mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
refactor: cleanup and consolidate tool packages
- Delete deprecated files (agent_tool_v1_deprecated.py, swarm_tool_v1_deprecated.py) - Consolidate duplicate LLM tools (merge llm_tool.py into llm_unified.py) - Rename agent CriticTool to CodeReviewTool to avoid conflict with reasoning critic - Fix stub classes in swarm_tool.py to raise ImportError instead of silent pass - Update CLAUDE.md with accurate tool counts (62 tools across 17 packages)
This commit is contained in:
@@ -26,4 +26,4 @@ if os.environ.get("HANZO_MCP_TRANSPORT") == "stdio":
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
__version__ = "0.10.8"
|
||||
__version__ = "0.10.9"
|
||||
|
||||
@@ -42,6 +42,8 @@ Action = Annotated[
|
||||
"close", # Close browser
|
||||
"tabs", # List/switch tabs
|
||||
"connect", # Connect to existing browser (CDP)
|
||||
"set_headless", # Switch headless/headed mode (restarts browser)
|
||||
"status", # Get current browser status
|
||||
],
|
||||
Field(description="Browser action to perform"),
|
||||
]
|
||||
@@ -381,6 +383,49 @@ class BrowserTool:
|
||||
"tabs": [{"index": i, "url": p.url} for i, p in enumerate(pool.pages)],
|
||||
}
|
||||
|
||||
elif action == "set_headless":
|
||||
# Toggle or set headless mode (requires browser restart)
|
||||
# Use the 'text' param as "true"/"false" or toggle if not provided
|
||||
if text is not None:
|
||||
new_headless = text.lower() in ("true", "1", "yes", "on")
|
||||
else:
|
||||
# Toggle current mode
|
||||
new_headless = not pool._headless
|
||||
|
||||
current_url = page.url if page else None
|
||||
old_mode = "headless" if pool._headless else "headed"
|
||||
|
||||
# Close and relaunch with new mode
|
||||
await pool.close()
|
||||
self.headless = new_headless
|
||||
page = await pool.ensure_browser(headless=new_headless)
|
||||
|
||||
# Restore URL if we had one
|
||||
if current_url and current_url != "about:blank":
|
||||
await page.goto(current_url)
|
||||
|
||||
new_mode = "headless" if new_headless else "headed"
|
||||
return {
|
||||
"success": True,
|
||||
"previous_mode": old_mode,
|
||||
"current_mode": new_mode,
|
||||
"headless": new_headless,
|
||||
"url": page.url,
|
||||
}
|
||||
|
||||
elif action == "status":
|
||||
# Get current browser status
|
||||
return {
|
||||
"success": True,
|
||||
"initialized": pool._initialized,
|
||||
"headless": pool._headless,
|
||||
"mode": "headless" if pool._headless else "headed",
|
||||
"cdp_endpoint": pool._cdp_endpoint,
|
||||
"tabs": len(pool.pages),
|
||||
"current_url": page.url if page else None,
|
||||
"playwright_available": PLAYWRIGHT_AVAILABLE,
|
||||
}
|
||||
|
||||
else:
|
||||
return {"error": f"Unknown action: {action}"}
|
||||
|
||||
@@ -416,6 +461,12 @@ class BrowserTool:
|
||||
Uses shared browser instance for low latency.
|
||||
Supports CDP connection for cross-MCP browser sharing.
|
||||
|
||||
DISPLAY INSTRUCTIONS: Show the action result as a bullet point.
|
||||
For navigate: • Navigated to [url] (status: [status])
|
||||
For click/type/fill: • [action] completed on [selector]
|
||||
For screenshot: • Screenshot captured, [format] ([size] bytes)
|
||||
For errors: • Error: [message]
|
||||
|
||||
Actions:
|
||||
- navigate: Go to URL (requires url)
|
||||
- click: Click element (requires selector)
|
||||
@@ -429,6 +480,8 @@ class BrowserTool:
|
||||
- close: Close browser
|
||||
- tabs: List or switch tabs
|
||||
- connect: Connect to existing browser via CDP
|
||||
- set_headless: Switch headless/headed mode (text="true"/"false" or omit to toggle)
|
||||
- status: Get current browser status (headless, tabs, url)
|
||||
|
||||
Examples:
|
||||
browser(action="navigate", url="https://example.com")
|
||||
@@ -436,6 +489,9 @@ class BrowserTool:
|
||||
browser(action="type", selector="input[name=q]", text="hello")
|
||||
browser(action="screenshot", full_page=True)
|
||||
browser(action="connect", cdp_endpoint="http://localhost:9222")
|
||||
browser(action="set_headless", text="false") # Switch to visible
|
||||
browser(action="set_headless") # Toggle mode
|
||||
browser(action="status") # Check current state
|
||||
"""
|
||||
return await tool_instance.execute(
|
||||
action=action,
|
||||
|
||||
@@ -12,7 +12,7 @@ def register_version_tool(mcp: FastMCP) -> None:
|
||||
"""Register the version tool with the MCP server."""
|
||||
|
||||
@mcp.tool()
|
||||
async def version() -> dict[str, Any]:
|
||||
async def version() -> str:
|
||||
"""Get hanzo-mcp version and environment information.
|
||||
|
||||
Returns version info including:
|
||||
@@ -22,9 +22,17 @@ def register_version_tool(mcp: FastMCP) -> None:
|
||||
"""
|
||||
import platform
|
||||
|
||||
return {
|
||||
info = {
|
||||
"hanzo_mcp": hanzo_mcp.__version__,
|
||||
"python": sys.version.split()[0],
|
||||
"platform": platform.system(),
|
||||
"arch": platform.machine(),
|
||||
}
|
||||
|
||||
# Return formatted string for clean display
|
||||
lines = [
|
||||
f"hanzo-mcp: v{info['hanzo_mcp']}",
|
||||
f"python: {info['python']}",
|
||||
f"platform: {info['platform']} ({info['arch']})",
|
||||
]
|
||||
return "\n".join(lines)
|
||||
|
||||
@@ -108,6 +108,10 @@ class TodoTool(TodoBaseTool):
|
||||
"""Get the tool description."""
|
||||
return """Manage todos. Actions: list (default), add, update, remove, clear.
|
||||
|
||||
DISPLAY INSTRUCTIONS: Output is plain text with formatted todo list.
|
||||
Show directly without modification. The list uses === headers ===
|
||||
and checkmarks for completed items.
|
||||
|
||||
Usage:
|
||||
todo
|
||||
todo "Fix the bug in authentication"
|
||||
|
||||
@@ -25,12 +25,11 @@ class MCPResourceDocument:
|
||||
return result
|
||||
|
||||
def to_json_string(self) -> str:
|
||||
"""Convert to clean JSON string - just the data, not wrapped."""
|
||||
"""Convert to clean JSON string - just the data, no wrapper."""
|
||||
import json
|
||||
|
||||
# Return just the data content, not wrapped in {"data": ...}
|
||||
# This makes output cleaner when MCP adds its own wrapper
|
||||
return json.dumps({"result": self.data}, indent=2)
|
||||
# Return just the data content directly, no wrapper
|
||||
return json.dumps(self.data, indent=2)
|
||||
|
||||
def to_readable_string(self) -> str:
|
||||
"""Convert to human-readable formatted string for display."""
|
||||
|
||||
@@ -4,7 +4,7 @@ This package provides agent orchestration and multi-agent tools:
|
||||
- agent: Agent spawning and orchestration
|
||||
- swarm: Multi-agent swarm coordination
|
||||
- CLI agents: Claude, Codex, Gemini, Grok CLI tools
|
||||
- critic: Critical analysis tool
|
||||
- code_review: Critical code review tool
|
||||
- review: Code review tool
|
||||
|
||||
Install:
|
||||
@@ -25,7 +25,7 @@ import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Core tools that don't have external dependencies
|
||||
from .critic_tool import CriticTool
|
||||
from .critic_tool import CodeReviewTool
|
||||
from .iching_tool import IChingTool
|
||||
from .review_tool import ReviewTool
|
||||
from .network_tool import NetworkTool
|
||||
@@ -58,7 +58,7 @@ except ImportError as e:
|
||||
|
||||
# Export list for tool discovery
|
||||
TOOLS = [
|
||||
CriticTool,
|
||||
CodeReviewTool,
|
||||
IChingTool,
|
||||
ReviewTool,
|
||||
ClarificationTool,
|
||||
@@ -73,7 +73,7 @@ TOOLS = [
|
||||
__all__ = [
|
||||
"register_tools",
|
||||
"TOOLS",
|
||||
"CriticTool",
|
||||
"CodeReviewTool",
|
||||
"IChingTool",
|
||||
"ReviewTool",
|
||||
"ClarificationTool",
|
||||
|
||||
@@ -1,660 +0,0 @@
|
||||
"""Agent tool implementation for Hanzo AI.
|
||||
|
||||
This module implements the AgentTool that allows Claude to delegate tasks to sub-agents,
|
||||
enabling concurrent execution of multiple operations and specialized processing.
|
||||
"""
|
||||
|
||||
import re
|
||||
import json
|
||||
import time
|
||||
import asyncio
|
||||
|
||||
# Import litellm with warnings suppressed
|
||||
import warnings
|
||||
from typing import Unpack, Annotated, TypedDict, final, override
|
||||
from collections.abc import Iterable
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", DeprecationWarning)
|
||||
import litellm
|
||||
from pydantic import Field
|
||||
from mcp.server import FastMCP
|
||||
from openai.types.chat import ChatCompletionToolParam, ChatCompletionMessageParam
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
|
||||
from hanzo_mcp.tools.filesystem import Edit, MultiEdit, get_read_only_filesystem_tools
|
||||
from hanzo_mcp.tools.common.context import (
|
||||
ToolContext,
|
||||
create_tool_context,
|
||||
)
|
||||
from hanzo_mcp.tools.common.batch_tool import BatchTool
|
||||
|
||||
from hanzo_tools.core import BaseTool, PermissionManager, auto_timeout
|
||||
|
||||
from .prompt import (
|
||||
get_default_model,
|
||||
get_system_prompt,
|
||||
get_model_parameters,
|
||||
get_allowed_agent_tools,
|
||||
)
|
||||
from .critic_tool import CriticTool, CriticProtocol
|
||||
from .iching_tool import IChingTool
|
||||
from .review_tool import ReviewTool, ReviewProtocol
|
||||
from .tool_adapter import (
|
||||
convert_tools_to_openai_functions,
|
||||
)
|
||||
from .clarification_tool import ClarificationTool
|
||||
from .clarification_protocol import (
|
||||
ClarificationType,
|
||||
AgentClarificationMixin,
|
||||
)
|
||||
|
||||
Prompt = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Task for the agent to perform (must include absolute paths starting with /)",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class AgentToolParams(TypedDict, total=False):
|
||||
"""Parameters for the AgentTool.
|
||||
|
||||
Attributes:
|
||||
prompts: Task(s) for the agent to perform (must include absolute paths starting with /)
|
||||
"""
|
||||
|
||||
prompts: str | list[str]
|
||||
|
||||
|
||||
@final
|
||||
class AgentTool(AgentClarificationMixin, BaseTool):
|
||||
"""Tool for delegating tasks to sub-agents.
|
||||
|
||||
The AgentTool allows Claude to create and manage sub-agents for performing
|
||||
specialized tasks concurrently, such as code search, analysis, and more.
|
||||
|
||||
Agents can request clarification from the main loop up to once per task.
|
||||
"""
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name.
|
||||
|
||||
Returns:
|
||||
Tool name
|
||||
"""
|
||||
return "agent"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description.
|
||||
|
||||
Returns:
|
||||
Tool description
|
||||
"""
|
||||
# Note: Consider adding glob pattern support in future
|
||||
at = [t.name for t in self.available_tools]
|
||||
|
||||
return f"""Launch a new agent that has access to the following tools: {at}. When you are searching for a keyword or file and are not confident that you will find the right match in the first few tries, use the Agent tool to perform the search for you.
|
||||
|
||||
When to use the Agent tool:
|
||||
- If you are searching for a keyword like \"config\" or \"logger\", or for questions like \"which file does X?\", the Agent tool is strongly recommended
|
||||
- When you need to perform edits across multiple files based on search results
|
||||
- When you need to delegate complex file modification tasks
|
||||
|
||||
When NOT to use the Agent tool:
|
||||
- If you want to read a specific file path, use the read or glob tool instead of the Agent tool, to find the match more quickly
|
||||
- If you are searching for a specific class definition like \"class Foo\", use the glob tool instead, to find the match more quickly
|
||||
- If you are searching for code within a specific file or set of 2-3 files, use the read tool instead of the Agent tool, to find the match more quickly
|
||||
- Writing code and running bash commands (use other tools for that)
|
||||
- Other tasks that are not related to searching for a keyword or file
|
||||
|
||||
Usage notes:
|
||||
1. Launch multiple agents concurrently whenever possible, to maximize performance; to do that, use a single message with multiple tool uses
|
||||
2. When the agent is done, it will return a single message back to you. The result returned by the agent is not visible to the user. To show the user the result, you should send a text message back to the user with a concise summary of the result.
|
||||
3. Each agent invocation is stateless. You will not be able to send additional messages to the agent, nor will the agent be able to communicate with you outside of its final report. Therefore, your prompt should contain a highly detailed task description for the agent to perform autonomously and you should specify exactly what information the agent should return back to you in its final and only message to you.
|
||||
4. The agent's outputs should generally be trusted
|
||||
5. Clearly tell the agent whether you expect it to write code or just to do research (search, file reads, web fetches, etc.), since it is not aware of the user's intent"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
model: str | None = None,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
max_iterations: int = 10,
|
||||
max_tool_uses: int = 30,
|
||||
) -> None:
|
||||
"""Initialize the agent tool.
|
||||
|
||||
Args:
|
||||
|
||||
permission_manager: Permission manager for access control
|
||||
model: Optional model name override in LiteLLM format (e.g., "openai/gpt-4o")
|
||||
api_key: Optional API key for the model provider
|
||||
base_url: Optional base URL for the model provider API endpoint
|
||||
max_tokens: Optional maximum tokens for model responses
|
||||
max_iterations: Maximum number of iterations for agent (default: 10)
|
||||
max_tool_uses: Maximum number of total tool uses for agent (default: 30)
|
||||
"""
|
||||
|
||||
self.permission_manager = permission_manager
|
||||
self.model_override = model
|
||||
self.api_key_override = api_key
|
||||
self.base_url_override = base_url
|
||||
self.max_tokens_override = max_tokens
|
||||
self.max_iterations = max_iterations
|
||||
self.max_tool_uses = max_tool_uses
|
||||
self.available_tools: list[BaseTool] = []
|
||||
self.available_tools.extend(get_read_only_filesystem_tools(self.permission_manager))
|
||||
self.available_tools.extend(get_read_only_jupyter_tools(self.permission_manager))
|
||||
|
||||
# Always add edit tools - agents should have edit access
|
||||
self.available_tools.append(Edit(self.permission_manager))
|
||||
self.available_tools.append(MultiEdit(self.permission_manager))
|
||||
|
||||
# Add clarification tool for agents
|
||||
self.available_tools.append(ClarificationTool())
|
||||
|
||||
# Add critic tool for agents (devil's advocate)
|
||||
self.available_tools.append(CriticTool())
|
||||
|
||||
# Add review tool for agents (balanced review)
|
||||
self.available_tools.append(ReviewTool())
|
||||
|
||||
# Add I Ching tool for creative guidance
|
||||
self.available_tools.append(IChingTool())
|
||||
|
||||
self.available_tools.append(BatchTool({t.name: t for t in self.available_tools}))
|
||||
|
||||
# Initialize protocols
|
||||
self.critic_protocol = CriticProtocol()
|
||||
self.review_protocol = ReviewProtocol()
|
||||
|
||||
@override
|
||||
@auto_timeout("agent_tool_v1_deprecated")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[AgentToolParams],
|
||||
) -> str:
|
||||
"""Execute the tool with the given parameters.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Tool execution result
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Create tool context
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract and validate parameters
|
||||
prompts = params.get("prompts")
|
||||
|
||||
if prompts is None:
|
||||
await tool_ctx.error("No prompts provided")
|
||||
return """Error: At least one prompt must be provided.
|
||||
|
||||
IMPORTANT REMINDER FOR CLAUDE:
|
||||
The dispatch_agent tool requires prompts parameter. Please provide either:
|
||||
- A single prompt as a string
|
||||
- Multiple prompts as an array of strings
|
||||
|
||||
Each prompt must contain absolute paths starting with /.
|
||||
Example of correct usage:
|
||||
- prompts: "Search for all instances of the 'config' variable in /Users/bytedance/project/hanzo-mcp"
|
||||
- prompts: ["Find files in /path/to/project", "Search code in /path/to/src"]"""
|
||||
|
||||
# Handle both string and list inputs
|
||||
if isinstance(prompts, str):
|
||||
prompt_list = [prompts]
|
||||
elif isinstance(prompts, list):
|
||||
if not prompts:
|
||||
await tool_ctx.error("Empty prompts list provided")
|
||||
return "Error: At least one prompt must be provided when using a list."
|
||||
if not all(isinstance(p, str) for p in prompts):
|
||||
await tool_ctx.error("All prompts must be strings")
|
||||
return "Error: All prompts in the list must be strings."
|
||||
prompt_list = prompts
|
||||
else:
|
||||
await tool_ctx.error("Invalid prompts parameter type")
|
||||
return "Error: Parameter 'prompts' must be a string or an array of strings."
|
||||
|
||||
# Validate absolute paths in all prompts
|
||||
absolute_path_pattern = r"/(?:[^/\s]+/)*[^/\s]+"
|
||||
for prompt in prompt_list:
|
||||
if not re.search(absolute_path_pattern, prompt):
|
||||
await tool_ctx.error(f"Prompt does not contain absolute path: {prompt[:50]}...")
|
||||
return """Error: All prompts must contain at least one absolute path.
|
||||
|
||||
IMPORTANT REMINDER FOR CLAUDE:
|
||||
When using the dispatch_agent tool, always include absolute paths in your prompts.
|
||||
Example of correct usage:
|
||||
- "Search for all instances of the 'config' variable in /Users/bytedance/project/hanzo-mcp"
|
||||
- "Find files that import the database module in /Users/bytedance/project/hanzo-mcp/src"
|
||||
|
||||
The agent cannot access files without knowing their absolute locations."""
|
||||
|
||||
# Execute agent(s) - always use _execute_multiple_agents for list inputs
|
||||
if len(prompt_list) == 1:
|
||||
await tool_ctx.info("Launching agent")
|
||||
result = await self._execute_multiple_agents(prompt_list, tool_ctx)
|
||||
execution_time = time.time() - start_time
|
||||
formatted_result = f"""Agent execution completed in {execution_time:.2f} seconds.
|
||||
|
||||
AGENT RESPONSE:
|
||||
{result}"""
|
||||
await tool_ctx.info(f"Agent execution completed in {execution_time:.2f}s")
|
||||
return formatted_result
|
||||
else:
|
||||
await tool_ctx.info(f"Launching {len(prompt_list)} agents in parallel")
|
||||
result = await self._execute_multiple_agents(prompt_list, tool_ctx)
|
||||
execution_time = time.time() - start_time
|
||||
formatted_result = f"""Multi-agent execution completed in {execution_time:.2f} seconds ({len(prompt_list)} agents).
|
||||
|
||||
AGENT RESPONSES:
|
||||
{result}"""
|
||||
await tool_ctx.info(f"Multi-agent execution completed in {execution_time:.2f}s")
|
||||
return formatted_result
|
||||
|
||||
async def _execute_agent(self, prompt: str, tool_ctx: ToolContext) -> str:
|
||||
"""Execute a single agent with the given prompt.
|
||||
|
||||
Args:
|
||||
prompt: The task prompt for the agent
|
||||
tool_ctx: Tool context for logging
|
||||
|
||||
Returns:
|
||||
Agent execution result
|
||||
"""
|
||||
# Get available tools for the agent
|
||||
agent_tools = get_allowed_agent_tools(
|
||||
self.available_tools,
|
||||
self.permission_manager,
|
||||
)
|
||||
|
||||
# Convert tools to OpenAI format
|
||||
openai_tools = convert_tools_to_openai_functions(agent_tools)
|
||||
|
||||
# Log execution start
|
||||
await tool_ctx.info("Starting agent execution")
|
||||
|
||||
# Create a result container
|
||||
result = ""
|
||||
|
||||
try:
|
||||
# Create system prompt for this agent
|
||||
system_prompt = get_system_prompt(
|
||||
agent_tools,
|
||||
self.permission_manager,
|
||||
)
|
||||
|
||||
# Execute agent
|
||||
await tool_ctx.info(f"Executing agent task: {prompt[:50]}...")
|
||||
result = await self._execute_agent_with_tools(system_prompt, prompt, agent_tools, openai_tools, tool_ctx)
|
||||
except Exception as e:
|
||||
# Log and return error result
|
||||
error_message = f"Error executing agent: {str(e)}"
|
||||
await tool_ctx.error(error_message)
|
||||
return f"Error: {error_message}"
|
||||
|
||||
return result if result else "No results returned from agent"
|
||||
|
||||
async def _execute_multiple_agents(self, prompts: list[str], tool_ctx: ToolContext) -> str:
|
||||
"""Execute multiple agents concurrently.
|
||||
|
||||
Args:
|
||||
prompts: List of prompts for the agents
|
||||
tool_ctx: Tool context for logging
|
||||
|
||||
Returns:
|
||||
Combined results from all agents
|
||||
"""
|
||||
# Get available tools for the agents
|
||||
agent_tools = get_allowed_agent_tools(
|
||||
self.available_tools,
|
||||
self.permission_manager,
|
||||
)
|
||||
|
||||
# Convert tools to OpenAI format
|
||||
openai_tools = convert_tools_to_openai_functions(agent_tools)
|
||||
|
||||
# Create system prompt for the agents
|
||||
system_prompt = get_system_prompt(
|
||||
agent_tools,
|
||||
self.permission_manager,
|
||||
)
|
||||
|
||||
# Create tasks for parallel execution
|
||||
tasks = []
|
||||
for i, prompt in enumerate(prompts):
|
||||
await tool_ctx.info(f"Creating agent task {i + 1}: {prompt[:50]}...")
|
||||
task = self._execute_agent_with_tools(system_prompt, prompt, agent_tools, openai_tools, tool_ctx)
|
||||
tasks.append(task)
|
||||
|
||||
# Execute all agents concurrently
|
||||
await tool_ctx.info(f"Executing {len(tasks)} agents in parallel")
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
# Handle single agent case
|
||||
if len(results) == 1:
|
||||
if isinstance(results[0], Exception):
|
||||
await tool_ctx.error(f"Agent execution failed: {str(results[0])}")
|
||||
return f"Error: {str(results[0])}"
|
||||
return results[0]
|
||||
|
||||
# Format results for multiple agents
|
||||
formatted_results = []
|
||||
for i, result in enumerate(results):
|
||||
if isinstance(result, Exception):
|
||||
formatted_results.append(f"Agent {i + 1} Error:\n{str(result)}")
|
||||
await tool_ctx.error(f"Agent {i + 1} failed: {str(result)}")
|
||||
else:
|
||||
formatted_results.append(f"Agent {i + 1} Result:\n{result}")
|
||||
|
||||
return "\n\n---\n\n".join(formatted_results)
|
||||
|
||||
async def _execute_agent_with_tools(
|
||||
self,
|
||||
system_prompt: str,
|
||||
user_prompt: str,
|
||||
available_tools: list[BaseTool],
|
||||
openai_tools: list[ChatCompletionToolParam],
|
||||
tool_ctx: ToolContext,
|
||||
) -> str:
|
||||
"""Execute agent with tool handling.
|
||||
|
||||
Args:
|
||||
system_prompt: System prompt for the agent
|
||||
user_prompt: User prompt for the agent
|
||||
available_tools: List of available tools
|
||||
openai_tools: List of tools in OpenAI format
|
||||
tool_ctx: Tool context for logging
|
||||
|
||||
Returns:
|
||||
Agent execution result
|
||||
"""
|
||||
# Get model parameters and name
|
||||
model = get_default_model(self.model_override)
|
||||
params = get_model_parameters(max_tokens=self.max_tokens_override)
|
||||
|
||||
# Initialize messages
|
||||
messages: Iterable[ChatCompletionMessageParam] = []
|
||||
messages.append({"role": "system", "content": system_prompt})
|
||||
messages.append({"role": "user", "content": user_prompt})
|
||||
|
||||
# Track tool usage for metrics
|
||||
tool_usage = {}
|
||||
total_tool_use_count = 0
|
||||
iteration_count = 0
|
||||
max_tool_uses = self.max_tool_uses # Safety limit to prevent infinite loops
|
||||
max_iterations = self.max_iterations # Add a maximum number of iterations for safety
|
||||
|
||||
# Execute until the agent completes or reaches the limit
|
||||
while total_tool_use_count < max_tool_uses and iteration_count < max_iterations:
|
||||
iteration_count += 1
|
||||
await tool_ctx.info(f"Calling model (iteration {iteration_count})...")
|
||||
|
||||
try:
|
||||
# Configure model parameters based on capabilities
|
||||
completion_params = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"tools": openai_tools,
|
||||
"tool_choice": "auto",
|
||||
"temperature": params["temperature"],
|
||||
"timeout": params["timeout"],
|
||||
}
|
||||
|
||||
if self.api_key_override:
|
||||
completion_params["api_key"] = self.api_key_override
|
||||
|
||||
# Add max_tokens if provided
|
||||
if params.get("max_tokens"):
|
||||
completion_params["max_tokens"] = params.get("max_tokens")
|
||||
|
||||
# Add base_url if provided
|
||||
if self.base_url_override:
|
||||
completion_params["base_url"] = self.base_url_override
|
||||
|
||||
# Make the model call
|
||||
response = litellm.completion(**completion_params) # pyright: ignore
|
||||
|
||||
if len(response.choices) == 0: # pyright: ignore
|
||||
raise ValueError("No response choices returned")
|
||||
|
||||
message = response.choices[0].message # pyright: ignore
|
||||
|
||||
# Add message to conversation history
|
||||
messages.append(message) # pyright: ignore
|
||||
|
||||
# If no tool calls, we're done
|
||||
if not message.tool_calls:
|
||||
return message.content or "Agent completed with no response."
|
||||
|
||||
# Process tool calls
|
||||
tool_call_count = len(message.tool_calls)
|
||||
await tool_ctx.info(f"Processing {tool_call_count} tool calls")
|
||||
|
||||
# Track which tool calls we've processed to ensure all get results
|
||||
tool_results_to_add = []
|
||||
|
||||
for tool_call in message.tool_calls:
|
||||
total_tool_use_count += 1
|
||||
function_name = tool_call.function.name
|
||||
|
||||
# Track usage
|
||||
tool_usage[function_name] = tool_usage.get(function_name, 0) + 1
|
||||
|
||||
# Log tool usage
|
||||
await tool_ctx.info(f"Agent using tool: {function_name}")
|
||||
|
||||
# Parse the arguments
|
||||
try:
|
||||
function_args = json.loads(tool_call.function.arguments)
|
||||
except json.JSONDecodeError:
|
||||
function_args = {}
|
||||
|
||||
# Find the matching tool
|
||||
tool = next((t for t in available_tools if t.name == function_name), None)
|
||||
if not tool:
|
||||
tool_result = f"Error: Tool '{function_name}' not found"
|
||||
# Special handling for clarification requests
|
||||
elif function_name == "request_clarification":
|
||||
try:
|
||||
# Extract clarification parameters
|
||||
request_type = function_args.get("type", "ADDITIONAL_INFO")
|
||||
question = function_args.get("question", "")
|
||||
context = function_args.get("context", {})
|
||||
options = function_args.get("options")
|
||||
|
||||
# Convert string type to enum
|
||||
clarification_type = ClarificationType[request_type]
|
||||
|
||||
# Request clarification
|
||||
answer = await self.request_clarification(
|
||||
request_type=clarification_type,
|
||||
question=question,
|
||||
context=context,
|
||||
options=options,
|
||||
)
|
||||
|
||||
tool_result = self.format_clarification_in_output(question, answer)
|
||||
except Exception as e:
|
||||
tool_result = f"Error processing clarification: {str(e)}"
|
||||
# Special handling for critic requests
|
||||
elif function_name == "critic":
|
||||
try:
|
||||
# Extract critic parameters
|
||||
review_type = function_args.get("review_type", "GENERAL")
|
||||
work_description = function_args.get("work_description", "")
|
||||
code_snippets = function_args.get("code_snippets")
|
||||
file_paths = function_args.get("file_paths")
|
||||
specific_concerns = function_args.get("specific_concerns")
|
||||
|
||||
# Request critical review
|
||||
tool_result = self.critic_protocol.request_review(
|
||||
review_type=review_type,
|
||||
work_description=work_description,
|
||||
code_snippets=code_snippets,
|
||||
file_paths=file_paths,
|
||||
specific_concerns=specific_concerns,
|
||||
)
|
||||
except Exception as e:
|
||||
tool_result = f"Error processing critic review: {str(e)}"
|
||||
# Special handling for review requests
|
||||
elif function_name == "review":
|
||||
try:
|
||||
# Extract review parameters
|
||||
focus = function_args.get("focus", "GENERAL")
|
||||
work_description = function_args.get("work_description", "")
|
||||
code_snippets = function_args.get("code_snippets")
|
||||
file_paths = function_args.get("file_paths")
|
||||
context = function_args.get("context")
|
||||
|
||||
# Request balanced review
|
||||
tool_result = self.review_protocol.request_review(
|
||||
focus=focus,
|
||||
work_description=work_description,
|
||||
code_snippets=code_snippets,
|
||||
file_paths=file_paths,
|
||||
context=context,
|
||||
)
|
||||
except Exception as e:
|
||||
tool_result = f"Error processing review: {str(e)}"
|
||||
else:
|
||||
try:
|
||||
tool_result = await tool.call(ctx=tool_ctx.mcp_context, **function_args)
|
||||
except Exception as e:
|
||||
tool_result = f"Error executing {function_name}: {str(e)}"
|
||||
|
||||
await tool_ctx.info(
|
||||
f"tool {function_name} run with args {function_args} and return {tool_result[: min(100, len(tool_result))]}"
|
||||
)
|
||||
# Add the tool result to the list
|
||||
tool_results_to_add.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"name": function_name,
|
||||
"content": tool_result,
|
||||
}
|
||||
)
|
||||
|
||||
# Add all tool results to messages atomically
|
||||
messages.extend(tool_results_to_add)
|
||||
|
||||
# Log progress
|
||||
await tool_ctx.info(f"Processed {len(message.tool_calls)} tool calls. Total: {total_tool_use_count}")
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Error in model call: {str(e)}")
|
||||
# Avoid trying to JSON serialize message objects
|
||||
await tool_ctx.error(f"Message count: {len(messages)}")
|
||||
|
||||
# CRITICAL FIX: Ensure we add tool_result messages for any tool_use blocks that were added
|
||||
# Get the last assistant message (which has the tool_calls)
|
||||
last_assistant_msg = None
|
||||
for msg in reversed(messages):
|
||||
if msg.get("role") == "assistant" and hasattr(msg, "tool_calls"):
|
||||
last_assistant_msg = msg
|
||||
break
|
||||
|
||||
# If there are orphaned tool calls, add error results for them
|
||||
if last_assistant_msg and hasattr(last_assistant_msg, "tool_calls"):
|
||||
# Count how many tool results we already have after this message
|
||||
tool_results_count = 0
|
||||
found_assistant = False
|
||||
for msg in reversed(messages):
|
||||
if msg == last_assistant_msg:
|
||||
found_assistant = True
|
||||
break
|
||||
if found_assistant and msg.get("role") == "tool":
|
||||
tool_results_count += 1
|
||||
|
||||
# Add error results for any missing tool calls
|
||||
expected_results = len(last_assistant_msg.tool_calls)
|
||||
if tool_results_count < expected_results:
|
||||
for i in range(tool_results_count, expected_results):
|
||||
tool_call = last_assistant_msg.tool_calls[i]
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"name": tool_call.function.name,
|
||||
"content": f"Error: Tool execution interrupted - {str(e)}",
|
||||
}
|
||||
)
|
||||
|
||||
return f"Error in agent execution: {str(e)}"
|
||||
|
||||
# If we've reached the limit, add a warning and get final response
|
||||
if total_tool_use_count >= max_tool_uses or iteration_count >= max_iterations:
|
||||
messages.append(
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You have reached the maximum iteration. Please provide your final response.",
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
# Make a final call to get the result
|
||||
final_response = litellm.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=params["temperature"],
|
||||
timeout=params["timeout"],
|
||||
max_tokens=params.get("max_tokens"),
|
||||
)
|
||||
|
||||
return (
|
||||
final_response.choices[0].message.content or "Agent reached max iteration limit without a response."
|
||||
) # pyright: ignore
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Error in final model call: {str(e)}")
|
||||
return f"Error in final response: {str(e)}"
|
||||
|
||||
# Should not reach here but just in case
|
||||
return "Agent execution completed after maximum iterations."
|
||||
|
||||
def _format_result(self, result: str, execution_time: float) -> str:
|
||||
"""Format agent result with metrics.
|
||||
|
||||
Args:
|
||||
result: Raw result from agent
|
||||
execution_time: Execution time in seconds
|
||||
|
||||
Returns:
|
||||
Formatted result with metrics
|
||||
"""
|
||||
return f"""Agent execution completed in {execution_time:.2f} seconds.
|
||||
|
||||
AGENT RESPONSE:
|
||||
{result}
|
||||
"""
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this agent tool with the MCP server.
|
||||
|
||||
Creates a wrapper function with explicitly defined parameters that match
|
||||
the tool's parameter schema and registers it with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
tool_self = self # Create a reference to self for use in the closure
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def dispatch_agent(prompts: str | list[str], ctx: MCPContext) -> str:
|
||||
return await tool_self.call(ctx, prompts=prompts)
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Critic tool for agents to request critical review from main loop."""
|
||||
"""Code review tool for agents to request critical review from main loop."""
|
||||
|
||||
from enum import Enum
|
||||
from typing import List, Optional, override
|
||||
@@ -21,10 +21,10 @@ class ReviewType(Enum):
|
||||
GENERAL = "general"
|
||||
|
||||
|
||||
class CriticTool(BaseTool):
|
||||
"""Tool for agents to request critical review from the main loop."""
|
||||
class CodeReviewTool(BaseTool):
|
||||
"""Tool for agents to request critical code review from the main loop."""
|
||||
|
||||
name = "critic"
|
||||
name = "code_review"
|
||||
|
||||
@property
|
||||
@override
|
||||
|
||||
@@ -37,35 +37,44 @@ try:
|
||||
|
||||
HANZO_AGENTS_AVAILABLE = True
|
||||
except ImportError:
|
||||
# Define minimal stubs if hanzo-agents is not available
|
||||
# hanzo-agents not available - provide stub classes that raise on use
|
||||
HANZO_AGENTS_AVAILABLE = False
|
||||
|
||||
class Agent:
|
||||
pass
|
||||
class _StubMeta(type):
|
||||
"""Metaclass that raises ImportError when stub class is instantiated."""
|
||||
|
||||
class State:
|
||||
pass
|
||||
def __call__(cls, *args, **kwargs):
|
||||
raise ImportError(
|
||||
f"{cls.__name__} requires hanzo-agents package. "
|
||||
"Install with: pip install hanzo-agents"
|
||||
)
|
||||
|
||||
class Network:
|
||||
pass
|
||||
class Agent(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class Tool:
|
||||
pass
|
||||
class State(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class History:
|
||||
pass
|
||||
class Network(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class ModelRegistry:
|
||||
pass
|
||||
class Tool(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class InferenceResult:
|
||||
pass
|
||||
class History(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class ToolCall:
|
||||
pass
|
||||
class ModelRegistry(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class Router:
|
||||
pass
|
||||
class InferenceResult(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class ToolCall(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class Router(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
|
||||
# Import optional components with fallbacks
|
||||
@@ -80,52 +89,53 @@ except ImportError:
|
||||
DeterministicRouter,
|
||||
)
|
||||
except ImportError:
|
||||
# Define stubs if not available
|
||||
class DeterministicRouter:
|
||||
pass
|
||||
# Stubs that raise on use
|
||||
class DeterministicRouter(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class LLMRouter:
|
||||
pass
|
||||
class LLMRouter(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class HybridRouter:
|
||||
pass
|
||||
class HybridRouter(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
|
||||
def _stub_factory(name: str):
|
||||
"""Create a stub function that raises ImportError."""
|
||||
|
||||
def stub(*args, **kwargs):
|
||||
raise ImportError(
|
||||
f"{name} requires hanzo-agents package. "
|
||||
"Install with: pip install hanzo-agents"
|
||||
)
|
||||
|
||||
stub.__name__ = name
|
||||
return stub
|
||||
|
||||
|
||||
try:
|
||||
from hanzo_agents import create_memory_kv, create_memory_vector
|
||||
except ImportError:
|
||||
try:
|
||||
# Try core module import
|
||||
from hanzo_agents.core.memory import create_memory_kv, create_memory_vector
|
||||
except ImportError:
|
||||
# Define stubs if not available
|
||||
def create_memory_kv(*args, **kwargs):
|
||||
pass
|
||||
|
||||
def create_memory_vector(*args, **kwargs):
|
||||
pass
|
||||
create_memory_kv = _stub_factory("create_memory_kv")
|
||||
create_memory_vector = _stub_factory("create_memory_vector")
|
||||
|
||||
|
||||
try:
|
||||
from hanzo_agents import sequential_router, conditional_router, state_based_router
|
||||
except ImportError:
|
||||
try:
|
||||
# Try core module import
|
||||
from hanzo_agents.core.router import (
|
||||
sequential_router,
|
||||
conditional_router,
|
||||
state_based_router,
|
||||
)
|
||||
except ImportError:
|
||||
# Define stubs if not available
|
||||
def sequential_router(*args, **kwargs):
|
||||
pass
|
||||
|
||||
def conditional_router(*args, **kwargs):
|
||||
pass
|
||||
|
||||
def state_based_router(*args, **kwargs):
|
||||
pass
|
||||
sequential_router = _stub_factory("sequential_router")
|
||||
conditional_router = _stub_factory("conditional_router")
|
||||
state_based_router = _stub_factory("state_based_router")
|
||||
|
||||
|
||||
try:
|
||||
@@ -136,18 +146,18 @@ try:
|
||||
OpenAICodexAgent,
|
||||
)
|
||||
except ImportError:
|
||||
# Define stub classes if not available
|
||||
class ClaudeCodeAgent(Agent):
|
||||
pass
|
||||
# Stubs that raise on use
|
||||
class ClaudeCodeAgent(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class OpenAICodexAgent(Agent):
|
||||
pass
|
||||
class OpenAICodexAgent(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class GeminiAgent(Agent):
|
||||
pass
|
||||
class GeminiAgent(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
class GrokAgent(Agent):
|
||||
pass
|
||||
class GrokAgent(metaclass=_StubMeta):
|
||||
"""Stub - requires hanzo-agents."""
|
||||
|
||||
|
||||
from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
|
||||
|
||||
@@ -1,550 +0,0 @@
|
||||
"""Swarm tool implementation for parallel and hierarchical agent execution.
|
||||
|
||||
This module implements the SwarmTool that enables both parallel execution of multiple
|
||||
agent instances and hierarchical workflows with specialized roles.
|
||||
"""
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
List,
|
||||
Unpack,
|
||||
Optional,
|
||||
TypedDict,
|
||||
final,
|
||||
override,
|
||||
)
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, PermissionManager, auto_timeout, create_tool_context
|
||||
|
||||
from .agent_tool import AgentTool
|
||||
|
||||
|
||||
class AgentNode(TypedDict):
|
||||
"""Node in the agent network.
|
||||
|
||||
Attributes:
|
||||
id: Unique identifier for this agent
|
||||
query: The specific query/task for this agent
|
||||
model: Optional model override (e.g., 'claude-3-5-sonnet', 'gpt-4o')
|
||||
role: Optional role description (e.g., 'architect', 'frontend', 'reviewer')
|
||||
connections: List of agent IDs this agent connects to (sends results to)
|
||||
receives_from: Optional list of agent IDs this agent receives input from
|
||||
file_path: Optional specific file for the agent to work on
|
||||
"""
|
||||
|
||||
id: str
|
||||
query: str
|
||||
model: Optional[str]
|
||||
role: Optional[str]
|
||||
connections: Optional[List[str]]
|
||||
receives_from: Optional[List[str]]
|
||||
file_path: Optional[str]
|
||||
|
||||
|
||||
class SwarmConfig(TypedDict):
|
||||
"""Configuration for an agent network.
|
||||
|
||||
Attributes:
|
||||
agents: Dictionary of agent configurations keyed by ID
|
||||
entry_point: ID of the first agent to execute (optional, defaults to finding roots)
|
||||
topology: Optional topology type (tree, dag, pipeline, star, mesh)
|
||||
"""
|
||||
|
||||
agents: Dict[str, AgentNode]
|
||||
entry_point: Optional[str]
|
||||
topology: Optional[str]
|
||||
|
||||
|
||||
class SwarmToolParams(TypedDict):
|
||||
"""Parameters for the SwarmTool.
|
||||
|
||||
Attributes:
|
||||
config: Agent network configuration
|
||||
query: Initial query to send to entry point agent(s)
|
||||
context: Optional context shared by all agents
|
||||
max_concurrent: Maximum number of concurrent agents (default: 10)
|
||||
"""
|
||||
|
||||
config: SwarmConfig
|
||||
query: str
|
||||
context: Optional[str]
|
||||
max_concurrent: Optional[int]
|
||||
|
||||
|
||||
@final
|
||||
class SwarmTool(BaseTool):
|
||||
"""Tool for executing multiple agent tasks in parallel.
|
||||
|
||||
The SwarmTool enables efficient parallel processing of multiple files or tasks
|
||||
by spawning independent agent instances for each task.
|
||||
"""
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "swarm"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Execute a network of AI agents with flexible connection topologies.
|
||||
|
||||
This tool enables sophisticated agent orchestration where agents can be connected
|
||||
in various network patterns. Each agent can pass results to connected agents,
|
||||
enabling complex workflows.
|
||||
|
||||
Features:
|
||||
- Flexible agent networks (tree, DAG, pipeline, star, mesh)
|
||||
- Each agent can use different models (Claude, GPT-4, Gemini, etc.)
|
||||
- Agents automatically pass results to connected agents
|
||||
- Parallel execution with dependency management
|
||||
- Full editing capabilities for each agent
|
||||
|
||||
Common Topologies:
|
||||
|
||||
1. Tree (Architect pattern):
|
||||
architect → [frontend, backend, database] → reviewer
|
||||
|
||||
2. Pipeline (Sequential processing):
|
||||
analyzer → planner → implementer → tester → reviewer
|
||||
|
||||
3. Star (Central coordinator):
|
||||
coordinator ← → [agent1, agent2, agent3, agent4]
|
||||
|
||||
4. DAG (Complex dependencies):
|
||||
Multiple agents with custom connections
|
||||
|
||||
Usage Example:
|
||||
|
||||
swarm(
|
||||
config={
|
||||
"agents": {
|
||||
"architect": {
|
||||
"id": "architect",
|
||||
"query": "Analyze codebase and create refactoring plan",
|
||||
"model": "claude-3-5-sonnet",
|
||||
"connections": ["frontend", "backend", "database"]
|
||||
},
|
||||
"frontend": {
|
||||
"id": "frontend",
|
||||
"query": "Refactor UI components based on architect's plan",
|
||||
"role": "Frontend Developer",
|
||||
"connections": ["reviewer"]
|
||||
},
|
||||
"backend": {
|
||||
"id": "backend",
|
||||
"query": "Refactor API endpoints based on architect's plan",
|
||||
"role": "Backend Developer",
|
||||
"connections": ["reviewer"]
|
||||
},
|
||||
"database": {
|
||||
"id": "database",
|
||||
"query": "Optimize database schema based on architect's plan",
|
||||
"role": "Database Expert",
|
||||
"connections": ["reviewer"]
|
||||
},
|
||||
"reviewer": {
|
||||
"id": "reviewer",
|
||||
"query": "Review all changes and ensure consistency",
|
||||
"model": "gpt-4o",
|
||||
"receives_from": ["frontend", "backend", "database"]
|
||||
}
|
||||
},
|
||||
"entry_point": "architect"
|
||||
},
|
||||
query="Refactor the authentication system for better security and performance"
|
||||
)
|
||||
|
||||
Models can be specified as:
|
||||
- Full: 'anthropic/claude-3-5-sonnet-20241022'
|
||||
- Short: 'claude-3-5-sonnet', 'gpt-4o', 'gemini-1.5-pro'
|
||||
- CLI tools: 'claude_cli', 'codex_cli', 'gemini_cli', 'grok_cli'
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
model: str | None = None,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
agent_max_iterations: int = 10,
|
||||
agent_max_tool_uses: int = 30,
|
||||
):
|
||||
"""Initialize the swarm tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
model: Optional model name override (defaults to Claude Sonnet)
|
||||
api_key: Optional API key for the model provider
|
||||
base_url: Optional base URL for the model provider
|
||||
max_tokens: Optional maximum tokens for model responses
|
||||
agent_max_iterations: Max iterations per agent (default: 10)
|
||||
agent_max_tool_uses: Max tool uses per agent (default: 30)
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
# Default to latest Claude Sonnet if no model specified
|
||||
from .code_auth import get_latest_claude_model
|
||||
|
||||
self.model = model or f"anthropic/{get_latest_claude_model()}"
|
||||
self.api_key = api_key or os.environ.get("ANTHROPIC_API_KEY") or os.environ.get("CLAUDE_API_KEY")
|
||||
self.base_url = base_url
|
||||
self.max_tokens = max_tokens
|
||||
self.agent_max_iterations = agent_max_iterations
|
||||
self.agent_max_tool_uses = agent_max_tool_uses
|
||||
|
||||
@override
|
||||
@auto_timeout("swarm_tool_v1_deprecated")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[SwarmToolParams],
|
||||
) -> str:
|
||||
"""Execute the swarm tool.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Combined results from all agents
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
agents = params.get("agents", [])
|
||||
manager_query = params.get("manager_query")
|
||||
reviewer_query = params.get("reviewer_query")
|
||||
common_context = params.get("common_context", "")
|
||||
max_concurrent = params.get("max_concurrent", 10)
|
||||
semaphore = asyncio.Semaphore(max_concurrent)
|
||||
|
||||
if not agents:
|
||||
await tool_ctx.error("No agents provided")
|
||||
return "Error: At least one agent must be provided."
|
||||
|
||||
# Extract parameters
|
||||
config = params.get("config", {})
|
||||
initial_query = params.get("query", "")
|
||||
context = params.get("context", "")
|
||||
|
||||
agents_config = config.get("agents", {})
|
||||
entry_point = config.get("entry_point")
|
||||
|
||||
await tool_ctx.info(f"Starting swarm execution with {len(agents_config)} agents")
|
||||
|
||||
# Build agent network
|
||||
agent_instances = {}
|
||||
agent_results = {}
|
||||
execution_queue = asyncio.Queue()
|
||||
completed_agents = set()
|
||||
|
||||
# Create agent instances
|
||||
for agent_id, agent_config in agents_config.items():
|
||||
model = agent_config.get("model", self.model)
|
||||
|
||||
# Support CLI tools
|
||||
cli_tools = {
|
||||
"claude_cli": self._get_cli_tool("claude_cli"),
|
||||
"codex_cli": self._get_cli_tool("codex_cli"),
|
||||
"gemini_cli": self._get_cli_tool("gemini_cli"),
|
||||
"grok_cli": self._get_cli_tool("grok_cli"),
|
||||
}
|
||||
|
||||
if model in cli_tools:
|
||||
agent = cli_tools[model]
|
||||
else:
|
||||
# Regular agent with model
|
||||
agent = AgentTool(
|
||||
permission_manager=self.permission_manager,
|
||||
model=self._normalize_model(model),
|
||||
api_key=self.api_key,
|
||||
base_url=self.base_url,
|
||||
max_tokens=self.max_tokens,
|
||||
max_iterations=self.agent_max_iterations,
|
||||
max_tool_uses=self.agent_max_tool_uses,
|
||||
)
|
||||
|
||||
agent_instances[agent_id] = agent
|
||||
|
||||
# Find entry points (agents with no incoming connections)
|
||||
if entry_point:
|
||||
await execution_queue.put((entry_point, initial_query, {}))
|
||||
else:
|
||||
# Find root agents (no receives_from)
|
||||
roots = []
|
||||
for agent_id, agent_config in agents_config.items():
|
||||
if not agent_config.get("receives_from"):
|
||||
# Check if any other agent connects to this one
|
||||
has_incoming = False
|
||||
for other_config in agents_config.values():
|
||||
if other_config.get("connections") and agent_id in other_config["connections"]:
|
||||
has_incoming = True
|
||||
break
|
||||
if not has_incoming:
|
||||
roots.append(agent_id)
|
||||
|
||||
if not roots:
|
||||
await tool_ctx.error("No entry point found in agent network")
|
||||
return "Error: Could not determine entry point for agent network"
|
||||
|
||||
for root in roots:
|
||||
await execution_queue.put((root, initial_query, {}))
|
||||
|
||||
# Execute agents in network order
|
||||
async def execute_agent(agent_id: str, query: str, inputs: Dict[str, str]) -> str:
|
||||
"""Execute a single agent in the network."""
|
||||
async with semaphore:
|
||||
try:
|
||||
agent_config = agents_config[agent_id]
|
||||
agent = agent_instances[agent_id]
|
||||
|
||||
await tool_ctx.info(f"Executing agent: {agent_id} ({agent_config.get('role', 'Agent')})")
|
||||
|
||||
# Build prompt with context and inputs
|
||||
prompt_parts = []
|
||||
|
||||
# Add role context
|
||||
if agent_config.get("role"):
|
||||
prompt_parts.append(f"Your role: {agent_config['role']}")
|
||||
|
||||
# Add shared context
|
||||
if context:
|
||||
prompt_parts.append(f"Context:\n{context}")
|
||||
|
||||
# Add inputs from connected agents
|
||||
if inputs:
|
||||
prompt_parts.append("Input from previous agents:")
|
||||
for input_agent, input_result in inputs.items():
|
||||
prompt_parts.append(f"\n--- From {input_agent} ---\n{input_result}")
|
||||
|
||||
# Add file context if specified
|
||||
if agent_config.get("file_path"):
|
||||
prompt_parts.append(f"\nFile to work on: {agent_config['file_path']}")
|
||||
|
||||
# Add the main query
|
||||
prompt_parts.append(f"\nTask: {agent_config['query']}")
|
||||
|
||||
# Combine query with initial query if this is entry point
|
||||
if query and query != agent_config["query"]:
|
||||
prompt_parts.append(f"\nMain objective: {query}")
|
||||
|
||||
full_prompt = "\n\n".join(prompt_parts)
|
||||
|
||||
# Execute the agent
|
||||
result = await agent.call(ctx, prompts=full_prompt)
|
||||
|
||||
await tool_ctx.info(f"Agent {agent_id} completed")
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Agent {agent_id} failed: {str(e)}"
|
||||
await tool_ctx.error(error_msg)
|
||||
return f"Error: {error_msg}"
|
||||
|
||||
# Process agent network
|
||||
running_tasks = set()
|
||||
|
||||
while not execution_queue.empty() or running_tasks:
|
||||
# Start new tasks up to concurrency limit
|
||||
while not execution_queue.empty() and len(running_tasks) < max_concurrent:
|
||||
agent_id, query, inputs = await execution_queue.get()
|
||||
|
||||
if agent_id not in completed_agents:
|
||||
# Check if all dependencies are met
|
||||
agent_config = agents_config[agent_id]
|
||||
receives_from = agent_config.get("receives_from", [])
|
||||
|
||||
# Collect inputs from dependencies
|
||||
ready = True
|
||||
for dep in receives_from:
|
||||
if dep not in agent_results:
|
||||
ready = False
|
||||
# Re-queue for later
|
||||
await execution_queue.put((agent_id, query, inputs))
|
||||
break
|
||||
else:
|
||||
inputs[dep] = agent_results[dep]
|
||||
|
||||
if ready:
|
||||
# Execute agent
|
||||
task = asyncio.create_task(execute_agent(agent_id, query, inputs))
|
||||
running_tasks.add(task)
|
||||
|
||||
async def handle_completion(task, agent_id=agent_id):
|
||||
result = await task
|
||||
agent_results[agent_id] = result
|
||||
completed_agents.add(agent_id)
|
||||
running_tasks.discard(task)
|
||||
|
||||
# Queue connected agents
|
||||
agent_config = agents_config[agent_id]
|
||||
connections = agent_config.get("connections", [])
|
||||
for next_agent in connections:
|
||||
if next_agent in agents_config:
|
||||
await execution_queue.put((next_agent, "", {agent_id: result}))
|
||||
|
||||
asyncio.create_task(handle_completion(task))
|
||||
|
||||
# Wait a bit if we're at capacity
|
||||
if running_tasks:
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Wait for all tasks to complete
|
||||
if running_tasks:
|
||||
await asyncio.gather(*running_tasks, return_exceptions=True)
|
||||
|
||||
# Format results
|
||||
return self._format_network_results(agents_config, agent_results, entry_point)
|
||||
|
||||
def _normalize_model(self, model: str) -> str:
|
||||
"""Normalize model names to full format."""
|
||||
model_map = {
|
||||
"claude-3-5-sonnet": "anthropic/claude-3-5-sonnet-20241022",
|
||||
"claude-3-opus": "anthropic/claude-3-opus-20240229",
|
||||
"gpt-4o": "openai/gpt-4o",
|
||||
"gpt-4": "openai/gpt-4",
|
||||
"gemini-1.5-pro": "google/gemini-1.5-pro",
|
||||
"gemini-1.5-flash": "google/gemini-1.5-flash",
|
||||
}
|
||||
return model_map.get(model, model)
|
||||
|
||||
def _get_cli_tool(self, tool_name: str):
|
||||
"""Get CLI tool instance."""
|
||||
# Import here to avoid circular imports
|
||||
if tool_name == "claude_cli":
|
||||
from .claude_cli_tool import ClaudeCLITool
|
||||
|
||||
return ClaudeCLITool(self.permission_manager)
|
||||
elif tool_name == "codex_cli":
|
||||
from .codex_cli_tool import CodexCLITool
|
||||
|
||||
return CodexCLITool(self.permission_manager)
|
||||
elif tool_name == "gemini_cli":
|
||||
from .gemini_cli_tool import GeminiCLITool
|
||||
|
||||
return GeminiCLITool(self.permission_manager)
|
||||
elif tool_name == "grok_cli":
|
||||
from .grok_cli_tool import GrokCLITool
|
||||
|
||||
return GrokCLITool(self.permission_manager)
|
||||
return None
|
||||
|
||||
def _format_network_results(
|
||||
self,
|
||||
agents_config: Dict[str, Any],
|
||||
results: Dict[str, str],
|
||||
entry_point: Optional[str],
|
||||
) -> str:
|
||||
"""Format results from agent network execution."""
|
||||
output = ["Agent Network Execution Results"]
|
||||
output.append("=" * 80)
|
||||
output.append(f"Total agents: {len(agents_config)}")
|
||||
output.append(f"Completed: {len(results)}")
|
||||
output.append(f"Failed: {len([r for r in results.values() if r.startswith('Error:')])}")
|
||||
|
||||
if entry_point:
|
||||
output.append(f"Entry point: {entry_point}")
|
||||
|
||||
output.append("\nExecution Flow:")
|
||||
output.append("-" * 40)
|
||||
|
||||
# Show results in execution order
|
||||
def format_agent_tree(agent_id: str, level: int = 0) -> List[str]:
|
||||
lines = []
|
||||
indent = " " * level
|
||||
|
||||
if agent_id in agents_config:
|
||||
config = agents_config[agent_id]
|
||||
role = config.get("role", "Agent")
|
||||
model = config.get("model", "default")
|
||||
|
||||
status = "✅" if agent_id in results and not results[agent_id].startswith("Error:") else "❌"
|
||||
lines.append(f"{indent}{status} {agent_id} ({role}) [{model}]")
|
||||
|
||||
# Show connections
|
||||
connections = config.get("connections", [])
|
||||
for conn in connections:
|
||||
if conn in agents_config:
|
||||
lines.extend(format_agent_tree(conn, level + 1))
|
||||
|
||||
return lines
|
||||
|
||||
# Start from entry point or roots
|
||||
if entry_point:
|
||||
output.extend(format_agent_tree(entry_point))
|
||||
else:
|
||||
# Find roots
|
||||
roots = []
|
||||
for agent_id in agents_config:
|
||||
has_incoming = False
|
||||
for config in agents_config.values():
|
||||
if config.get("connections") and agent_id in config["connections"]:
|
||||
has_incoming = True
|
||||
break
|
||||
if not has_incoming:
|
||||
roots.append(agent_id)
|
||||
|
||||
for root in roots:
|
||||
output.extend(format_agent_tree(root))
|
||||
|
||||
# Detailed results
|
||||
output.append("\n\nDetailed Results:")
|
||||
output.append("=" * 80)
|
||||
|
||||
for agent_id, result in results.items():
|
||||
config = agents_config.get(agent_id, {})
|
||||
role = config.get("role", "Agent")
|
||||
|
||||
output.append(f"\n### {agent_id} ({role})")
|
||||
output.append("-" * 40)
|
||||
|
||||
if result.startswith("Error:"):
|
||||
output.append(result)
|
||||
else:
|
||||
# Show first part of result
|
||||
lines = result.split("\n")
|
||||
preview_lines = lines[:10]
|
||||
output.extend(preview_lines)
|
||||
|
||||
if len(lines) > 10:
|
||||
output.append(f"... ({len(lines) - 10} more lines)")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this swarm tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def swarm(
|
||||
ctx: MCPContext,
|
||||
config: dict[str, Any],
|
||||
query: str,
|
||||
context: Optional[str] = None,
|
||||
max_concurrent: int = 10,
|
||||
) -> str:
|
||||
# Convert to typed format
|
||||
typed_config = SwarmConfig(
|
||||
agents=config.get("agents", {}),
|
||||
entry_point=config.get("entry_point"),
|
||||
topology=config.get("topology"),
|
||||
)
|
||||
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
config=typed_config,
|
||||
query=query,
|
||||
context=context,
|
||||
max_concurrent=max_concurrent,
|
||||
)
|
||||
@@ -16,21 +16,16 @@ logger = logging.getLogger(__name__)
|
||||
# Lazy imports for heavy dependencies
|
||||
_tools = []
|
||||
|
||||
try:
|
||||
from .llm_tool import LLMTool
|
||||
|
||||
_tools.append(LLMTool)
|
||||
except ImportError as e:
|
||||
logger.debug(f"LLMTool not available: {e}")
|
||||
LLMTool = None
|
||||
|
||||
try:
|
||||
from .llm_unified import UnifiedLLMTool
|
||||
|
||||
# Alias for backwards compatibility
|
||||
LLMTool = UnifiedLLMTool
|
||||
_tools.append(UnifiedLLMTool)
|
||||
except ImportError as e:
|
||||
logger.debug(f"UnifiedLLMTool not available: {e}")
|
||||
UnifiedLLMTool = None
|
||||
LLMTool = None
|
||||
|
||||
try:
|
||||
from .consensus_tool import ConsensusTool
|
||||
|
||||
@@ -17,7 +17,7 @@ from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, ToolContext, auto_timeout, create_tool_context
|
||||
|
||||
from .llm_tool import LLMTool
|
||||
from .llm_unified import LLMTool
|
||||
|
||||
Prompt = Annotated[
|
||||
str,
|
||||
|
||||
@@ -9,7 +9,7 @@ from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, ToolContext, auto_timeout, create_tool_context
|
||||
|
||||
from .llm_tool import LLMTool
|
||||
from .llm_unified import LLMTool
|
||||
|
||||
Action = Annotated[
|
||||
str,
|
||||
@@ -275,7 +275,7 @@ Providers are automatically detected based on environment variables:
|
||||
def _list_all_models(self, provider: Optional[str] = None) -> str:
|
||||
"""List all available models from LiteLLM."""
|
||||
try:
|
||||
from .llm_tool import LLMTool
|
||||
from .llm_unified import LLMTool
|
||||
|
||||
all_models = LLMTool.get_all_models()
|
||||
|
||||
|
||||
@@ -1,883 +0,0 @@
|
||||
"""Unified LLM tool with multiple actions including consensus mode."""
|
||||
|
||||
import os
|
||||
import json
|
||||
import asyncio
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
List,
|
||||
Unpack,
|
||||
Optional,
|
||||
Annotated,
|
||||
TypedDict,
|
||||
final,
|
||||
override,
|
||||
)
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, ToolContext, auto_timeout, create_tool_context
|
||||
|
||||
# Check if litellm is available
|
||||
try:
|
||||
import litellm
|
||||
|
||||
LITELLM_AVAILABLE = True
|
||||
except ImportError:
|
||||
LITELLM_AVAILABLE = False
|
||||
|
||||
|
||||
# Parameter types
|
||||
Action = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Action to perform: query, consensus, list, models, enable, disable, test",
|
||||
default="query",
|
||||
),
|
||||
]
|
||||
|
||||
Model = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Model name (e.g., gpt-4, claude-3-opus-20240229)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Models = Annotated[
|
||||
Optional[List[str]],
|
||||
Field(
|
||||
description="List of models for consensus mode",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Prompt = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="The prompt to send to the LLM",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
SystemPrompt = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="System prompt to set context",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Temperature = Annotated[
|
||||
float,
|
||||
Field(
|
||||
description="Temperature for response randomness (0-2)",
|
||||
default=0.7,
|
||||
),
|
||||
]
|
||||
|
||||
MaxTokens = Annotated[
|
||||
Optional[int],
|
||||
Field(
|
||||
description="Maximum tokens in response",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
JsonMode = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Request JSON formatted response",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
Stream = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Stream the response",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
Provider = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Provider name for list/enable/disable actions",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
IncludeRaw = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Include raw responses in consensus mode",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
JudgeModel = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Model to use as judge/aggregator in consensus",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
DevilsAdvocate = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Enable devil's advocate mode (10th model critiques others)",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
ConsensusSize = Annotated[
|
||||
Optional[int],
|
||||
Field(
|
||||
description="Number of models to use in consensus (default: 3)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class LLMParams(TypedDict, total=False):
|
||||
"""Parameters for LLM tool."""
|
||||
|
||||
action: str
|
||||
model: Optional[str]
|
||||
models: Optional[List[str]]
|
||||
prompt: Optional[str]
|
||||
system_prompt: Optional[str]
|
||||
temperature: float
|
||||
max_tokens: Optional[int]
|
||||
json_mode: bool
|
||||
stream: bool
|
||||
provider: Optional[str]
|
||||
include_raw: bool
|
||||
judge_model: Optional[str]
|
||||
devils_advocate: bool
|
||||
consensus_size: Optional[int]
|
||||
|
||||
|
||||
@final
|
||||
class LLMTool(BaseTool):
|
||||
"""Unified LLM tool with multiple actions."""
|
||||
|
||||
# Config file for settings
|
||||
CONFIG_FILE = Path.home() / ".hanzo" / "mcp" / "llm_config.json"
|
||||
|
||||
# Default consensus models in order of preference
|
||||
DEFAULT_CONSENSUS_MODELS = [
|
||||
"gpt-4o", # OpenAI's latest
|
||||
"claude-3-opus-20240229", # Claude's most capable
|
||||
"gemini/gemini-1.5-pro", # Google's best
|
||||
"groq/llama3-70b-8192", # Fast Groq
|
||||
"mistral/mistral-large-latest", # Mistral's best
|
||||
"perplexity/llama-3.1-sonar-large-128k-chat", # Perplexity with search
|
||||
]
|
||||
|
||||
# API key environment variables
|
||||
API_KEY_ENV_VARS = {
|
||||
"openai": ["OPENAI_API_KEY"],
|
||||
"anthropic": ["ANTHROPIC_API_KEY", "CLAUDE_API_KEY"],
|
||||
"google": ["GOOGLE_API_KEY", "GEMINI_API_KEY"],
|
||||
"groq": ["GROQ_API_KEY"],
|
||||
"mistral": ["MISTRAL_API_KEY"],
|
||||
"perplexity": ["PERPLEXITY_API_KEY", "PERPLEXITYAI_API_KEY"],
|
||||
"together": ["TOGETHER_API_KEY", "TOGETHERAI_API_KEY"],
|
||||
"cohere": ["COHERE_API_KEY"],
|
||||
"replicate": ["REPLICATE_API_KEY"],
|
||||
"huggingface": ["HUGGINGFACE_API_KEY", "HF_TOKEN"],
|
||||
"bedrock": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"],
|
||||
"vertex": ["GOOGLE_APPLICATION_CREDENTIALS"],
|
||||
"azure": ["AZURE_API_KEY"],
|
||||
"voyage": ["VOYAGE_API_KEY"],
|
||||
"deepseek": ["DEEPSEEK_API_KEY"],
|
||||
}
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the unified LLM tool."""
|
||||
self.available_providers = self._detect_available_providers()
|
||||
self.config = self._load_config()
|
||||
|
||||
def _detect_available_providers(self) -> Dict[str, List[str]]:
|
||||
"""Detect which providers have API keys configured."""
|
||||
available = {}
|
||||
|
||||
for provider, env_vars in self.API_KEY_ENV_VARS.items():
|
||||
for var in env_vars:
|
||||
if os.getenv(var):
|
||||
available[provider] = env_vars
|
||||
break
|
||||
|
||||
return available
|
||||
|
||||
def _load_config(self) -> Dict[str, Any]:
|
||||
"""Load configuration from file."""
|
||||
if self.CONFIG_FILE.exists():
|
||||
try:
|
||||
with open(self.CONFIG_FILE, "r") as f:
|
||||
return json.load(f)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Default config
|
||||
return {
|
||||
"disabled_providers": [],
|
||||
"consensus_models": None, # Use defaults if None
|
||||
"default_judge_model": "gpt-4o",
|
||||
"consensus_size": 3,
|
||||
}
|
||||
|
||||
def _save_config(self):
|
||||
"""Save configuration to file."""
|
||||
self.CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(self.CONFIG_FILE, "w") as f:
|
||||
json.dump(self.config, f, indent=2)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "llm"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
available = list(self.available_providers.keys())
|
||||
|
||||
return f"""Query LLMs. Default: single query. Actions: consensus, list, models, test.
|
||||
|
||||
Usage:
|
||||
llm "What is the capital of France?"
|
||||
llm "Explain this code" --model gpt-4o
|
||||
llm --action consensus "Is this approach correct?" --devils-advocate
|
||||
llm --action models --provider openai
|
||||
|
||||
Available: {", ".join(available) if available else "None"}"""
|
||||
|
||||
@override
|
||||
@auto_timeout("llm")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[LLMParams],
|
||||
) -> str:
|
||||
"""Execute LLM action."""
|
||||
# Create tool context only if we have a proper MCP context
|
||||
tool_ctx = None
|
||||
try:
|
||||
if hasattr(ctx, "client") and ctx.client and hasattr(ctx.client, "server"):
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
if tool_ctx:
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
except Exception:
|
||||
# Running in test mode without MCP context
|
||||
pass
|
||||
|
||||
# Fast test path: allow offline deterministic responses
|
||||
if os.getenv("HANZO_MCP_FAST_TESTS") == "1":
|
||||
action = params.get("action", "query")
|
||||
if action == "query":
|
||||
prompt = params.get("prompt", "") or ""
|
||||
# Simple keyword routing for tests
|
||||
if "generate tasks" in prompt.lower():
|
||||
return '{"project": "Demo", "tasks": [{"title": "Init repo", "slug": "init-repo"}, {"title": "Add CI", "slug": "add-ci"}]}'
|
||||
return "Architecture Proposal: minimal service layers and tests"
|
||||
elif action == "consensus":
|
||||
return "Consensus: aligned"
|
||||
# Other actions can fall through to regular path if available
|
||||
|
||||
if not LITELLM_AVAILABLE:
|
||||
return "Error: LiteLLM is not installed. Install it with: pip install litellm"
|
||||
|
||||
# Extract action
|
||||
action = params.get("action", "query")
|
||||
|
||||
# Route to appropriate handler
|
||||
if action == "query":
|
||||
return await self._handle_query(tool_ctx, params)
|
||||
elif action == "consensus":
|
||||
return await self._handle_consensus(tool_ctx, params)
|
||||
elif action == "list":
|
||||
return self._handle_list()
|
||||
elif action == "models":
|
||||
return self._handle_models(params.get("provider"))
|
||||
elif action == "enable":
|
||||
return self._handle_enable(params.get("provider"))
|
||||
elif action == "disable":
|
||||
return self._handle_disable(params.get("provider"))
|
||||
elif action == "test":
|
||||
return await self._handle_test(tool_ctx, params.get("model"), params.get("provider"))
|
||||
else:
|
||||
return f"Error: Unknown action '{action}'. Valid actions: query, consensus, list, models, enable, disable, test"
|
||||
|
||||
async def _handle_query(self, tool_ctx, params: Dict[str, Any]) -> str:
|
||||
"""Handle single model query."""
|
||||
model = params.get("model")
|
||||
prompt = params.get("prompt")
|
||||
|
||||
if not prompt:
|
||||
return "Error: prompt is required for query action"
|
||||
|
||||
# Auto-select model if not specified
|
||||
if not model:
|
||||
if self.available_providers:
|
||||
# Use first available model
|
||||
if "openai" in self.available_providers:
|
||||
model = "gpt-4o-mini"
|
||||
elif "anthropic" in self.available_providers:
|
||||
model = "claude-3-haiku-20240307"
|
||||
elif "google" in self.available_providers:
|
||||
model = "gemini/gemini-1.5-flash"
|
||||
else:
|
||||
# Use first provider's default
|
||||
provider = list(self.available_providers.keys())[0]
|
||||
model = f"{provider}/default"
|
||||
else:
|
||||
return "Error: No model specified and no API keys found"
|
||||
|
||||
# Check if we have API key for this model
|
||||
provider = self._get_provider_for_model(model)
|
||||
if provider and provider not in self.available_providers:
|
||||
env_vars = self.API_KEY_ENV_VARS.get(provider, [])
|
||||
return f"Error: No API key found for {provider}. Set one of: {', '.join(env_vars)}"
|
||||
|
||||
# Build messages
|
||||
messages = []
|
||||
if params.get("system_prompt"):
|
||||
messages.append({"role": "system", "content": params["system_prompt"]})
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
|
||||
# Build kwargs
|
||||
kwargs = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"temperature": params.get("temperature", 0.7),
|
||||
}
|
||||
|
||||
if params.get("max_tokens"):
|
||||
kwargs["max_tokens"] = params["max_tokens"]
|
||||
|
||||
if params.get("json_mode"):
|
||||
kwargs["response_format"] = {"type": "json_object"}
|
||||
|
||||
if params.get("stream"):
|
||||
kwargs["stream"] = True
|
||||
|
||||
try:
|
||||
if tool_ctx:
|
||||
await tool_ctx.info(f"Querying {model}...")
|
||||
|
||||
if kwargs.get("stream"):
|
||||
# Handle streaming response
|
||||
response_text = ""
|
||||
async for chunk in await litellm.acompletion(**kwargs):
|
||||
if chunk.choices[0].delta.content:
|
||||
response_text += chunk.choices[0].delta.content
|
||||
return response_text
|
||||
else:
|
||||
# Regular response
|
||||
response = await litellm.acompletion(**kwargs)
|
||||
return response.choices[0].message.content
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
if "model_not_found" in error_msg or "does not exist" in error_msg:
|
||||
return f"Error: Model '{model}' not found. Use 'llm --action models' to see available models."
|
||||
else:
|
||||
return f"Error calling LLM: {error_msg}"
|
||||
|
||||
async def _handle_consensus(self, tool_ctx, params: Dict[str, Any]) -> str:
|
||||
"""Handle consensus mode with multiple models."""
|
||||
prompt = params.get("prompt")
|
||||
if not prompt:
|
||||
return "Error: prompt is required for consensus action"
|
||||
|
||||
# Determine models to use
|
||||
models = params.get("models")
|
||||
if not models:
|
||||
# Use configured or default models
|
||||
consensus_size = params.get("consensus_size") or self.config.get("consensus_size", 3)
|
||||
models = self._get_consensus_models(consensus_size)
|
||||
|
||||
if not models:
|
||||
return "Error: No models available for consensus. Set API keys for at least 2 providers."
|
||||
|
||||
if len(models) < 2:
|
||||
return "Error: Consensus requires at least 2 models"
|
||||
|
||||
# Check for devil's advocate mode
|
||||
devils_advocate = params.get("devils_advocate", False)
|
||||
if devils_advocate and len(models) < 3:
|
||||
return "Error: Devil's advocate mode requires at least 3 models"
|
||||
|
||||
if tool_ctx:
|
||||
await tool_ctx.info(f"Running consensus with {len(models)} models...")
|
||||
|
||||
# Query models in parallel
|
||||
system_prompt = params.get("system_prompt")
|
||||
temperature = params.get("temperature", 0.7)
|
||||
max_tokens = params.get("max_tokens")
|
||||
|
||||
# Split models if using devil's advocate
|
||||
if devils_advocate:
|
||||
consensus_models = models[:-1]
|
||||
devil_model = models[-1]
|
||||
else:
|
||||
consensus_models = models
|
||||
devil_model = None
|
||||
|
||||
# Query consensus models
|
||||
responses = await self._query_models_parallel(
|
||||
consensus_models, prompt, system_prompt, temperature, max_tokens, tool_ctx
|
||||
)
|
||||
|
||||
# Get devil's advocate response if enabled
|
||||
devil_response = None
|
||||
if devil_model:
|
||||
# Create devil's advocate prompt
|
||||
responses_text = "\n\n".join(
|
||||
[f"Model {i + 1}: {resp['response']}" for i, resp in enumerate(responses) if resp["response"]]
|
||||
)
|
||||
|
||||
devil_prompt = f"""You are a critical analyst. Review these responses to the question below and provide a devil's advocate perspective. Challenge assumptions, point out weaknesses, and suggest alternative viewpoints.
|
||||
|
||||
Original Question: {prompt}
|
||||
|
||||
Responses from other models:
|
||||
{responses_text}
|
||||
|
||||
Provide your critical analysis:"""
|
||||
|
||||
devil_result = await self._query_single_model(
|
||||
devil_model, devil_prompt, system_prompt, temperature, max_tokens
|
||||
)
|
||||
|
||||
if devil_result["success"]:
|
||||
devil_response = {
|
||||
"model": devil_model,
|
||||
"response": devil_result["response"],
|
||||
"time_ms": devil_result["time_ms"],
|
||||
}
|
||||
|
||||
# Aggregate responses
|
||||
judge_model = params.get("judge_model") or self.config.get("default_judge_model", "gpt-4o")
|
||||
include_raw = params.get("include_raw", False)
|
||||
|
||||
return await self._aggregate_consensus(responses, prompt, judge_model, include_raw, devil_response, tool_ctx)
|
||||
|
||||
def _handle_list(self) -> str:
|
||||
"""List available providers."""
|
||||
output = ["=== LLM Providers ==="]
|
||||
|
||||
# Get all possible providers
|
||||
all_providers = sorted(self.API_KEY_ENV_VARS.keys())
|
||||
disabled = self.config.get("disabled_providers", [])
|
||||
|
||||
output.append(f"Total providers: {len(all_providers)}")
|
||||
output.append(f"Available: {len(self.available_providers)}")
|
||||
output.append(f"Disabled: {len(disabled)}\n")
|
||||
|
||||
for provider in all_providers:
|
||||
status_parts = []
|
||||
|
||||
# Check if API key exists
|
||||
if provider in self.available_providers:
|
||||
status_parts.append("✅ API key found")
|
||||
else:
|
||||
status_parts.append("❌ No API key")
|
||||
|
||||
# Check if disabled
|
||||
if provider in disabled:
|
||||
status_parts.append("🚫 Disabled")
|
||||
|
||||
# Show environment variables
|
||||
env_vars = self.API_KEY_ENV_VARS.get(provider, [])
|
||||
status = " | ".join(status_parts)
|
||||
|
||||
output.append(f"{provider}: {status}")
|
||||
output.append(f" Environment variables: {', '.join(env_vars)}")
|
||||
|
||||
output.append("\nUse 'llm --action enable/disable --provider <name>' to manage providers")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _handle_models(self, provider: Optional[str] = None) -> str:
|
||||
"""List available models."""
|
||||
try:
|
||||
all_models = self._get_all_models()
|
||||
|
||||
if not all_models:
|
||||
return "No models available or LiteLLM not properly initialized"
|
||||
|
||||
output = ["=== Available LLM Models ==="]
|
||||
|
||||
if provider:
|
||||
# Show models for specific provider
|
||||
provider_lower = provider.lower()
|
||||
models = all_models.get(provider_lower, [])
|
||||
|
||||
if not models:
|
||||
return f"No models found for provider '{provider}'"
|
||||
|
||||
output.append(f"\n{provider.upper()} ({len(models)} models):")
|
||||
output.append("-" * 40)
|
||||
|
||||
# Show first 50 models
|
||||
for model in models[:50]:
|
||||
output.append(f" {model}")
|
||||
|
||||
if len(models) > 50:
|
||||
output.append(f" ... and {len(models) - 50} more")
|
||||
else:
|
||||
# Show summary of all providers
|
||||
total_models = sum(len(models) for models in all_models.values())
|
||||
output.append(f"Total models available: {total_models}")
|
||||
output.append("")
|
||||
|
||||
# Show providers with counts
|
||||
for provider_name, models in sorted(all_models.items()):
|
||||
if models:
|
||||
available = "✅" if provider_name in self.available_providers else "❌"
|
||||
output.append(f"{available} {provider_name}: {len(models)} models")
|
||||
|
||||
output.append("\nUse 'llm --action models --provider <name>' to see specific models")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
return f"Error listing models: {str(e)}"
|
||||
|
||||
def _handle_enable(self, provider: Optional[str]) -> str:
|
||||
"""Enable a provider."""
|
||||
if not provider:
|
||||
return "Error: provider is required for enable action"
|
||||
|
||||
provider = provider.lower()
|
||||
disabled = self.config.get("disabled_providers", [])
|
||||
|
||||
if provider in disabled:
|
||||
disabled.remove(provider)
|
||||
self.config["disabled_providers"] = disabled
|
||||
self._save_config()
|
||||
return f"Successfully enabled {provider}"
|
||||
else:
|
||||
return f"{provider} is already enabled"
|
||||
|
||||
def _handle_disable(self, provider: Optional[str]) -> str:
|
||||
"""Disable a provider."""
|
||||
if not provider:
|
||||
return "Error: provider is required for disable action"
|
||||
|
||||
provider = provider.lower()
|
||||
disabled = self.config.get("disabled_providers", [])
|
||||
|
||||
if provider not in disabled:
|
||||
disabled.append(provider)
|
||||
self.config["disabled_providers"] = disabled
|
||||
self._save_config()
|
||||
return f"Successfully disabled {provider}"
|
||||
else:
|
||||
return f"{provider} is already disabled"
|
||||
|
||||
async def _handle_test(self, tool_ctx, model: Optional[str], provider: Optional[str]) -> str:
|
||||
"""Test a model or provider."""
|
||||
if not model and not provider:
|
||||
return "Error: Either model or provider is required for test action"
|
||||
|
||||
# If provider specified, test its default model
|
||||
if provider and not model:
|
||||
provider = provider.lower()
|
||||
if provider == "openai":
|
||||
model = "gpt-3.5-turbo"
|
||||
elif provider == "anthropic":
|
||||
model = "claude-3-haiku-20240307"
|
||||
elif provider == "google":
|
||||
model = "gemini/gemini-1.5-flash"
|
||||
elif provider == "groq":
|
||||
model = "groq/llama3-8b-8192"
|
||||
else:
|
||||
model = f"{provider}/default"
|
||||
|
||||
# Test the model
|
||||
test_prompt = "Say 'Hello from Hanzo AI!' in exactly 5 words."
|
||||
|
||||
try:
|
||||
if tool_ctx:
|
||||
await tool_ctx.info(f"Testing {model}...")
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": test_prompt}],
|
||||
temperature=0,
|
||||
max_tokens=20,
|
||||
)
|
||||
|
||||
result = response.choices[0].message.content
|
||||
return f"✅ {model} is working!\nResponse: {result}"
|
||||
|
||||
except Exception as e:
|
||||
return f"❌ {model} failed: {str(e)}"
|
||||
|
||||
def _get_consensus_models(self, size: int) -> List[str]:
|
||||
"""Get models for consensus based on availability."""
|
||||
# Use configured models if set
|
||||
configured = self.config.get("consensus_models")
|
||||
if configured:
|
||||
return configured[:size]
|
||||
|
||||
# Otherwise, build list from available providers
|
||||
models = []
|
||||
disabled = self.config.get("disabled_providers", [])
|
||||
|
||||
# Try default models first
|
||||
for model in self.DEFAULT_CONSENSUS_MODELS:
|
||||
if len(models) >= size:
|
||||
break
|
||||
|
||||
provider = self._get_provider_for_model(model)
|
||||
if provider and provider in self.available_providers and provider not in disabled:
|
||||
models.append(model)
|
||||
|
||||
# If still need more, add from available providers
|
||||
if len(models) < size:
|
||||
for provider in self.available_providers:
|
||||
if provider in disabled:
|
||||
continue
|
||||
|
||||
if provider == "openai" and "gpt-4o" not in models:
|
||||
models.append("gpt-4o")
|
||||
elif provider == "anthropic" and "claude-3-opus-20240229" not in models:
|
||||
models.append("claude-3-opus-20240229")
|
||||
elif provider == "google" and "gemini/gemini-1.5-pro" not in models:
|
||||
models.append("gemini/gemini-1.5-pro")
|
||||
|
||||
if len(models) >= size:
|
||||
break
|
||||
|
||||
return models
|
||||
|
||||
async def _query_models_parallel(
|
||||
self,
|
||||
models: List[str],
|
||||
prompt: str,
|
||||
system_prompt: Optional[str],
|
||||
temperature: float,
|
||||
max_tokens: Optional[int],
|
||||
tool_ctx,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Query multiple models in parallel."""
|
||||
|
||||
async def query_with_info(model: str) -> Dict[str, Any]:
|
||||
result = await self._query_single_model(model, prompt, system_prompt, temperature, max_tokens)
|
||||
return {
|
||||
"model": model,
|
||||
"response": result.get("response"),
|
||||
"success": result.get("success", False),
|
||||
"error": result.get("error"),
|
||||
"time_ms": result.get("time_ms", 0),
|
||||
}
|
||||
|
||||
# Run all queries in parallel
|
||||
tasks = [query_with_info(model) for model in models]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
# Report results
|
||||
successful = sum(1 for r in results if r["success"])
|
||||
if tool_ctx:
|
||||
await tool_ctx.info(f"Completed {successful}/{len(models)} model queries")
|
||||
|
||||
return results
|
||||
|
||||
async def _query_single_model(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
system_prompt: Optional[str],
|
||||
temperature: float,
|
||||
max_tokens: Optional[int],
|
||||
) -> Dict[str, Any]:
|
||||
"""Query a single model and return result with metadata."""
|
||||
import time
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
messages = []
|
||||
if system_prompt:
|
||||
messages.append({"role": "system", "content": system_prompt})
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
|
||||
kwargs = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
}
|
||||
if max_tokens:
|
||||
kwargs["max_tokens"] = max_tokens
|
||||
|
||||
response = await litellm.acompletion(**kwargs)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"response": response.choices[0].message.content,
|
||||
"time_ms": int((time.time() - start_time) * 1000),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"success": False,
|
||||
"error": str(e),
|
||||
"time_ms": int((time.time() - start_time) * 1000),
|
||||
}
|
||||
|
||||
async def _aggregate_consensus(
|
||||
self,
|
||||
responses: List[Dict[str, Any]],
|
||||
original_prompt: str,
|
||||
judge_model: str,
|
||||
include_raw: bool,
|
||||
devil_response: Optional[Dict[str, Any]],
|
||||
tool_ctx,
|
||||
) -> str:
|
||||
"""Aggregate consensus responses using a judge model."""
|
||||
# Prepare response data
|
||||
successful_responses = [r for r in responses if r["success"]]
|
||||
|
||||
if not successful_responses:
|
||||
return "Error: All models failed to respond"
|
||||
|
||||
# Format responses for aggregation
|
||||
responses_text = "\n\n".join([f"Model: {r['model']}\nResponse: {r['response']}" for r in successful_responses])
|
||||
|
||||
if devil_response:
|
||||
responses_text += f"\n\nDevil's Advocate ({devil_response['model']}):\n{devil_response['response']}"
|
||||
|
||||
# Create aggregation prompt
|
||||
aggregation_prompt = f"""Analyze the following responses from multiple AI models to this question:
|
||||
|
||||
<original_question>
|
||||
{original_prompt}
|
||||
</original_question>
|
||||
|
||||
<model_responses>
|
||||
{responses_text}
|
||||
</model_responses>
|
||||
|
||||
Please provide:
|
||||
1. A synthesis of the key points where models agree
|
||||
2. Notable differences or disagreements between responses
|
||||
3. A balanced conclusion incorporating the best insights
|
||||
{f"4. Evaluation of the devil's advocate critique" if devil_response else ""}
|
||||
|
||||
Be concise and highlight the most important findings."""
|
||||
|
||||
# Get aggregation
|
||||
try:
|
||||
if tool_ctx:
|
||||
await tool_ctx.info(f"Aggregating responses with {judge_model}...")
|
||||
|
||||
judge_result = await self._query_single_model(judge_model, aggregation_prompt, None, 0.3, None)
|
||||
|
||||
if not judge_result["success"]:
|
||||
return f"Error: Judge model failed: {judge_result.get('error', 'Unknown error')}"
|
||||
|
||||
# Format output
|
||||
output = [f"=== Consensus Analysis ({len(successful_responses)} models) ===\n"]
|
||||
output.append(judge_result["response"])
|
||||
|
||||
# Add model list
|
||||
output.append(f"\nModels consulted: {', '.join([r['model'] for r in successful_responses])}")
|
||||
if devil_response:
|
||||
output.append(f"Devil's Advocate: {devil_response['model']}")
|
||||
|
||||
# Add timing info
|
||||
avg_time = sum(r["time_ms"] for r in responses) / len(responses)
|
||||
output.append(f"\nAverage response time: {avg_time:.0f}ms")
|
||||
|
||||
# Include raw responses if requested
|
||||
if include_raw:
|
||||
output.append("\n\n=== Raw Responses ===")
|
||||
for r in successful_responses:
|
||||
output.append(f"\n{r['model']}:")
|
||||
output.append("-" * 40)
|
||||
output.append(r["response"])
|
||||
|
||||
if devil_response:
|
||||
output.append(f"\nDevil's Advocate ({devil_response['model']}):")
|
||||
output.append("-" * 40)
|
||||
output.append(devil_response["response"])
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
return f"Error during aggregation: {str(e)}"
|
||||
|
||||
def _get_provider_for_model(self, model: str) -> Optional[str]:
|
||||
"""Determine the provider for a given model."""
|
||||
model_lower = model.lower()
|
||||
|
||||
# Check explicit provider prefix
|
||||
if "/" in model:
|
||||
return model.split("/")[0]
|
||||
|
||||
# Check model prefixes
|
||||
if model_lower.startswith("gpt"):
|
||||
return "openai"
|
||||
elif model_lower.startswith("claude"):
|
||||
return "anthropic"
|
||||
elif model_lower.startswith("gemini"):
|
||||
return "google"
|
||||
elif model_lower.startswith("command"):
|
||||
return "cohere"
|
||||
|
||||
# Default to OpenAI
|
||||
return "openai"
|
||||
|
||||
def _get_all_models(self) -> Dict[str, List[str]]:
|
||||
"""Get all available models from LiteLLM."""
|
||||
try:
|
||||
import litellm
|
||||
|
||||
# Get all models
|
||||
all_models = litellm.model_list
|
||||
|
||||
# Organize by provider
|
||||
providers = {}
|
||||
|
||||
for model in all_models:
|
||||
# Extract provider
|
||||
if "/" in model:
|
||||
provider = model.split("/")[0]
|
||||
elif model.startswith("gpt"):
|
||||
provider = "openai"
|
||||
elif model.startswith("claude"):
|
||||
provider = "anthropic"
|
||||
elif model.startswith("gemini"):
|
||||
provider = "google"
|
||||
elif model.startswith("command"):
|
||||
provider = "cohere"
|
||||
else:
|
||||
provider = "other"
|
||||
|
||||
if provider not in providers:
|
||||
providers[provider] = []
|
||||
providers[provider].append(model)
|
||||
|
||||
# Sort models within each provider
|
||||
for provider in providers:
|
||||
providers[provider] = sorted(providers[provider])
|
||||
|
||||
return providers
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -868,3 +868,7 @@ Be concise and highlight the most important findings."""
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
|
||||
|
||||
# Alias for backwards compatibility
|
||||
LLMTool = UnifiedLLMTool
|
||||
|
||||
@@ -7,7 +7,7 @@ from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, auto_timeout
|
||||
|
||||
from .llm_tool import LLMTool
|
||||
from .llm_unified import LLMTool
|
||||
|
||||
Prompt = Annotated[
|
||||
str,
|
||||
|
||||
Reference in New Issue
Block a user