MCP with Gemini Integration

""" Tools module for Gemini LLM Integration. Provides dynamic tool registration and discovery with enhanced error handling. """ import os import importlib import logging from typing import List, Callable, Any, Dict, Optional from dataclasses import dataclass logger = logging.getLogger(__name__) @dataclass class ToolDefinition: """Definition of a tool that can be registered.""" name: str description: str handler: Callable input_schema: Dict[str, Any] examples: Optional[List[str]] = None # Example queries for this tool def __post_init__(self): """Validate tool definition after initialization.""" if not self.name or not isinstance(self.name, str): raise ValueError("Tool name must be a non-empty string") if not self.description or not isinstance(self.description, str): raise ValueError("Tool description must be a non-empty string") if not callable(self.handler): raise ValueError("Tool handler must be callable") if not isinstance(self.input_schema, dict): raise ValueError("Tool input_schema must be a dictionary") class ToolRegistry: """Enhanced registry for managing and discovering tools with better error handling.""" def __init__(self): self.tools: List[ToolDefinition] = [] self._discovered = False def register_tool(self, tool_def: ToolDefinition): """Register a tool definition with validation.""" if not isinstance(tool_def, ToolDefinition): raise ValueError("tool_def must be a ToolDefinition instance") # Check for duplicate tool names existing_tool = self.get_tool(tool_def.name) if existing_tool: raise ValueError(f"Tool '{tool_def.name}' is already registered") self.tools.append(tool_def) logger.info(f"Registered tool: {tool_def.name} - {tool_def.description}") def discover_and_register_tools(self) -> List[ToolDefinition]: """Discover and register all available tools with enhanced error handling.""" if self._discovered: return self.tools try: # Get the tools directory tools_dir = os.path.dirname(__file__) if not os.path.exists(tools_dir): raise RuntimeError(f"Tools directory not found: {tools_dir}") # Discover all Python files in the tools directory discovered_count = 0 for filename in os.listdir(tools_dir): if filename.endswith('.py') and filename != '__init__.py': module_name = filename[:-3] # Remove .py extension try: # Import the module - try different import paths module = None import_paths = [ f'src.tools.{module_name}', f'tools.{module_name}', module_name ] for import_path in import_paths: try: module = importlib.import_module(import_path) break except ImportError: continue if module is None: logger.error(f"Could not import {filename} with any import path") continue # Look for register_tool function if hasattr(module, 'register_tool'): tool_def = module.register_tool() if tool_def: self.register_tool(tool_def) discovered_count += 1 logger.info(f"Successfully discovered tool: {tool_def.name}") else: logger.warning(f"No register_tool function found in {filename}") except Exception as e: logger.error(f"Failed to load tool from {filename}: {e}") continue self._discovered = True logger.info(f"Tool discovery completed. Found {discovered_count} tools out of {len(self.tools)} total.") if not self.tools: logger.warning("No tools were discovered. Check that tools exist and have register_tool() functions.") except Exception as e: logger.error(f"Error during tool discovery: {e}") raise RuntimeError(f"Tool discovery failed: {str(e)}") return self.tools def get_tool(self, name: str) -> Optional[ToolDefinition]: """Get a tool by name.""" for tool in self.tools: if tool.name == name: return tool return None def list_tools(self) -> List[str]: """List all registered tool names.""" return [tool.name for tool in self.tools] def get_tool_descriptions(self) -> Dict[str, str]: """Get a dictionary of tool names to descriptions.""" return {tool.name: tool.description for tool in self.tools} def validate_tool_call(self, tool_name: str, arguments: Dict[str, Any]) -> bool: """Validate that a tool call has the correct arguments.""" tool = self.get_tool(tool_name) if not tool: return False # Check required properties required_props = tool.input_schema.get("required", []) for prop in required_props: if prop not in arguments: return False return True