Create mcp_agent.py

mcp_agent.py

@@ -0,0 +1,117 @@
+# mcp_agent.py
+import asyncio
+from typing import List, Dict, Optional
+from fastmcp import Client
+from fastmcp.client.transports import SSETransport
+from llama_index.core.workflow import Context
+from llama_index.tools.mcp import McpToolSpec
+from llama_index.core.agent.workflow import FunctionAgent, ToolCall
+from mcp.client.sse import sse_client as original_sse_client
+from mcp.client.session import ClientSession
+from datetime import timedelta
+from contextlib import asynccontextmanager
+
+class CustomMCPClient:
+    def __init__(self, url: str, headers: dict = None, timeout: int = 30):
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timeout
+
+    async def call_tool(self, tool_name: str, arguments: dict):
+        async with self._create_session() as session:
+            return await session.call_tool(tool_name, arguments)
+
+    async def list_tools(self):
+        async with self._create_session() as session:
+            return await session.list_tools()
+
+    @asynccontextmanager
+    async def _create_session(self):
+        async with original_sse_client(
+            url=self.url,
+            headers=self.headers,
+            timeout=self.timeout
+        ) as streams:
+            async with ClientSession(
+                *streams,
+                read_timeout_seconds=timedelta(seconds=self.timeout)
+            ) as session:
+                await session.initialize()
+                yield session
+
+class MCPAgent:
+    def __init__(self, sse_url: str, hf_token: str, llm):
+        self.sse_url = sse_url
+        self.hf_token = hf_token
+        self.llm = llm
+        self.agent = None
+        self.agent_context = None
+        
+        self.SYSTEM_PROMPT = """
+        Eres un **Asistente de IA especializado en llamadas a herramientas (Tool Calling)**.
+        Tu función es responder EXCLUSIVAMENTE usando las herramientas asignadas, sin usar conocimientos previos.
+        
+        ### 🔹 Herramientas Disponibles:
+        1. **query_space**: Consulta información técnica en espacios especializados (PRINCIPAL/PREFERENTE)
+        2. **list_spaces_names**: Lista espacios disponibles. (Contiene los 'space_name', útil para usar query_space)
+        3. **search_tavily**: Recuperador de información y contenido web que usa el motor de búsqueda tavily
+        
+        ### ⚠️ Reglas Estrictas:
+        1. **OBLIGATORIO**: Usar `query_space` para cualquier consulta técnica.
+        2. **PROHIBIDO**: Inventar respuestas o usar conocimientos no verificados.
+        3. **PRIORIDAD**: Siempre validar información con herramientas antes de responder.
+        4. **PRECISIÓN**: Si la herramienta falla, NO intentes adivinar la respuesta.
+        """
+    
+    async def initialize(self):
+        """Inicializa el cliente MCP y el agente"""
+        mcp_client = CustomMCPClient(
+            self.sse_url,
+            headers={"Authorization": f"Bearer {self.hf_token}"}
+        )
+        mcp_tool = McpToolSpec(client=mcp_client)
+        
+        tools = await mcp_tool.to_tool_list_async()
+        self.agent = FunctionAgent(
+            name="MCP_Agent",
+            description="Agente que utiliza herramientas MCP para responder.",
+            tools=tools,
+            llm=self.llm,
+            system_prompt=self.SYSTEM_PROMPT,
+        )
+        self.agent_context = Context(self.agent)
+        
+        return await mcp_client.list_tools()
+    
+    async def process_message(self, message: str) -> str:
+        """Procesa un mensaje del usuario y devuelve la respuesta del agente"""
+        if not self.agent:
+            await self.initialize()
+            
+        handler = self.agent.run(message, ctx=self.agent_context)
+        return str(await handler)
+    
+    async def stream_response(self, message: str):
+        """Genera la respuesta del agente en formato de streaming"""
+        if not self.agent:
+            await self.initialize()
+            
+        handler = self.agent.run(message, ctx=self.agent_context)
+        async for event in handler.stream_events():
+            if isinstance(event, ToolCall):
+                yield {
+                    "type": "tool_call",
+                    "tool_name": event.tool_name,
+                    "arguments": event.tool_kwargs
+                }
+            elif isinstance(event, ToolCallResult):
+                yield {
+                    "type": "tool_result",
+                    "tool_name": event.tool_name,
+                    "result": event.tool_output
+                }
+        
+        yield {
+            "type": "final_response",
+            "content": str(await handler)
+        }