305 lines
9.5 KiB
Python
305 lines
9.5 KiB
Python
"""Tool registry for auto-discovering and managing chatbot tools.
|
|
|
|
This module provides a central registry that automatically discovers all tools
|
|
in the chatbotTools directory structure and provides methods to query them.
|
|
The registry is built in-memory at startup and does not require a database.
|
|
"""
|
|
|
|
import importlib
|
|
import inspect
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional
|
|
|
|
from langchain_core.tools import BaseTool
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ToolMetadata:
|
|
"""Metadata about a discovered chatbot tool.
|
|
|
|
Attributes:
|
|
tool_id: Unique identifier (e.g., 'shared.tavily_search')
|
|
name: Function name of the tool
|
|
category: Category of the tool ('shared' or 'customer')
|
|
description: Tool description from docstring
|
|
tool_instance: The actual LangChain tool instance
|
|
module_path: Full Python module path
|
|
"""
|
|
|
|
tool_id: str
|
|
name: str
|
|
category: str
|
|
description: str
|
|
tool_instance: BaseTool
|
|
module_path: str
|
|
|
|
def __str__(self) -> str:
|
|
"""Return a pretty-printed string representation for logging."""
|
|
return (
|
|
f"ToolMetadata(\n"
|
|
f" tool_id='{self.tool_id}',\n"
|
|
f" name='{self.name}',\n"
|
|
f" category='{self.category}',\n"
|
|
f" description='{self.description}',\n"
|
|
f" module_path='{self.module_path}'\n"
|
|
f")"
|
|
)
|
|
|
|
|
|
class ToolRegistry:
|
|
"""Central registry for all chatbot tools.
|
|
|
|
This class discovers and catalogs all tools decorated with @tool in the
|
|
chatbotTools directory structure. Tools are automatically discovered at
|
|
initialization by scanning the filesystem.
|
|
|
|
The registry provides methods to query tools by ID, category, or get all tools.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize an empty tool registry."""
|
|
self._tools: Dict[str, ToolMetadata] = {}
|
|
self._initialized: bool = False
|
|
|
|
def initialize(self) -> None:
|
|
"""Discover and register all tools from the chatbotTools directory.
|
|
|
|
This method scans both sharedTools and customerTools directories,
|
|
imports all tool*.py modules, and extracts functions decorated with @tool.
|
|
|
|
This method is idempotent - calling it multiple times has no effect
|
|
after the first initialization.
|
|
"""
|
|
if self._initialized:
|
|
logger.debug("Tool registry already initialized, skipping")
|
|
return
|
|
|
|
logger.info("Initializing tool registry...")
|
|
|
|
# Get base path to chatbotTools directory
|
|
base_path = Path(__file__).parent.parent / "chatbotTools"
|
|
|
|
if not base_path.exists():
|
|
logger.warning(f"chatbotTools directory not found at {base_path}")
|
|
self._initialized = True
|
|
return
|
|
|
|
# Discover tools in each category
|
|
self._discover_category(
|
|
category_path=base_path / "sharedTools", category="shared"
|
|
)
|
|
self._discover_category(
|
|
category_path=base_path / "customerTools", category="customer"
|
|
)
|
|
|
|
self._initialized = True
|
|
logger.info(f"Tool registry initialized with {len(self._tools)} tools")
|
|
|
|
def _discover_category(self, *, category_path: Path, category: str) -> None:
|
|
"""Discover all tools in a specific category directory.
|
|
|
|
Args:
|
|
category_path: Path to the category directory (sharedTools or customerTools)
|
|
category: Category name ('shared' or 'customer')
|
|
"""
|
|
if not category_path.exists():
|
|
logger.warning(f"Category directory not found: {category_path}")
|
|
return
|
|
|
|
logger.debug(f"Discovering tools in category: {category}")
|
|
|
|
# Find all tool*.py files (excluding __init__.py)
|
|
tool_files = [
|
|
f for f in category_path.glob("tool*.py") if f.name != "__init__.py"
|
|
]
|
|
|
|
for tool_file in tool_files:
|
|
self._import_and_register_tools(
|
|
tool_file=tool_file, category=category, category_path=category_path
|
|
)
|
|
|
|
logger.debug(f"Discovered {len(tool_files)} tool files in {category}")
|
|
|
|
def _import_and_register_tools(
|
|
self, *, tool_file: Path, category: str, category_path: Path
|
|
) -> None:
|
|
"""Import a tool module and register all discovered tools.
|
|
|
|
Args:
|
|
tool_file: Path to the tool Python file
|
|
category: Category name ('shared' or 'customer')
|
|
category_path: Path to the category directory
|
|
"""
|
|
# Construct module name
|
|
module_name = (
|
|
f"modules.features.chatBot.chatbotTools.{category}Tools.{tool_file.stem}"
|
|
)
|
|
|
|
try:
|
|
# Import the module
|
|
module = importlib.import_module(module_name)
|
|
|
|
# Find all BaseTool instances in the module
|
|
tools_found = 0
|
|
for name, obj in inspect.getmembers(module):
|
|
if isinstance(obj, BaseTool):
|
|
self._register_tool(
|
|
tool_instance=obj,
|
|
name=name,
|
|
category=category,
|
|
module_path=module_name,
|
|
)
|
|
tools_found += 1
|
|
|
|
if tools_found == 0:
|
|
logger.warning(f"No tools found in {module_name}")
|
|
else:
|
|
logger.debug(f"Loaded {tools_found} tool(s) from {module_name}")
|
|
|
|
except ImportError as e:
|
|
logger.error(
|
|
f"Import error loading tools from {module_name}: {str(e)}. "
|
|
f"This tool will not be available."
|
|
)
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Unexpected error loading tools from {module_name}: {type(e).__name__}: {str(e)}"
|
|
)
|
|
|
|
def _register_tool(
|
|
self, *, tool_instance: BaseTool, name: str, category: str, module_path: str
|
|
) -> None:
|
|
"""Register a single tool in the registry.
|
|
|
|
Args:
|
|
tool_instance: The LangChain tool instance
|
|
name: Function name of the tool
|
|
category: Category name ('shared' or 'customer')
|
|
module_path: Full Python module path
|
|
"""
|
|
tool_id = f"{category}.{name}"
|
|
|
|
# Check for duplicate tool IDs
|
|
if tool_id in self._tools:
|
|
logger.warning(f"Duplicate tool ID detected: {tool_id}, overwriting")
|
|
|
|
metadata = ToolMetadata(
|
|
tool_id=tool_id,
|
|
name=name,
|
|
category=category,
|
|
description=tool_instance.description or "",
|
|
tool_instance=tool_instance,
|
|
module_path=module_path,
|
|
)
|
|
|
|
self._tools[tool_id] = metadata
|
|
logger.debug(f"Registered tool: {tool_id}")
|
|
|
|
def get_all_tools(self) -> List[ToolMetadata]:
|
|
"""Get all registered tools.
|
|
|
|
Returns:
|
|
List of all tool metadata objects
|
|
"""
|
|
return list(self._tools.values())
|
|
|
|
def get_tool(self, *, tool_id: str) -> Optional[ToolMetadata]:
|
|
"""Get a specific tool by its ID.
|
|
|
|
Args:
|
|
tool_id: The tool identifier (e.g., 'shared.tavily_search')
|
|
|
|
Returns:
|
|
Tool metadata if found, None otherwise
|
|
"""
|
|
return self._tools.get(tool_id)
|
|
|
|
def get_tools_by_category(self, *, category: str) -> List[ToolMetadata]:
|
|
"""Get all tools in a specific category.
|
|
|
|
Args:
|
|
category: Category name ('shared' or 'customer')
|
|
|
|
Returns:
|
|
List of tool metadata for the specified category
|
|
"""
|
|
return [t for t in self._tools.values() if t.category == category]
|
|
|
|
def list_tool_ids(self) -> List[str]:
|
|
"""Get a list of all registered tool IDs.
|
|
|
|
Returns:
|
|
List of tool ID strings
|
|
"""
|
|
return list(self._tools.keys())
|
|
|
|
def get_tool_instances(self, *, tool_ids: List[str]) -> List[BaseTool]:
|
|
"""Get actual tool instances for a list of tool IDs.
|
|
|
|
This is useful for filtering tools based on user permissions.
|
|
|
|
Args:
|
|
tool_ids: List of tool IDs to retrieve
|
|
|
|
Returns:
|
|
List of BaseTool instances for the specified IDs
|
|
"""
|
|
instances = []
|
|
for tool_id in tool_ids:
|
|
metadata = self.get_tool(tool_id=tool_id)
|
|
if metadata:
|
|
instances.append(metadata.tool_instance)
|
|
else:
|
|
logger.warning(f"Tool ID not found in registry: {tool_id}")
|
|
return instances
|
|
|
|
@property
|
|
def is_initialized(self) -> bool:
|
|
"""Check if the registry has been initialized.
|
|
|
|
Returns:
|
|
True if initialized, False otherwise
|
|
"""
|
|
return self._initialized
|
|
|
|
|
|
# Global registry instance
|
|
_registry: Optional[ToolRegistry] = None
|
|
|
|
|
|
def get_registry() -> ToolRegistry:
|
|
"""Get the global tool registry instance.
|
|
|
|
This function ensures the registry is initialized on first access.
|
|
Subsequent calls return the same instance.
|
|
|
|
Returns:
|
|
The global ToolRegistry instance
|
|
"""
|
|
global _registry
|
|
|
|
if _registry is None:
|
|
_registry = ToolRegistry()
|
|
|
|
if not _registry.is_initialized:
|
|
_registry.initialize()
|
|
|
|
return _registry
|
|
|
|
|
|
def reinitialize_registry() -> ToolRegistry:
|
|
"""Force reinitialize the tool registry.
|
|
|
|
This is useful for testing or when tools are added dynamically.
|
|
|
|
Returns:
|
|
The reinitialized ToolRegistry instance
|
|
"""
|
|
global _registry
|
|
_registry = ToolRegistry()
|
|
_registry.initialize()
|
|
return _registry
|