698 lines
29 KiB
Python
698 lines
29 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
MCP client that uses Ollama for inference and LangChain create_agent with
|
||
runtime-registered MCP tools (see https://docs.langchain.com/oss/python/langchain/agents#runtime-tool-registration).
|
||
"""
|
||
|
||
import json
|
||
import sys
|
||
import os
|
||
import asyncio
|
||
from pathlib import Path
|
||
from typing import Optional, Dict, Any, List, Callable, Awaitable
|
||
|
||
import requests
|
||
from fastmcp import Client as FastMcpClient
|
||
from ollama import ResponseError as OllamaResponseError
|
||
from pydantic import BaseModel, ConfigDict, Field, create_model
|
||
|
||
# LangChain agent and middleware
|
||
try:
|
||
from langchain.agents import create_agent
|
||
from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse, ToolCallRequest
|
||
from langchain_core.tools import StructuredTool, tool
|
||
from langchain_ollama import ChatOllama
|
||
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, ToolMessage
|
||
except ImportError as e:
|
||
print(f"Missing dependency: {e}. Install with: pip install langchain langgraph langchain-community langchain-ollama", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
|
||
@tool
|
||
def getTime() -> str:
|
||
"""Get the current time in ISO format."""
|
||
from datetime import datetime
|
||
return datetime.now().isoformat()
|
||
|
||
|
||
@tool
|
||
def countWords(text: str) -> int:
|
||
"""Count the number of words in a text."""
|
||
return len(text.split())
|
||
|
||
|
||
def loadMcpConfig(configPath: Optional[str] = None) -> Dict[str, str]:
|
||
"""Load MCP server URLs from mcp.json. Returns dict serverName -> url."""
|
||
if configPath is None:
|
||
# Default: mcpServer/mcp.json relative to project root or cwd
|
||
base = Path(__file__).resolve().parent.parent
|
||
configPath = str(base / "mcpServer" / "mcp.json")
|
||
path = Path(configPath)
|
||
if not path.exists():
|
||
return {}
|
||
try:
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
except (json.JSONDecodeError, OSError) as e:
|
||
print(f"Warning: Could not load MCP config from {path}: {e}", file=sys.stderr)
|
||
return {}
|
||
servers = data.get("mcpServers") or data.get("mcp_servers") or {}
|
||
return {name: info.get("url", "") for name, info in servers.items() if isinstance(info, dict) and info.get("url")}
|
||
|
||
|
||
class GenericToolArgs(BaseModel):
|
||
"""Accept any keyword arguments for MCP tool calls (fallback when schema is missing)."""
|
||
model_config = ConfigDict(extra="allow")
|
||
|
||
|
||
def _jsonSchemaTypeToPython(jsonType: str) -> type:
|
||
"""Map JSON schema type to Python type."""
|
||
return {"string": str, "integer": int, "number": float, "boolean": bool, "array": list, "object": dict}.get(jsonType, str)
|
||
|
||
|
||
def _defaultForJsonType(jsonType: str) -> Any:
|
||
"""Sensible default for optional MCP params so server does not receive null."""
|
||
return {"string": "", "integer": 0, "number": 0.0, "boolean": False, "array": [], "object": {}}.get(jsonType, "")
|
||
|
||
|
||
def _defaultsFromInputSchema(inputSchema: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""Build default values for all params so we never send null to the MCP server (LLM may omit required params)."""
|
||
if not inputSchema:
|
||
return {}
|
||
properties = inputSchema.get("properties") or {}
|
||
out: Dict[str, Any] = {}
|
||
for name, spec in properties.items():
|
||
if not isinstance(spec, dict):
|
||
continue
|
||
if "default" in spec:
|
||
out[name] = spec["default"]
|
||
else:
|
||
out[name] = _defaultForJsonType(spec.get("type", "string"))
|
||
return out
|
||
|
||
|
||
def buildArgsSchemaFromMcpInputSchema(toolName: str, inputSchema: Dict[str, Any]) -> type[BaseModel]:
|
||
"""Build a Pydantic model from MCP tool inputSchema so the LLM gets exact parameter names (path, content, etc.)."""
|
||
if not inputSchema:
|
||
return GenericToolArgs
|
||
properties = inputSchema.get("properties") or {}
|
||
required = set(inputSchema.get("required") or [])
|
||
if not properties:
|
||
return GenericToolArgs
|
||
fields: Dict[str, Any] = {}
|
||
for name, spec in properties.items():
|
||
if not isinstance(spec, dict):
|
||
continue
|
||
desc = spec.get("description", "")
|
||
jsonType = spec.get("type", "string")
|
||
pyType = _jsonSchemaTypeToPython(jsonType)
|
||
if name in required:
|
||
fields[name] = (pyType, Field(..., description=desc))
|
||
else:
|
||
fields[name] = (Optional[pyType], Field(None, description=desc))
|
||
if not fields:
|
||
return GenericToolArgs
|
||
return create_model(f"McpArgs_{toolName}", **fields)
|
||
|
||
|
||
class OllamaClient:
|
||
"""Client for interacting with Ollama API."""
|
||
|
||
def __init__(self, baseUrl: str = "http://localhost:11434", model: str = "gpt-oss:20b"):
|
||
self.baseUrl = baseUrl
|
||
self.model = model
|
||
|
||
def listModels(self) -> List[str]:
|
||
"""List available Ollama models."""
|
||
try:
|
||
response = requests.get(f"{self.baseUrl}/api/tags", timeout=10)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
return [model["name"] for model in data.get("models", [])]
|
||
except requests.RequestException as e:
|
||
print(f"Error listing models: {e}", file=sys.stderr)
|
||
return []
|
||
|
||
def chat(self, messages: List[Dict[str, str]], options: Optional[Dict[str, Any]] = None) -> str:
|
||
"""Send chat messages to Ollama and get response."""
|
||
payload = {
|
||
"model": self.model,
|
||
"messages": messages,
|
||
"stream": False,
|
||
}
|
||
|
||
if options:
|
||
payload["options"] = options
|
||
|
||
try:
|
||
response = requests.post(
|
||
f"{self.baseUrl}/api/chat",
|
||
json=payload,
|
||
timeout=60*60*60
|
||
)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
return data.get("message", {}).get("content", "")
|
||
except requests.RequestException as e:
|
||
print(f"Error in chat request: {e}", file=sys.stderr)
|
||
raise
|
||
|
||
def generate(self, prompt: str, options: Optional[Dict[str, Any]] = None) -> str:
|
||
"""Generate text from a prompt using Ollama."""
|
||
payload = {
|
||
"model": self.model,
|
||
"prompt": prompt,
|
||
"stream": False,
|
||
}
|
||
|
||
if options:
|
||
payload["options"] = options
|
||
|
||
try:
|
||
response = requests.post(
|
||
f"{self.baseUrl}/api/generate",
|
||
json=payload,
|
||
timeout=120
|
||
)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
return data.get("response", "")
|
||
except requests.RequestException as e:
|
||
print(f"Error in generate request: {e}", file=sys.stderr)
|
||
raise
|
||
|
||
def checkHealth(self) -> bool:
|
||
"""Check if Ollama server is accessible."""
|
||
try:
|
||
response = requests.get(f"{self.baseUrl}/api/tags", timeout=5)
|
||
return response.status_code == 200
|
||
except requests.RequestException:
|
||
return False
|
||
|
||
|
||
class McpServerWrapper:
|
||
"""Wrapper around FastMCP Client for easier use."""
|
||
|
||
def __init__(self, httpUrl: str, headers: Optional[Dict[str, str]] = None):
|
||
self.httpUrl = httpUrl.rstrip("/")
|
||
self.headers = headers or {}
|
||
self.client: Optional[FastMcpClient] = None
|
||
self.serverTools: List[Dict[str, Any]] = []
|
||
|
||
async def connect(self) -> bool:
|
||
"""Connect and initialize with MCP server via HTTP."""
|
||
try:
|
||
# FastMcpClient doesn't support headers parameter directly
|
||
# Headers would need to be passed via custom transport or auth
|
||
# For now, we initialize without headers
|
||
self.client = FastMcpClient(self.httpUrl)
|
||
await self.client.__aenter__()
|
||
# Load tools after connection
|
||
tools = await self.listServerTools()
|
||
return True
|
||
except Exception as e:
|
||
print(f"Error connecting to MCP server: {e}", file=sys.stderr)
|
||
return False
|
||
|
||
async def disconnect(self):
|
||
"""Disconnect from MCP server."""
|
||
if self.client:
|
||
await self.client.__aexit__(None, None, None)
|
||
self.client = None
|
||
|
||
async def listServerTools(self) -> List[Dict[str, Any]]:
|
||
"""List tools available from MCP server."""
|
||
if not self.client:
|
||
return []
|
||
|
||
try:
|
||
tools = await self.client.list_tools()
|
||
self.serverTools = tools
|
||
return tools
|
||
except Exception as e:
|
||
print(f"Error listing tools: {e}", file=sys.stderr)
|
||
return []
|
||
|
||
async def callServerTool(self, name: str, arguments: Dict[str, Any]) -> Any:
|
||
"""Call a tool on the MCP server."""
|
||
if not self.client:
|
||
raise RuntimeError("Not connected to MCP server")
|
||
|
||
try:
|
||
result = await self.client.call_tool(name, arguments)
|
||
# FastMCP call_tool returns a result object with .content
|
||
if hasattr(result, 'content'):
|
||
content = result.content
|
||
# If content is a list, return it as is (will be serialized later)
|
||
if isinstance(content, list):
|
||
return content
|
||
return content
|
||
elif isinstance(result, list):
|
||
# Handle list of results
|
||
if len(result) > 0:
|
||
# Extract content from each item if it exists
|
||
contents = []
|
||
for item in result:
|
||
if hasattr(item, 'content'):
|
||
contents.append(item.content)
|
||
else:
|
||
contents.append(item)
|
||
return contents if len(contents) > 1 else contents[0] if contents else None
|
||
return result
|
||
return result
|
||
except Exception as e:
|
||
raise RuntimeError(f"Tool call failed: {str(e)}")
|
||
|
||
async def listServerResources(self) -> List[Dict[str, Any]]:
|
||
"""List resources available from MCP server."""
|
||
if not self.client:
|
||
return []
|
||
|
||
try:
|
||
resources = await self.client.list_resources()
|
||
return resources
|
||
except Exception as e:
|
||
print(f"Error listing resources: {e}", file=sys.stderr)
|
||
return []
|
||
|
||
|
||
def _serializeToolResult(result: Any) -> Any:
|
||
"""Serialize tool result to JSON-serializable format."""
|
||
if hasattr(result, "text"):
|
||
return result.text
|
||
if hasattr(result, "content"):
|
||
content = result.content
|
||
if hasattr(content, "text"):
|
||
return content.text
|
||
return content
|
||
if isinstance(result, list):
|
||
return [_serializeToolResult(item) for item in result]
|
||
if isinstance(result, dict):
|
||
return {k: _serializeToolResult(v) for k, v in result.items()}
|
||
return result
|
||
|
||
|
||
def _makeMcpToolCoroutine(
|
||
toolName: str,
|
||
server: McpServerWrapper,
|
||
defaultArgs: Dict[str, Any],
|
||
toolTimeout: Optional[float] = None,
|
||
) -> Callable[..., Awaitable[Any]]:
|
||
async def _invoke(**kwargs: Any) -> Any:
|
||
merged = {**defaultArgs, **kwargs}
|
||
# Strip None values - MCP server Zod schemas often reject null for optional params (expect number | undefined, not number | null)
|
||
merged = {k: v for k, v in merged.items() if v is not None}
|
||
try:
|
||
if toolTimeout is not None and toolTimeout > 0:
|
||
result = await asyncio.wait_for(
|
||
server.callServerTool(toolName, merged),
|
||
timeout=toolTimeout,
|
||
)
|
||
else:
|
||
result = await server.callServerTool(toolName, merged)
|
||
except asyncio.TimeoutError:
|
||
return (
|
||
f"[Tool timeout] '{toolName}' exceeded {toolTimeout}s. "
|
||
"The operation may have hung (e.g. command not found, subprocess blocking). "
|
||
"Try an alternative (e.g. 'python' instead of 'python3') or increase --tool-timeout."
|
||
)
|
||
return _serializeToolResult(result)
|
||
return _invoke
|
||
|
||
|
||
async def buildMcpLangChainTools(
|
||
mcpServers: List[McpServerWrapper],
|
||
toolTimeout: Optional[float] = None,
|
||
) -> List[StructuredTool]:
|
||
"""Build LangChain StructuredTools from connected MCP servers (runtime tool registration)."""
|
||
tools: List[StructuredTool] = []
|
||
for server in mcpServers:
|
||
rawTools = await server.listServerTools()
|
||
for raw in rawTools:
|
||
name = getattr(raw, "name", None) or (raw.get("name") if isinstance(raw, dict) else None)
|
||
description = getattr(raw, "description", None) or (raw.get("description", "") if isinstance(raw, dict) else "")
|
||
inputSchema = getattr(raw, "inputSchema", None) or getattr(raw, "input_schema", None) or (raw.get("inputSchema") or raw.get("input_schema") if isinstance(raw, dict) else None)
|
||
if not name:
|
||
continue
|
||
description = description or f"MCP tool: {name}"
|
||
schemaDict = inputSchema or {}
|
||
argsSchema = buildArgsSchemaFromMcpInputSchema(name, schemaDict)
|
||
defaultArgs = _defaultsFromInputSchema(schemaDict)
|
||
tool = StructuredTool.from_function(
|
||
name=name,
|
||
description=description,
|
||
args_schema=argsSchema,
|
||
coroutine=_makeMcpToolCoroutine(name, server, defaultArgs, toolTimeout),
|
||
)
|
||
tools.append(tool)
|
||
return tools
|
||
|
||
|
||
class LogToolCallsMiddleware(AgentMiddleware):
|
||
"""Middleware that logs every tool call (name and args)."""
|
||
|
||
def wrap_tool_call(self, request: ToolCallRequest, handler: Callable):
|
||
_logToolCallRequest(request)
|
||
return handler(request)
|
||
|
||
async def awrap_tool_call(self, request: ToolCallRequest, handler: Callable):
|
||
_logToolCallRequest(request)
|
||
return await handler(request)
|
||
|
||
|
||
def _extractTextFromAIMessageContent(content: Any) -> str:
|
||
"""Extract plain text from AIMessage.content (str or list of content blocks)."""
|
||
if content is None:
|
||
return ""
|
||
if isinstance(content, str):
|
||
return content.strip()
|
||
if isinstance(content, list):
|
||
parts: List[str] = []
|
||
for block in content:
|
||
if isinstance(block, dict) and "text" in block:
|
||
parts.append(str(block["text"]))
|
||
elif isinstance(block, str):
|
||
parts.append(block)
|
||
return "\n".join(parts).strip() if parts else ""
|
||
return str(content).strip()
|
||
|
||
|
||
def _extractFinalResponse(result: Dict[str, Any]) -> str:
|
||
"""Extract the final assistant text from agent result; handle recursion limit / no final message."""
|
||
messages = result.get("messages") or []
|
||
for msg in reversed(messages):
|
||
if isinstance(msg, AIMessage) and hasattr(msg, "content"):
|
||
text = _extractTextFromAIMessageContent(msg.content)
|
||
if text:
|
||
return text
|
||
return (
|
||
"Agent stopped without a final text response (e.g. hit step limit after tool calls). "
|
||
"Try again or increase --recursion-limit."
|
||
)
|
||
|
||
|
||
def _logToolCallRequest(request: ToolCallRequest) -> None:
|
||
tc = request.tool_call
|
||
name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", None)
|
||
args = tc.get("args", tc.get("arguments", {})) if isinstance(tc, dict) else getattr(tc, "args", getattr(tc, "arguments", {}))
|
||
argsStr = json.dumps(args, ensure_ascii=False)
|
||
if len(argsStr) > 500:
|
||
argsStr = argsStr[:497] + "..."
|
||
print(f"[Tool Call] {name} args={argsStr}", file=sys.stderr)
|
||
|
||
|
||
class McpToolsMiddleware(AgentMiddleware):
|
||
"""Middleware that adds MCP tools at runtime and handles their execution (runtime tool registration)."""
|
||
|
||
def __init__(self, mcpTools: List[StructuredTool], staticToolNames: Optional[List[str]] = None):
|
||
self.mcpTools = mcpTools
|
||
self.mcpToolsByName = {t.name: t for t in mcpTools}
|
||
staticNames = set(staticToolNames or [])
|
||
self.validToolNames = staticNames | set(self.mcpToolsByName.keys())
|
||
|
||
def wrap_model_call(self, request: ModelRequest, handler: Callable) -> ModelResponse:
|
||
updated = request.override(tools=[*request.tools, *self.mcpTools])
|
||
return handler(updated)
|
||
|
||
async def awrap_model_call(self, request: ModelRequest, handler: Callable):
|
||
updated = request.override(tools=[*request.tools, *self.mcpTools])
|
||
return await handler(updated)
|
||
|
||
def _toolExists(self, name: Optional[str]) -> bool:
|
||
return bool(name and name in self.validToolNames)
|
||
|
||
def _unknownToolErrorToolMessage(self, request: ToolCallRequest, name: str) -> ToolMessage:
|
||
available = ", ".join(sorted(self.validToolNames))
|
||
content = (
|
||
f"[Error] Tool '{name}' does not exist. "
|
||
f"Only the following tools are available: {available}. "
|
||
"Do not call tools that are not in this list."
|
||
)
|
||
tc = request.tool_call
|
||
toolCallId = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)
|
||
return ToolMessage(
|
||
content=content,
|
||
tool_call_id=toolCallId or "unknown",
|
||
name=name or "unknown",
|
||
status="error",
|
||
)
|
||
|
||
def wrap_tool_call(self, request: ToolCallRequest, handler: Callable):
|
||
name = request.tool_call.get("name") if isinstance(request.tool_call, dict) else getattr(request.tool_call, "name", None)
|
||
if not self._toolExists(name):
|
||
return self._unknownToolErrorToolMessage(request, name or "<unknown>")
|
||
if name and name in self.mcpToolsByName:
|
||
return handler(request.override(tool=self.mcpToolsByName[name]))
|
||
return handler(request)
|
||
|
||
async def awrap_tool_call(self, request: ToolCallRequest, handler: Callable):
|
||
name = request.tool_call.get("name") if isinstance(request.tool_call, dict) else getattr(request.tool_call, "name", None)
|
||
if not self._toolExists(name):
|
||
return self._unknownToolErrorToolMessage(request, name or "<unknown>")
|
||
if name and name in self.mcpToolsByName:
|
||
return await handler(request.override(tool=self.mcpToolsByName[name]))
|
||
return await handler(request)
|
||
|
||
''' TODO Use this if you want sequential thinking
|
||
SYSTEM_PROMPT = """
|
||
ROLE:
|
||
Sei un esperto Analista di Cybersecurity specializzato in CTF (Capture The Flag) e analisi di vulnerabilità. Operi in un ambiente Linux sandbox dove la tua unica area di lavoro è la directory /tmp.
|
||
|
||
WORKSPACE CONSTRAINT: IL "SINGLE SOURCE OF TRUTH"
|
||
- Obbligo Assoluto: Tutte le operazioni di lettura, scrittura, download e analisi devono avvenire esclusivamente all'interno di /tmp.
|
||
- Percorsi: Ogni file deve essere referenziato con il percorso assoluto (es. /tmp/binary.bin). Non usare mai directory come ~/, /home o altre al di fuori di /tmp.
|
||
- Condivisione: Ricorda che /tmp è montata su tutti i container MCP (fetch, filesystem, ecc.). Se scarichi un file con fetch in /tmp, il tool filesystem lo troverà immediatamente lì.
|
||
|
||
TOOLSET & WORKFLOW:
|
||
Utilizza i tuoi tool secondo questa logica:
|
||
1. sequentialthinking (Pianificazione): Usa questo tool PRIMA di ogni azione complessa. Suddividi la sfida in step logici (es. 1. Download, 2. Analisi Header, 3. Estrazione Flag). Ti aiuta a non perdere il filo durante task lunghi.
|
||
2. fetch (Ingestion): Usalo per recuperare binari, exploit o dati remoti. Salva l'output sempre in /tmp.
|
||
3. filesystem (Manipolazione): Usalo per ispezionare i file scaricati, creare script di exploit o leggere file di log e flag direttamente in /tmp.
|
||
4. memory (Stato): Utilizza questo tool per memorizzare scoperte chiave, indirizzi di memoria, offset o password trovate durante la sfida. Ti serve per mantenere il contesto tra diverse fasi del ragionamento.
|
||
|
||
METODOLOGIA DI ANALISI:
|
||
- Ipotesi e Test: Prima di agire, formula un'ipotesi basata sui dati presenti in /tmp.
|
||
- Verifica Intermedia: Dopo ogni comando o modifica ai file, verifica il risultato usando il tool filesystem. Non dare mai per scontato che un'operazione sia riuscita senza controllare.
|
||
- Pulizia Mentale: Se una strategia fallisce, usa sequentialthinking per rivedere il piano e aggiorna il tool memory con il motivo del fallimento per non ripetere lo stesso errore.
|
||
|
||
REGOLE DI COMUNICAZIONE:
|
||
- Sii estremamente tecnico, sintetico e preciso.
|
||
- Se un file non è presente in /tmp, non provare a indovinarne il contenuto; usa fetch per ottenerlo o filesystem per cercarlo.
|
||
- Rispondi con l'output delle tue analisi e l'eventuale flag trovata nel formato richiesto dalla sfida.
|
||
"""
|
||
'''
|
||
|
||
SYSTEM_PROMPT = "ROLE:\nSei un esperto Analista di Cybersecurity specializzato in CTF (Capture The Flag) e analisi di vulnerabilità. Operi in un ambiente Linux sandbox dove la tua unica area di lavoro è la directory /tmp.\n\nWORKSPACE CONSTRAINT: IL \"SINGLE SOURCE OF TRUTH\"\n- Obbligo Assoluto: Tutte le operazioni di lettura, scrittura e analisi devono avvenire esclusivamente all'interno di /tmp.\n- Percorsi: Ogni file deve essere referenziato con il percorso assoluto (es. /tmp/binary.bin). Non usare mai directory esterne a /tmp.\n- Condivisione: /tmp è montata su tutti i container MCP. I file creati o modificati da un tool sono immediatamente visibili agli altri.\n\nSTRETTO DIVIETO DI ALLUCINAZIONE TOOL:\n- USA ESCLUSIVAMENTE I TOOL MCP FORNITI: 'memory', 'filesystem'.\n\nREGOLE DI COMUNICAZIONE:\n- Sii estremamente tecnico, sintetico e preciso.\n- Non fare mai riferimento a tool che non siano 'memory' o 'filesystem'."
|
||
|
||
class OllamaMcpClient:
|
||
"""MCP client that uses Ollama and LangChain create_agent with optional runtime MCP tools."""
|
||
|
||
def __init__(
|
||
self,
|
||
ollamaClient: OllamaClient,
|
||
mcpTools: Optional[List[StructuredTool]] = None,
|
||
systemPrompt: Optional[str] = None,
|
||
):
|
||
self.ollamaClient = ollamaClient
|
||
self.mcpTools = mcpTools or []
|
||
self.systemPrompt = systemPrompt or SYSTEM_PROMPT
|
||
staticTools: List[Any] = [getTime, countWords]
|
||
staticToolNames = [getTime.name, countWords.name]
|
||
middleware: List[AgentMiddleware] = [LogToolCallsMiddleware()]
|
||
if self.mcpTools:
|
||
middleware.append(McpToolsMiddleware(self.mcpTools, staticToolNames=staticToolNames))
|
||
model = ChatOllama(
|
||
base_url=ollamaClient.baseUrl,
|
||
model=ollamaClient.model,
|
||
temperature=0.1,
|
||
)
|
||
self.agent = create_agent(
|
||
model,
|
||
tools=staticTools,
|
||
middleware=middleware,
|
||
system_prompt=self.systemPrompt,
|
||
)
|
||
|
||
async def processRequest(self, prompt: str, context: Optional[List[str]] = None, recursionLimit: int = 50) -> str:
|
||
"""Process a request using the LangChain agent (ReAct loop with tools)."""
|
||
messages: List[Any] = [HumanMessage(content=prompt)]
|
||
if context:
|
||
messages.insert(0, SystemMessage(content=f"Context:\n{chr(10).join(context)}"))
|
||
config: Dict[str, Any] = {"recursion_limit": recursionLimit}
|
||
toolParseRetryPrompt = (
|
||
"ATTENZIONE: Una chiamata write_file ha prodotto JSON non valido. "
|
||
"Quando scrivi file con codice Python: usa \\n per le newline nel JSON, escapa le virgolette con \\. "
|
||
"Non aggiungere parametri extra (es. overwrite). Usa edit_file per modifiche incrementali se il contenuto è lungo."
|
||
)
|
||
try:
|
||
result = await self.agent.ainvoke({"messages": messages}, config=config)
|
||
except OllamaResponseError as e:
|
||
errStr = str(e)
|
||
if "error parsing tool call" in errStr:
|
||
print(f"[Agent Error]: Tool call parse error, retrying with guidance: {errStr[:200]}...", file=sys.stderr)
|
||
retryMessages: List[Any] = [SystemMessage(content=toolParseRetryPrompt)]
|
||
retryMessages.extend(messages)
|
||
result = await self.agent.ainvoke({"messages": retryMessages}, config=config)
|
||
else:
|
||
print(f"[Agent Error]: {e}", file=sys.stderr)
|
||
raise
|
||
except Exception as e:
|
||
print(f"[Agent Error]: {e}", file=sys.stderr)
|
||
raise
|
||
return _extractFinalResponse(result)
|
||
|
||
def listTools(self) -> List[str]:
|
||
"""List tool names (static + MCP)."""
|
||
names = [getTime.name, countWords.name]
|
||
names.extend(t.name for t in self.mcpTools)
|
||
return names
|
||
|
||
|
||
async def async_main(args, ollamaClient: OllamaClient):
|
||
"""Async main: MCP tools come only from mcp.json (Docker containers exposing SSE). Ollama is used only as LLM."""
|
||
mcpTools: List[StructuredTool] = []
|
||
mcpServers: List[McpServerWrapper] = []
|
||
|
||
# MCP servers from config file (mcp.json) – Docker containers with SSE endpoints
|
||
serverUrls: Dict[str, str] = loadMcpConfig(args.mcp_config)
|
||
if args.mcp_server:
|
||
serverUrls["default"] = args.mcp_server.rstrip("/")
|
||
|
||
# Which servers to use: default = all from mcp.json; or --mcp-tools fetch,filesystem to pick a subset
|
||
wantServers = [s.strip() for s in (args.mcp_tools or "").split(",") if s.strip()]
|
||
if not wantServers and serverUrls:
|
||
wantServers = list(serverUrls.keys())
|
||
print(f"MCP tools from config (all SSE servers): {wantServers}", file=sys.stderr)
|
||
for name in wantServers:
|
||
url = serverUrls.get(name)
|
||
if not url:
|
||
print(f"Warning: MCP server '{name}' not in config (known: {list(serverUrls.keys())})", file=sys.stderr)
|
||
continue
|
||
wrapper = McpServerWrapper(httpUrl=url)
|
||
if await wrapper.connect():
|
||
mcpServers.append(wrapper)
|
||
print(f"Connected to MCP server '{name}' at {url}", file=sys.stderr)
|
||
else:
|
||
print(f"Error: Failed to connect to MCP server '{name}' at {url}", file=sys.stderr)
|
||
|
||
if mcpServers:
|
||
mcpTools = await buildMcpLangChainTools(mcpServers, toolTimeout=getattr(args, "tool_timeout", None))
|
||
#print(f"Loaded {len(mcpTools)} MCP tools: {[t.name for t in mcpTools]}", file=sys.stderr)
|
||
|
||
mcpClient = OllamaMcpClient(ollamaClient, mcpTools=mcpTools)
|
||
print(f"Agent tools: {mcpClient.listTools()}", file=sys.stderr)
|
||
|
||
if args.prompt:
|
||
response = await mcpClient.processRequest(args.prompt, recursionLimit=args.recursion_limit)
|
||
print(response)
|
||
elif args.interactive:
|
||
print("MCP Client with Ollama (LangChain agent) - Interactive Mode")
|
||
print("Type 'quit' or 'exit' to exit\n")
|
||
while True:
|
||
try:
|
||
prompt = input("You: ").strip()
|
||
if prompt.lower() in ["quit", "exit"]:
|
||
break
|
||
if not prompt:
|
||
continue
|
||
response = await mcpClient.processRequest(prompt, recursionLimit=args.recursion_limit)
|
||
print(f"Assistant: {response}\n")
|
||
except KeyboardInterrupt:
|
||
print("\nGoodbye!")
|
||
break
|
||
except Exception as e:
|
||
print(f"Error: {e}", file=sys.stderr)
|
||
|
||
for wrapper in mcpServers:
|
||
await wrapper.disconnect()
|
||
|
||
|
||
def main():
|
||
"""Main function to run the MCP client."""
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="MCP client using Ollama")
|
||
parser.add_argument(
|
||
"--base-url",
|
||
default="http://localhost:11434",
|
||
help="Ollama base URL (default: http://localhost:11434)"
|
||
)
|
||
parser.add_argument(
|
||
"--model",
|
||
default="gpt-oss:20b",
|
||
help="Ollama model to use (default: ministral-3)"
|
||
)
|
||
parser.add_argument(
|
||
"--list-models",
|
||
action="store_true",
|
||
help="List available Ollama models and exit"
|
||
)
|
||
parser.add_argument(
|
||
"--prompt",
|
||
help="Prompt to send to the model"
|
||
)
|
||
parser.add_argument(
|
||
"--interactive",
|
||
"-i",
|
||
action="store_true",
|
||
help="Run in interactive mode"
|
||
)
|
||
parser.add_argument(
|
||
"--mcp-config",
|
||
default=None,
|
||
help="Path to mcp.json (default: mcpServer/mcp.json relative to project)"
|
||
)
|
||
parser.add_argument(
|
||
"--mcp-tools",
|
||
default="",
|
||
help="Comma-separated MCP server names from mcp.json (default: all servers in config). E.g. fetch,filesystem"
|
||
)
|
||
parser.add_argument(
|
||
"--mcp-server",
|
||
help="Override: single MCP SSE URL (e.g. http://localhost:3000/sse). Added as server 'default' in addition to mcp.json."
|
||
)
|
||
parser.add_argument(
|
||
"--mcp-headers",
|
||
help="Additional headers for MCP server as JSON string (e.g. '{\"Authorization\": \"Bearer token\"}')"
|
||
)
|
||
parser.add_argument(
|
||
"--recursion-limit",
|
||
type=int,
|
||
default=5000,
|
||
help="Max agent steps (model + tool calls) before stopping (default: 50)"
|
||
)
|
||
parser.add_argument(
|
||
"--tool-timeout",
|
||
type=float,
|
||
default=60,
|
||
help="Timeout in seconds for each MCP tool call. Prevents agent from freezing when a tool hangs (e.g. run with missing executable). Default: 60"
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
# Initialize Ollama client
|
||
ollamaClient = OllamaClient(baseUrl=args.base_url, model=args.model)
|
||
|
||
# Check health
|
||
if not ollamaClient.checkHealth():
|
||
print(f"Error: Cannot connect to Ollama at {args.base_url}", file=sys.stderr)
|
||
print("Make sure Ollama is running and accessible.", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
# List models if requested
|
||
if args.list_models:
|
||
models = ollamaClient.listModels()
|
||
print("Available models:")
|
||
for model in models:
|
||
print(f" - {model}")
|
||
sys.exit(0)
|
||
|
||
# Run async main
|
||
asyncio.run(async_main(args, ollamaClient))
|
||
|
||
if not args.prompt and not args.interactive:
|
||
parser.print_help()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |