File size: 6,741 Bytes
8f51ef2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 |
"""
Base Tool Manager
This module provides a general framework for managing different types of tools
in the agentic pipeline. It can be extended for specific tool types.
"""
import os
import sys
from abc import ABC, abstractmethod
from typing import Dict, List, Any, Optional, Type, Callable
from pathlib import Path
import json
import tempfile
from dataclasses import dataclass
from enum import Enum
from langchain_core.tools import BaseTool, tool
from pydantic import BaseModel, Field
# Add parent directory to path for imports
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
# Define base configuration classes here since they might not exist in new_agent
@dataclass
class ToolConfig:
"""Configuration for a tool."""
name: str
tool_type: str
description: str
fallback_enabled: bool = True
class ToolStatus(Enum):
"""Status of a tool."""
NOT_AVAILABLE = "not_available"
AVAILABLE = "available"
FALLBACK = "fallback"
ERROR = "error"
@dataclass
class ToolConfigExtended(ToolConfig):
"""Extended configuration for a tool."""
input_schema: Type[BaseModel] = None
cache_results: bool = True
class BaseToolManager(ABC):
"""
Abstract base class for tool managers.
Provides common functionality for tool initialization and management.
"""
def __init__(self, config: ToolConfigExtended):
"""
Initialize the tool manager.
Args:
config: Tool configuration
"""
self.config = config
self.status = ToolStatus.NOT_AVAILABLE
self.tool = None
self.temp_dir = Path(tempfile.mkdtemp())
self.temp_dir.mkdir(parents=True, exist_ok=True)
# Initialize tool
self._initialize_tool()
@abstractmethod
def _initialize_tool(self):
"""Initialize the specific tool. Must be implemented by subclasses."""
pass
@abstractmethod
def _create_tool(self) -> BaseTool:
"""Create the tool instance. Must be implemented by subclasses."""
pass
@abstractmethod
def _create_fallback_tool(self) -> BaseTool:
"""Create fallback tool. Must be implemented by subclasses."""
pass
def is_available(self) -> bool:
"""Check if the tool is available."""
return self.status in [ToolStatus.AVAILABLE, ToolStatus.FALLBACK]
def get_status(self) -> ToolStatus:
"""Get current tool status."""
return self.status
def get_tool(self) -> Optional[BaseTool]:
"""Get the tool instance."""
return self.tool
def get_tool_info(self) -> Dict[str, Any]:
"""Get information about the tool."""
return {
"name": self.config.name,
"type": self.config.tool_type,
"status": self.status.value,
"description": self.config.description,
"temp_dir": str(self.temp_dir)
}
def _set_status(self, status: ToolStatus):
"""Set the tool status."""
self.status = status
def _handle_error(self, error: Exception, fallback: bool = True):
"""Handle errors during tool initialization."""
print(f"Error in {self.config.name}: {error}")
self._set_status(ToolStatus.ERROR)
if fallback and self.config.fallback_enabled:
self._initialize_fallback()
def _initialize_fallback(self):
"""Initialize fallback tool when main tool fails."""
print(f"Initializing fallback for {self.config.name}")
self._set_status(ToolStatus.FALLBACK)
self.tool = self._create_fallback_tool()
def cleanup(self):
"""Clean up resources."""
self.tool = None
self._set_status(ToolStatus.NOT_AVAILABLE)
class ToolRegistry:
"""
Registry for managing multiple tool managers.
Provides a centralized way to access different tools.
"""
def __init__(self):
self._managers: Dict[str, BaseToolManager] = {}
self._configs: Dict[str, ToolConfigExtended] = {}
self._tools: Dict[str, BaseTool] = {}
def register_tool(self, name: str, manager: BaseToolManager, config: ToolConfigExtended):
"""Register a tool manager."""
self._managers[name] = manager
self._configs[name] = config
self._tools[name] = manager.get_tool()
print(f"Registered tool: {name}")
def get_tool(self, name: str) -> Optional[BaseTool]:
"""Get a tool by name."""
return self._tools.get(name)
def get_manager(self, name: str) -> Optional[BaseToolManager]:
"""Get a tool manager by name."""
return self._managers.get(name)
def get_available_tools(self) -> List[str]:
"""Get list of available tool names."""
return [name for name, manager in self._managers.items() if manager.is_available()]
def get_tool_info(self, name: str) -> Optional[Dict[str, Any]]:
"""Get information about a specific tool."""
manager = self._managers.get(name)
if manager:
return manager.get_tool_info()
return None
def get_all_tools(self) -> Dict[str, BaseTool]:
"""Get all available tools."""
return {name: tool for name, tool in self._tools.items() if tool is not None}
def cleanup_all(self):
"""Clean up all registered tools."""
for manager in self._managers.values():
manager.cleanup()
self._managers.clear()
self._configs.clear()
self._tools.clear()
# Global tool registry
tool_registry = ToolRegistry()
def create_tool_wrapper(tool_name: str, description: str, input_schema: Type[BaseModel]) -> Callable:
"""
Create a tool wrapper function for LangChain integration.
Args:
tool_name: Name of the tool
description: Description of the tool
input_schema: Input schema for the tool
Returns:
Tool wrapper function
"""
def tool_wrapper(**kwargs) -> Dict[str, Any]:
"""Tool wrapper function."""
tool = tool_registry.get_tool(tool_name)
if tool is None:
return {"error": f"Tool {tool_name} not available"}
try:
# Convert kwargs to input schema
input_data = input_schema(**kwargs)
result = tool.run(input_data.dict())
return result
except Exception as e:
return {"error": f"Tool {tool_name} failed: {str(e)}"}
# Set function attributes for LangChain
tool_wrapper.__name__ = tool_name
tool_wrapper.__doc__ = description
return tool_wrapper
|