"""LlamaIndex toolkit for Go Micro services.""" import json import re from typing import Any, Dict, List, Optional from dataclasses import dataclass import requests from llama_index.core.tools import FunctionTool, ToolMetadata from pydantic import BaseModel, Field from go_micro_llamaindex.exceptions import ( GoMicroConnectionError, GoMicroAuthError, GoMicroToolError, ) @dataclass class GoMicroConfig: """Configuration for Go Micro MCP gateway connection. Attributes: gateway_url: URL of the MCP gateway (e.g., http://localhost:3000) auth_token: Optional bearer authentication token timeout: Request timeout in seconds retry_count: Number of retries on failure retry_delay: Delay between retries in seconds verify_ssl: Whether to verify SSL certificates """ gateway_url: str auth_token: Optional[str] = None timeout: int = 30 retry_count: int = 3 retry_delay: float = 1.0 verify_ssl: bool = True class GoMicroTool(BaseModel): """Represents a Go Micro service tool. Attributes: name: Tool name (e.g., "users.Users.Get") service: Service name (e.g., "users") endpoint: Endpoint name (e.g., "Users.Get") description: Tool description example: Example input JSON scopes: Required auth scopes metadata: Additional metadata from service """ name: str service: str endpoint: str description: str example: Optional[str] = None scopes: Optional[List[str]] = None metadata: Dict[str, str] = Field(default_factory=dict) class GoMicroToolkit: """LlamaIndex toolkit for Go Micro services. This class provides integration between LlamaIndex and Go Micro services via the Model Context Protocol (MCP) gateway. Example: >>> toolkit = GoMicroToolkit.from_gateway("http://localhost:3000") >>> tools = toolkit.get_tools() >>> for tool in tools: ... print(f"Tool: {tool.metadata.name}") """ def __init__(self, config: GoMicroConfig): """Initialize the toolkit. Args: config: Configuration for MCP gateway connection """ self.config = config self._tools: Optional[List[GoMicroTool]] = None self._session = requests.Session() if config.auth_token: self._session.headers.update({ "Authorization": f"Bearer {config.auth_token}" }) @classmethod def from_gateway( cls, gateway_url: str, auth_token: Optional[str] = None, **kwargs: Any ) -> "GoMicroToolkit": """Create toolkit from MCP gateway URL. Args: gateway_url: URL of the MCP gateway auth_token: Optional bearer authentication token **kwargs: Additional configuration options Returns: GoMicroToolkit instance Example: >>> toolkit = GoMicroToolkit.from_gateway("http://localhost:3000") """ config = GoMicroConfig( gateway_url=gateway_url, auth_token=auth_token, **kwargs ) return cls(config) def _make_request( self, method: str, path: str, **kwargs: Any ) -> requests.Response: """Make HTTP request to MCP gateway. Args: method: HTTP method (GET, POST, etc.) path: API path **kwargs: Additional request arguments Returns: Response object Raises: GoMicroConnectionError: If connection fails GoMicroAuthError: If authentication fails """ url = f"{self.config.gateway_url}{path}" kwargs.setdefault("timeout", self.config.timeout) kwargs.setdefault("verify", self.config.verify_ssl) try: response = self._session.request(method, url, **kwargs) if response.status_code == 401: raise GoMicroAuthError("Authentication failed") elif response.status_code == 403: raise GoMicroAuthError("Forbidden: insufficient permissions") response.raise_for_status() return response except requests.ConnectionError as e: raise GoMicroConnectionError( f"Failed to connect to MCP gateway at {url}: {e}" ) except requests.Timeout as e: raise GoMicroConnectionError( f"Request to MCP gateway timed out: {e}" ) except requests.RequestException as e: if isinstance(e, (GoMicroConnectionError, GoMicroAuthError)): raise raise GoMicroConnectionError(f"Request failed: {e}") def refresh(self) -> None: """Refresh tool list from MCP gateway. Raises: GoMicroConnectionError: If unable to connect to gateway """ response = self._make_request("GET", "/mcp/tools") data = response.json() tools_data = data.get("tools", []) self._tools = [ GoMicroTool( name=tool["name"], service=tool["service"], endpoint=tool["endpoint"], description=tool.get("description", ""), example=tool.get("example"), scopes=tool.get("scopes"), metadata=tool.get("metadata", {}) ) for tool in tools_data ] def get_tools( self, service_filter: Optional[str] = None, name_pattern: Optional[str] = None, include: Optional[List[str]] = None, exclude: Optional[List[str]] = None, ) -> List[FunctionTool]: """Get LlamaIndex tools from Go Micro services. Args: service_filter: Filter tools by service name name_pattern: Filter tools by name pattern (regex) include: List of tool names to include exclude: List of tool names to exclude Returns: List of LlamaIndex FunctionTool objects Example: >>> toolkit = GoMicroToolkit.from_gateway("http://localhost:3000") >>> all_tools = toolkit.get_tools() >>> user_tools = toolkit.get_tools(service_filter="users") """ if self._tools is None: self.refresh() tools = self._tools or [] if service_filter: tools = [t for t in tools if t.service == service_filter] if name_pattern: pattern = re.compile(name_pattern) tools = [t for t in tools if pattern.match(t.name)] if include: tools = [t for t in tools if t.name in include] if exclude: tools = [t for t in tools if t.name not in exclude] return [self._create_llamaindex_tool(tool) for tool in tools] def _create_llamaindex_tool(self, tool: GoMicroTool) -> FunctionTool: """Create a LlamaIndex FunctionTool from a GoMicroTool. Args: tool: GoMicroTool to convert Returns: LlamaIndex FunctionTool object """ toolkit = self def tool_func(arguments: str) -> str: """Execute the tool. Args: arguments: JSON string with tool arguments Returns: JSON string with tool result """ return toolkit.call_tool(tool.name, arguments) description = tool.description if tool.example: description += f"\n\nExample input: {tool.example}" return FunctionTool.from_defaults( fn=tool_func, name=tool.name, description=description, ) def call_tool(self, tool_name: str, arguments: str) -> str: """Call a specific tool directly. Args: tool_name: Name of the tool to call arguments: JSON string with tool arguments Returns: JSON string with tool result Raises: GoMicroToolError: If tool execution fails Example: >>> toolkit = GoMicroToolkit.from_gateway("http://localhost:3000") >>> result = toolkit.call_tool( ... "users.Users.Get", ... '{"id": "user-123"}' ... ) """ try: args = json.loads(arguments) if isinstance(arguments, str) else arguments except json.JSONDecodeError as e: raise GoMicroToolError(f"Invalid JSON arguments: {e}") try: response = self._make_request( "POST", "/mcp/call", json={"name": tool_name, "arguments": args} ) return json.dumps(response.json()) except requests.RequestException as e: raise GoMicroToolError(f"Tool execution failed: {e}") def list_tools(self) -> List[GoMicroTool]: """Get raw list of available tools. Returns: List of GoMicroTool objects Example: >>> toolkit = GoMicroToolkit.from_gateway("http://localhost:3000") >>> for tool in toolkit.list_tools(): ... print(f"{tool.name}: {tool.description}") """ if self._tools is None: self.refresh() return self._tools or []