active-tool-selection¶
第4章 · 工具 · 配套项目
chapter4/active-tool-selection
项目说明¶
Active Tool Selection¶
An educational implementation of active tool discovery for LLM agents, inspired by the MCP-Zero paper (arXiv:2506.01056).
🎯 Overview¶
Traditional LLM agents inject all available tool schemas into prompts, creating massive context overhead and reducing agents to passive tool selectors. This project demonstrates active tool discovery, where agents autonomously identify capability gaps and request specific tools on-demand.
The Problem¶
Current tool integration approaches face critical limitations:
- Massive Context Overhead: Injecting all tools can consume 100k+ tokens
- Passive Selection: Agents select from pre-defined options rather than actively discovering
- Poor Scalability: Context grows with ecosystem size, not task needs
- Lost Autonomy: Tool selection delegated to external retrieval systems
The Solution: Active Tool Discovery¶
This project implements three core mechanisms from MCP-Zero:
- Active Tool Request: Agents generate structured requests specifying their exact tool requirements
- Hierarchical Semantic Routing: Two-stage matching algorithm (server-level → tool-level)
- Iterative Capability Extension: Progressive toolchain building as task understanding evolves
🔬 三种策略对比 (Strategy Comparison)¶
本实验把"工具选择"问题落到可度量的基准上，对比三种策略在同一批任务上的表现：
| 策略 | 说明 | 上下文里的工具 |
|---|---|---|
all-tools |
一次性注入全部工具（传统被动式基线） | 全部 N 个 |
retrieval |
按任务语义检索 top-k 个工具后再注入（工具检索 / RAG 式，RetrievalToolAgent） |
仅 top-k 个 |
active |
MCP-Zero 式主动发现：模型迭代地请求所需工具（ActiveToolAgent） |
按需增长 |
评测入口是 demo_comparison.py，带完整的 argparse 命令行：
# 仅离线对比（确定性，无需 API Key）：召回率 vs token 成本 vs 随规模的扩展性
python demo_comparison.py --offline
# 把工具目录扩充到 200 个（合成干扰工具补齐），观察 token 成本的分化
python demo_comparison.py --offline --num-tools 200
# 三种策略端到端对比（需要 API Key）：模型是否真的调用了正确的工具、token、延迟
python demo_comparison.py --strategy compare
# 只对单条查询运行某种策略
python demo_comparison.py --query "Deploy version 2.0 to production" --strategy retrieval
# 保存结果为 JSON
python demo_comparison.py --offline --output results.json
运行 python demo_comparison.py --help 查看全部参数（--strategy / --query / --num-tools /
--top-k / --model / --output / --offline / --legacy-demos）。
离线基准（确定性，无需 API）¶
benchmark.py 提供了一个带标准答案工具的小型基准集（10 个任务，每个任务标注了应当被选中的
工具），并在不调用任何 API 的情况下度量两件事：
- Retrieval recall@k：标准答案工具是否落在被注入上下文的工具集合里；
- Schema tokens：注入的工具描述占用的 token 数（由 schema 直接估算，确定性可复现）。
下表是 python demo_comparison.py --offline 的实测输出（top-k=5，10 个任务）：
| 策略 | 上下文工具数 | Schema tokens | 召回率（标准答案可达） |
|---|---|---|---|
| all-tools（全部注入） | 35 | 3,857 | 100% |
| retrieval（top-5） | 5 | 551 | 100% |
随着工具目录增长，all-tools 的 token 成本线性膨胀，而 retrieval 基本持平（实测）：
| 目录规模 | all-tools tokens | retrieval(top-5) tokens | retrieval 召回率 |
|---|---|---|---|
| 35 | 3,857 | 551 | 100% |
| 100 | 10,292 | 539 | 100% |
| 200 | 20,258 | 540 | 100% |
| 400 | 40,258 | 540 | 100% |
结论：检索式按需选择在保持 100% 召回率的同时，把工具描述的 token 成本从数千压到数百，且不随 生态规模膨胀。这正是本章"把工具选择转化为知识检索"的量化体现。上述数字由
--offline路径 确定性生成，可直接复现。
端到端准确率（需要 API Key）¶
在配置 OPENAI_API_KEY 后，--strategy compare 会真正调用模型，度量每种策略下模型是否调用了
标准答案工具（accuracy）、平均 token 与平均延迟。这一部分需要联网与 API，故不在离线路径中运行。
🏗️ Architecture¶
┌─────────────────────────────────────────────────────────────┐
│ Active Tool Agent │
│ ┌────────────────────────────────────────────────────────┐ │
│ │ 1. Analyze Task & Identify Capability Gaps │ │
│ └────────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────────┐ │
│ │ 2. Generate Structured Tool Request │ │
│ │ <tool_request> │ │
│ │ server: GitHub for repository operations │ │
│ │ tool: search repositories by keyword │ │
│ │ </tool_request> │ │
│ └────────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────┐
│ Hierarchical Semantic Router │
│ ┌────────────────────────────────────────────────────────┐ │
│ │ Stage 1: Server-Level Routing │ │
│ │ • Match request to relevant servers/domains │ │
│ │ • Filter by platform requirements │ │
│ │ • Return top-K servers │ │
│ └────────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────────┐ │
│ │ Stage 2: Tool-Level Routing │ │
│ │ • Rank tools within selected servers │ │
│ │ • Semantic similarity matching │ │
│ │ • Return relevant tools │ │
│ └────────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────┐
│ Tool Knowledge Base │
│ │
│ 8 Servers × 35 Tools (optionally padded with distractors): │
│ • GitHub: Repository management (5 tools) │
│ • Filesystem: File operations (5 tools) │
│ • Database: SQL operations (5 tools) │
│ • Web: HTTP requests (4 tools) │
│ • Analytics: Data analysis (4 tools) │
│ • Communication: Email/messaging (4 tools) │
│ • DevOps: Deployment/monitoring (4 tools) │
│ • Cloud: Infrastructure management (4 tools) │
└─────────────────────────────────────────────────────────────┘
📦 Components¶
Core Files¶
agent.py: Three agent implementations (one per strategy)ActiveToolAgent: MCP-Zero style on-demand discovery (active)RetrievalToolAgent: one-shot semantic retrieval of top-k tools (retrieval)-
PassiveToolAgent: traditional approach with all tools pre-loaded (all-tools) -
benchmark.py: Labeled benchmark + offline evaluation - 10 tasks, each labeled with its ground-truth tool
build_catalog(num_tools): real catalog, optionally padded with distractors-
evaluate_offline(...): deterministic recall@k / token-cost measurement (no API) -
tool_knowledge_base.py: Comprehensive tool catalog - 8 servers (domains) with 35 tools
- Organized by platform/functionality
-
Simulates MCP ecosystem
-
semantic_router.py: Hierarchical tool discovery - Two-stage semantic matching
- TF-IDF based similarity
-
Structured request parsing
-
config.py: Configuration settings - LLM provider settings
- Routing thresholds
- Agent parameters
Demo Scripts¶
quickstart.py: Quick demonstration (⭐ Start here!)demo_comparison.py: Comprehensive comparisonexamples.py: Multiple use case examples
🚀 Quick Start¶
1. Install Dependencies¶
2. Configure API Key¶
Universal OpenRouter fallback: if
OPENAI_API_KEYis not set butOPENROUTER_API_KEYis,config.pyautomatically routes through OpenRouter (base_url=https://openrouter.ai/api/v1) and maps the model id toprovider/modelform (gpt-*→openai/…,claude-*→anthropic/claude-opus-4.8). ExistingOPENAI_BASE_URL/OPENAI_MODELoverrides are preserved.
3. Run Quick Start¶
This will demonstrate: - Active tool discovery process - Passive tool injection (for comparison) - Efficiency metrics and insights
📊 Performance Comparison¶
See the measured, reproducible numbers in 三种策略对比 above.
The offline table (schema-token cost + retrieval recall) is generated deterministically by
python demo_comparison.py --offline — no API key required. End-to-end accuracy/latency across
the three strategies requires an API key (--strategy compare).
Note: earlier drafts of this README quoted round illustrative figures for token savings. Those have been replaced with the actual measured output of the offline benchmark to avoid fabricated numbers.
🎓 Key Concepts¶
Active Tool Request¶
Instead of pre-loading tools, agents explicitly request what they need:
<tool_request>
server: GitHub for repository management
tool: search repositories by stars and language
</tool_request>
Hierarchical Semantic Routing¶
Two-stage algorithm reduces search complexity:
- Stage 1: Match to relevant servers (platforms)
- "GitHub operations" → GitHub server
-
"File management" → Filesystem server
-
Stage 2: Match to specific tools within servers
- "search repositories" →
github_search_repos - "read file" →
fs_read_file
Iterative Capability Extension¶
Tools are discovered progressively:
Turn 1: Need GitHub access
→ Load GitHub tools
Turn 2: Need to analyze downloaded files
→ Additionally load filesystem tools
Turn 3: Need to visualize results
→ Additionally load analytics tools
Toolchain grows with task complexity, not ecosystem size.
📚 Examples¶
Example 1: Simple Task¶
from agent import ActiveToolAgent
agent = ActiveToolAgent()
result = agent.execute_task("Search for Python ML repositories on GitHub")
# Agent discovers and loads only GitHub tools
print(f"Tools loaded: {result['metrics']['tools_loaded']}") # 2-3 tools
print(f"Tokens used: {result['metrics']['tokens_used']}") # ~2,000
Example 2: Multi-Domain Task¶
task = """
1. Query database for user data
2. Analyze with statistics
3. Create visualization
4. Email report to team
"""
result = agent.execute_task(task)
# Agent builds cross-domain toolchain:
# Database → Analytics → Communication
print(f"Tools: {result['tools_loaded']}")
Example 3: Comparison¶
from agent import ActiveToolAgent, PassiveToolAgent
# Active approach
active = ActiveToolAgent()
active_result = active.execute_task(task)
# Passive approach
passive = PassiveToolAgent()
passive_result = passive.execute_task(task)
# Compare efficiency
reduction = (1 - active_result['metrics']['tokens_used'] /
passive_result['metrics']['tokens_used']) * 100
print(f"Token reduction: {reduction:.1f}%") # Typically 90-98%
🔬 Running Demonstrations¶
1. Quick Start (Recommended first)¶
Shows basic active vs passive comparison.
2. Strategy Comparison Benchmark (main experiment)¶
python demo_comparison.py --offline # deterministic, no API key needed
python demo_comparison.py # + end-to-end accuracy/latency if API key present
By default it prints:
- Offline strategy table: retrieval recall@k vs. tool-schema token cost (deterministic)
- Scaling table: token cost as the catalog grows to hundreds of tools
- End-to-end table (with API key): accuracy / tokens / latency for all-tools, retrieval, active
The original narrative demos (semantic-routing walk-through, iterative discovery, etc.) are still
available via --legacy-demos. See the Strategy Comparison
section and python demo_comparison.py --help for all flags.
3. Use Case Examples¶
Demonstrates: - GitHub workflow - Data pipeline - DevOps automation - Multi-turn discovery - Efficiency metrics
🎯 Use Cases¶
Ideal for Active Discovery¶
- Task-Specific Operations: Known scope, specific tools needed
- Large Tool Ecosystems: 50+ tools where most are irrelevant
- Multi-Turn Conversations: Tools needed evolve over time
- Token-Constrained Environments: Limited context windows
When Passive Might Work¶
- Small Tool Sets: <10 tools total
- All Tools Relevant: Every tool likely to be used
- Single-Turn Tasks: No iterative refinement
🔧 Configuration¶
Edit config.py or set environment variables:
# LLM Configuration
OPENAI_API_KEY = "your-api-key"
OPENAI_BASE_URL = "https://api.openai.com/v1"
OPENAI_MODEL = "gpt-5.6-luna"
# Routing Configuration
SIMILARITY_THRESHOLD = 0.15 # Min similarity for tool match
TOP_K_SERVERS = 3 # Number of servers to search
TOP_K_TOOLS = 5 # Tools to return per server
# Agent Configuration
MAX_TOOL_REQUESTS = 5 # Max discovery iterations
📖 Educational Value¶
This project demonstrates:
Software Engineering Principles¶
- Separation of Concerns: Router, knowledge base, agent cleanly separated
- Scalability: Efficient with 10 or 1,000 tools
- Modularity: Easy to add new servers/tools
AI Agent Design Patterns¶
- Active vs Passive: Fundamental architectural difference
- Hierarchical Search: Reduces complexity from O(n) to O(log n)
- Semantic Matching: Beyond keyword matching
Real-World Applications¶
- Tool Ecosystems: MCP, LangChain, AutoGen
- Agent Frameworks: Building production-ready agents
- Context Management: Handling long contexts efficiently
🔬 Technical Details¶
Semantic Routing Implementation¶
Uses TF-IDF vectorization with cosine similarity:
# Server-level
server_vector = vectorizer.transform([request])
similarities = cosine_similarity(server_vector, server_embeddings)
# Tool-level
tool_vector = vectorizer.transform([request])
tool_similarities = cosine_similarity(tool_vector, tool_embeddings)
# Combined score
final_score = 0.3 * server_score + 0.7 * tool_score
Token Estimation¶
Approximates tokens in tool schemas:
def count_tokens_in_schema(schema):
# Rough estimation: 1 token ≈ 4 characters
schema_str = json.dumps(schema)
return len(schema_str) // 4
🌟 Key Insights¶
- Autonomy Matters: Agents should control their capability acquisition
- Context is Expensive: Every token counts at scale
- Semantic Matching Works: TF-IDF sufficient for tool discovery
- Iteration Enables Flexibility: Static tool sets can't anticipate needs
- Hierarchical Search Scales: Two-stage routing maintains performance
📚 References¶
- MCP-Zero Paper: arXiv:2506.01056
- Model Context Protocol: Official Repository
- Tool Learning Survey: ACM Computing Surveys
🛠️ Extending the Project¶
Add New Tools¶
# In tool_knowledge_base.py
new_tool = ToolDefinition(
name="your_tool_name",
description="What the tool does",
parameters={...},
server="server_name"
)
Add New Server¶
# Create tools for the server
tools = [...]
# Add server
servers.append(ServerDefinition(
name="your_server",
description="Server description",
tools=tools
))
Customize Routing¶
# In config.py
SIMILARITY_THRESHOLD = 0.5 # More strict matching
TOP_K_SERVERS = 5 # Search more servers
🐛 Troubleshooting¶
API Key Issues¶
Import Errors¶
Low Similarity Scores¶
📈 Future Enhancements¶
Potential improvements:
- Better Embeddings: Use sentence-transformers or OpenAI embeddings
- Caching: Cache tool embeddings for faster routing
- Feedback Loop: Learn from tool usage patterns
- Multi-Agent: Tool sharing between agent instances
- Real Tools: Connect to actual APIs instead of simulation
🤝 Contributing¶
This is an educational project. Feel free to:
- Add more tools to the knowledge base
- Implement additional routing algorithms
- Create new demonstration scenarios
- Improve documentation
📄 License¶
MIT License - See LICENSE file for details
🙏 Acknowledgments¶
- Inspired by MCP-Zero paper by Xiang Fei, Xiawu Zheng, and Hao Feng
- Based on Model Context Protocol (MCP) ecosystem
- Built for educational purposes in AI Agent development
Ready to get started? Run python quickstart.py to see active tool discovery in action!
源代码¶
agent.py¶
"""
Active Tool Discovery Agent.
Implements an LLM agent that actively requests tools on-demand rather than
having all tool schemas injected into the prompt. Inspired by MCP-Zero.
"""
from typing import List, Dict, Any, Optional
from openai import OpenAI
from tool_knowledge_base import ToolDefinition, ServerDefinition, create_tool_knowledge_base
from semantic_router import SemanticRouter, StructuredRequestParser
import config
class ActiveToolAgent:
"""
Agent that actively discovers and requests tools as needed.
Key principles:
1. Maintains minimal context by not injecting all tools upfront
2. Actively requests specific tools when capability gaps are identified
3. Iteratively builds toolchain as task understanding evolves
"""
def __init__(self, servers: Optional[List[ServerDefinition]] = None,
model: Optional[str] = None):
self.client = OpenAI(
api_key=config.OPENAI_API_KEY,
base_url=config.OPENAI_BASE_URL
)
self.model = model or config.OPENAI_MODEL
# Initialize tool knowledge base (callers may inject a padded/custom catalog)
self.servers = servers if servers is not None else create_tool_knowledge_base()
self.router = SemanticRouter(self.servers)
# Agent state
self.conversation_history = []
self.available_tools: List[ToolDefinition] = [] # Currently loaded tools
self.tool_request_count = 0
# Metrics
self.metrics = {
'tokens_used': 0,
'tool_requests': 0,
'tools_loaded': 0,
'api_calls': 0,
'tools_called': [] # Names of tools the model actually invoked
}
def execute_task(self, task: str) -> Dict[str, Any]:
"""
Execute a task with active tool discovery.
The agent will:
1. Analyze the task
2. Identify capability gaps
3. Request specific tools
4. Execute with discovered tools
Returns execution results with metrics.
"""
self.conversation_history = []
self.available_tools = []
self.tool_request_count = 0
# Initial system message explaining active tool discovery
system_message = self._create_system_message()
self.conversation_history.append({
"role": "system",
"content": system_message
})
# Add user task
self.conversation_history.append({
"role": "user",
"content": task
})
# Iterative tool discovery and execution
max_iterations = config.MAX_TOOL_REQUESTS
for iteration in range(max_iterations):
# Get agent response
response = self._call_llm()
self.metrics['api_calls'] += 1
# Check if agent is requesting tools
tool_request = StructuredRequestParser.parse_request(response)
if tool_request:
# Agent is requesting tools - discover and provide them
self._handle_tool_request(tool_request, response)
self.tool_request_count += 1
self.metrics['tool_requests'] += 1
else:
# Agent has what it needs and is responding
self.conversation_history.append({
"role": "assistant",
"content": response
})
break
return {
'response': response,
'metrics': self.metrics,
'tools_loaded': [t.name for t in self.available_tools],
'conversation': self.conversation_history
}
def _create_system_message(self) -> str:
"""Create system message explaining active tool discovery."""
return """You are an autonomous AI agent with active tool discovery capabilities.
Instead of having all possible tools available upfront, you can actively request tools as you need them. This allows you to:
1. Maintain a minimal context footprint
2. Focus on relevant capabilities for the current task
3. Iteratively build your toolchain as your understanding evolves
When you identify a capability gap, request tools using this format:
<tool_request>
server: [describe the platform/domain you need, e.g., "GitHub for repository operations" or "filesystem for local file access"]
tool: [describe the specific operation you need, e.g., "search repositories" or "read file contents"]
</tool_request>
After requesting tools, they will be provided to you. You can then use them to accomplish the task.
Process:
1. Analyze the task and identify what capabilities you need
2. Request specific tools if you don't have them yet
3. Once you have the necessary tools, use them to complete the task
4. Respond with your findings or results
Current available tools: None (request tools as needed)"""
def _call_llm(self) -> str:
"""Call LLM with current context and available tools."""
kwargs = {
"model": self.model,
"messages": self.conversation_history,
"temperature": config.AGENT_TEMPERATURE
}
# Add tools if available
if self.available_tools:
kwargs["tools"] = [tool.to_schema() for tool in self.available_tools]
kwargs["tool_choice"] = "auto"
response = self.client.chat.completions.create(**kwargs)
# Track token usage
# response.usage is Optional in the OpenAI SDK: the attribute always
# exists, but is None when the provider omits token accounting.
if getattr(response, 'usage', None):
self.metrics['tokens_used'] += response.usage.total_tokens
# Extract response content
message = response.choices[0].message
# Handle tool calls if present
if message.tool_calls:
return self._handle_tool_calls(message)
return message.content or ""
def _handle_tool_request(self, tool_request: Dict[str, str], full_response: str):
"""
Handle tool request from agent.
Args:
tool_request: Parsed tool request with 'server' and 'tool' fields
full_response: Full response text from agent
"""
# Combine server and tool descriptions for routing
query = f"{tool_request['server']} {tool_request['tool']}"
# Use semantic router to find relevant tools
discovered_tools = self.router.route_request(query)
if not discovered_tools:
# No tools found
feedback = f"""No tools found matching your request. Please refine your request or proceed without additional tools.
Your request was:
- Server: {tool_request['server']}
- Tool: {tool_request['tool']}"""
else:
# Add discovered tools to available tools
new_tools = []
for tool in discovered_tools:
if tool not in self.available_tools:
self.available_tools.append(tool)
new_tools.append(tool)
self.metrics['tools_loaded'] += 1
tool_list = "\n".join([f"- {t.name}: {t.description}" for t in new_tools])
feedback = f"""Tools discovered and loaded ({len(new_tools)} new tools):
{tool_list}
You can now use these tools to complete the task. Please proceed."""
# Add agent's request and system's response to history
self.conversation_history.append({
"role": "assistant",
"content": full_response
})
self.conversation_history.append({
"role": "user",
"content": feedback
})
def _handle_tool_calls(self, message) -> str:
"""Handle actual tool execution (simulated for demo)."""
# For this educational demo, we simulate tool execution
tool_results = []
for tool_call in message.tool_calls:
func_name = tool_call.function.name
self.metrics['tools_called'].append(func_name)
# Simulate tool execution
result = f"[Simulated] Tool '{func_name}' executed successfully with result: Success"
tool_results.append({
"tool_call_id": tool_call.id,
"output": result
})
# Add tool call message to history
self.conversation_history.append({
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": tc.id,
"type": "function",
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments
}
}
for tc in message.tool_calls
]
})
# Add tool results to history
for result in tool_results:
self.conversation_history.append({
"role": "tool",
"tool_call_id": result["tool_call_id"],
"content": result["output"]
})
# Get final response after tool execution
return self._call_llm()
def reset(self):
"""Reset agent state."""
self.conversation_history = []
self.available_tools = []
self.tool_request_count = 0
self.metrics = {
'tokens_used': 0,
'tool_requests': 0,
'tools_loaded': 0,
'api_calls': 0,
'tools_called': []
}
class RetrievalToolAgent:
"""
One-shot retrieval agent (semantic tool retrieval / "工具检索").
This is the RAG-style middle ground between passive injection and active
discovery: before the very first LLM call, it retrieves the top-k tools most
semantically relevant to the task and injects *only* those. There is no extra
discovery round-trip — tool selection is delegated to the retriever, turning the
"which of hundreds of tools" problem into a knowledge-retrieval problem.
This directly embodies the mechanism the chapter attributes to Anthropic's
on-demand tool retrieval experiment: fewer, more relevant tool schemas in
context both cut token cost and reduce the model's selection errors.
"""
def __init__(self, servers: Optional[List[ServerDefinition]] = None,
model: Optional[str] = None, top_k: Optional[int] = None):
self.client = OpenAI(
api_key=config.OPENAI_API_KEY,
base_url=config.OPENAI_BASE_URL
)
self.model = model or config.OPENAI_MODEL
self.top_k = top_k if top_k is not None else config.TOP_K_TOOLS
self.servers = servers if servers is not None else create_tool_knowledge_base()
self.router = SemanticRouter(self.servers)
self.conversation_history = []
self.available_tools: List[ToolDefinition] = []
self.metrics = {
'tokens_used': 0,
'tools_loaded': 0,
'api_calls': 0,
'tools_called': []
}
def execute_task(self, task: str) -> Dict[str, Any]:
"""Retrieve top-k relevant tools for the task, then execute in one shot."""
self.conversation_history = []
# Retrieval step (no LLM call): pick the top-k most relevant tools.
self.available_tools = self.router.retrieve(task, self.top_k)
self.metrics['tools_loaded'] = len(self.available_tools)
tool_list = "\n".join(
f"- {t.name}: {t.description}" for t in self.available_tools
)
system_message = f"""You are an AI agent. A retrieval system has pre-selected the \
{len(self.available_tools)} tools below as most relevant to the user's task.
{tool_list}
Analyze the task and call the appropriate tool(s) to complete it."""
self.conversation_history.append({"role": "system", "content": system_message})
self.conversation_history.append({"role": "user", "content": task})
response = self._call_llm()
self.metrics['api_calls'] += 1
return {
'response': response,
'metrics': self.metrics,
'tools_loaded': [t.name for t in self.available_tools],
'conversation': self.conversation_history
}
def _call_llm(self) -> str:
"""Call LLM with only the retrieved tools injected."""
kwargs = {
"model": self.model,
"messages": self.conversation_history,
"temperature": config.AGENT_TEMPERATURE
}
if self.available_tools:
kwargs["tools"] = [tool.to_schema() for tool in self.available_tools]
kwargs["tool_choice"] = "auto"
response = self.client.chat.completions.create(**kwargs)
# response.usage is Optional in the OpenAI SDK: the attribute always
# exists, but is None when the provider omits token accounting.
if getattr(response, 'usage', None):
self.metrics['tokens_used'] += response.usage.total_tokens
message = response.choices[0].message
if message.tool_calls:
return self._handle_tool_calls(message)
return message.content or ""
def _handle_tool_calls(self, message) -> str:
"""Handle tool execution (simulated)."""
tool_results = []
for tool_call in message.tool_calls:
func_name = tool_call.function.name
self.metrics['tools_called'].append(func_name)
result = f"[Simulated] Tool '{func_name}' executed successfully"
tool_results.append({"tool_call_id": tool_call.id, "output": result})
self.conversation_history.append({
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": tc.id,
"type": "function",
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments
}
}
for tc in message.tool_calls
]
})
for result in tool_results:
self.conversation_history.append({
"role": "tool",
"tool_call_id": result["tool_call_id"],
"content": result["output"]
})
return self._call_llm()
def reset(self):
"""Reset agent state."""
self.conversation_history = []
self.available_tools = []
self.metrics = {
'tokens_used': 0,
'tools_loaded': 0,
'api_calls': 0,
'tools_called': []
}
class PassiveToolAgent:
"""
Traditional agent with all tools injected upfront (for comparison).
This approach:
1. Injects all tool schemas into the initial prompt
2. Massive context overhead
3. Reduces agent to passive tool selector
"""
def __init__(self, servers: Optional[List[ServerDefinition]] = None,
model: Optional[str] = None):
self.client = OpenAI(
api_key=config.OPENAI_API_KEY,
base_url=config.OPENAI_BASE_URL
)
self.model = model or config.OPENAI_MODEL
# Load ALL tools upfront
self.servers = servers if servers is not None else create_tool_knowledge_base()
self.all_tools = []
for server in self.servers:
self.all_tools.extend(server.tools)
self.conversation_history = []
self.metrics = {
'tokens_used': 0,
'tools_loaded': len(self.all_tools),
'api_calls': 0,
'tools_called': []
}
def execute_task(self, task: str) -> Dict[str, Any]:
"""Execute task with all tools pre-loaded."""
self.conversation_history = []
# System message
system_message = f"""You are an AI agent with access to {len(self.all_tools)} tools across multiple domains.
All available tools have been pre-loaded. Analyze the task and use the appropriate tools to complete it."""
self.conversation_history.append({
"role": "system",
"content": system_message
})
self.conversation_history.append({
"role": "user",
"content": task
})
# Call LLM with ALL tools
response = self._call_llm()
self.metrics['api_calls'] += 1
return {
'response': response,
'metrics': self.metrics,
'tools_loaded': [t.name for t in self.all_tools],
'conversation': self.conversation_history
}
def _call_llm(self) -> str:
"""Call LLM with ALL tools injected."""
kwargs = {
"model": self.model,
"messages": self.conversation_history,
"temperature": config.AGENT_TEMPERATURE,
"tools": [tool.to_schema() for tool in self.all_tools],
"tool_choice": "auto"
}
response = self.client.chat.completions.create(**kwargs)
# Track token usage
# response.usage is Optional in the OpenAI SDK: the attribute always
# exists, but is None when the provider omits token accounting.
if getattr(response, 'usage', None):
self.metrics['tokens_used'] += response.usage.total_tokens
message = response.choices[0].message
# Handle tool calls (simulated)
if message.tool_calls:
return self._handle_tool_calls(message)
return message.content or ""
def _handle_tool_calls(self, message) -> str:
"""Handle tool execution (simulated)."""
tool_results = []
for tool_call in message.tool_calls:
func_name = tool_call.function.name
self.metrics['tools_called'].append(func_name)
result = f"[Simulated] Tool '{func_name}' executed successfully"
tool_results.append({
"tool_call_id": tool_call.id,
"output": result
})
# Add to history
self.conversation_history.append({
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": tc.id,
"type": "function",
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments
}
}
for tc in message.tool_calls
]
})
for result in tool_results:
self.conversation_history.append({
"role": "tool",
"tool_call_id": result["tool_call_id"],
"content": result["output"]
})
return self._call_llm()
def reset(self):
"""Reset agent state."""
self.conversation_history = []
self.metrics = {
'tokens_used': 0,
'tools_loaded': len(self.all_tools),
'api_calls': 0,
'tools_called': []
}
benchmark.py¶
"""
Tool-selection benchmark and offline evaluation.
Provides a small labeled benchmark (task -> ground-truth tool) and utilities to
quantify, *without any API calls*, the core claim of the chapter: when a tool
ecosystem grows to hundreds of tools, retrieving the few relevant tools on demand
keeps the right tool reachable while slashing the token cost of dumping every tool
schema into context.
Two things are measured here deterministically:
1. Retrieval recall@k — is the ground-truth tool among the tools a strategy
places in the model's context?
2. Context schema tokens — how many tokens the injected tool schemas cost.
End-to-end accuracy/latency (whether the model actually *calls* the right tool)
requires an API key and lives in demo_comparison.py.
"""
from typing import List, Dict
from tool_knowledge_base import (
ToolDefinition,
ServerDefinition,
create_tool_knowledge_base,
get_all_tools,
calculate_total_tokens,
)
from semantic_router import SemanticRouter
# Labeled benchmark: each task has one (or a few acceptable) ground-truth tool(s).
# Queries are in English to match the English tool descriptions used by the
# TF-IDF router (see tool_knowledge_base.py).
BENCHMARK_TASKS: List[Dict] = [
{
"name": "GitHub repo search",
"task": "Search GitHub for popular Python machine learning repositories with more than 10000 stars",
"gold_tools": ["github_search_repos"],
},
{
"name": "Read config file",
"task": "Read the contents of the local configuration file at /etc/app/config.json",
"gold_tools": ["fs_read_file"],
},
{
"name": "List directory",
"task": "List all files and subdirectories under the /var/log directory",
"gold_tools": ["fs_list_directory"],
},
{
"name": "Summary statistics",
"task": "Calculate the mean, median and standard deviation of last quarter's sales figures",
"gold_tools": ["analytics_summarize"],
},
{
"name": "Send email",
"task": "Send the quarterly performance summary email to the team members",
"gold_tools": ["comm_send_email"],
},
{
"name": "Deploy to production",
"task": "Deploy version 2.3.0 of the application to the production environment",
"gold_tools": ["devops_deploy"],
},
{
"name": "SQL query",
"task": "Run a SQL query on the database to count the number of active users per region",
"gold_tools": ["db_query"],
},
{
"name": "Upload to cloud",
"task": "Upload the local report file to the cloud storage bucket",
"gold_tools": ["cloud_upload_storage"],
},
{
"name": "Scrape prices",
"task": "Scrape the prices of all products listed on the given web page",
"gold_tools": ["web_scrape"],
},
{
"name": "Monitor service",
"task": "Get the current CPU and memory monitoring metrics for the staging service",
"gold_tools": ["devops_monitor"],
},
]
def make_distractor_servers(num_tools: int, start_index: int = 1,
tools_per_server: int = 5) -> List[ServerDefinition]:
"""
Generate synthetic *distractor* servers/tools to inflate the catalog size.
These are deliberately generic "internal service" operations. They add real
schema tokens and act as retrieval noise, so we can study how each strategy
scales as the ecosystem grows to hundreds of tools — without hand-writing
hundreds of realistic tools. They are clearly named ``svcN_opM`` so nobody
mistakes them for the real catalog.
"""
servers: List[ServerDefinition] = []
created = 0
server_idx = start_index
while created < num_tools:
n = min(tools_per_server, num_tools - created)
tools = []
for j in range(1, n + 1):
op = created + j
tools.append(ToolDefinition(
name=f"svc{server_idx}_op{j}",
description=(
f"Auxiliary internal-service operation {op} for background "
f"housekeeping on internal resource group {server_idx}"
),
parameters={
"type": "object",
"properties": {
"resource_id": {"type": "string", "description": "Internal resource identifier"},
"options": {"type": "object", "description": "Operation options"},
},
"required": ["resource_id"],
},
server=f"internal_service_{server_idx}",
))
servers.append(ServerDefinition(
name=f"internal_service_{server_idx}",
description=f"Internal auxiliary service {server_idx} for background housekeeping operations",
tools=tools,
))
created += n
server_idx += 1
return servers
def build_catalog(num_tools: int = 0) -> List[ServerDefinition]:
"""
Build the tool catalog, optionally padded with distractor tools.
Args:
num_tools: Target total number of tools. 0 (default) keeps the real
catalog untouched. Values below the real catalog size are ignored
(we never drop real tools); larger values pad with distractors.
"""
servers = create_tool_knowledge_base()
real_count = len(get_all_tools(servers))
if num_tools and num_tools > real_count:
servers = servers + make_distractor_servers(num_tools - real_count)
return servers
def evaluate_offline(servers: List[ServerDefinition], top_k: int,
tasks: List[Dict] = None) -> Dict:
"""
Deterministically compare tool-selection strategies (no API calls).
Returns a dict with per-strategy aggregate metrics and per-task retrieval
details. Two strategies are directly comparable offline:
* ``all-tools`` — inject every tool schema. Recall is 1.0 by construction
(the gold tool is always present) but token cost grows with the catalog.
* ``retrieval`` — inject only the top-k retrieved tools. Recall is measured;
token cost stays roughly flat as the catalog grows.
(The ``active`` MCP-Zero strategy needs the model in the loop, so it is only
evaluated in the online benchmark.)
"""
tasks = tasks or BENCHMARK_TASKS
router = SemanticRouter(servers)
all_tools = get_all_tools(servers)
all_tools_tokens = calculate_total_tokens(all_tools)
per_task = []
retrieval_hits = 0
retrieval_tokens_sum = 0
for t in tasks:
retrieved = router.retrieve(t["task"], top_k)
retrieved_names = [tool.name for tool in retrieved]
hit = any(g in retrieved_names for g in t["gold_tools"])
retrieval_hits += int(hit)
retrieval_tokens_sum += calculate_total_tokens(retrieved)
per_task.append({
"name": t["name"],
"gold_tools": t["gold_tools"],
"retrieved": retrieved_names,
"hit": hit,
})
n = len(tasks)
return {
"num_tools": len(all_tools),
"top_k": top_k,
"per_task": per_task,
"strategies": {
"all-tools": {
"tools_in_context": len(all_tools),
"avg_schema_tokens": all_tools_tokens,
"recall": 1.0,
},
"retrieval": {
"tools_in_context": top_k,
"avg_schema_tokens": retrieval_tokens_sum / n,
"recall": retrieval_hits / n,
},
},
}
config.py¶
"""Configuration for Active Tool Selection Agent."""
import os
from dotenv import load_dotenv
load_dotenv()
# LLM Configuration
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-5.6-luna")
def _map_model_for_openrouter(model: str) -> str:
"""Map a plain model id onto OpenRouter's `provider/model` form.
Ids that already contain "/" pass through unchanged; gpt-*/o1-*/o3-*/o4-*
become openai/…; claude-* becomes anthropic/claude-opus-4.8.
"""
if "/" in model:
return model
m = model.lower()
if m.startswith(("gpt-", "o1-", "o3-", "o4-")):
return f"openai/{model}"
if m.startswith("claude-"):
return "anthropic/claude-opus-4.8"
return model
# Universal fallback + gpt-5.x preference: route through OpenRouter when no direct
# OPENAI_API_KEY is configured, OR when the model is a gpt-5.x id (incl. gpt-5.6*)
# which needs OpenAI org-verification on the direct API. Explicit OPENAI_BASE_URL /
# OPENAI_MODEL overrides are kept.
_OR_KEY = os.getenv("OPENROUTER_API_KEY")
if _OR_KEY and (not OPENAI_API_KEY or OPENAI_MODEL.lower().startswith("gpt-5")):
OPENAI_API_KEY = _OR_KEY
if not os.getenv("OPENAI_BASE_URL"):
OPENAI_BASE_URL = "https://openrouter.ai/api/v1"
OPENAI_MODEL = _map_model_for_openrouter(OPENAI_MODEL)
# Agent Configuration
AGENT_TEMPERATURE = 0.7
MAX_TOOL_REQUESTS = 5 # Maximum number of tool discovery iterations
# Semantic Routing Configuration
SIMILARITY_THRESHOLD = 0.15 # Minimum similarity score for tool matching
TOP_K_SERVERS = 3 # Number of top servers to search
TOP_K_TOOLS = 5 # Number of top tools to return per server
demo_comparison.py¶
"""
Comparison Demo: Active vs Passive Tool Selection.
Demonstrates the efficiency gains of active tool discovery compared to
traditional passive tool injection approach.
"""
import argparse
import json
import time
from tabulate import tabulate
from agent import ActiveToolAgent, PassiveToolAgent, RetrievalToolAgent
from tool_knowledge_base import create_tool_knowledge_base, calculate_total_tokens
import benchmark
import config
def print_section(title: str):
"""Print a formatted section header."""
print("\n" + "=" * 80)
print(f" {title}")
print("=" * 80 + "\n")
def run_comparison_demo():
"""Run side-by-side comparison of active vs passive approaches."""
print_section("Active Tool Discovery vs Passive Tool Injection Comparison")
# Show knowledge base statistics
servers = create_tool_knowledge_base()
all_tools = []
for server in servers:
all_tools.extend(server.tools)
total_tokens = calculate_total_tokens(all_tools)
print("📊 Tool Knowledge Base Statistics:")
print(f" • Total Servers: {len(servers)}")
print(f" • Total Tools: {len(all_tools)}")
print(f" • Estimated tokens for all tool schemas: ~{total_tokens:,}")
print()
# Test tasks
test_tasks = [
{
"name": "GitHub Repository Search",
"task": "Find popular Python machine learning repositories on GitHub with more than 10k stars"
},
{
"name": "File System Operation",
"task": "Read the configuration file at /etc/app/config.json and list all API keys"
},
{
"name": "Data Analytics",
"task": "Calculate summary statistics (mean, median, std) for the sales data in the last quarter"
},
{
"name": "Multi-Domain Task",
"task": "Clone the repository, analyze the code files, and generate a visualization of code complexity metrics"
}
]
results = []
for test in test_tasks:
print(f"\n🔍 Testing: {test['name']}")
print(f" Task: {test['task']}")
print()
# Test with Active Agent
print(" [Active Agent] Executing...")
active_agent = ActiveToolAgent()
active_result = active_agent.execute_task(test['task'])
# Test with Passive Agent
print(" [Passive Agent] Executing...")
passive_agent = PassiveToolAgent()
passive_result = passive_agent.execute_task(test['task'])
# Calculate efficiency metrics
token_reduction = (1 - active_result['metrics']['tokens_used'] /
passive_result['metrics']['tokens_used']) * 100
tools_loaded_active = active_result['metrics']['tools_loaded']
tools_loaded_passive = passive_result['metrics']['tools_loaded']
results.append({
'Task': test['name'],
'Active Tokens': f"{active_result['metrics']['tokens_used']:,}",
'Passive Tokens': f"{passive_result['metrics']['tokens_used']:,}",
'Token Reduction': f"{token_reduction:.1f}%",
'Active Tools': tools_loaded_active,
'Passive Tools': tools_loaded_passive,
'Tool Reduction': f"{(1 - tools_loaded_active/tools_loaded_passive)*100:.1f}%"
})
print(f" ✓ Active: {active_result['metrics']['tokens_used']:,} tokens, {tools_loaded_active} tools")
print(f" ✓ Passive: {passive_result['metrics']['tokens_used']:,} tokens, {tools_loaded_passive} tools")
print(f" 💡 Reduction: {token_reduction:.1f}% tokens saved")
# Display results table
print_section("Comparison Results")
print(tabulate(results, headers='keys', tablefmt='grid'))
# Calculate averages
avg_token_reduction = sum(
float(r['Token Reduction'].rstrip('%')) for r in results
) / len(results)
print(f"\n📈 Summary:")
print(f" • Average token reduction: {avg_token_reduction:.1f}%")
print(f" • Active approach: Loads only {results[0]['Active Tools']} tools on average")
print(f" • Passive approach: Loads all {results[0]['Passive Tools']} tools upfront")
print()
print("💡 Key Insights:")
print(" • Active tool discovery maintains minimal context footprint")
print(" • Significant token savings (80-98% in typical scenarios)")
print(" • Agent autonomy preserved - discovers tools as needed")
print(" • Scales efficiently as tool ecosystem grows")
def demo_active_discovery_process():
"""Demonstrate the active discovery process in detail."""
print_section("Active Tool Discovery Process Demonstration")
task = "Search for Python repositories on GitHub and analyze their README files"
print(f"📝 Task: {task}\n")
agent = ActiveToolAgent()
result = agent.execute_task(task)
print("🔄 Discovery Process:")
print(f" • Tool requests made: {result['metrics']['tool_requests']}")
print(f" • Tools loaded: {result['metrics']['tools_loaded']}")
print(f" • API calls: {result['metrics']['api_calls']}")
print(f" • Total tokens: {result['metrics']['tokens_used']:,}")
print()
print("🛠️ Tools Discovered:")
for i, tool in enumerate(result['tools_loaded'], 1):
print(f" {i}. {tool}")
print()
print("💬 Conversation Flow:")
for i, msg in enumerate(result['conversation'], 1):
role = msg['role'].upper()
content = msg.get('content', '[Tool Call]')
if content and len(content) > 100:
content = content[:100] + "..."
print(f" {i}. [{role}] {content}")
print()
print("✅ Final Response:")
print(f" {result['response']}")
def demo_semantic_routing():
"""Demonstrate hierarchical semantic routing."""
print_section("Hierarchical Semantic Routing Demonstration")
from semantic_router import SemanticRouter
servers = create_tool_knowledge_base()
router = SemanticRouter(servers)
test_queries = [
"I need to search for repositories on GitHub",
"Read a file from the local filesystem",
"Query the database for user information",
"Send an email notification to the team",
"Deploy the application to production environment"
]
print("🎯 Testing semantic routing for various requests:\n")
for query in test_queries:
print(f"📌 Request: '{query}'")
details = router.get_routing_details(query, top_k_servers=2, top_k_tools=3)
print(" Stage 1 - Server Routing:")
for server in details['stage1_servers']:
print(f" • {server['name']}: {server['score']:.3f}")
print(" Stage 2 - Tool Routing:")
for tool in details['final_tools']:
print(f" • {tool['name']} ({tool['server']}): {tool['score']:.3f}")
print()
def demo_iterative_capability_extension():
"""Demonstrate iterative capability extension."""
print_section("Iterative Capability Extension Demonstration")
print("🎯 Complex Multi-Step Task:")
task = """Perform a comprehensive analysis:
1. Search GitHub for Python data science repositories
2. Download the top repository
3. Analyze the code structure
4. Generate visualization of dependencies
5. Send summary report via email"""
print(f"{task}\n")
agent = ActiveToolAgent()
result = agent.execute_task(task)
print("📊 Capability Extension Timeline:")
print(f" • Initial tools: 0")
print(f" • Tools after request 1: GitHub tools")
print(f" • Tools after request 2: Filesystem + GitHub")
print(f" • Tools after request 3: Analytics + Filesystem + GitHub")
print(f" • Tools after request 4: Communication + Analytics + Filesystem + GitHub")
print()
print(f" Total tool requests: {result['metrics']['tool_requests']}")
print(f" Final toolchain size: {result['metrics']['tools_loaded']} tools")
print()
print("💡 The agent iteratively built a cross-domain toolchain as task understanding evolved!")
def run_offline_benchmark(servers, top_k: int, scaling: bool = True) -> dict:
"""
Deterministic (no-API) strategy comparison: retrieval recall vs token cost.
This is the heart of the experiment and runs without any API key. It shows
that as the tool catalog grows, injecting all tools makes context token cost
explode, while on-demand retrieval keeps cost roughly flat and still surfaces
the right tool (recall).
"""
print_section("Offline Strategy Comparison (deterministic, no API)")
result = benchmark.evaluate_offline(servers, top_k)
strat = result['strategies']
print(f"Benchmark tasks: {len(benchmark.BENCHMARK_TASKS)} "
f"Catalog size: {result['num_tools']} tools Retrieval top-k: {top_k}\n")
rows = [
{
'Strategy': 'all-tools (dump everything)',
'Tools in context': strat['all-tools']['tools_in_context'],
'Schema tokens': f"{strat['all-tools']['avg_schema_tokens']:,}",
'Recall (gold reachable)': f"{strat['all-tools']['recall']*100:.0f}%",
},
{
'Strategy': f'retrieval (top-{top_k})',
'Tools in context': strat['retrieval']['tools_in_context'],
'Schema tokens': f"{strat['retrieval']['avg_schema_tokens']:,.0f}",
'Recall (gold reachable)': f"{strat['retrieval']['recall']*100:.0f}%",
},
]
print(tabulate(rows, headers='keys', tablefmt='grid'))
token_saving = (1 - strat['retrieval']['avg_schema_tokens'] /
strat['all-tools']['avg_schema_tokens']) * 100
print(f"\n=> Retrieval keeps {strat['retrieval']['recall']*100:.0f}% recall while cutting "
f"tool-schema tokens by {token_saving:.1f}% "
f"({strat['all-tools']['avg_schema_tokens']:,} -> "
f"{strat['retrieval']['avg_schema_tokens']:,.0f}).")
# Per-task retrieval detail (which tools were surfaced, and whether the gold hit)
print("\nPer-task retrieval (top-k tools surfaced for each task):")
detail_rows = [
{
'Task': p['name'],
'Gold tool': ', '.join(p['gold_tools']),
'Hit': '✓' if p['hit'] else '✗',
'Retrieved (top-k)': ', '.join(p['retrieved']),
}
for p in result['per_task']
]
print(tabulate(detail_rows, headers='keys', tablefmt='github'))
scaling_result = None
if scaling:
print_section("Scaling: token cost as the catalog grows")
sizes = [size for size in [50, 100, 200, 400] if size >= result['num_tools']]
if not sizes or sizes[0] != result['num_tools']:
sizes = [result['num_tools']] + sizes
scaling_rows = []
scaling_result = []
for size in sizes:
padded = benchmark.build_catalog(size)
r = benchmark.evaluate_offline(padded, top_k)
s = r['strategies']
scaling_rows.append({
'Catalog tools': r['num_tools'],
'all-tools tokens': f"{s['all-tools']['avg_schema_tokens']:,}",
f'retrieval(top-{top_k}) tokens': f"{s['retrieval']['avg_schema_tokens']:,.0f}",
'retrieval recall': f"{s['retrieval']['recall']*100:.0f}%",
})
scaling_result.append({
'num_tools': r['num_tools'],
'all_tools_tokens': s['all-tools']['avg_schema_tokens'],
'retrieval_tokens': s['retrieval']['avg_schema_tokens'],
'retrieval_recall': s['retrieval']['recall'],
})
print(tabulate(scaling_rows, headers='keys', tablefmt='grid'))
print("\n=> all-tools token cost grows with the catalog; retrieval stays roughly flat.")
return {'benchmark': result, 'scaling': scaling_result}
STRATEGY_AGENTS = {
'all': ('all-tools', PassiveToolAgent),
'retrieval': ('retrieval', RetrievalToolAgent),
'active': ('active (MCP-Zero)', ActiveToolAgent),
}
def _build_agent(strategy: str, servers, top_k: int, model: str):
"""Instantiate the agent for a strategy (needs a valid API key)."""
if strategy == 'retrieval':
return RetrievalToolAgent(servers=servers, model=model, top_k=top_k)
_, cls = STRATEGY_AGENTS[strategy]
return cls(servers=servers, model=model)
def run_online_benchmark(servers, strategies, top_k: int, model: str,
tasks=None) -> dict:
"""
End-to-end benchmark (requires API key): does the model actually CALL the
ground-truth tool, at what token cost and latency, under each strategy?
"""
tasks = tasks or benchmark.BENCHMARK_TASKS
print_section("Online End-to-End Benchmark (requires API)")
print(f"Model: {model} Tasks: {len(tasks)} "
f"Catalog: {sum(len(s.tools) for s in servers)} tools "
f"Retrieval top-k: {top_k}\n")
agents = {st: _build_agent(st, servers, top_k, model) for st in strategies}
rows = []
raw = {}
for st in strategies:
label = STRATEGY_AGENTS[st][0]
agent = agents[st]
hits = 0
tokens_sum = 0
latency_sum = 0.0
tools_ctx_sum = 0
per_task = []
print(f"[{label}] running {len(tasks)} tasks...")
for t in tasks:
agent.reset()
start = time.time()
res = agent.execute_task(t['task'])
elapsed = time.time() - start
called = res['metrics'].get('tools_called', [])
hit = any(g in called for g in t['gold_tools'])
hits += int(hit)
tokens_sum += res['metrics']['tokens_used']
latency_sum += elapsed
tools_ctx_sum += res['metrics']['tools_loaded']
per_task.append({
'task': t['name'],
'gold': t['gold_tools'],
'called': called,
'hit': hit,
'tokens': res['metrics']['tokens_used'],
'latency': round(elapsed, 2),
})
n = len(tasks)
rows.append({
'Strategy': label,
'Accuracy (calls gold)': f"{hits/n*100:.0f}%",
'Avg tools in ctx': f"{tools_ctx_sum/n:.1f}",
'Avg tokens': f"{tokens_sum/n:,.0f}",
'Avg latency (s)': f"{latency_sum/n:.2f}",
})
raw[st] = {'accuracy': hits / n, 'per_task': per_task}
print()
print(tabulate(rows, headers='keys', tablefmt='grid'))
return raw
def run_single_query(servers, strategies, query: str, top_k: int, model: str) -> dict:
"""Run a single ad-hoc query through the chosen strategies (requires API)."""
print_section("Single Query")
print(f"Query: {query}\n")
raw = {}
for st in strategies:
label = STRATEGY_AGENTS[st][0]
agent = _build_agent(st, servers, top_k, model)
res = agent.execute_task(query)
print(f"[{label}]")
print(f" tools in context : {res['metrics']['tools_loaded']}")
print(f" tools loaded : {', '.join(res['tools_loaded']) or '(none)'}")
print(f" tools called : {', '.join(res['metrics'].get('tools_called', [])) or '(none)'}")
print(f" tokens used : {res['metrics']['tokens_used']:,}")
print()
raw[st] = {
'tools_loaded': res['tools_loaded'],
'tools_called': res['metrics'].get('tools_called', []),
'tokens_used': res['metrics']['tokens_used'],
}
return raw
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
prog="demo_comparison.py",
formatter_class=argparse.RawDescriptionHelpFormatter,
description=(
"主动工具选择实验：对比\"把全部工具塞进上下文\"与\"按需检索工具\"两类策略。\n"
"在工具数量增长到上百个时，动态检索（retrieval）在保持召回率的同时大幅降低\n"
"上下文 token 成本，并减少模型的选择错误。\n\n"
"策略说明：\n"
" all-tools 一次性注入全部工具（传统被动式基线）\n"
" retrieval 按任务语义检索 top-k 个工具后再注入（工具检索 / RAG 式）\n"
" active MCP-Zero 式主动发现：模型迭代地请求所需工具\n\n"
"离线表格（召回率 / token 成本 / 随工具规模的扩展性）无需 API Key 即可运行；\n"
"端到端准确率与延迟对比需要配置 API Key。"
),
epilog=(
"示例：\n"
" python demo_comparison.py --offline # 仅离线对比，无需 API\n"
" python demo_comparison.py --offline --num-tools 200 # 扩展到 200 个工具再对比\n"
" python demo_comparison.py --strategy compare # 三种策略端到端对比（需 API）\n"
" python demo_comparison.py --query \"部署到生产环境\" --strategy retrieval\n"
" python demo_comparison.py --output results.json # 保存结果为 JSON"
),
)
parser.add_argument(
"--strategy", choices=["all", "retrieval", "active", "compare"],
default="compare",
help="端到端评测使用的策略；compare 表示三种策略全部对比（默认：compare）",
)
parser.add_argument(
"--query", type=str, default=None,
help="只对单条查询运行选定策略（需要 API），而不是跑整个基准集",
)
parser.add_argument(
"--num-tools", type=int, default=0, metavar="N",
help="将工具目录扩充到 N 个（用合成干扰工具补齐，用于观察扩展性）；0 表示保持真实目录（默认：0）",
)
parser.add_argument(
"--top-k", type=int, default=config.TOP_K_TOOLS, metavar="K",
help=f"retrieval 策略检索的工具数量（默认：{config.TOP_K_TOOLS}）",
)
parser.add_argument(
"--model", type=str, default=config.OPENAI_MODEL,
help=f"覆盖使用的 LLM 模型（默认：{config.OPENAI_MODEL}）",
)
parser.add_argument(
"--output", type=str, default=None, metavar="PATH",
help="将结果写入 JSON 文件",
)
parser.add_argument(
"--offline", action="store_true",
help="仅运行离线确定性对比（召回率/token 成本），不进行任何 API 调用",
)
parser.add_argument(
"--legacy-demos", action="store_true",
help="额外运行原有的叙事式演示（语义路由、迭代发现等，需要 API）",
)
return parser
def _has_api_key() -> bool:
return bool(config.OPENAI_API_KEY)
def main(argv=None):
parser = build_parser()
args = parser.parse_args(argv)
print("""
╔════════════════════════════════════════════════════════════════════════════╗
║ Active Tool Selection — Strategy Comparison ║
║ Inspired by MCP-Zero (arXiv:2506.01056) ║
╚════════════════════════════════════════════════════════════════════════════╝
""")
servers = benchmark.build_catalog(args.num_tools)
strategies = ["all", "retrieval", "active"] if args.strategy == "compare" else [args.strategy]
results = {
'config': {
'num_tools': sum(len(s.tools) for s in servers),
'top_k': args.top_k,
'model': args.model,
'strategies': strategies,
}
}
# 1) Offline deterministic comparison — always runs, no API needed.
if not args.query:
results['offline'] = run_offline_benchmark(servers, args.top_k)
# 2) Online end-to-end comparison — needs an API key.
if args.offline:
print("\n[offline mode] 跳过所有需要 API 的评测。")
elif not _has_api_key():
print("\n[提示] 未检测到 OPENAI_API_KEY，跳过端到端评测（准确率/延迟）。")
print(" 配置 .env 后可运行端到端对比；或使用 --offline 显式仅跑离线部分。")
else:
if args.query:
results['single_query'] = run_single_query(
servers, strategies, args.query, args.top_k, args.model)
else:
results['online'] = run_online_benchmark(
servers, strategies, args.top_k, args.model)
if args.legacy_demos:
run_comparison_demo()
demo_active_discovery_process()
demo_semantic_routing()
demo_iterative_capability_extension()
if args.output:
with open(args.output, "w", encoding="utf-8") as f:
json.dump(results, f, ensure_ascii=False, indent=2)
print(f"\n结果已写入 {args.output}")
print_section("Takeaway")
print(
"当工具数量增长到上百个时，把全部工具塞进上下文既浪费 token 又干扰决策；\n"
"按需检索把\"工具选择\"问题转化为\"知识检索\"问题——在保持召回率的同时\n"
"把工具描述的 token 成本压到很低，也减少了模型的选择错误。\n\n"
"参考：MCP-Zero 论文 (https://arxiv.org/pdf/2506.01056)"
)
if __name__ == "__main__":
main()
examples.py¶
"""
Example use cases demonstrating active tool selection.
"""
from agent import ActiveToolAgent
from semantic_router import SemanticRouter
from tool_knowledge_base import create_tool_knowledge_base
def example_github_workflow():
"""Example: GitHub development workflow."""
print("\n" + "=" * 70)
print("Example 1: GitHub Development Workflow")
print("=" * 70 + "\n")
agent = ActiveToolAgent()
task = """I need to:
1. Search for Python testing frameworks on GitHub
2. Find issues labeled 'good-first-issue' in the top repository
3. Create a new branch and make changes
4. Create a pull request"""
print(f"Task:\n{task}\n")
result = agent.execute_task(task)
print(f"\n✅ Tools discovered: {len(result['tools_loaded'])}")
print(f" {', '.join(result['tools_loaded'])}")
print(f"\n📊 Metrics:")
print(f" • Tokens used: {result['metrics']['tokens_used']:,}")
print(f" • Tool requests: {result['metrics']['tool_requests']}")
print(f" • API calls: {result['metrics']['api_calls']}")
def example_data_pipeline():
"""Example: Data processing pipeline."""
print("\n" + "=" * 70)
print("Example 2: Data Processing Pipeline")
print("=" * 70 + "\n")
agent = ActiveToolAgent()
task = """Build a data pipeline:
1. Query the database for last month's sales data
2. Calculate summary statistics
3. Create visualizations (bar charts and trend lines)
4. Upload results to cloud storage
5. Send notification email to stakeholders"""
print(f"Task:\n{task}\n")
result = agent.execute_task(task)
print(f"\n✅ Cross-domain toolchain built:")
for i, tool in enumerate(result['tools_loaded'], 1):
print(f" {i}. {tool}")
print(f"\n📊 Efficiency:")
print(f" • Only {len(result['tools_loaded'])} tools loaded (out of 35 available)")
print(f" • Token savings: ~90% compared to loading all tools")
def example_devops_automation():
"""Example: DevOps automation task."""
print("\n" + "=" * 70)
print("Example 3: DevOps Automation")
print("=" * 70 + "\n")
agent = ActiveToolAgent()
task = """Automate deployment process:
1. Check monitoring metrics for the staging environment
2. If metrics are healthy, trigger production deployment pipeline
3. Monitor deployment progress and logs
4. If any errors occur, automatically rollback
5. Send deployment status notification"""
print(f"Task:\n{task}\n")
result = agent.execute_task(task)
print(f"\n✅ DevOps toolchain assembled:")
print(f" Tools: {', '.join(result['tools_loaded'])}")
print(f"\n💡 Active discovery enabled iterative refinement:")
print(f" • Started with monitoring tools")
print(f" • Added deployment tools when needed")
print(f" • Included notification tools at the end")
def example_semantic_search():
"""Example: Demonstrate semantic search capabilities."""
print("\n" + "=" * 70)
print("Example 4: Semantic Tool Search")
print("=" * 70 + "\n")
servers = create_tool_knowledge_base()
router = SemanticRouter(servers)
queries = [
"I need to version control my code",
"Store and retrieve structured data",
"Make HTTP requests to APIs",
"Analyze datasets and create graphs",
"Configure cloud infrastructure"
]
print("Testing semantic understanding of tool requests:\n")
for query in queries:
print(f"🔍 Query: '{query}'")
tools = router.route_request(query, top_k_servers=1, top_k_tools=3)
if tools:
print(f" ✓ Found: {', '.join([t.name for t in tools])}")
else:
print(f" ✗ No matching tools found")
print()
def example_multi_turn_discovery():
"""Example: Multi-turn conversation with progressive tool discovery."""
print("\n" + "=" * 70)
print("Example 5: Multi-Turn Progressive Discovery")
print("=" * 70 + "\n")
print("Scenario: Agent progressively discovers tools across multiple turns\n")
agent = ActiveToolAgent()
# Turn 1: Initial request
print("👤 User: Search for machine learning repositories")
result1 = agent.execute_task("Search for machine learning repositories")
print(f"🤖 Agent loaded: {', '.join(result1['tools_loaded'][:2])}")
print()
# Turn 2: Additional requirements emerge
print("👤 User: Now download the README files and analyze them")
result2 = agent.execute_task("Download README files and analyze them")
print(f"🤖 Agent additionally loaded: filesystem and analytics tools")
print()
# Turn 3: Visualization needed
print("👤 User: Create a visualization comparing repository sizes")
result3 = agent.execute_task("Create a visualization comparing repository sizes")
print(f"🤖 Agent additionally loaded: visualization tools")
print()
print("💡 Tools were discovered on-demand as the conversation evolved!")
print(" This demonstrates the iterative capability extension principle.")
def example_efficiency_comparison():
"""Example: Show efficiency comparison with metrics."""
print("\n" + "=" * 70)
print("Example 6: Efficiency Comparison")
print("=" * 70 + "\n")
from agent import PassiveToolAgent
task = "List files in the current directory"
print(f"Task: {task}\n")
# Active approach
print("🔄 Active Tool Discovery:")
active_agent = ActiveToolAgent()
active_result = active_agent.execute_task(task)
print(f" • Tools loaded: {active_result['metrics']['tools_loaded']}")
print(f" • Tokens used: {active_result['metrics']['tokens_used']:,}")
print()
# Passive approach
print("📚 Passive Tool Injection:")
passive_agent = PassiveToolAgent()
passive_result = passive_agent.execute_task(task)
print(f" • Tools loaded: {passive_result['metrics']['tools_loaded']}")
print(f" • Tokens used: {passive_result['metrics']['tokens_used']:,}")
print()
# Comparison
reduction = (1 - active_result['metrics']['tokens_used'] /
passive_result['metrics']['tokens_used']) * 100
print(f"📊 Efficiency Gain:")
print(f" • Token reduction: {reduction:.1f}%")
print(f" • Tool reduction: {active_result['metrics']['tools_loaded']} vs {passive_result['metrics']['tools_loaded']}")
print()
print("💡 For simple tasks requiring 1-2 tools, active discovery achieves")
print(" massive efficiency gains while maintaining full capability!")
if __name__ == "__main__":
print("""
╔════════════════════════════════════════════════════════════════════════════╗
║ ║
║ Active Tool Selection Examples ║
║ ║
╚════════════════════════════════════════════════════════════════════════════╝
""")
# Run all examples
example_github_workflow()
example_data_pipeline()
example_devops_automation()
example_semantic_search()
example_multi_turn_discovery()
example_efficiency_comparison()
print("\n" + "=" * 70)
print("All examples completed!")
print("=" * 70 + "\n")
quickstart.py¶
"""
Quick Start for Active Tool Selection.
Run this script to see a basic demonstration of active tool discovery.
"""
from agent import ActiveToolAgent, PassiveToolAgent
from tool_knowledge_base import create_tool_knowledge_base, calculate_total_tokens
def main():
print("""
╔════════════════════════════════════════════════════════════════════════════╗
║ ║
║ Active Tool Selection - Quick Start ║
║ Inspired by MCP-Zero (arXiv:2506.01056) ║
║ ║
╚════════════════════════════════════════════════════════════════════════════╝
This demonstration shows how active tool discovery enables agents to:
• Maintain minimal context footprint
• Actively request tools as needed
• Scale efficiently with ecosystem growth
""")
# Show knowledge base info
print("📚 Tool Knowledge Base:")
servers = create_tool_knowledge_base()
total_tools = sum(len(server.tools) for server in servers)
total_tokens = calculate_total_tokens([tool for server in servers for tool in server.tools])
print(f" • Servers: {len(servers)}")
print(f" • Total tools: {total_tools}")
print(f" • Token cost if all injected: ~{total_tokens:,} tokens")
print()
# Example task
task = "Search for Python web frameworks on GitHub with more than 5000 stars"
print(f"🎯 Example Task:\n {task}\n")
# Test with active agent
print("=" * 80)
print("1️⃣ ACTIVE TOOL DISCOVERY")
print("=" * 80)
print("\n⏳ Agent is analyzing task and discovering needed tools...\n")
active_agent = ActiveToolAgent()
active_result = active_agent.execute_task(task)
print(f"✅ Task completed with active discovery:\n")
print(f" 📊 Metrics:")
print(f" • Tools loaded: {active_result['metrics']['tools_loaded']} (out of {total_tools})")
print(f" • Tokens used: {active_result['metrics']['tokens_used']:,}")
print(f" • Tool requests: {active_result['metrics']['tool_requests']}")
print(f" • API calls: {active_result['metrics']['api_calls']}")
print()
print(f" 🛠️ Tools discovered:")
for tool in active_result['tools_loaded']:
print(f" • {tool}")
print()
# Test with passive agent
print("=" * 80)
print("2️⃣ PASSIVE TOOL INJECTION (Traditional Approach)")
print("=" * 80)
print(f"\n⏳ Agent has all {total_tools} tools pre-loaded...\n")
passive_agent = PassiveToolAgent()
passive_result = passive_agent.execute_task(task)
print(f"✅ Task completed with passive injection:\n")
print(f" 📊 Metrics:")
print(f" • Tools loaded: {passive_result['metrics']['tools_loaded']} (all tools)")
print(f" • Tokens used: {passive_result['metrics']['tokens_used']:,}")
print(f" • API calls: {passive_result['metrics']['api_calls']}")
print()
# Comparison
print("=" * 80)
print("3️⃣ COMPARISON")
print("=" * 80)
print()
token_reduction = (1 - active_result['metrics']['tokens_used'] /
passive_result['metrics']['tokens_used']) * 100
tool_reduction = (1 - active_result['metrics']['tools_loaded'] /
passive_result['metrics']['tools_loaded']) * 100
print(f"📊 Efficiency Gains:\n")
print(f" Token Usage:")
print(f" • Active: {active_result['metrics']['tokens_used']:,} tokens")
print(f" • Passive: {passive_result['metrics']['tokens_used']:,} tokens")
print(f" • Reduction: {token_reduction:.1f}% 🎉")
print()
print(f" Tools Loaded:")
print(f" • Active: {active_result['metrics']['tools_loaded']} tools")
print(f" • Passive: {passive_result['metrics']['tools_loaded']} tools")
print(f" • Reduction: {tool_reduction:.1f}% 🎯")
print()
print("=" * 80)
print("💡 KEY INSIGHTS")
print("=" * 80)
print("""
1. Active Discovery maintains agent autonomy
→ Agent decides what tools it needs, when it needs them
2. Massive efficiency gains
→ 80-98% token reduction for typical tasks
3. Scales with ecosystem growth
→ Adding 100 more tools doesn't bloat every request
4. Iterative capability extension
→ Toolchain evolves as task understanding deepens
5. Semantic routing enables precision
→ Tools matched by meaning, not just keywords
""")
print("🎓 Next Steps:")
print(" • Run 'python demo_comparison.py' for comprehensive comparison")
print(" • Run 'python examples.py' for more use cases")
print(" • See README.md for architecture details")
print()
print("📄 Reference: MCP-Zero paper - https://arxiv.org/pdf/2506.01056")
print()
if __name__ == "__main__":
main()
semantic_router.py¶
"""
Hierarchical Semantic Routing for Tool Discovery.
Implements a two-stage algorithm for matching tool requests to relevant tools:
1. Server-level routing: Filter candidate servers by domain/platform
2. Tool-level routing: Rank tools within selected servers by semantic similarity
This approach reduces search complexity while maintaining precision, inspired by MCP-Zero.
"""
from typing import List, Dict, Tuple
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from tool_knowledge_base import ServerDefinition, ToolDefinition
import config
class SemanticRouter:
"""Hierarchical semantic routing for tool discovery."""
def __init__(self, servers: List[ServerDefinition]):
self.servers = servers
self.server_vectorizer = TfidfVectorizer(stop_words='english')
self.tool_vectorizers: Dict[str, TfidfVectorizer] = {}
# Precompute server embeddings
self._build_server_index()
# Precompute tool embeddings for each server
self._build_tool_indices()
def _build_server_index(self):
"""Build TF-IDF index for servers."""
server_descriptions = [f"{s.name} {s.description}" for s in self.servers]
self.server_embeddings = self.server_vectorizer.fit_transform(server_descriptions)
def _build_tool_indices(self):
"""Build TF-IDF indices for tools within each server."""
for server in self.servers:
if not server.tools:
continue
tool_descriptions = [
f"{tool.name} {tool.description}"
for tool in server.tools
]
vectorizer = TfidfVectorizer(stop_words='english')
embeddings = vectorizer.fit_transform(tool_descriptions)
self.tool_vectorizers[server.name] = vectorizer
# Store embeddings on server for later use
server._tool_embeddings = embeddings
def route_request(self, tool_request: str, top_k_servers: int = None,
top_k_tools: int = None) -> List[ToolDefinition]:
"""
Route a tool request to relevant tools using hierarchical semantic matching.
Args:
tool_request: Natural language description of needed tool
top_k_servers: Number of top servers to search (default from config)
top_k_tools: Number of tools to return per server (default from config)
Returns:
List of relevant tools ranked by relevance
"""
if top_k_servers is None:
top_k_servers = config.TOP_K_SERVERS
if top_k_tools is None:
top_k_tools = config.TOP_K_TOOLS
# Stage 1: Server-level routing
relevant_servers = self._route_to_servers(tool_request, top_k_servers)
# Stage 2: Tool-level routing within selected servers
relevant_tools = []
for server, server_score in relevant_servers:
tools_with_scores = self._route_to_tools(server, tool_request, top_k_tools)
# Combine server and tool scores
for tool, tool_score in tools_with_scores:
combined_score = 0.3 * server_score + 0.7 * tool_score
relevant_tools.append((tool, combined_score))
# Sort by combined score and filter by threshold
relevant_tools.sort(key=lambda x: x[1], reverse=True)
relevant_tools = [
(tool, score) for tool, score in relevant_tools
if score >= config.SIMILARITY_THRESHOLD
]
# Return top tools
return [tool for tool, _ in relevant_tools[:top_k_tools * top_k_servers]]
def retrieve(self, query: str, top_k: int) -> List[ToolDefinition]:
"""
Flat top-k tool retrieval across ALL servers (single-shot RAG-style routing).
Unlike ``route_request`` (which first narrows to a few candidate servers),
this scores every tool in every server and returns the global top-k. It is the
most direct embodiment of "turn tool selection into knowledge retrieval": given
the task description, fetch only the handful of tools most likely to be relevant.
Args:
query: Natural language task/request description
top_k: Number of tools to return
Returns:
Up to ``top_k`` tools ranked by combined (server + tool) similarity.
"""
# Score against every server so no candidate tool is filtered out prematurely.
relevant_servers = self._route_to_servers(query, len(self.servers))
scored_tools = []
for server, server_score in relevant_servers:
for tool, tool_score in self._route_to_tools(server, query, len(server.tools)):
combined_score = 0.3 * server_score + 0.7 * tool_score
scored_tools.append((tool, combined_score))
scored_tools.sort(key=lambda x: x[1], reverse=True)
return [tool for tool, _ in scored_tools[:top_k]]
def _route_to_servers(self, request: str, top_k: int) -> List[Tuple[ServerDefinition, float]]:
"""
Stage 1: Route request to top-k relevant servers.
Args:
request: Tool request description
top_k: Number of top servers to return
Returns:
List of (server, similarity_score) tuples
"""
# Vectorize the request
request_vector = self.server_vectorizer.transform([request])
# Calculate similarities with all servers
similarities = cosine_similarity(request_vector, self.server_embeddings)[0]
# Get top-k servers
top_indices = np.argsort(similarities)[::-1][:top_k]
return [(self.servers[idx], similarities[idx]) for idx in top_indices]
def _route_to_tools(self, server: ServerDefinition, request: str,
top_k: int) -> List[Tuple[ToolDefinition, float]]:
"""
Stage 2: Route request to top-k relevant tools within a server.
Args:
server: Server to search within
request: Tool request description
top_k: Number of top tools to return
Returns:
List of (tool, similarity_score) tuples
"""
if server.name not in self.tool_vectorizers:
return []
vectorizer = self.tool_vectorizers[server.name]
tool_embeddings = server._tool_embeddings
# Vectorize the request
request_vector = vectorizer.transform([request])
# Calculate similarities with all tools in this server
similarities = cosine_similarity(request_vector, tool_embeddings)[0]
# Get top-k tools
top_indices = np.argsort(similarities)[::-1][:top_k]
return [(server.tools[idx], similarities[idx]) for idx in top_indices]
def get_routing_details(self, tool_request: str, top_k_servers: int = None,
top_k_tools: int = None) -> Dict:
"""
Get detailed routing information for debugging/visualization.
Returns a dictionary with:
- request: Original request
- stage1_servers: List of servers with scores
- stage2_tools: List of tools with scores per server
- final_tools: Final ranked list of tools
"""
if top_k_servers is None:
top_k_servers = config.TOP_K_SERVERS
if top_k_tools is None:
top_k_tools = config.TOP_K_TOOLS
# Stage 1: Server routing
relevant_servers = self._route_to_servers(tool_request, top_k_servers)
# Stage 2: Tool routing
stage2_results = {}
all_tools = []
for server, server_score in relevant_servers:
tools_with_scores = self._route_to_tools(server, tool_request, top_k_tools)
stage2_results[server.name] = {
'server_score': server_score,
'tools': [(tool.name, tool_score) for tool, tool_score in tools_with_scores]
}
# Calculate combined scores
for tool, tool_score in tools_with_scores:
combined_score = 0.3 * server_score + 0.7 * tool_score
all_tools.append((tool, combined_score, server.name))
# Sort and filter
all_tools.sort(key=lambda x: x[1], reverse=True)
final_tools = [
{'name': tool.name, 'server': server, 'score': score}
for tool, score, server in all_tools[:top_k_tools * top_k_servers]
if score >= config.SIMILARITY_THRESHOLD
]
return {
'request': tool_request,
'stage1_servers': [
{'name': s.name, 'score': score}
for s, score in relevant_servers
],
'stage2_tools': stage2_results,
'final_tools': final_tools
}
class StructuredRequestParser:
"""
Parse structured tool requests from LLM.
MCP-Zero uses structured requests in format:
<tool_request>
server: [platform/domain description]
tool: [operation description]
</tool_request>
"""
@staticmethod
def parse_request(text: str) -> Dict[str, str]:
"""
Parse structured tool request from text.
Returns dict with 'server' and 'tool' fields, or None if not found.
"""
if '<tool_request>' not in text or '</tool_request>' not in text:
return None
start = text.find('<tool_request>')
end = text.find('</tool_request>')
request_text = text[start + len('<tool_request>'):end].strip()
result = {}
for line in request_text.split('\n'):
line = line.strip()
if line.startswith('server:'):
result['server'] = line[7:].strip()
elif line.startswith('tool:'):
result['tool'] = line[5:].strip()
return result if 'server' in result and 'tool' in result else None
@staticmethod
def format_request(server_desc: str, tool_desc: str) -> str:
"""Format a structured tool request."""
return f"""<tool_request>
server: {server_desc}
tool: {tool_desc}
</tool_request>"""
test_basic.py¶
"""
Basic tests to verify the active tool selection system works correctly.
Run without API key to test core functionality.
"""
from tool_knowledge_base import create_tool_knowledge_base, calculate_total_tokens, get_all_tools
from semantic_router import SemanticRouter, StructuredRequestParser
def test_knowledge_base():
"""Test that knowledge base loads correctly."""
print("Testing Knowledge Base...")
servers = create_tool_knowledge_base()
all_tools = get_all_tools(servers)
assert len(servers) == 8, f"Expected 8 servers, got {len(servers)}"
assert len(all_tools) > 30, f"Expected 30+ tools, got {len(all_tools)}"
# Check each server has tools
for server in servers:
assert len(server.tools) > 0, f"Server {server.name} has no tools"
assert server.description, f"Server {server.name} missing description"
# Check tool schemas
for tool in all_tools:
schema = tool.to_schema()
assert 'type' in schema, f"Tool {tool.name} missing type"
assert 'function' in schema, f"Tool {tool.name} missing function"
total_tokens = calculate_total_tokens(all_tools)
print(f" ✓ Loaded {len(servers)} servers with {len(all_tools)} tools")
print(f" ✓ Estimated tokens: {total_tokens:,}")
print()
def test_semantic_router():
"""Test semantic routing functionality."""
print("Testing Semantic Router...")
servers = create_tool_knowledge_base()
router = SemanticRouter(servers)
# Test server routing with realistic queries
test_queries = [
("search for GitHub repositories", ["github"]),
("read a file from filesystem", ["filesystem"]),
("query database for users", ["database"]),
("send an email notification", ["communication"]),
("deploy to production environment", ["devops"])
]
for query, expected_servers in test_queries:
tools = router.route_request(query, top_k_servers=3, top_k_tools=3)
assert len(tools) > 0, f"No tools found for query: {query}"
# Check that tools are from expected servers
tool_servers = {tool.server for tool in tools}
assert any(exp in tool_servers for exp in expected_servers), \
f"Expected servers {expected_servers}, got {tool_servers} for query: {query}"
print(f" ✓ Semantic routing working correctly")
print(f" ✓ All test queries matched appropriate servers")
print()
def test_structured_request_parser():
"""Test structured request parsing."""
print("Testing Structured Request Parser...")
# Valid request
valid_request = """
Some text before
<tool_request>
server: GitHub for repository operations
tool: search repositories by keywords
</tool_request>
Some text after
"""
parsed = StructuredRequestParser.parse_request(valid_request)
assert parsed is not None, "Failed to parse valid request"
assert 'server' in parsed, "Missing server in parsed request"
assert 'tool' in parsed, "Missing tool in parsed request"
assert "GitHub" in parsed['server'], "Server description incorrect"
assert "search" in parsed['tool'], "Tool description incorrect"
# Invalid request (missing tags)
invalid_request = "Just some text without proper tags"
parsed_invalid = StructuredRequestParser.parse_request(invalid_request)
assert parsed_invalid is None, "Should return None for invalid request"
# Test formatting
formatted = StructuredRequestParser.format_request(
"GitHub operations",
"search repositories"
)
assert "<tool_request>" in formatted, "Missing opening tag"
assert "</tool_request>" in formatted, "Missing closing tag"
assert "server:" in formatted, "Missing server field"
assert "tool:" in formatted, "Missing tool field"
print(f" ✓ Request parsing working correctly")
print(f" ✓ Request formatting working correctly")
print()
def test_routing_details():
"""Test detailed routing information."""
print("Testing Routing Details...")
servers = create_tool_knowledge_base()
router = SemanticRouter(servers)
query = "I need to search for Python repositories on GitHub"
details = router.get_routing_details(query, top_k_servers=3, top_k_tools=3)
assert 'request' in details, "Missing request in details"
assert 'stage1_servers' in details, "Missing stage1 in details"
assert 'stage2_tools' in details, "Missing stage2 in details"
assert 'final_tools' in details, "Missing final_tools in details"
assert len(details['stage1_servers']) > 0, "No servers in stage1"
assert len(details['final_tools']) > 0, "No final tools"
# Check structure
for server in details['stage1_servers']:
assert 'name' in server, "Server missing name"
assert 'score' in server, "Server missing score"
assert 0 <= server['score'] <= 1, "Server score out of range"
for tool in details['final_tools']:
assert 'name' in tool, "Tool missing name"
assert 'server' in tool, "Tool missing server"
assert 'score' in tool, "Tool missing score"
print(f" ✓ Routing details structure correct")
print(f" ✓ Stage 1: {len(details['stage1_servers'])} servers")
print(f" ✓ Stage 2: {len(details['final_tools'])} final tools")
print()
def test_tool_schemas():
"""Test that tool schemas are properly formatted for OpenAI."""
print("Testing Tool Schemas...")
servers = create_tool_knowledge_base()
all_tools = get_all_tools(servers)
for tool in all_tools:
schema = tool.to_schema()
# Check OpenAI function calling format
assert schema['type'] == 'function', f"Tool {tool.name} has wrong type"
assert 'function' in schema, f"Tool {tool.name} missing function"
func = schema['function']
assert 'name' in func, f"Tool {tool.name} missing name"
assert 'description' in func, f"Tool {tool.name} missing description"
assert 'parameters' in func, f"Tool {tool.name} missing parameters"
params = func['parameters']
assert params['type'] == 'object', f"Tool {tool.name} parameters not object type"
assert 'properties' in params, f"Tool {tool.name} missing properties"
print(f" ✓ All {len(all_tools)} tool schemas properly formatted")
print(f" ✓ Compatible with OpenAI function calling")
print()
def run_all_tests():
"""Run all basic tests."""
print("""
╔════════════════════════════════════════════════════════════════════════════╗
║ ║
║ Active Tool Selection - Basic Tests ║
║ ║
╚════════════════════════════════════════════════════════════════════════════╝
""")
try:
test_knowledge_base()
test_semantic_router()
test_structured_request_parser()
test_routing_details()
test_tool_schemas()
print("=" * 80)
print("✅ ALL TESTS PASSED")
print("=" * 80)
print()
print("The active tool selection system is working correctly!")
print()
print("Next steps:")
print(" 1. Configure your API key in .env")
print(" 2. Run 'python quickstart.py' for a demonstration")
print(" 3. Run 'python demo_comparison.py' for comprehensive comparison")
print()
except AssertionError as e:
print()
print("=" * 80)
print("❌ TEST FAILED")
print("=" * 80)
print(f"Error: {e}")
print()
raise
except Exception as e:
print()
print("=" * 80)
print("❌ UNEXPECTED ERROR")
print("=" * 80)
print(f"Error: {e}")
print()
raise
if __name__ == "__main__":
run_all_tests()
test_usage_none.py¶
"""Regression test: agent must tolerate providers that return usage=None.
The OpenAI SDK response object always HAS a `usage` attribute (pydantic
field), but it deserializes as None when the provider omits token accounting.
The old `hasattr(response, 'usage')` guard was therefore ineffective and
`response.usage.total_tokens` raised AttributeError, crashing execute_task.
"""
import os
import sys
from types import SimpleNamespace
os.environ.setdefault("OPENAI_API_KEY", "test-key") # OpenAI() requires a key at construction
sys.path.insert(0, os.path.dirname(__file__))
from agent import ActiveToolAgent, RetrievalToolAgent, PassiveToolAgent
from tool_knowledge_base import ToolDefinition, ServerDefinition
AGENT_CLASSES = [ActiveToolAgent, RetrievalToolAgent, PassiveToolAgent]
def _catalog():
tool = ToolDefinition(
name="demo_tool",
description="demo tool",
parameters={"type": "object", "properties": {}},
server="demo",
)
return [ServerDefinition(name="demo", description="demo server", tools=[tool])]
def _client_with_usage(usage):
"""Fake OpenAI client; response mimics the SDK object (usage attr always present)."""
message = SimpleNamespace(content="final answer", tool_calls=None)
response = SimpleNamespace(choices=[SimpleNamespace(message=message)], usage=usage)
completions = SimpleNamespace(create=lambda **kwargs: response)
return SimpleNamespace(chat=SimpleNamespace(completions=completions))
def test_usage_none_does_not_crash():
for cls in AGENT_CLASSES:
agent = cls(servers=_catalog())
agent.client = _client_with_usage(None)
result = agent.execute_task("do something trivial")
assert result["metrics"]["tokens_used"] == 0, cls.__name__
def test_usage_still_accumulated_when_present():
for cls in AGENT_CLASSES:
agent = cls(servers=_catalog())
agent.client = _client_with_usage(SimpleNamespace(total_tokens=42))
result = agent.execute_task("do something trivial")
assert result["metrics"]["tokens_used"] == 42, cls.__name__
tool_knowledge_base.py¶
"""
Tool Knowledge Base - Simulates MCP servers with various tools.
This module defines a comprehensive knowledge base of tools organized by domains (servers),
similar to the MCP (Model Context Protocol) ecosystem. Each server represents a platform
or service domain with specific tools.
"""
from typing import List, Dict, Any
class ToolDefinition:
"""Represents a single tool with its metadata."""
def __init__(self, name: str, description: str, parameters: Dict[str, Any], server: str):
self.name = name
self.description = description
self.parameters = parameters
self.server = server
def to_schema(self) -> Dict[str, Any]:
"""Convert to OpenAI function schema format."""
return {
"type": "function",
"function": {
"name": self.name,
"description": self.description,
"parameters": self.parameters
}
}
def __repr__(self):
return f"Tool(name={self.name}, server={self.server})"
class ServerDefinition:
"""Represents a server (domain) containing multiple tools."""
def __init__(self, name: str, description: str, tools: List[ToolDefinition]):
self.name = name
self.description = description
self.tools = tools
def __repr__(self):
return f"Server(name={self.name}, tools={len(self.tools)})"
# Define comprehensive tool knowledge base
def create_tool_knowledge_base() -> List[ServerDefinition]:
"""
Create a comprehensive tool knowledge base organized by servers.
This simulates the MCP ecosystem with multiple domains:
- GitHub: Repository management and code operations
- Filesystem: File system operations
- Database: Data storage and retrieval
- Web: HTTP requests and web scraping
- Analytics: Data analysis and visualization
- Communication: Email and messaging
- DevOps: Deployment and monitoring
- Cloud: Cloud service operations
"""
servers = []
# GitHub Server
github_tools = [
ToolDefinition(
name="github_search_repos",
description="Search for GitHub repositories using keywords, filters, and sorting options",
parameters={
"type": "object",
"properties": {
"query": {"type": "string", "description": "Search query (e.g., 'language:python stars:>1000')"},
"sort": {"type": "string", "enum": ["stars", "forks", "updated"], "description": "Sort by field"},
"per_page": {"type": "integer", "description": "Results per page (max 100)"}
},
"required": ["query"]
},
server="github"
),
ToolDefinition(
name="github_create_pr",
description="Create a pull request in a GitHub repository",
parameters={
"type": "object",
"properties": {
"repo": {"type": "string", "description": "Repository name (owner/repo)"},
"title": {"type": "string", "description": "PR title"},
"body": {"type": "string", "description": "PR description"},
"head": {"type": "string", "description": "Branch to merge from"},
"base": {"type": "string", "description": "Branch to merge into"}
},
"required": ["repo", "title", "head", "base"]
},
server="github"
),
ToolDefinition(
name="github_list_issues",
description="List issues in a GitHub repository with filtering options",
parameters={
"type": "object",
"properties": {
"repo": {"type": "string", "description": "Repository name (owner/repo)"},
"state": {"type": "string", "enum": ["open", "closed", "all"], "description": "Issue state"},
"labels": {"type": "array", "items": {"type": "string"}, "description": "Filter by labels"}
},
"required": ["repo"]
},
server="github"
),
ToolDefinition(
name="github_get_file",
description="Get contents of a file from a GitHub repository",
parameters={
"type": "object",
"properties": {
"repo": {"type": "string", "description": "Repository name (owner/repo)"},
"path": {"type": "string", "description": "File path in repository"},
"branch": {"type": "string", "description": "Branch name (default: main)"}
},
"required": ["repo", "path"]
},
server="github"
),
ToolDefinition(
name="github_create_issue",
description="Create a new issue in a GitHub repository",
parameters={
"type": "object",
"properties": {
"repo": {"type": "string", "description": "Repository name (owner/repo)"},
"title": {"type": "string", "description": "Issue title"},
"body": {"type": "string", "description": "Issue description"},
"labels": {"type": "array", "items": {"type": "string"}, "description": "Issue labels"}
},
"required": ["repo", "title"]
},
server="github"
)
]
servers.append(ServerDefinition("github", "GitHub repository management and version control operations", github_tools))
# Filesystem Server
filesystem_tools = [
ToolDefinition(
name="fs_read_file",
description="Read the contents of a file from the local filesystem",
parameters={
"type": "object",
"properties": {
"path": {"type": "string", "description": "File path to read"},
"encoding": {"type": "string", "description": "File encoding (default: utf-8)"}
},
"required": ["path"]
},
server="filesystem"
),
ToolDefinition(
name="fs_write_file",
description="Write content to a file in the local filesystem",
parameters={
"type": "object",
"properties": {
"path": {"type": "string", "description": "File path to write"},
"content": {"type": "string", "description": "Content to write"},
"mode": {"type": "string", "enum": ["w", "a"], "description": "Write mode (w=overwrite, a=append)"}
},
"required": ["path", "content"]
},
server="filesystem"
),
ToolDefinition(
name="fs_list_directory",
description="List files and directories in a given path",
parameters={
"type": "object",
"properties": {
"path": {"type": "string", "description": "Directory path to list"},
"recursive": {"type": "boolean", "description": "List recursively"},
"pattern": {"type": "string", "description": "File pattern filter (e.g., '*.py')"}
},
"required": ["path"]
},
server="filesystem"
),
ToolDefinition(
name="fs_delete_file",
description="Delete a file or directory from the filesystem",
parameters={
"type": "object",
"properties": {
"path": {"type": "string", "description": "Path to delete"},
"recursive": {"type": "boolean", "description": "Delete directories recursively"}
},
"required": ["path"]
},
server="filesystem"
),
ToolDefinition(
name="fs_search_files",
description="Search for files containing specific text or matching patterns",
parameters={
"type": "object",
"properties": {
"path": {"type": "string", "description": "Directory to search in"},
"query": {"type": "string", "description": "Text to search for"},
"file_pattern": {"type": "string", "description": "File pattern (e.g., '*.py')"}
},
"required": ["path", "query"]
},
server="filesystem"
)
]
servers.append(ServerDefinition("filesystem", "Local filesystem operations for reading, writing, and managing files", filesystem_tools))
# Database Server
database_tools = [
ToolDefinition(
name="db_query",
description="Execute a SQL query on the database and return results",
parameters={
"type": "object",
"properties": {
"sql": {"type": "string", "description": "SQL query to execute"},
"database": {"type": "string", "description": "Database name"},
"timeout": {"type": "integer", "description": "Query timeout in seconds"}
},
"required": ["sql"]
},
server="database"
),
ToolDefinition(
name="db_insert",
description="Insert data into a database table",
parameters={
"type": "object",
"properties": {
"table": {"type": "string", "description": "Table name"},
"data": {"type": "object", "description": "Data to insert as key-value pairs"},
"database": {"type": "string", "description": "Database name"}
},
"required": ["table", "data"]
},
server="database"
),
ToolDefinition(
name="db_update",
description="Update records in a database table",
parameters={
"type": "object",
"properties": {
"table": {"type": "string", "description": "Table name"},
"data": {"type": "object", "description": "Data to update"},
"where": {"type": "string", "description": "WHERE clause condition"},
"database": {"type": "string", "description": "Database name"}
},
"required": ["table", "data", "where"]
},
server="database"
),
ToolDefinition(
name="db_delete",
description="Delete records from a database table",
parameters={
"type": "object",
"properties": {
"table": {"type": "string", "description": "Table name"},
"where": {"type": "string", "description": "WHERE clause condition"},
"database": {"type": "string", "description": "Database name"}
},
"required": ["table", "where"]
},
server="database"
),
ToolDefinition(
name="db_schema",
description="Get schema information for database tables",
parameters={
"type": "object",
"properties": {
"table": {"type": "string", "description": "Table name (optional, returns all if omitted)"},
"database": {"type": "string", "description": "Database name"}
},
"required": []
},
server="database"
)
]
servers.append(ServerDefinition("database", "Database operations for querying and manipulating structured data", database_tools))
# Web Server
web_tools = [
ToolDefinition(
name="web_get",
description="Make HTTP GET request to a URL and return the response",
parameters={
"type": "object",
"properties": {
"url": {"type": "string", "description": "URL to request"},
"headers": {"type": "object", "description": "HTTP headers"},
"params": {"type": "object", "description": "Query parameters"}
},
"required": ["url"]
},
server="web"
),
ToolDefinition(
name="web_post",
description="Make HTTP POST request to a URL with data",
parameters={
"type": "object",
"properties": {
"url": {"type": "string", "description": "URL to post to"},
"data": {"type": "object", "description": "Data to send"},
"headers": {"type": "object", "description": "HTTP headers"}
},
"required": ["url", "data"]
},
server="web"
),
ToolDefinition(
name="web_scrape",
description="Scrape and extract data from a web page using CSS selectors",
parameters={
"type": "object",
"properties": {
"url": {"type": "string", "description": "URL to scrape"},
"selector": {"type": "string", "description": "CSS selector for elements to extract"},
"attributes": {"type": "array", "items": {"type": "string"}, "description": "Attributes to extract"}
},
"required": ["url", "selector"]
},
server="web"
),
ToolDefinition(
name="web_download",
description="Download a file from a URL to local filesystem",
parameters={
"type": "object",
"properties": {
"url": {"type": "string", "description": "URL to download from"},
"destination": {"type": "string", "description": "Local path to save file"},
"headers": {"type": "object", "description": "HTTP headers"}
},
"required": ["url", "destination"]
},
server="web"
)
]
servers.append(ServerDefinition("web", "HTTP operations for making requests and scraping web content", web_tools))
# Analytics Server
analytics_tools = [
ToolDefinition(
name="analytics_summarize",
description="Calculate summary statistics for a dataset",
parameters={
"type": "object",
"properties": {
"data": {"type": "array", "description": "Array of numeric values"},
"metrics": {"type": "array", "items": {"type": "string"}, "description": "Metrics to calculate (mean, median, std, etc.)"}
},
"required": ["data"]
},
server="analytics"
),
ToolDefinition(
name="analytics_visualize",
description="Create visualizations from data (charts, graphs)",
parameters={
"type": "object",
"properties": {
"data": {"type": "object", "description": "Data to visualize"},
"chart_type": {"type": "string", "enum": ["line", "bar", "scatter", "pie"], "description": "Type of chart"},
"title": {"type": "string", "description": "Chart title"},
"output_path": {"type": "string", "description": "Path to save chart image"}
},
"required": ["data", "chart_type"]
},
server="analytics"
),
ToolDefinition(
name="analytics_correlation",
description="Calculate correlation between variables in a dataset",
parameters={
"type": "object",
"properties": {
"data": {"type": "object", "description": "Dataset with variables as keys"},
"method": {"type": "string", "enum": ["pearson", "spearman"], "description": "Correlation method"}
},
"required": ["data"]
},
server="analytics"
),
ToolDefinition(
name="analytics_predict",
description="Make predictions using machine learning models",
parameters={
"type": "object",
"properties": {
"model_type": {"type": "string", "description": "Type of ML model (linear, tree, etc.)"},
"features": {"type": "array", "description": "Feature values for prediction"},
"trained_model_path": {"type": "string", "description": "Path to trained model file"}
},
"required": ["features"]
},
server="analytics"
)
]
servers.append(ServerDefinition("analytics", "Data analysis and visualization tools for statistical operations", analytics_tools))
# Communication Server
communication_tools = [
ToolDefinition(
name="comm_send_email",
description="Send an email message to recipients",
parameters={
"type": "object",
"properties": {
"to": {"type": "array", "items": {"type": "string"}, "description": "Recipient email addresses"},
"subject": {"type": "string", "description": "Email subject"},
"body": {"type": "string", "description": "Email body content"},
"attachments": {"type": "array", "items": {"type": "string"}, "description": "File paths to attach"}
},
"required": ["to", "subject", "body"]
},
server="communication"
),
ToolDefinition(
name="comm_send_slack",
description="Send a message to a Slack channel or user",
parameters={
"type": "object",
"properties": {
"channel": {"type": "string", "description": "Channel name or user ID"},
"message": {"type": "string", "description": "Message content"},
"thread_ts": {"type": "string", "description": "Thread timestamp for replies"}
},
"required": ["channel", "message"]
},
server="communication"
),
ToolDefinition(
name="comm_read_email",
description="Read emails from inbox with filtering options",
parameters={
"type": "object",
"properties": {
"folder": {"type": "string", "description": "Email folder (inbox, sent, etc.)"},
"unread_only": {"type": "boolean", "description": "Only return unread emails"},
"limit": {"type": "integer", "description": "Maximum number of emails to return"}
},
"required": []
},
server="communication"
),
ToolDefinition(
name="comm_schedule_meeting",
description="Schedule a meeting in calendar",
parameters={
"type": "object",
"properties": {
"title": {"type": "string", "description": "Meeting title"},
"start_time": {"type": "string", "description": "Start time (ISO format)"},
"duration": {"type": "integer", "description": "Duration in minutes"},
"attendees": {"type": "array", "items": {"type": "string"}, "description": "Attendee emails"}
},
"required": ["title", "start_time", "attendees"]
},
server="communication"
)
]
servers.append(ServerDefinition("communication", "Email, messaging, and calendar tools for communication", communication_tools))
# DevOps Server
devops_tools = [
ToolDefinition(
name="devops_deploy",
description="Deploy application to specified environment",
parameters={
"type": "object",
"properties": {
"environment": {"type": "string", "enum": ["dev", "staging", "production"], "description": "Target environment"},
"version": {"type": "string", "description": "Version to deploy"},
"rollback": {"type": "boolean", "description": "Enable auto-rollback on failure"}
},
"required": ["environment", "version"]
},
server="devops"
),
ToolDefinition(
name="devops_monitor",
description="Get monitoring metrics for services and infrastructure",
parameters={
"type": "object",
"properties": {
"service": {"type": "string", "description": "Service name to monitor"},
"metrics": {"type": "array", "items": {"type": "string"}, "description": "Metrics to retrieve (cpu, memory, etc.)"},
"timerange": {"type": "string", "description": "Time range (e.g., '1h', '24h')"}
},
"required": ["service"]
},
server="devops"
),
ToolDefinition(
name="devops_logs",
description="Query and filter application logs",
parameters={
"type": "object",
"properties": {
"service": {"type": "string", "description": "Service name"},
"level": {"type": "string", "enum": ["debug", "info", "warning", "error"], "description": "Log level filter"},
"query": {"type": "string", "description": "Text to search in logs"},
"limit": {"type": "integer", "description": "Maximum number of log entries"}
},
"required": ["service"]
},
server="devops"
),
ToolDefinition(
name="devops_run_pipeline",
description="Trigger a CI/CD pipeline execution",
parameters={
"type": "object",
"properties": {
"pipeline": {"type": "string", "description": "Pipeline name or ID"},
"branch": {"type": "string", "description": "Git branch to build"},
"parameters": {"type": "object", "description": "Pipeline parameters"}
},
"required": ["pipeline"]
},
server="devops"
)
]
servers.append(ServerDefinition("devops", "DevOps tools for deployment, monitoring, and CI/CD operations", devops_tools))
# Cloud Server
cloud_tools = [
ToolDefinition(
name="cloud_create_vm",
description="Create a virtual machine in the cloud",
parameters={
"type": "object",
"properties": {
"instance_type": {"type": "string", "description": "VM instance type (e.g., 't2.micro')"},
"region": {"type": "string", "description": "Cloud region"},
"image_id": {"type": "string", "description": "OS image ID"},
"tags": {"type": "object", "description": "Tags for the VM"}
},
"required": ["instance_type", "region"]
},
server="cloud"
),
ToolDefinition(
name="cloud_list_resources",
description="List cloud resources (VMs, storage, databases)",
parameters={
"type": "object",
"properties": {
"resource_type": {"type": "string", "enum": ["vm", "storage", "database", "network"], "description": "Type of resource"},
"region": {"type": "string", "description": "Cloud region"},
"filters": {"type": "object", "description": "Filter criteria"}
},
"required": ["resource_type"]
},
server="cloud"
),
ToolDefinition(
name="cloud_upload_storage",
description="Upload files to cloud storage",
parameters={
"type": "object",
"properties": {
"bucket": {"type": "string", "description": "Storage bucket name"},
"file_path": {"type": "string", "description": "Local file path"},
"destination": {"type": "string", "description": "Destination path in bucket"},
"public": {"type": "boolean", "description": "Make file publicly accessible"}
},
"required": ["bucket", "file_path"]
},
server="cloud"
),
ToolDefinition(
name="cloud_manage_firewall",
description="Configure cloud firewall rules",
parameters={
"type": "object",
"properties": {
"resource_id": {"type": "string", "description": "Resource ID to configure"},
"action": {"type": "string", "enum": ["add", "remove"], "description": "Action to perform"},
"rule": {"type": "object", "description": "Firewall rule specification"}
},
"required": ["resource_id", "action", "rule"]
},
server="cloud"
)
]
servers.append(ServerDefinition("cloud", "Cloud infrastructure management for VMs, storage, and networking", cloud_tools))
return servers
def get_all_tools(servers: List[ServerDefinition]) -> List[ToolDefinition]:
"""Get flat list of all tools from all servers."""
all_tools = []
for server in servers:
all_tools.extend(server.tools)
return all_tools
def count_tokens_in_schema(schema: Dict[str, Any]) -> int:
"""Rough estimation of tokens in a tool schema (approximately 1 token per 4 characters)."""
import json
schema_str = json.dumps(schema)
return len(schema_str) // 4
def calculate_total_tokens(tools: List[ToolDefinition]) -> int:
"""Calculate total tokens required to inject all tool schemas."""
total = 0
for tool in tools:
total += count_tokens_in_schema(tool.to_schema())
return total