Spaces:
Running
Running
| import asyncio | |
| import json | |
| from typing import Any | |
| from src.services.mcp.prompt import MCP_TOOL_SPECS | |
| _MCP_TOOL_NAMES = {spec["name"] for spec in MCP_TOOL_SPECS} | |
| def _normalize_tool_input(tool_input: Any) -> dict[str, Any]: | |
| if tool_input is None: | |
| return {} | |
| if isinstance(tool_input, dict): | |
| return tool_input | |
| if isinstance(tool_input, str): | |
| raw = tool_input.strip() | |
| if not raw: | |
| return {} | |
| try: | |
| parsed = json.loads(raw) | |
| return parsed if isinstance(parsed, dict) else {"input": parsed} | |
| except json.JSONDecodeError: | |
| return {"message": raw} | |
| return {"input": tool_input} | |
| async def execute_mcp_tool( | |
| tool_name: str, | |
| tool_input: Any, | |
| *, | |
| brain_pat: str | None, | |
| ) -> str: | |
| if tool_name not in _MCP_TOOL_NAMES: | |
| return f"Error: unknown tool '{tool_name}'" | |
| from src.services.mcp.main import ( | |
| auth_token_var, | |
| get_search_operation_instructions, | |
| list_brains, | |
| search_memory, | |
| search_semantically, | |
| traverse_graph, | |
| ) | |
| params = _normalize_tool_input(tool_input) | |
| token = auth_token_var.set(brain_pat) | |
| try: | |
| if tool_name == "get_search_operation_instructions": | |
| message = str(params.get("message") or "") | |
| result = get_search_operation_instructions(message) | |
| return result if isinstance(result, str) else json.dumps(result, default=str) | |
| if tool_name == "list_brains": | |
| result = await list_brains() | |
| return result if isinstance(result, str) else json.dumps(result, default=str) | |
| if tool_name == "search_semantically": | |
| result = await search_semantically( | |
| query=str(params.get("query") or ""), | |
| brain_id=str(params.get("brain_id") or ""), | |
| ) | |
| return result if isinstance(result, str) else json.dumps(result, default=str) | |
| if tool_name == "traverse_graph": | |
| result = await traverse_graph( | |
| brain_id=str(params.get("brain_id") or ""), | |
| start_uuid=str(params.get("start_uuid") or ""), | |
| start_name=str(params.get("start_name") or ""), | |
| start_labels=params.get("start_labels"), | |
| depth=int(params.get("depth") or 2), | |
| rel_types=params.get("rel_types"), | |
| node_labels=params.get("node_labels"), | |
| direction=str(params.get("direction") or "both"), | |
| limit=int(params.get("limit") or 50), | |
| ) | |
| return result if isinstance(result, str) else json.dumps(result, default=str) | |
| if tool_name == "search_memory": | |
| result = await search_memory( | |
| db_query=str(params.get("db_query") or ""), | |
| brain_id=str(params.get("brain_id") or ""), | |
| ) | |
| return result if isinstance(result, str) else json.dumps(result, default=str) | |
| return f"Error: tool '{tool_name}' is not implemented" | |
| except Exception as exc: | |
| return f"Error executing tool '{tool_name}': {exc}" | |
| finally: | |
| auth_token_var.reset(token) | |