Merge branch 'test'

agent_test.py

@@ -0,0 +1,240 @@
+import asyncio
+import os
+import json
+import base64
+from typing import List, Dict, Any, Union
+from contextlib import AsyncExitStack
+from io import BytesIO
+from PIL import Image
+import gradio as gr
+from gradio.components.chatbot import ChatMessage
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+from dotenv import load_dotenv
+from langchain_openai import ChatOpenAI
+
+load_dotenv()
+
+loop = asyncio.new_event_loop()
+asyncio.set_event_loop(loop)
+
+class MCPClientWrapper:
+    def __init__(self):
+        self.session = None
+        self.exit_stack = None
+        self.mistral = ChatOpenAI(model_name="mistralai/mistral-small", temperature=0.7, openai_api_key=os.getenv("OPENROUTER_API_KEY"))
+        self.tools = []
+    
+    def connect(self, server_path: str) -> str:
+        return loop.run_until_complete(self._connect(server_path))
+    
+    async def _connect(self, server_path: str) -> str:
+        if self.exit_stack:
+            await self.exit_stack.aclose()
+        
+        self.exit_stack = AsyncExitStack()
+        
+        is_python = server_path.endswith('.py')
+        command = "python" if is_python else "node"
+        
+        server_params = StdioServerParameters(
+            command=command,
+            args=[server_path],
+            env={"PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"}
+        )
+        
+        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
+        self.stdio, self.write = stdio_transport
+        
+        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
+        await self.session.initialize()
+        
+        response = await self.session.list_tools()
+        self.tools = [{ 
+            "name": tool.name,
+            "description": tool.description,
+            "input_schema": tool.inputSchema
+        } for tool in response.tools]
+        
+        tool_names = [tool["name"] for tool in self.tools]
+        return f"Connected to MCP server. Available tools: {', '.join(tool_names)}"
+    
+    def process_message(self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]) -> tuple:
+        if not self.session:
+            return history + [
+                {"role": "user", "content": message}, 
+                {"role": "assistant", "content": "Please connect to an MCP server first."}
+            ], gr.Textbox(value="")
+        
+        new_messages = loop.run_until_complete(self._process_query(message, history))
+        return history + [{"role": "user", "content": message}] + new_messages, gr.Textbox(value="")
+    
+    async def _process_query(self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]):
+        claude_messages = []
+        for msg in history:
+            if isinstance(msg, ChatMessage):
+                role, content = msg.role, msg.content
+            else:
+                role, content = msg.get("role"), msg.get("content")
+            
+            if role in ["user", "assistant", "system"]:
+                claude_messages.append({"role": role, "content": content})
+        
+        claude_messages.append({"role": "user", "content": message})
+        
+        response = self.mistral.messages.create(
+            model="claude-3-5-sonnet-20241022",
+            max_tokens=1000,
+            messages=claude_messages,
+            tools=self.tools
+        )
+
+        result_messages = []
+        
+        for content in response.content:
+            if content.type == 'text':
+                result_messages.append({
+                    "role": "assistant", 
+                    "content": content.text
+                })
+                
+            elif content.type == 'tool_use':
+                tool_name = content.name
+                tool_args = content.input
+                
+                result_messages.append({
+                    "role": "assistant",
+                    "content": f"I'll use the {tool_name} tool to help answer your question.",
+                    "metadata": {
+                        "title": f"Using tool: {tool_name}",
+                        "log": f"Parameters: {json.dumps(tool_args, ensure_ascii=True)}",
+                        "status": "pending",
+                        "id": f"tool_call_{tool_name}"
+                    }
+                })
+                
+                result_messages.append({
+                    "role": "assistant",
+                    "content": "```json\n" + json.dumps(tool_args, indent=2, ensure_ascii=True) + "\n```",
+                    "metadata": {
+                        "parent_id": f"tool_call_{tool_name}",
+                        "id": f"params_{tool_name}",
+                        "title": "Tool Parameters"
+                    }
+                })
+                
+                result = await self.session.call_tool(tool_name, tool_args)
+                
+                if result_messages and "metadata" in result_messages[-2]:
+                    result_messages[-2]["metadata"]["status"] = "done"
+                
+                result_messages.append({
+                    "role": "assistant",
+                    "content": "Here are the results from the tool:",
+                    "metadata": {
+                        "title": f"Tool Result for {tool_name}",
+                        "status": "done",
+                        "id": f"result_{tool_name}"
+                    }
+                })
+                
+                result_content = result.content
+                if isinstance(result_content, list):
+                    result_content = "\n".join(str(item) for item in result_content)
+                
+                try:
+                    result_json = json.loads(result_content)
+                    if isinstance(result_json, dict) and "type" in result_json:
+                        if result_json["type"] == "image" and "url" in result_json:
+                            result_messages.append({
+                                "role": "assistant",
+                                "content": {"path": result_json["url"], "alt_text": result_json.get("message", "Generated image")},
+                                "metadata": {
+                                    "parent_id": f"result_{tool_name}",
+                                    "id": f"image_{tool_name}",
+                                    "title": "Generated Image"
+                                }
+                            })
+                        else:
+                            result_messages.append({
+                                "role": "assistant",
+                                "content": "```\n" + result_content + "\n```",
+                                "metadata": {
+                                    "parent_id": f"result_{tool_name}",
+                                    "id": f"raw_result_{tool_name}",
+                                    "title": "Raw Output"
+                                }
+                            })
+                except:
+                    result_messages.append({
+                        "role": "assistant",
+                        "content": "```\n" + result_content + "\n```",
+                        "metadata": {
+                            "parent_id": f"result_{tool_name}",
+                            "id": f"raw_result_{tool_name}",
+                            "title": "Raw Output"
+                        }
+                    })
+                
+                claude_messages.append({"role": "user", "content": f"Tool result for {tool_name}: {result_content}"})
+                next_response = self.mistral.messages.create(
+                    model="claude-3-5-sonnet-20241022",
+                    max_tokens=1000,
+                    messages=claude_messages,
+                )
+                
+                if next_response.content and next_response.content[0].type == 'text':
+                    result_messages.append({
+                        "role": "assistant",
+                        "content": next_response.content[0].text
+                    })
+
+        return result_messages
+
+client = MCPClientWrapper()
+
+def gradio_interface():
+    with gr.Blocks(title="MCP Weather Client") as demo:
+        gr.Markdown("# MCP Weather Assistant")
+        gr.Markdown("Connect to your MCP weather server and chat with the assistant")
+        
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=4):
+                server_path = gr.Textbox(
+                    label="Server Script Path",
+                    placeholder="Enter path to server script (e.g., weather.py)",
+                    value="gradio_mcp_server.py"
+                )
+            with gr.Column(scale=1):
+                connect_btn = gr.Button("Connect")
+        
+        status = gr.Textbox(label="Connection Status", interactive=False)
+        
+        chatbot = gr.Chatbot(
+            value=[], 
+            height=500,
+            type="messages",
+            show_copy_button=True,
+            avatar_images=("👤", "🤖")
+        )
+        
+        with gr.Row(equal_height=True):
+            msg = gr.Textbox(
+                label="Your Question",
+                placeholder="Ask about weather or alerts (e.g., What's the weather in New York?)",
+                scale=4
+            )
+            clear_btn = gr.Button("Clear Chat", scale=1)
+        
+        connect_btn.click(client.connect, inputs=server_path, outputs=status)
+        msg.submit(client.process_message, [msg, chatbot], [chatbot, msg])
+        clear_btn.click(lambda: [], None, chatbot)
+        
+    return demo
+
+if __name__ == "__main__":
+    if not os.getenv("OPENROUTER_API_KEY"):
+        print("Warning: OPENROUTER_API_KEY not found in environment. Please set it in your .env file.")
+    
+    interface = gradio_interface()
+    interface.launch(debug=True)

gradio_interface/app.py

@@ -1,7 +1,204 @@
+import os
 import gradio as gr
+from os import getenv
+import base64
+from io import BytesIO
+from dotenv import load_dotenv
+import requests
+import socket
+import logging
+import json
 
-def greet(name):
-    return "Hello " + name + "!!"
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.callbacks import StreamingStdOutCallbackHandler
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# Load environment
+dotenv_path = os.path.join(os.path.dirname(__file__), '.env')
+load_dotenv(dotenv_path=dotenv_path)
+
+# Connectivity test
+def test_connectivity(url="https://openrouter.helicone.ai/api/v1"):
+    try:
+        return requests.get(url, timeout=5).status_code == 200
+    except (requests.RequestException, socket.error):
+        return False
+
+# Helper to make direct API calls to OpenRouter when LangChain fails
+def direct_api_call(messages, api_key, base_url):
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {api_key}",
+        "HTTP-Referer": "https://your-app-domain.com",  # Add your domain
+        "X-Title": "Image Analysis App"
+    }
+    
+    if getenv("HELICONE_API_KEY"):
+        headers["Helicone-Auth"] = f"Bearer {getenv('HELICONE_API_KEY')}"
+    
+    payload = {
+        "model": "google/gemini-flash-1.5",
+        "messages": messages,
+        "stream": False,
+    }
+    
+    try:
+        response = requests.post(
+            f"{base_url}/chat/completions",
+            headers=headers,
+            json=payload,
+            timeout=30
+        )
+        response.raise_for_status()
+        return response.json()["choices"][0]["message"]["content"]
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Initialize LLM with streaming and retry logic
+def init_llm():
+    if not test_connectivity():
+        raise RuntimeError("No hay conexión a OpenRouter. Verifica red y claves.")
+    return ChatOpenAI(
+        openai_api_key=getenv("OPENROUTER_API_KEY"),
+        openai_api_base=getenv("OPENROUTER_BASE_URL"),
+        model_name="google/gemini-flash-1.5",
+        streaming=True,
+        callbacks=[StreamingStdOutCallbackHandler()],
+        model_kwargs={
+            "extra_headers": {"Helicone-Auth": f"Bearer {getenv('HELICONE_API_KEY')}"}
+        },
+    )
+
+# Try to initialize LLM but handle failures gracefully
+try:
+    llm = init_llm()
+except Exception as e:
+    llm = None
+
+# Helpers
+def encode_image_to_base64(pil_image):
+    buffer = BytesIO()
+    pil_image.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode()
+
+# Core logic
+def generate_response(message, chat_history, image):
+    # Convert chat history to standard format
+    formatted_history = []
+    for msg in chat_history:
+        role = msg.get('role')
+        content = msg.get('content')
+        if role == 'user':
+            formatted_history.append({"role": "user", "content": content})
+        else:
+            formatted_history.append({"role": "assistant", "content": content})
+    
+    # Prepare system message
+    system_msg = {"role": "system", "content": "You are an expert image analysis assistant. Answer succinctly."}
+    
+    # Prepare the latest message with image if provided
+    if image:
+        base64_image = encode_image_to_base64(image)
+        
+        # Format for direct API call (OpenRouter/OpenAI format)
+        api_messages = [system_msg] + formatted_history + [{
+            "role": "user",
+            "content": [
+                {"type": "text", "text": message},
+                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_image}"}}
+            ]
+        }]
+        
+        # For LangChain format
+        content_for_langchain = [
+            {"type": "text", "text": message},
+            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_image}"}}
+        ]
+    else:
+        api_messages = [system_msg] + formatted_history + [{"role": "user", "content": message}]
+        content_for_langchain = message
+
+    # Build LangChain messages
+    lc_messages = [HumanMessage(content="You are an expert image analysis assistant. Answer succinctly.")]
+    for msg in chat_history:
+        role = msg.get('role')
+        content = msg.get('content')
+        if role == 'user':
+            lc_messages.append(HumanMessage(content=content))
+        else:
+            lc_messages.append(AIMessage(content=content))
+    
+    lc_messages.append(HumanMessage(content=content_for_langchain))
+
+    try:
+        # First try with LangChain
+        if llm:
+            try:
+                try:
+                    stream_iter = llm.stream(lc_messages)
+                    partial = ""
+                    for chunk in stream_iter:
+                        if chunk is None:
+                            continue
+                        content = getattr(chunk, 'content', None)
+                        if content is None:
+                            continue
+                        partial += content
+                        yield partial
+                    
+                    # If we got this far, streaming worked
+                    return
+                except Exception as e:
+                    print(f"Streaming failed: {e}. Falling back to non-streaming mode")
+                    
+                # Try non-streaming
+                try:
+                    response = llm.invoke(lc_messages)
+                    yield response.content
+                    return
+                except Exception as e:
+                    raise e
+            except Exception as e:
+                raise e
+        
+        response_text = direct_api_call(
+            api_messages,
+            getenv("OPENROUTER_API_KEY"),
+            getenv("OPENROUTER_BASE_URL")
+        )
+        yield response_text
+        
+    except Exception as e:
+        import traceback
+        error_trace = traceback.format_exc()
+        yield f"⚠️ Error al generar respuesta: {str(e)}. Intenta más tarde."
+
+# Gradio interface
+def process_message(message, chat_history, image):
+    if chat_history is None:
+        chat_history = []
+    if image is None:
+        chat_history.append({'role':'assistant','content':'Por favor sube una imagen.'})
+        return "", chat_history
+    chat_history.append({'role':'user','content':message})
+    chat_history.append({'role':'assistant','content':'⏳ Procesando...'})
+    yield "", chat_history
+    for chunk in generate_response(message, chat_history, image):
+        chat_history[-1]['content'] = chunk
+        yield "", chat_history
+    return "", chat_history
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(type='messages', height=600)
+            msg = gr.Textbox(label="Mensaje", placeholder="Escribe tu pregunta...")
+            clear = gr.ClearButton([msg, chatbot])
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="pil", label="Sube Imagen")
+            info = gr.Textbox(label="Info Imagen", interactive=False)
+
+    msg.submit(process_message, [msg, chatbot, image_input], [msg, chatbot])
+    image_input.change(lambda img: f"Tamaño: {img.size}" if img else "Sin imagen.", [image_input], [info])
+
+demo.launch()