CHKIM79/scalable-ai-agent-system
0
1"""2Tool Management System3Handles web search, code execution, database operations, file operations, mathematical calculations4"""5import asyncio6import logging7import json8import subprocess9import tempfile10import os11import sqlite312import requests13from typing import Dict, List, Any, Optional, Callable, Union14from dataclasses import dataclass, field15from enum import Enum16import aiohttp17import math18import numpy as np19from datetime import datetime20 21 22class ToolCategory(Enum):23 WEB_SEARCH = "web_search"24 CODE_EXECUTION = "code_execution"25 DATABASE = "database"26 FILE_OPERATIONS = "file_operations"27 MATHEMATICAL = "mathematical"28 API_INTEGRATION = "api_integration"29 SYSTEM = "system"30 31 32@dataclass33class ToolDefinition:34 name: str35 category: ToolCategory36 function: Callable37 description: str38 parameters: Dict[str, Any] = field(default_factory=dict)39 timeout: int = 3040 requires_auth: bool = False41 42 43@dataclass44class ToolResult:45 tool_name: str46 success: bool47 result: Any48 execution_time: float49 error_message: Optional[str] = None50 metadata: Dict[str, Any] = field(default_factory=dict)51 52 53class WebSearchTool:54 """Web search capabilities"""55 56 def __init__(self, api_key: str = None):57 self.api_key = api_key58 self.session = None59 60 async def search(self, query: str, num_results: int = 5) -> Dict[str, Any]:61 """Perform web search"""62 if not self.session:63 self.session = aiohttp.ClientSession()64 65 # Simulate web search (replace with actual search API)66 try:67 # Using DuckDuckGo Instant Answer API as example68 url = f"https://api.duckduckgo.com/"69 params = {70 'q': query,71 'format': 'json',72 'no_html': '1',73 'skip_disambig': '1'74 }75 76 async with self.session.get(url, params=params) as response:77 data = await response.json()78 79 results = []80 if data.get('AbstractText'):81 results.append({82 'title': data.get('Heading', 'Search Result'),83 'snippet': data.get('AbstractText'),84 'url': data.get('AbstractURL', ''),85 'source': data.get('AbstractSource', 'DuckDuckGo')86 })87 88 # Add related topics89 for topic in data.get('RelatedTopics', [])[:num_results-1]:90 if isinstance(topic, dict) and 'Text' in topic:91 results.append({92 'title': topic.get('Text', '').split(' - ')[0],93 'snippet': topic.get('Text', ''),94 'url': topic.get('FirstURL', ''),95 'source': 'DuckDuckGo'96 })97 98 return {99 'query': query,100 'results': results[:num_results],101 'total_results': len(results)102 }103 104 except Exception as e:105 return {106 'query': query,107 'results': [],108 'error': str(e)109 }110 111 112class CodeExecutionTool:113 """Safe code execution in sandboxed environment"""114 115 def __init__(self):116 self.allowed_imports = {117 'math', 'json', 'datetime', 'random', 'string', 'collections',118 'itertools', 'functools', 'operator', 're', 'urllib.parse'119 }120 121 async def execute_python(self, code: str, timeout: int = 10) -> Dict[str, Any]:122 """Execute Python code safely"""123 try:124 # Create temporary file125 with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:126 # Add safety restrictions127 safe_code = f"""128import sys129import os130import subprocess131import importlib132 133# Restrict dangerous operations134def restricted_import(name, *args, **kwargs):135 allowed = {self.allowed_imports}136 if name not in allowed:137 raise ImportError(f"Import of {{name}} is not allowed")138 return original_import(name, *args, **kwargs)139 140original_import = __builtins__.__import__141__builtins__.__import__ = restricted_import142 143# Disable dangerous functions144__builtins__.open = None145__builtins__.exec = None146__builtins__.eval = None147 148# User code149{code}150"""151 f.write(safe_code)152 f.flush()153 154 # Execute with timeout155 process = await asyncio.create_subprocess_exec(156 'python', f.name,157 stdout=asyncio.subprocess.PIPE,158 stderr=asyncio.subprocess.PIPE159 )160 161 try:162 stdout, stderr = await asyncio.wait_for(163 process.communicate(), timeout=timeout164 )165 166 return {167 'success': process.returncode == 0,168 'stdout': stdout.decode('utf-8'),169 'stderr': stderr.decode('utf-8'),170 'return_code': process.returncode171 }172 173 except asyncio.TimeoutError:174 process.kill()175 return {176 'success': False,177 'stdout': '',178 'stderr': 'Execution timed out',179 'return_code': -1180 }181 finally:182 os.unlink(f.name)183 184 except Exception as e:185 return {186 'success': False,187 'stdout': '',188 'stderr': str(e),189 'return_code': -1190 }191 192 193class DatabaseTool:194 """Database operations"""195 196 def __init__(self, db_path: str = ":memory:"):197 self.db_path = db_path198 self.connection = None199 200 async def connect(self):201 """Connect to database"""202 self.connection = sqlite3.connect(self.db_path)203 self.connection.row_factory = sqlite3.Row204 205 async def execute_query(self, query: str, params: tuple = None) -> Dict[str, Any]:206 """Execute SQL query safely"""207 if not self.connection:208 await self.connect()209 210 try:211 cursor = self.connection.cursor()212 213 # Basic SQL injection protection214 if any(dangerous in query.upper() for dangerous in ['DROP', 'DELETE', 'UPDATE', 'INSERT'] 215 if not query.upper().strip().startswith(('SELECT', 'WITH'))):216 return {217 'success': False,218 'error': 'Only SELECT queries are allowed for safety'219 }220 221 if params:222 cursor.execute(query, params)223 else:224 cursor.execute(query)225 226 if query.upper().strip().startswith('SELECT'):227 results = [dict(row) for row in cursor.fetchall()]228 return {229 'success': True,230 'results': results,231 'row_count': len(results)232 }233 else:234 self.connection.commit()235 return {236 'success': True,237 'affected_rows': cursor.rowcount238 }239 240 except Exception as e:241 return {242 'success': False,243 'error': str(e)244 }245 246 247class FileOperationsTool:248 """Safe file operations"""249 250 def __init__(self, base_path: str = "/tmp/agent_files"):251 self.base_path = base_path252 os.makedirs(base_path, exist_ok=True)253 254 def _safe_path(self, path: str) -> str:255 """Ensure path is within allowed directory"""256 abs_path = os.path.abspath(os.path.join(self.base_path, path))257 if not abs_path.startswith(os.path.abspath(self.base_path)):258 raise ValueError("Path outside allowed directory")259 return abs_path260 261 async def read_file(self, filename: str) -> Dict[str, Any]:262 """Read file content"""263 try:264 safe_path = self._safe_path(filename)265 with open(safe_path, 'r', encoding='utf-8') as f:266 content = f.read()267 268 return {269 'success': True,270 'content': content,271 'size': len(content)272 }273 except Exception as e:274 return {275 'success': False,276 'error': str(e)277 }278 279 async def write_file(self, filename: str, content: str) -> Dict[str, Any]:280 """Write content to file"""281 try:282 safe_path = self._safe_path(filename)283 os.makedirs(os.path.dirname(safe_path), exist_ok=True)284 285 with open(safe_path, 'w', encoding='utf-8') as f:286 f.write(content)287 288 return {289 'success': True,290 'bytes_written': len(content.encode('utf-8'))291 }292 except Exception as e:293 return {294 'success': False,295 'error': str(e)296 }297 298 async def list_files(self, directory: str = "") -> Dict[str, Any]:299 """List files in directory"""300 try:301 safe_path = self._safe_path(directory)302 files = []303 304 for item in os.listdir(safe_path):305 item_path = os.path.join(safe_path, item)306 files.append({307 'name': item,308 'is_directory': os.path.isdir(item_path),309 'size': os.path.getsize(item_path) if os.path.isfile(item_path) else 0,310 'modified': datetime.fromtimestamp(os.path.getmtime(item_path)).isoformat()311 })312 313 return {314 'success': True,315 'files': files,316 'count': len(files)317 }318 except Exception as e:319 return {320 'success': False,321 'error': str(e)322 }323 324 325class MathematicalTool:326 """Mathematical calculations and operations"""327 328 async def calculate(self, expression: str) -> Dict[str, Any]:329 """Safely evaluate mathematical expressions"""330 try:331 # Allow only safe mathematical operations332 allowed_names = {333 'abs': abs, 'round': round, 'min': min, 'max': max, 'sum': sum, 'pow': pow,334 'math': math, 'pi': math.pi, 'e': math.e,335 'sin': math.sin, 'cos': math.cos, 'tan': math.tan,336 'sqrt': math.sqrt, 'log': math.log, 'exp': math.exp337 }338 339 # Remove dangerous functions340 safe_dict = {"__builtins__": {}}341 safe_dict.update(allowed_names)342 343 result = eval(expression, safe_dict)344 345 return {346 'success': True,347 'result': result,348 'expression': expression349 }350 351 except Exception as e:352 return {353 'success': False,354 'error': str(e),355 'expression': expression356 }357 358 async def statistics(self, numbers: List[float]) -> Dict[str, Any]:359 """Calculate statistical measures"""360 try:361 if not numbers:362 return {'success': False, 'error': 'Empty list provided'}363 364 arr = np.array(numbers)365 366 return {367 'success': True,368 'count': len(numbers),369 'mean': float(np.mean(arr)),370 'median': float(np.median(arr)),371 'std_dev': float(np.std(arr)),372 'variance': float(np.var(arr)),373 'min': float(np.min(arr)),374 'max': float(np.max(arr)),375 'sum': float(np.sum(arr))376 }377 378 except Exception as e:379 return {380 'success': False,381 'error': str(e)382 }383 384 385class ToolManager:386 """Main tool management system"""387 388 def __init__(self):389 self.tools: Dict[str, ToolDefinition] = {}390 self.web_search = WebSearchTool()391 self.code_executor = CodeExecutionTool()392 self.database = DatabaseTool()393 self.file_ops = FileOperationsTool()394 self.math_tool = MathematicalTool()395 self.logger = logging.getLogger(__name__)396 397 # Register built-in tools398 self._register_builtin_tools()399 400 def _register_builtin_tools(self):401 """Register all built-in tools"""402 403 # Web search tools404 self.register_tool(405 "web_search",406 self.web_search.search,407 "Search the web for information",408 ToolCategory.WEB_SEARCH,409 {"query": str, "num_results": int}410 )411 412 # Code execution tools413 self.register_tool(414 "execute_python",415 self.code_executor.execute_python,416 "Execute Python code safely",417 ToolCategory.CODE_EXECUTION,418 {"code": str, "timeout": int}419 )420 421 # Database tools422 self.register_tool(423 "db_query",424 self.database.execute_query,425 "Execute SQL query",426 ToolCategory.DATABASE,427 {"query": str, "params": tuple}428 )429 430 # File operation tools431 self.register_tool(432 "read_file",433 self.file_ops.read_file,434 "Read file content",435 ToolCategory.FILE_OPERATIONS,436 {"filename": str}437 )438 439 self.register_tool(440 "write_file",441 self.file_ops.write_file,442 "Write content to file",443 ToolCategory.FILE_OPERATIONS,444 {"filename": str, "content": str}445 )446 447 self.register_tool(448 "list_files",449 self.file_ops.list_files,450 "List files in directory",451 ToolCategory.FILE_OPERATIONS,452 {"directory": str}453 )454 455 # Mathematical tools456 self.register_tool(457 "calculate",458 self.math_tool.calculate,459 "Perform mathematical calculations",460 ToolCategory.MATHEMATICAL,461 {"expression": str}462 )463 464 self.register_tool(465 "statistics",466 self.math_tool.statistics,467 "Calculate statistical measures",468 ToolCategory.MATHEMATICAL,469 {"numbers": List[float]}470 )471 472 def register_tool(self, 473 name: str, 474 function: Callable, 475 description: str,476 category: ToolCategory = ToolCategory.SYSTEM,477 parameters: Dict[str, Any] = None,478 timeout: int = 30,479 requires_auth: bool = False):480 """Register a new tool"""481 482 tool_def = ToolDefinition(483 name=name,484 category=category,485 function=function,486 description=description,487 parameters=parameters or {},488 timeout=timeout,489 requires_auth=requires_auth490 )491 492 self.tools[name] = tool_def493 self.logger.info(f"Registered tool: {name}")494 495 async def execute_tool(self, tool_name: str, parameters: Dict[str, Any] = None, **kwargs) -> ToolResult:496 """Execute a tool with given parameters"""497 start_time = asyncio.get_event_loop().time()498 499 if tool_name not in self.tools:500 return ToolResult(501 tool_name=tool_name,502 success=False,503 result=None,504 execution_time=0,505 error_message=f"Tool '{tool_name}' not found"506 )507 508 tool_def = self.tools[tool_name]509 510 try:511 # Merge parameters and kwargs512 all_params = {}513 if parameters:514 all_params.update(parameters)515 all_params.update(kwargs)516 517 # Execute with timeout518 result = await asyncio.wait_for(519 tool_def.function(**all_params),520 timeout=tool_def.timeout521 )522 523 execution_time = asyncio.get_event_loop().time() - start_time524 525 return ToolResult(526 tool_name=tool_name,527 success=True,528 result=result,529 execution_time=execution_time530 )531 532 except asyncio.TimeoutError:533 return ToolResult(534 tool_name=tool_name,535 success=False,536 result=None,537 execution_time=tool_def.timeout,538 error_message=f"Tool execution timed out after {tool_def.timeout}s"539 )540 541 except Exception as e:542 execution_time = asyncio.get_event_loop().time() - start_time543 return ToolResult(544 tool_name=tool_name,545 success=False,546 result=None,547 execution_time=execution_time,548 error_message=str(e)549 )550 551 def get_available_tools(self, category: ToolCategory = None) -> List[Dict[str, Any]]:552 """Get list of available tools"""553 tools_list = []554 555 for name, tool_def in self.tools.items():556 if category is None or tool_def.category == category:557 tools_list.append({558 'name': name,559 'category': tool_def.category.value,560 'description': tool_def.description,561 'parameters': tool_def.parameters,562 'timeout': tool_def.timeout563 })564 565 return tools_list566 567 def get_tool_schema(self, tool_name: str) -> Optional[Dict[str, Any]]:568 """Get OpenAI Functions compatible schema for a tool"""569 if tool_name not in self.tools:570 return None571 572 tool_def = self.tools[tool_name]573 574 # Convert to OpenAI Functions format575 properties = {}576 required = []577 578 for param_name, param_type in tool_def.parameters.items():579 if param_type == str:580 properties[param_name] = {"type": "string"}581 elif param_type == int:582 properties[param_name] = {"type": "integer"}583 elif param_type == float:584 properties[param_name] = {"type": "number"}585 elif param_type == bool:586 properties[param_name] = {"type": "boolean"}587 elif param_type == list:588 properties[param_name] = {"type": "array"}589 590 required.append(param_name)591 592 return {593 "name": tool_name,594 "description": tool_def.description,595 "parameters": {596 "type": "object",597 "properties": properties,598 "required": required599 }600 }601 602 async def initialize(self):603 """Initialize tool manager"""604 await self.database.connect()605 self.logger.info("Tool manager initialized")606 607 async def shutdown(self):608 """Shutdown tool manager"""609 if self.web_search.session:610 await self.web_search.session.close()611 if self.database.connection:612 self.database.connection.close()613 self.logger.info("Tool manager shutdown")614