geronimo-pericoli commited on
Commit
1aaf234
verified
1 Parent(s): d4bc3ee

Delete mcp_agent.py

Browse files
Files changed (1) hide show
  1. mcp_agent.py +0 -117
mcp_agent.py DELETED
@@ -1,117 +0,0 @@
1
- # mcp_agent.py
2
- import asyncio
3
- from typing import List, Dict, Optional, Generator, AsyncGenerator
4
- from fastmcp import Client
5
- from fastmcp.client.transports import SSETransport
6
- from llama_index.core.workflow import Context
7
- from llama_index.tools.mcp import McpToolSpec
8
- from llama_index.core.agent.workflow import FunctionAgent, ToolCall, ToolCallResult
9
- from mcp.client.sse import sse_client as original_sse_client
10
- from mcp.client.session import ClientSession
11
- from datetime import timedelta
12
- from contextlib import asynccontextmanager
13
-
14
- class CustomMCPClient:
15
- def __init__(self, url: str, headers: dict = None, timeout: int = 30):
16
- self.url = url
17
- self.headers = headers or {}
18
- self.timeout = timeout
19
-
20
- async def call_tool(self, tool_name: str, arguments: dict):
21
- async with self._create_session() as session:
22
- return await session.call_tool(tool_name, arguments)
23
-
24
- async def list_tools(self):
25
- async with self._create_session() as session:
26
- return await session.list_tools()
27
-
28
- @asynccontextmanager
29
- async def _create_session(self):
30
- async with original_sse_client(
31
- url=self.url,
32
- headers=self.headers,
33
- timeout=self.timeout
34
- ) as streams:
35
- async with ClientSession(
36
- *streams,
37
- read_timeout_seconds=timedelta(seconds=self.timeout)
38
- ) as session:
39
- await session.initialize()
40
- yield session
41
-
42
- class MCPAgent:
43
- def __init__(self, sse_url: str, hf_token: str, llm):
44
- self.sse_url = sse_url
45
- self.hf_token = hf_token
46
- self.llm = llm
47
- self.agent = None
48
- self.agent_context = None
49
-
50
- self.SYSTEM_PROMPT = """
51
- Eres un **Asistente de IA especializado en llamadas a herramientas (Tool Calling)**.
52
- Tu funci贸n es responder EXCLUSIVAMENTE usando las herramientas asignadas, sin usar conocimientos previos.
53
-
54
- ### 馃敼 Herramientas Disponibles:
55
- 1. **query_space**: Consulta informaci贸n t茅cnica en espacios especializados (PRINCIPAL/PREFERENTE)
56
- 2. **list_spaces_names**: Lista espacios disponibles. (Contiene los 'space_name', 煤til para usar query_space)
57
- 3. **search_tavily**: Recuperador de informaci贸n y contenido web que usa el motor de b煤squeda tavily
58
-
59
- ### 鈿狅笍 Reglas Estrictas:
60
- 1. **OBLIGATORIO**: Usar `query_space` para cualquier consulta t茅cnica.
61
- 2. **PROHIBIDO**: Inventar respuestas o usar conocimientos no verificados.
62
- 3. **PRIORIDAD**: Siempre validar informaci贸n con herramientas antes de responder.
63
- 4. **PRECISI脫N**: Si la herramienta falla, NO intentes adivinar la respuesta.
64
- """
65
-
66
- async def initialize(self):
67
- """Inicializa el cliente MCP y el agente"""
68
- mcp_client = CustomMCPClient(
69
- self.sse_url,
70
- headers={"Authorization": f"Bearer {self.hf_token}"}
71
- )
72
- mcp_tool = McpToolSpec(client=mcp_client)
73
-
74
- tools = await mcp_tool.to_tool_list_async()
75
- self.agent = FunctionAgent(
76
- name="MCP_Agent",
77
- description="Agente que utiliza herramientas MCP para responder.",
78
- tools=tools,
79
- llm=self.llm,
80
- system_prompt=self.SYSTEM_PROMPT,
81
- )
82
- self.agent_context = Context(self.agent)
83
-
84
- return await mcp_client.list_tools()
85
-
86
- async def process_message(self, message: str) -> str:
87
- """Procesa un mensaje del usuario y devuelve la respuesta del agente"""
88
- if not self.agent:
89
- await self.initialize()
90
-
91
- handler = self.agent.run(message, ctx=self.agent_context)
92
- return str(await handler)
93
-
94
- async def stream_response(self, message: str) -> AsyncGenerator[Dict, None]:
95
- """Genera la respuesta del agente en formato de streaming"""
96
- if not self.agent:
97
- await self.initialize()
98
-
99
- handler = self.agent.run(message, ctx=self.agent_context)
100
- async for event in handler.stream_events():
101
- if isinstance(event, ToolCall):
102
- yield {
103
- "type": "tool_call",
104
- "tool_name": event.tool_name,
105
- "arguments": event.tool_kwargs
106
- }
107
- elif isinstance(event, ToolCallResult):
108
- yield {
109
- "type": "tool_result",
110
- "tool_name": event.tool_name,
111
- "result": event.tool_output
112
- }
113
-
114
- yield {
115
- "type": "final_response",
116
- "content": str(await handler)
117
- }