Coverage for agentos/swarm/tool_registry.py: 31%
319 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-10 01:20 +0800
1"""
2v1.9.8: Dynamic Tool Registry + Intelligent Tool Router.
4ToolRegistry: schema-based tool catalog with versioning, capability tags, and dependency tracking.
5ToolRouter: LLM-driven tool selection with semantic matching, confidence scoring, and fallback chains.
6"""
8from __future__ import annotations
10import json
11import time
12from collections import defaultdict
13from collections.abc import Callable
14from dataclasses import dataclass, field
15from enum import StrEnum
16from typing import Any
18# ── Tool Schema ───────────────────────────────────────────────────
21class ToolCategory(StrEnum):
22 """Top-level tool category for coarse-grained routing."""
24 FILE = "file"
25 NETWORK = "network"
26 CODE = "code"
27 SYSTEM = "system"
28 DATA = "data"
29 AGENT = "agent"
30 CUSTOM = "custom"
33@dataclass
34class ToolParam:
35 """Parameter definition for a tool."""
37 name: str
38 type: str # str, int, float, bool, list, dict
39 description: str = ""
40 required: bool = False
41 default: Any = None
42 enum_values: list[str] | None = None
43 min_value: float | None = None
44 max_value: float | None = None
45 pattern: str = "" # Regex validation pattern
48@dataclass
49class ToolSchema:
50 """Complete tool schema definition."""
52 name: str # Unique tool name
53 description: str # Human-readable description
54 category: ToolCategory = ToolCategory.CUSTOM
55 params: list[ToolParam] = field(default_factory=list)
56 returns: str = "any" # Return type description
57 version: str = "1.0.0"
58 capabilities: list[str] = field(default_factory=list) # e.g. ["read", "text", "file"]
59 tags: list[str] = field(default_factory=list) # Searchable tags
60 dependencies: list[str] = field(default_factory=list) # Required other tools
61 handler: Callable[..., Any] | None = None # Actual implementation
62 handler_ref: str = "" # String reference for serialization
63 cost_estimate: float = 0.0 # Relative cost (latency, tokens, etc)
64 is_destructive: bool = False # Data-modifying operations
65 requires_auth: bool = False # Needs authentication
66 rate_limit: int = 0 # Max calls per minute, 0 = unlimited
67 deprecated: bool = False
68 deprecated_message: str = ""
69 metadata: dict[str, Any] = field(default_factory=dict)
71 def to_openai_function(self) -> dict[str, Any]:
72 """Export schema as OpenAI function-calling format."""
73 properties = {}
74 required = []
75 for p in self.params:
76 prop: dict[str, Any] = {
77 "type": p.type,
78 "description": p.description,
79 }
80 if p.enum_values:
81 prop["enum"] = p.enum_values
82 properties[p.name] = prop
83 if p.required:
84 required.append(p.name)
86 return {
87 "type": "function",
88 "function": {
89 "name": self.name,
90 "description": self.description,
91 "parameters": {
92 "type": "object",
93 "properties": properties,
94 "required": required,
95 },
96 },
97 }
99 def match_score(self, query: str, keywords: list[str]) -> float:
100 """Compute relevance score for a natural language query."""
101 query_lower = query.lower()
102 score = 0.0
104 # Name exact match
105 if self.name.lower() == query_lower:
106 score += 10.0
107 elif self.name.lower() in query_lower:
108 score += 5.0
110 # Description match
111 desc_lower = self.description.lower()
112 if query_lower in desc_lower:
113 score += 3.0
114 for kw in keywords:
115 if kw in desc_lower:
116 score += 1.5
118 # Capability tag match
119 for cap in self.capabilities:
120 if cap.lower() in query_lower:
121 score += 2.0
123 # Tag match
124 for tag in self.tags:
125 if tag.lower() in query_lower or tag.lower() in keywords:
126 score += 1.0
128 # Parameter name match (user mentioned specific fields)
129 for p in self.params:
130 if p.name.lower() in query_lower:
131 score += 0.5
133 return score
136# ── Tool Registry ─────────────────────────────────────────────────
139class ToolRegistry:
140 """Central tool catalog with versioning, query, and lifecycle management.
142 Features:
143 - Schema-based registration with validation
144 - Semantic search over tool descriptions/capabilities/tags
145 - Version tracking and deprecation warnings
146 - Capability-based grouping
147 - Category-based organization
148 - Rate limiting enforcement
149 """
151 def __init__(self):
152 self._tools: dict[str, ToolSchema] = {}
153 self._by_category: dict[ToolCategory, list[str]] = defaultdict(list)
154 self._by_capability: dict[str, list[str]] = defaultdict(list)
155 self._by_tag: dict[str, list[str]] = defaultdict(list)
156 self._usage_counts: dict[str, int] = defaultdict(int)
157 self._rate_trackers: dict[str, list[float]] = defaultdict(list)
158 self._deprecation_log: list[dict[str, Any]] = []
160 def register(self, tool: ToolSchema) -> ToolSchema:
161 """Register a tool. Overwrites if same name (with warning)."""
162 if tool.name in self._tools:
163 existing = self._tools[tool.name]
164 if existing.version != tool.version:
165 # Version upgrade
166 pass
167 else:
168 pass # Overwrite silently
170 self._tools[tool.name] = tool
171 self._by_category[tool.category].append(tool.name)
172 for cap in tool.capabilities:
173 self._by_capability[cap].append(tool.name)
174 for tag in tool.tags:
175 self._by_tag[tag].append(tool.name)
177 return tool
179 def register_many(self, tools: list[ToolSchema]) -> list[ToolSchema]:
180 """Batch register tools."""
181 return [self.register(t) for t in tools]
183 def unregister(self, name: str) -> bool:
184 """Remove a tool from registry."""
185 if name not in self._tools:
186 return False
188 tool = self._tools.pop(name)
189 self._by_category[tool.category].remove(name)
190 for cap in tool.capabilities:
191 self._by_capability[cap].remove(name)
192 for tag in tool.tags:
193 self._by_tag[tag].remove(name)
194 return True
196 def get(self, name: str) -> ToolSchema | None:
197 """Get tool by name. Returns None and logs warning if deprecated."""
198 tool = self._tools.get(name)
199 if tool and tool.deprecated:
200 self._deprecation_log.append(
201 {
202 "tool": name,
203 "message": tool.deprecated_message,
204 "timestamp": time.time(),
205 }
206 )
207 return tool
209 def search(
210 self,
211 query: str,
212 top_k: int = 5,
213 category: ToolCategory | None = None,
214 exclude_deprecated: bool = True,
215 ) -> list[tuple[ToolSchema, float]]:
216 """Search tools by natural language query.
218 Returns ranked list of (ToolSchema, score).
219 """
220 keywords = query.lower().split()
221 candidates = []
223 tool_names = list(self._tools.keys())
224 if category:
225 tool_names = [n for n in tool_names if self._tools[n].category == category]
227 for name in tool_names:
228 tool = self._tools[name]
229 if exclude_deprecated and tool.deprecated:
230 continue
231 score = tool.match_score(query, keywords)
232 if score > 0:
233 # Boost frequently-used tools
234 usage_boost = min(self._usage_counts[name] * 0.1, 1.0)
235 candidates.append((tool, score + usage_boost))
237 candidates.sort(key=lambda x: -x[1])
238 return candidates[:top_k]
240 def search_by_capability(self, capability: str) -> list[ToolSchema]:
241 """Find all tools with a specific capability."""
242 names = self._by_capability.get(capability, [])
243 return [self._tools[n] for n in names if n in self._tools]
245 def search_by_tag(self, tag: str) -> list[ToolSchema]:
246 """Find all tools matching a tag."""
247 names = self._by_tag.get(tag, [])
248 return [self._tools[n] for n in names if n in self._tools]
250 def list_categories(self) -> dict[ToolCategory, int]:
251 """Count tools per category."""
252 return {cat: len(names) for cat, names in self._by_category.items() if names}
254 def list_capabilities(self) -> list[str]:
255 """List all registered capabilities."""
256 return sorted(self._by_capability.keys())
258 def list_tags(self) -> list[str]:
259 """List all registered tags."""
260 return sorted(self._by_tag.keys())
262 def export_openai_functions(
263 self,
264 category: ToolCategory | None = None,
265 exclude_deprecated: bool = True,
266 ) -> list[dict[str, Any]]:
267 """Export all tools as OpenAI function-calling format."""
268 result = []
269 for tool in self._tools.values():
270 if exclude_deprecated and tool.deprecated:
271 continue
272 if category and tool.category != category:
273 continue
274 result.append(tool.to_openai_function())
275 return result
277 def check_rate_limit(self, name: str) -> bool:
278 """Check if tool is within rate limit. Returns True if allowed."""
279 tool = self._tools.get(name)
280 if not tool or tool.rate_limit <= 0:
281 return True
283 now = time.time()
284 window_start = now - 60 # 1-minute window
285 calls = self._rate_trackers[name]
286 # Clean old entries
287 self._rate_trackers[name] = [t for t in calls if t > window_start]
289 return len(self._rate_trackers[name]) < tool.rate_limit
291 def record_usage(self, name: str) -> None:
292 """Record a tool usage for rate limiting and analytics."""
293 self._usage_counts[name] += 1
294 self._rate_trackers[name].append(time.time())
296 def get_stats(self) -> dict[str, Any]:
297 """Get registry statistics."""
298 return {
299 "total_tools": len(self._tools),
300 "categories": {str(k): len(v) for k, v in self._by_category.items()},
301 "capabilities": len(self._by_capability),
302 "tags": len(self._by_tag),
303 "deprecated": sum(1 for t in self._tools.values() if t.deprecated),
304 "top_used": sorted(
305 [(k, v) for k, v in self._usage_counts.items() if v > 0],
306 key=lambda x: -x[1],
307 )[:10],
308 }
311# ── Tool Router ───────────────────────────────────────────────────
314@dataclass
315class RoutingDecision:
316 """Result of tool routing decision."""
318 tool_name: str
319 tool_schema: ToolSchema | None
320 confidence: float # 0.0 - 1.0
321 reasoning: str # Why this tool was chosen
322 alternatives: list[str] # Fallback tool names
323 params: dict[str, Any] = field(default_factory=dict)
326@dataclass
327class RoutingContext:
328 """Context for tool routing decisions."""
330 task: str # User's task description
331 available_capabilities: list[str] = field(default_factory=list)
332 preferred_category: ToolCategory | None = None
333 exclude_destructive: bool = False
334 min_confidence: float = 0.3 # Minimum confidence threshold
335 max_alternatives: int = 3 # Max fallback alternatives
338class ToolRouter:
339 """Intelligent tool router with semantic matching and fallback chains.
341 Selects the best tool for a given task by:
342 1. Semantic matching via search query
343 2. LLM-driven selection (when available)
344 3. Rule-based fallback selection
345 4. Confidence scoring with threshold gating
346 """
348 def __init__(
349 self,
350 registry: ToolRegistry,
351 llm_selector: Callable[..., Any] | None = None,
352 ):
353 self.registry = registry
354 self.llm_selector = llm_selector # Optional LLM for smarter selection
356 def route(self, context: RoutingContext) -> RoutingDecision:
357 """Route a task to the best tool.
359 Priority:
360 1. LLM selector (if available) — best semantic understanding
361 2. Semantic search — keyword + capability matching
362 3. Default fallback
363 """
364 # Try LLM-based routing
365 if self.llm_selector and self._is_llm_worthwhile(context.task):
366 decision = self._llm_route(context)
367 if decision and decision.confidence >= context.min_confidence:
368 return decision
370 # Semantic search routing
371 return self._semantic_route(context)
373 def _is_llm_worthwhile(self, task: str) -> bool:
374 """Heuristic: LLM routing is worthwhile for complex tasks."""
375 # Simple one-word or obvious tool names don't need LLM
376 task_lower = task.lower().strip()
377 # If task is just a tool name, skip LLM
378 if task_lower in self.registry._tools:
379 return False
380 # If task is very short (~2 words), skip LLM
381 if len(task_lower.split()) <= 2:
382 return False
383 return True
385 def _llm_route(self, context: RoutingContext) -> RoutingDecision | None:
386 """Use LLM for intelligent tool selection."""
387 try:
388 tools_desc = self._build_tools_description(context)
389 prompt = (
390 f"Task: {context.task}\n\n"
391 f"Available tools:\n{tools_desc}\n\n"
392 "Select the best tool. Reply with JSON:\n"
393 '{"tool_name": "xxx", "confidence": 0.0-1.0, "reasoning": "why", '
394 '"alternatives": ["tool2", "tool3"]}'
395 )
396 result = self.llm_selector(prompt)
397 if isinstance(result, str):
398 result = json.loads(result)
400 tool_name = result.get("tool_name", "")
401 tool = self.registry.get(tool_name)
402 if not tool:
403 return None
405 return RoutingDecision(
406 tool_name=tool_name,
407 tool_schema=tool,
408 confidence=float(result.get("confidence", 0.5)),
409 reasoning=str(result.get("reasoning", "")),
410 alternatives=result.get("alternatives", []),
411 )
412 except Exception:
413 return None
415 def _semantic_route(self, context: RoutingContext) -> RoutingDecision:
416 """Semantic search-based routing with confidence scoring."""
417 candidates = self.registry.search(
418 query=context.task,
419 top_k=context.max_alternatives + 1,
420 category=context.preferred_category,
421 )
423 if not candidates:
424 return RoutingDecision(
425 tool_name="",
426 tool_schema=None,
427 confidence=0.0,
428 reasoning="No matching tool found",
429 alternatives=[],
430 )
432 # Filter destructive tools if excluded
433 if context.exclude_destructive:
434 candidates = [(t, s) for t, s in candidates if not t.is_destructive]
436 if not candidates:
437 return RoutingDecision(
438 tool_name="",
439 tool_schema=None,
440 confidence=0.0,
441 reasoning="All matching tools are destructive (excluded)",
442 alternatives=[],
443 )
445 # Normalize scores to 0-1 confidence
446 if len(candidates) == 1:
447 best_tool, raw_score = candidates[0]
448 confidence = min(raw_score / 10.0, 1.0)
449 alternatives = []
450 else:
451 scores = [s for _, s in candidates]
452 max_s = max(scores) if scores else 1
453 best_tool, raw_score = candidates[0]
454 confidence = min(raw_score / max_s, 1.0) if max_s > 0 else 0.5
455 alternatives = [t.name for t, _ in candidates[1 : context.max_alternatives + 1]]
457 return RoutingDecision(
458 tool_name=best_tool.name,
459 tool_schema=best_tool,
460 confidence=confidence,
461 reasoning=f"Best match: {best_tool.name} (score={raw_score:.1f})",
462 alternatives=alternatives,
463 )
465 def _build_tools_description(self, context: RoutingContext) -> str:
466 """Build a compact tool description for LLM prompt."""
467 tools = []
468 # Prioritize by category
469 names = list(self.registry._tools.keys())
470 if context.preferred_category:
471 cat_names = self.registry._by_category.get(context.preferred_category, [])
472 names = cat_names + [n for n in names if n not in cat_names]
474 for name in names[:20]: # Limit to avoid huge prompts
475 tool = self.registry._tools[name]
476 if tool.deprecated:
477 continue
478 if context.exclude_destructive and tool.is_destructive:
479 continue
480 params_desc = ", ".join(
481 f"{p.name}:{p.type}" + ("?" if not p.required else "") for p in tool.params[:5]
482 )
483 cap_tags = ", ".join(tool.capabilities[:3])
484 tools.append(
485 f"- {tool.name}: {tool.description[:100]}. "
486 f"Params: [{params_desc}]. Caps: [{cap_tags}]"
487 )
489 return "\n".join(tools)
492# ── Tool Execution Engine ─────────────────────────────────────────
495class ToolExecutionError(Exception):
496 """Raised when tool execution fails."""
498 def __init__(self, tool_name: str, message: str, recoverable: bool = True):
499 self.tool_name = tool_name
500 self.recoverable = recoverable
501 super().__init__(f"[{tool_name}] {message}")
504class ToolExecutor:
505 """Execution engine for registered tools with safety and error handling.
507 Features:
508 - Rate limit enforcement
509 - Parameter validation
510 - Destructive operation confirmation
511 - Timeout protection
512 - Error categorization (recoverable vs fatal)
513 """
515 def __init__(
516 self,
517 registry: ToolRegistry,
518 timeout: float = 30.0,
519 require_destructive_confirm: bool = True,
520 ):
521 self.registry = registry
522 self.timeout = timeout
523 self.require_destructive_confirm = require_destructive_confirm
524 self._pending_confirmations: dict[str, dict[str, Any]] = {}
526 def execute(
527 self,
528 tool_name: str,
529 params: dict[str, Any] | None = None,
530 force: bool = False,
531 ) -> Any:
532 """Execute a registered tool with safety checks.
534 Args:
535 tool_name: Registered tool name
536 params: Tool parameters
537 force: Skip destructive confirmation (use with caution)
539 Returns:
540 Tool execution result
542 Raises:
543 ToolExecutionError: On execution failure
544 """
545 params = params or {}
547 tool = self.registry.get(tool_name)
548 if not tool:
549 raise ToolExecutionError(
550 tool_name, f"Tool '{tool_name}' not registered", recoverable=False
551 )
553 if tool.deprecated:
554 raise ToolExecutionError(
555 tool_name,
556 f"Tool deprecated: {tool.deprecated_message}",
557 recoverable=False,
558 )
560 # Rate limit check
561 if not self.registry.check_rate_limit(tool_name):
562 raise ToolExecutionError(
563 tool_name,
564 f"Rate limit exceeded ({tool.rate_limit}/min)",
565 recoverable=True,
566 )
568 # Destructive check
569 if tool.is_destructive and self.require_destructive_confirm and not force:
570 self._pending_confirmations[tool_name] = params
571 raise ToolExecutionError(
572 tool_name,
573 "Destructive operation requires confirmation (pass force=True to skip)",
574 recoverable=False,
575 )
577 # Parameter validation
578 self._validate_params(tool, params)
580 # Record usage (before execution to prevent double-counting on retry)
581 self.registry.record_usage(tool_name)
583 # Execute
584 if not tool.handler:
585 raise ToolExecutionError(
586 tool_name,
587 "No handler registered for tool",
588 recoverable=False,
589 )
591 try:
592 result = tool.handler(**params)
593 except Exception as e:
594 raise ToolExecutionError(
595 tool_name,
596 f"Execution failed: {str(e)}",
597 recoverable=True,
598 ) from e
600 return result
602 def confirm_destructive(self, tool_name: str) -> Any:
603 """Confirm and execute a pending destructive operation."""
604 if tool_name not in self._pending_confirmations:
605 raise ToolExecutionError(tool_name, "No pending confirmation", recoverable=False)
606 params = self._pending_confirmations.pop(tool_name)
607 return self.execute(tool_name, params, force=True)
609 def cancel_destructive(self, tool_name: str) -> bool:
610 """Cancel a pending destructive operation."""
611 if tool_name in self._pending_confirmations:
612 del self._pending_confirmations[tool_name]
613 return True
614 return False
616 def _validate_params(self, tool: ToolSchema, params: dict[str, Any]) -> None:
617 """Validate parameters against schema."""
618 for p in tool.params:
619 if p.required and p.name not in params:
620 raise ToolExecutionError(
621 tool.name,
622 f"Missing required parameter: {p.name} ({p.description})",
623 recoverable=False,
624 )
626 if p.name in params:
627 value = params[p.name]
628 # Type check
629 if p.type == "str" and not isinstance(value, str):
630 raise ToolExecutionError(
631 tool.name, f"Parameter '{p.name}' must be string", recoverable=False
632 )
633 if p.type == "int" and not isinstance(value, int):
634 raise ToolExecutionError(
635 tool.name, f"Parameter '{p.name}' must be int", recoverable=False
636 )
637 if p.type == "float" and not isinstance(value, (int, float)):
638 raise ToolExecutionError(
639 tool.name, f"Parameter '{p.name}' must be number", recoverable=False
640 )
641 if p.type == "bool" and not isinstance(value, bool):
642 raise ToolExecutionError(
643 tool.name, f"Parameter '{p.name}' must be bool", recoverable=False
644 )
646 # Enum check
647 if p.enum_values and value not in p.enum_values:
648 raise ToolExecutionError(
649 tool.name,
650 f"Parameter '{p.name}' must be one of: {p.enum_values}",
651 recoverable=False,
652 )
654 # Range check
655 if isinstance(value, (int, float)):
656 if p.min_value is not None and value < p.min_value:
657 raise ToolExecutionError(
658 tool.name,
659 f"Parameter '{p.name}' minimum is {p.min_value}",
660 recoverable=False,
661 )
662 if p.max_value is not None and value > p.max_value:
663 raise ToolExecutionError(
664 tool.name,
665 f"Parameter '{p.name}' maximum is {p.max_value}",
666 recoverable=False,
667 )
669 def get_pending_confirmations(self) -> list[str]:
670 """List tools awaiting destructive confirmation."""
671 return list(self._pending_confirmations.keys())
674# ── Utility helpers ───────────────────────────────────────────────
677def create_tool(
678 name: str,
679 description: str,
680 handler: Callable,
681 category: ToolCategory = ToolCategory.CUSTOM,
682 params: list[ToolParam] | None = None,
683 capabilities: list[str] | None = None,
684 tags: list[str] | None = None,
685 is_destructive: bool = False,
686 rate_limit: int = 0,
687 **kwargs,
688) -> ToolSchema:
689 """Quick helper to create a tool schema."""
690 return ToolSchema(
691 name=name,
692 description=description,
693 category=category,
694 params=params or [],
695 capabilities=capabilities or [],
696 tags=tags or [],
697 handler=handler,
698 is_destructive=is_destructive,
699 rate_limit=rate_limit,
700 **kwargs,
701 )