mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
Route llm through api.hanzo.ai with Enso, and add kai_decide
The llm tool called upstream providers through the litellm fork, read provider keys, and never reached the MCP surface: its register() was a no-op and it skipped BaseTool's constructor. Kai had no Python tool. hanzo-tools-llm now carries two HIP-0300 tools over HanzoCloud: - llm (default action query): POST /v1/chat/completions with model enso-auto unless named, "auto" for Enso's pick, max_cost and max_latency_ms as X-Max-Cost and X-Max-Latency-Ms, returning id and the model that served; models: GET /v1/models, one row per model with the catalog's per-million prices, narrowed by family; feedback: POST /v1/ai/feedback with request_id, signal and a 1-3 rating. - kai_decide (default action decide): POST /v1/decisions with the TypeScript tool's schema and shape checks. The consensus, llm_manage and provider modules wrapped the upstream client and were never registered, so they go, with hanzo-consensus and the [full] extra. hanzo-tools-llm 0.2.1. Co-authored-by: Hanzo Dev <[email protected]>
This commit is contained in:
zooqueen
@@ -149,14 +149,10 @@ class TestToolPackages:
|
||||
reason="hanzo-tools-llm not installed",
|
||||
)
|
||||
def test_llm_tools(self):
|
||||
"""Test hanzo-tools-llm imports (tools depend on llm)."""
|
||||
from hanzo_tools.llm import TOOLS, LLM_AVAILABLE
|
||||
"""Test hanzo-tools-llm carries llm and kai_decide, both over HanzoCloud."""
|
||||
from hanzo_tools.llm import TOOLS
|
||||
|
||||
# LLM tools are optional, depend on llm
|
||||
if LLM_AVAILABLE:
|
||||
assert len(TOOLS) >= 1
|
||||
else:
|
||||
assert len(TOOLS) == 0
|
||||
assert [t.name for t in TOOLS] == ["llm", "kai_decide"]
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _module_installed("hanzo_tools.vector"),
|
||||
@@ -194,7 +190,7 @@ OPTIONAL_IMPORT_MODULES = [
|
||||
("hanzo_tools.mcp_tools", 1.0),
|
||||
("hanzo_tools.database", 1.0),
|
||||
("hanzo_tools.editor", 1.0),
|
||||
("hanzo_tools.llm", 2.0), # LLM has llm, allow more time
|
||||
("hanzo_tools.llm", 1.0),
|
||||
("hanzo_tools.vector", 2.0), # Vector has heavy deps
|
||||
]
|
||||
|
||||
|
||||
@@ -1,39 +1,28 @@
|
||||
# hanzo-tools-llm
|
||||
|
||||
LLM interaction tools for Hanzo AI MCP.
|
||||
Models through Hanzo as MCP tools. Every call goes to `api.hanzo.ai` with the
|
||||
Hanzo credential `hanzo_tools.core.HanzoCloud` resolves (`HANZO_API_KEY`, then
|
||||
`~/.hanzo/config.json`, then `hanzo auth token`).
|
||||
|
||||
## Tools
|
||||
|
||||
- `llm` - Core LLM interaction with multiple providers
|
||||
- `unified_llm` - Unified LLM interface
|
||||
- `consensus` - Multi-model consensus for higher accuracy
|
||||
- `llm_manage` - Model management and provider configuration
|
||||
- `llm` — action-routed, default action `query`:
|
||||
- `query`: `POST /v1/chat/completions`. `model` defaults to `enso-auto`;
|
||||
`"auto"` lets Enso pick across the models your org can serve. `max_cost`
|
||||
(USD per 1,000 tokens) and `max_latency_ms` ride as `X-Max-Cost` and
|
||||
`X-Max-Latency-Ms`. Returns `{id, model, content, finish_reason, usage}`.
|
||||
- `models`: `GET /v1/models`, one row per model with family, class and the
|
||||
catalog's per-million prices; `family` narrows it (`enso`, `kai`).
|
||||
- `feedback`: `POST /v1/ai/feedback` with `request_id`, `signal` and, for
|
||||
`signal: rating`, a `rating` of 1 to 3.
|
||||
- `kai_decide` — `POST /v1/decisions`: a `state` (text, object or array) and
|
||||
named typed `questions` (`choice`, `score`, `noul`), answered by Kai with
|
||||
calibrated probabilities.
|
||||
|
||||
## Installation
|
||||
## Install
|
||||
|
||||
```bash
|
||||
pip install hanzo-tools-llm
|
||||
|
||||
# With all LLM providers
|
||||
pip install hanzo-tools-llm[full]
|
||||
```
|
||||
|
||||
## Supported Providers
|
||||
|
||||
- OpenAI (GPT-4, GPT-4o, etc.)
|
||||
- Anthropic (Claude 4, Claude 3.5)
|
||||
- Together AI
|
||||
- Ollama (local models)
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from hanzo_tools.llm import TOOLS, LLM_AVAILABLE, register_tools
|
||||
|
||||
if LLM_AVAILABLE:
|
||||
register_tools(mcp_server)
|
||||
```
|
||||
|
||||
## Part of hanzo-tools
|
||||
|
||||
This package is part of the modular [hanzo-tools](../hanzo-tools) ecosystem.
|
||||
`hanzo-mcp` loads both tools through the `hanzo.tools` entry point.
|
||||
|
||||
@@ -1,69 +1,15 @@
|
||||
"""LLM tools for Hanzo AI.
|
||||
"""Models through Hanzo, as MCP tools.
|
||||
|
||||
Tools:
|
||||
- llm: Unified LLM interface with model management
|
||||
- consensus: Multi-model consensus
|
||||
- llm: completions routed by Enso, the model catalog with prices, feedback
|
||||
- kai_decide: typed questions to Kai, Hanzo's decision model
|
||||
|
||||
Install:
|
||||
pip install hanzo-tools-llm[full]
|
||||
Every call reaches api.hanzo.ai through hanzo_tools.core.HanzoCloud.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from .kai import KaiDecideTool
|
||||
from .llm_unified import UnifiedLLMTool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
TOOLS = [UnifiedLLMTool, KaiDecideTool]
|
||||
|
||||
# Lazy imports for heavy dependencies
|
||||
_tools = []
|
||||
|
||||
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
|
||||
|
||||
_tools.append(ConsensusTool)
|
||||
except ImportError as e:
|
||||
logger.debug(f"ConsensusTool not available: {e}")
|
||||
ConsensusTool = None
|
||||
|
||||
TOOLS = _tools
|
||||
LLM_AVAILABLE = len(_tools) > 0
|
||||
|
||||
__all__ = [
|
||||
"TOOLS",
|
||||
"LLM_AVAILABLE",
|
||||
"LLMTool",
|
||||
"UnifiedLLMTool",
|
||||
"ConsensusTool",
|
||||
"register_tools",
|
||||
]
|
||||
|
||||
|
||||
def register_tools(mcp_server, enabled_tools: dict[str, bool] | None = None):
|
||||
"""Register LLM tools with MCP server."""
|
||||
from hanzo_tools.core import ToolRegistry
|
||||
|
||||
enabled = enabled_tools or {}
|
||||
registered = []
|
||||
|
||||
for tool_class in TOOLS:
|
||||
if tool_class is None:
|
||||
continue
|
||||
tool_name = getattr(tool_class, "name", tool_class.__name__.lower())
|
||||
if enabled.get(tool_name, True):
|
||||
try:
|
||||
tool = tool_class()
|
||||
ToolRegistry.register_tool(mcp_server, tool)
|
||||
registered.append(tool)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to register {tool_name}: {e}")
|
||||
|
||||
return registered
|
||||
__all__ = ["TOOLS", "UnifiedLLMTool", "KaiDecideTool"]
|
||||
|
||||
@@ -1,131 +0,0 @@
|
||||
"""Consensus tool using Metastable protocol."""
|
||||
|
||||
import asyncio
|
||||
from typing import List, Optional, Annotated, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from hanzo_consensus import Result as ConsensusResult, run as run_consensus
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, auto_timeout, create_tool_context
|
||||
|
||||
from .llm_unified import LLMTool
|
||||
|
||||
|
||||
@final
|
||||
class ConsensusTool(BaseTool):
|
||||
"""Metastable consensus across multiple LLMs.
|
||||
|
||||
https://github.com/luxfi/consensus
|
||||
"""
|
||||
|
||||
name = "consensus"
|
||||
|
||||
DEFAULT_MODELS = [
|
||||
"gpt-4o-mini",
|
||||
"claude-3-5-sonnet-20241022",
|
||||
"gemini/gemini-1.5-pro",
|
||||
]
|
||||
|
||||
def __init__(self):
|
||||
self.llm_tool = LLMTool()
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
return """Metastable consensus across multiple LLMs.
|
||||
|
||||
Two-phase protocol:
|
||||
- Phase I (Sampling): Models propose, k-peer sampling, confidence
|
||||
- Phase II (Finality): Threshold aggregation, winner synthesis
|
||||
|
||||
Usage:
|
||||
consensus --prompt "Best approach?" --models '["gpt-4", "claude-3-5-sonnet"]'
|
||||
consensus --prompt "Review this" --rounds 2 --k 2
|
||||
|
||||
Reference: https://github.com/luxfi/consensus
|
||||
"""
|
||||
|
||||
@override
|
||||
@auto_timeout("consensus")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
prompt: Annotated[str, Field(description="Query")] = "",
|
||||
models: Annotated[Optional[List[str]], Field(description="Models")] = None,
|
||||
rounds: Annotated[int, Field(description="Sampling rounds")] = 3,
|
||||
k: Annotated[int, Field(description="Sample size")] = 3,
|
||||
alpha: Annotated[float, Field(description="Agreement")] = 0.6,
|
||||
beta_1: Annotated[float, Field(description="Preference")] = 0.5,
|
||||
beta_2: Annotated[float, Field(description="Decision")] = 0.8,
|
||||
temperature: Annotated[float, Field(description="Temperature")] = 0.7,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
"""Run Metastable consensus."""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
if not prompt:
|
||||
return "Error: prompt required"
|
||||
|
||||
model_list = models or self.DEFAULT_MODELS
|
||||
|
||||
# Filter to available models
|
||||
available = []
|
||||
for m in model_list:
|
||||
provider = self.llm_tool._get_provider_for_model(m)
|
||||
if provider in self.llm_tool.available_providers:
|
||||
available.append(m)
|
||||
|
||||
if not available:
|
||||
return "Error: No models available"
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Metastable consensus: {len(available)} models, {rounds} rounds"
|
||||
)
|
||||
|
||||
# Execute adapter
|
||||
async def execute(model: str, model_prompt: str) -> ConsensusResult:
|
||||
try:
|
||||
import time
|
||||
|
||||
start = time.time()
|
||||
result = await self.llm_tool.call(
|
||||
ctx,
|
||||
model=model,
|
||||
prompt=model_prompt,
|
||||
temperature=temperature,
|
||||
)
|
||||
ms = int((time.time() - start) * 1000)
|
||||
return ConsensusResult(id=model, output=result, ok=True, ms=ms)
|
||||
except Exception as e:
|
||||
return ConsensusResult(id=model, output="", ok=False, error=str(e))
|
||||
|
||||
# Run consensus
|
||||
state = await run_consensus(
|
||||
prompt=prompt,
|
||||
participants=available,
|
||||
execute=execute,
|
||||
rounds=rounds,
|
||||
k=min(k, len(available)),
|
||||
alpha=alpha,
|
||||
beta_1=beta_1,
|
||||
beta_2=beta_2,
|
||||
)
|
||||
|
||||
lines = [
|
||||
"=== Metastable Consensus ===",
|
||||
f"Models: {', '.join(available)}",
|
||||
f"Rounds: {rounds}, k={k}",
|
||||
f"Winner: {state.winner}",
|
||||
f"Finalized: {state.finalized}",
|
||||
"",
|
||||
"=== Synthesis ===",
|
||||
state.synthesis or "(no synthesis)",
|
||||
]
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register with MCP server."""
|
||||
pass
|
||||
@@ -0,0 +1,148 @@
|
||||
"""kai_decide — Kai, Hanzo's decision model, at POST /v1/decisions.
|
||||
|
||||
One state, any named typed questions, one call: the same tool, name and contract
|
||||
as the TypeScript (src/tools/kai.ts) and Rust (rust/src/tools/kai_tool.rs)
|
||||
runtimes. A question is {type, instructions, criteria}: ``choice`` picks a label,
|
||||
``score`` an ordinal level, ``noul`` the probability a statement holds. Answers
|
||||
come back with calibrated probabilities, and Kai writes no text.
|
||||
"""
|
||||
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, HanzoCloud, InvalidParamsError, ToolError
|
||||
from hanzo_tools.core.cloud import NO_KEY, CloudError
|
||||
|
||||
#: The model a decision is asked of when none is named.
|
||||
DEFAULT_MODEL = "kai"
|
||||
KINDS = ("choice", "score", "noul")
|
||||
|
||||
|
||||
def _content(v: Any) -> bool:
|
||||
"""Text, an object or an array: the wire's Content."""
|
||||
return isinstance(v, (str, dict, list))
|
||||
|
||||
|
||||
def fault(q: Any) -> str:
|
||||
"""What is wrong with one question under the wire contract, or ''."""
|
||||
if not isinstance(q, dict):
|
||||
return "must be {type, instructions, criteria}"
|
||||
kind = q.get("type")
|
||||
if kind not in KINDS:
|
||||
return f"type must be one of {', '.join(KINDS)}"
|
||||
if q.get("instructions") is not None and not _content(q["instructions"]):
|
||||
return "instructions must be text, an object or an array"
|
||||
c = q.get("criteria")
|
||||
if kind == "choice":
|
||||
if isinstance(c, list):
|
||||
if not all(isinstance(label, str) for label in c):
|
||||
return "criteria as a list must hold string labels"
|
||||
n = len(set(c))
|
||||
elif isinstance(c, dict):
|
||||
n = len(c)
|
||||
else:
|
||||
return "criteria must be {label: description} or [label, ...]"
|
||||
return "" if n >= 2 else f"criteria must name at least 2 labels, not {n}"
|
||||
if kind == "score":
|
||||
if not isinstance(c, list):
|
||||
return "criteria must be the levels, lowest first: [level0, level1, ...]"
|
||||
if not c:
|
||||
return "criteria must list at least 1 level"
|
||||
nulls = [i for i, level in enumerate(c) if level is None]
|
||||
return f"score level {nulls[0]} is null; describe every level" if nulls else ""
|
||||
if c is None:
|
||||
return ""
|
||||
if not isinstance(c, dict):
|
||||
return 'criteria must be {"true": ..., "false": ...}'
|
||||
odd = [k for k in c if str(k).lower() not in ("true", "false")]
|
||||
return f'criteria take only "true" and "false", not {", ".join(odd)}' if odd else ""
|
||||
|
||||
|
||||
def request(state: Any, questions: Any, model: Any = None) -> dict[str, Any]:
|
||||
"""Check a decision's shape and build its body; whether a state fits the model
|
||||
is the server's to say (422 state_too_long)."""
|
||||
if not _content(state):
|
||||
raise InvalidParamsError("state required: text, an object or an array", param="state")
|
||||
if not isinstance(questions, dict):
|
||||
raise InvalidParamsError("questions required: {name: {type, instructions, criteria}}", param="questions")
|
||||
if not 1 <= len(questions) <= 100:
|
||||
raise InvalidParamsError(f"questions must hold 1 to 100 questions, not {len(questions)}", param="questions")
|
||||
for name, q in questions.items():
|
||||
f = fault(q)
|
||||
if f:
|
||||
raise InvalidParamsError(f"question '{name}': {f}", param="questions")
|
||||
m = model or DEFAULT_MODEL
|
||||
if not isinstance(m, str):
|
||||
raise InvalidParamsError("model must be a model name, e.g. kai", param="model")
|
||||
return {"model": m, "state": state, "questions": questions}
|
||||
|
||||
|
||||
DECIDE_SCHEMA = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"state": {"type": ["string", "object", "array"], "items": {}, "description": "The case to decide about: text, a JSON object or an array"},
|
||||
"questions": {
|
||||
"type": "object",
|
||||
"description": "Question name → {type, instructions, criteria}; 1 to 100 questions",
|
||||
"additionalProperties": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"type": {"type": "string", "enum": list(KINDS)},
|
||||
"instructions": {"type": ["string", "object", "array"], "items": {}, "description": "Optional, recommended. What Kai answers"},
|
||||
"criteria": {
|
||||
"type": ["object", "array"],
|
||||
"items": {},
|
||||
"description": 'choice: {label: description} or [label, ...], 2 labels or more; score: [level0, level1, ...], lowest first, none null; noul: {"true": ..., "false": ...}, optional',
|
||||
},
|
||||
},
|
||||
"required": ["type"],
|
||||
},
|
||||
},
|
||||
"model": {"type": "string", "description": "Decision model (default kai)"},
|
||||
},
|
||||
"required": ["state", "questions"],
|
||||
}
|
||||
|
||||
|
||||
class KaiDecideTool(BaseTool):
|
||||
"""Typed questions to Kai about one case, answered with calibrated probabilities."""
|
||||
|
||||
name: ClassVar[str] = "kai_decide"
|
||||
VERSION: ClassVar[str] = "0.2.1"
|
||||
DEFAULT_ACTION: ClassVar[str] = "decide"
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._cloud: HanzoCloud | None = None
|
||||
|
||||
@self.action("decide", "Ask Kai typed questions about one case", schema=DECIDE_SCHEMA)
|
||||
async def decide(ctx: MCPContext, state: Any = None, questions: Any = None, model: str | None = None) -> dict:
|
||||
body = request(state, questions, model)
|
||||
if self._cloud is None:
|
||||
self._cloud = HanzoCloud()
|
||||
if not self._cloud.configured():
|
||||
raise ToolError(code="INVALID_PARAMS", message=NO_KEY)
|
||||
try:
|
||||
d = await self._cloud.post("/v1/decisions", body)
|
||||
except CloudError as e:
|
||||
raise ToolError(code="UPSTREAM", message=str(e))
|
||||
if not isinstance(d, dict) or not isinstance(d.get("answers"), dict):
|
||||
raise ToolError(code="UPSTREAM", message=f"not a decision: {str(d)[:200]}")
|
||||
return d
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return (
|
||||
"Ask Kai, Hanzo's decision model, typed questions about one case in one call (POST /v1/decisions). "
|
||||
"Reach for Kai when the answer is one of options you already know: classify, route, gate, rank, check. "
|
||||
"It writes no text. `questions` maps a name to {type, instructions, criteria}, 1 to 100 of them; "
|
||||
"`instructions` is optional but recommended, as the text Kai answers. choice picks one label: criteria "
|
||||
"{label: description} or [label, ...], 2 labels or more. score picks an ordinal level: criteria "
|
||||
"[level0, level1, ...], lowest first; act on the argmax of its probabilities, since score is the mean "
|
||||
'level index. noul gives the probability a statement holds: criteria {"true": ..., "false": ...}, '
|
||||
"optional; write it as a statement, not a bare yes/no question. Returns the decision {id, model, "
|
||||
"answers: {name: answer}, usage, routing, state_hash, latency_ms}. Probabilities are calibrated; "
|
||||
"confidence = (n·p_max − 1)/(n − 1) over n options. Billed on input tokens at the catalog's kai rate; "
|
||||
"output_tokens is 0."
|
||||
)
|
||||
@@ -1,455 +0,0 @@
|
||||
"""LLM management tool for enabling/disabling LLM providers."""
|
||||
|
||||
import json
|
||||
from typing import 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
|
||||
|
||||
from .llm_unified import LLMTool
|
||||
|
||||
Action = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Action to perform: list, enable, disable, test",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
Provider = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Provider name (for enable/disable/test actions)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Model = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Model to test (for test action)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class LLMManageParams(TypedDict, total=False):
|
||||
"""Parameters for LLM management tool."""
|
||||
|
||||
action: str
|
||||
provider: Optional[str]
|
||||
model: Optional[str]
|
||||
|
||||
|
||||
@final
|
||||
class LLMManageTool(BaseTool):
|
||||
"""Tool for managing LLM providers."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the LLM management tool."""
|
||||
self.llm_tool = LLMTool()
|
||||
self.config_file = Path.home() / ".hanzo" / "llm" / "providers.json"
|
||||
self.config_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._load_config()
|
||||
|
||||
def _load_config(self):
|
||||
"""Load provider configuration."""
|
||||
if self.config_file.exists():
|
||||
try:
|
||||
with open(self.config_file, "r") as f:
|
||||
self.config = json.load(f)
|
||||
except Exception:
|
||||
self.config = {"disabled_providers": []}
|
||||
else:
|
||||
self.config = {"disabled_providers": []}
|
||||
|
||||
def _save_config(self):
|
||||
"""Save provider configuration."""
|
||||
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_manage"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Manage LLM providers and test configurations.
|
||||
|
||||
Actions:
|
||||
- list: Show all providers and their status
|
||||
- models: List all available models (optionally filtered by provider)
|
||||
- enable: Enable a provider's tools
|
||||
- disable: Disable a provider's tools
|
||||
- test: Test a model to verify it works
|
||||
|
||||
Examples:
|
||||
- llm_manage --action list
|
||||
- llm_manage --action models
|
||||
- llm_manage --action models --provider openai
|
||||
- llm_manage --action enable --provider openai
|
||||
- llm_manage --action disable --provider perplexity
|
||||
- llm_manage --action test --model "gpt-4"
|
||||
- llm_manage --action test --provider groq --model "mixtral"
|
||||
|
||||
Providers are automatically detected based on environment variables:
|
||||
- OpenAI: OPENAI_API_KEY
|
||||
- Anthropic: ANTHROPIC_API_KEY or CLAUDE_API_KEY
|
||||
- Google: GOOGLE_API_KEY or GEMINI_API_KEY
|
||||
- Groq: GROQ_API_KEY
|
||||
- Mistral: MISTRAL_API_KEY
|
||||
- Perplexity: PERPLEXITY_API_KEY
|
||||
- And many more...
|
||||
"""
|
||||
|
||||
@override
|
||||
@auto_timeout("llm_manage")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[LLMManageParams],
|
||||
) -> str:
|
||||
"""Manage LLM providers.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Result of the management action
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
action = params.get("action")
|
||||
if not action:
|
||||
return "Error: action is required (list, enable, disable, test, models)"
|
||||
|
||||
provider = params.get("provider")
|
||||
model = params.get("model")
|
||||
|
||||
# Handle different actions
|
||||
if action == "list":
|
||||
return self._list_providers()
|
||||
elif action == "models":
|
||||
return self._list_all_models(provider)
|
||||
elif action == "enable":
|
||||
return self._enable_provider(provider)
|
||||
elif action == "disable":
|
||||
return self._disable_provider(provider)
|
||||
elif action == "test":
|
||||
return await self._test_model(ctx, provider, model)
|
||||
else:
|
||||
return f"Error: Invalid action '{action}'. Must be one of: list, models, enable, disable, test"
|
||||
|
||||
def _list_providers(self) -> str:
|
||||
"""List all providers and their status."""
|
||||
output = ["=== LLM Providers ==="]
|
||||
output.append("")
|
||||
|
||||
# Get all possible providers
|
||||
all_providers = sorted(LLMTool.API_KEY_ENV_VARS.keys())
|
||||
available_providers = self.llm_tool.available_providers
|
||||
disabled_providers = self.config.get("disabled_providers", [])
|
||||
|
||||
# Categorize providers
|
||||
active = []
|
||||
available_but_disabled = []
|
||||
no_api_key = []
|
||||
|
||||
for provider in all_providers:
|
||||
if provider in available_providers:
|
||||
if provider in disabled_providers:
|
||||
available_but_disabled.append(provider)
|
||||
else:
|
||||
active.append(provider)
|
||||
else:
|
||||
no_api_key.append(provider)
|
||||
|
||||
# Show active providers
|
||||
if active:
|
||||
output.append("✅ Active Providers (API key found, enabled):")
|
||||
for provider in active:
|
||||
env_vars = available_providers.get(provider, [])
|
||||
output.append(f" - {provider}: {', '.join(env_vars)}")
|
||||
|
||||
# Show example models
|
||||
examples = self._get_example_models(provider)
|
||||
if examples:
|
||||
output.append(f" Models: {', '.join(examples[:3])}")
|
||||
output.append("")
|
||||
|
||||
# Show disabled providers
|
||||
if available_but_disabled:
|
||||
output.append("⚠️ Available but Disabled (API key found, disabled):")
|
||||
for provider in available_but_disabled:
|
||||
env_vars = available_providers.get(provider, [])
|
||||
output.append(f" - {provider}: {', '.join(env_vars)}")
|
||||
output.append(
|
||||
f" Use: llm_manage --action enable --provider {provider}"
|
||||
)
|
||||
output.append("")
|
||||
|
||||
# Show providers without API keys
|
||||
if no_api_key:
|
||||
output.append("❌ No API Key Found:")
|
||||
for provider in no_api_key[:10]: # Show first 10
|
||||
env_vars = LLMTool.API_KEY_ENV_VARS.get(provider, [])
|
||||
output.append(f" - {provider}: Set one of {', '.join(env_vars)}")
|
||||
if len(no_api_key) > 10:
|
||||
output.append(f" ... and {len(no_api_key) - 10} more")
|
||||
output.append("")
|
||||
|
||||
# Summary
|
||||
output.append("=== Summary ===")
|
||||
output.append(f"Total providers: {len(all_providers)}")
|
||||
output.append(f"Active: {len(active)}")
|
||||
output.append(f"Disabled: {len(available_but_disabled)}")
|
||||
output.append(f"No API key: {len(no_api_key)}")
|
||||
|
||||
# Show available tools
|
||||
if active:
|
||||
output.append("\n=== Available LLM Tools ===")
|
||||
output.append("- llm: Universal LLM tool (all providers)")
|
||||
output.append("- consensus: Query multiple models in parallel")
|
||||
|
||||
provider_tools = []
|
||||
for provider in active:
|
||||
if provider in [
|
||||
"openai",
|
||||
"anthropic",
|
||||
"google",
|
||||
"groq",
|
||||
"mistral",
|
||||
"perplexity",
|
||||
]:
|
||||
tool_name = "gemini" if provider == "google" else provider
|
||||
provider_tools.append(tool_name)
|
||||
|
||||
if provider_tools:
|
||||
output.append(f"- Provider tools: {', '.join(provider_tools)}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _enable_provider(self, provider: Optional[str]) -> str:
|
||||
"""Enable a provider."""
|
||||
if not provider:
|
||||
return "Error: provider is required for enable action"
|
||||
|
||||
if provider not in self.llm_tool.available_providers:
|
||||
env_vars = LLMTool.API_KEY_ENV_VARS.get(provider, [])
|
||||
if env_vars:
|
||||
return f"Error: No API key found for {provider}. Set one of: {', '.join(env_vars)}"
|
||||
else:
|
||||
return f"Error: Unknown provider '{provider}'"
|
||||
|
||||
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 _disable_provider(self, provider: Optional[str]) -> str:
|
||||
"""Disable a provider."""
|
||||
if not provider:
|
||||
return "Error: provider is required for disable action"
|
||||
|
||||
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}. Its tools will no longer be available."
|
||||
else:
|
||||
return f"{provider} is already disabled"
|
||||
|
||||
def _list_all_models(self, provider: Optional[str] = None) -> str:
|
||||
"""List all available models from LLM."""
|
||||
try:
|
||||
from .llm_unified import LLMTool
|
||||
|
||||
all_models = LLMTool.get_all_models()
|
||||
|
||||
if not all_models:
|
||||
return "No models available or LLM not installed"
|
||||
|
||||
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:
|
||||
output.append(f"{provider_name}: {len(models)} models")
|
||||
|
||||
output.append(
|
||||
"\nUse 'llm_manage --action models --provider <name>' to see models for a specific provider"
|
||||
)
|
||||
|
||||
# Show recommended models
|
||||
output.append("\n=== Recommended Models ===")
|
||||
recommended = {
|
||||
"OpenAI": ["gpt-4o", "gpt-4o-mini", "gpt-3.5-turbo"],
|
||||
"Anthropic": [
|
||||
"claude-3-opus-20240229",
|
||||
"claude-3-5-sonnet-20241022",
|
||||
"claude-3-haiku-20240307",
|
||||
],
|
||||
"Google": ["gemini/gemini-1.5-pro", "gemini/gemini-1.5-flash"],
|
||||
"Groq": [
|
||||
"groq/llama3-70b-8192",
|
||||
"groq/llama3-8b-8192",
|
||||
"groq/gemma2-9b-it",
|
||||
],
|
||||
"Mistral": [
|
||||
"mistral/mistral-large-latest",
|
||||
"mistral/mistral-medium",
|
||||
],
|
||||
}
|
||||
|
||||
for provider_name, models in recommended.items():
|
||||
available = LLMTool().available_providers
|
||||
provider_key = provider_name.lower()
|
||||
|
||||
if provider_key in available:
|
||||
output.append(f"\n{provider_name} (✅ API key found):")
|
||||
for model in models:
|
||||
output.append(f" - {model}")
|
||||
else:
|
||||
output.append(f"\n{provider_name} (❌ No API key):")
|
||||
for model in models:
|
||||
output.append(f" - {model}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
return f"Error listing models: {str(e)}"
|
||||
|
||||
async def _test_model(
|
||||
self, ctx: MCPContext, provider: Optional[str], model: Optional[str]
|
||||
) -> str:
|
||||
"""Test a model to verify it works."""
|
||||
if not model and not provider:
|
||||
return "Error: Either model or provider is required for test action"
|
||||
|
||||
# Determine model to test
|
||||
if model:
|
||||
test_model = model
|
||||
else:
|
||||
# Use default model for provider
|
||||
default_models = {
|
||||
"openai": "gpt-3.5-turbo",
|
||||
"anthropic": "claude-3-haiku-20240307",
|
||||
"google": "gemini/gemini-pro",
|
||||
"groq": "groq/mixtral-8x7b-32768",
|
||||
"mistral": "mistral/mistral-tiny",
|
||||
"perplexity": "perplexity/sonar-small-online",
|
||||
}
|
||||
test_model = default_models.get(provider)
|
||||
if not test_model:
|
||||
return f"Error: No default model for provider '{provider}'. Please specify a model."
|
||||
|
||||
# Test the model
|
||||
test_prompt = "Hello! Please respond with 'OK' if you can hear me."
|
||||
|
||||
output = [f"Testing model: {test_model}"]
|
||||
output.append(f"Prompt: {test_prompt}")
|
||||
output.append("")
|
||||
|
||||
try:
|
||||
# Call the LLM
|
||||
params = {
|
||||
"model": test_model,
|
||||
"prompt": test_prompt,
|
||||
"max_tokens": 10,
|
||||
"temperature": 0,
|
||||
}
|
||||
|
||||
response = await self.llm_tool.call(ctx, **params)
|
||||
|
||||
if response.startswith("Error:"):
|
||||
output.append("❌ Test failed:")
|
||||
output.append(response)
|
||||
else:
|
||||
output.append("✅ Test successful!")
|
||||
output.append(f"Response: {response}")
|
||||
output.append("")
|
||||
output.append(f"Model '{test_model}' is working correctly.")
|
||||
|
||||
# Show provider info
|
||||
detected_provider = self.llm_tool._get_provider_for_model(test_model)
|
||||
if detected_provider:
|
||||
output.append(f"Provider: {detected_provider}")
|
||||
|
||||
except Exception as e:
|
||||
output.append("❌ Test failed with exception:")
|
||||
output.append(str(e))
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _get_example_models(self, provider: str) -> list[str]:
|
||||
"""Get example models for a provider."""
|
||||
examples = {
|
||||
"openai": ["gpt-4o", "gpt-4", "gpt-3.5-turbo", "o1-preview"],
|
||||
"anthropic": [
|
||||
"claude-3-opus-20240229",
|
||||
"claude-3-sonnet-20240229",
|
||||
"claude-3-haiku-20240307",
|
||||
],
|
||||
"google": [
|
||||
"gemini/gemini-pro",
|
||||
"gemini/gemini-1.5-pro",
|
||||
"gemini/gemini-1.5-flash",
|
||||
],
|
||||
"groq": [
|
||||
"groq/mixtral-8x7b-32768",
|
||||
"groq/llama3-70b-8192",
|
||||
"groq/llama3-8b-8192",
|
||||
],
|
||||
"mistral": [
|
||||
"mistral/mistral-large-latest",
|
||||
"mistral/mistral-medium",
|
||||
"mistral/mistral-small",
|
||||
],
|
||||
"perplexity": [
|
||||
"perplexity/sonar-medium-online",
|
||||
"perplexity/sonar-small-online",
|
||||
],
|
||||
}
|
||||
return examples.get(provider, [])
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,413 +0,0 @@
|
||||
"""Provider-specific LLM tools."""
|
||||
|
||||
from typing import Dict, Unpack, Optional, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, auto_timeout
|
||||
|
||||
from .llm_unified import LLMTool
|
||||
|
||||
Prompt = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="The prompt or question to send to the model",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
Model = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Specific model variant (defaults to provider's best model)",
|
||||
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.0-2.0)",
|
||||
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,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class ProviderToolParams(TypedDict, total=False):
|
||||
"""Parameters for provider-specific tools."""
|
||||
|
||||
prompt: str
|
||||
model: Optional[str]
|
||||
system_prompt: Optional[str]
|
||||
temperature: float
|
||||
max_tokens: Optional[int]
|
||||
json_mode: bool
|
||||
|
||||
|
||||
class BaseProviderTool(BaseTool):
|
||||
"""Base class for provider-specific LLM tools."""
|
||||
|
||||
def __init__(
|
||||
self, provider: str, default_model: str, model_variants: Dict[str, str]
|
||||
):
|
||||
"""Initialize provider tool.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
default_model: Default model to use
|
||||
model_variants: Map of short names to full model names
|
||||
"""
|
||||
self.provider = provider
|
||||
self.default_model = default_model
|
||||
self.model_variants = model_variants
|
||||
self.llm_tool = LLMTool()
|
||||
self.is_available = provider in self.llm_tool.available_providers
|
||||
|
||||
def get_full_model_name(self, model: Optional[str]) -> str:
|
||||
"""Get full model name from short name or default."""
|
||||
if not model:
|
||||
return self.default_model
|
||||
|
||||
# Check if it's a short name
|
||||
if model in self.model_variants:
|
||||
return self.model_variants[model]
|
||||
|
||||
# Return as-is if not found (assume full name)
|
||||
return model
|
||||
|
||||
@override
|
||||
@auto_timeout("provider_tools")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[ProviderToolParams],
|
||||
) -> str:
|
||||
"""Call the provider's LLM."""
|
||||
if not self.is_available:
|
||||
env_vars = LLMTool.API_KEY_ENV_VARS.get(self.provider, [])
|
||||
return f"Error: {self.provider.title()} API key not found. Set one of: {', '.join(env_vars)}"
|
||||
|
||||
# Get full model name
|
||||
model = self.get_full_model_name(params.get("model"))
|
||||
|
||||
# Prepare LLM tool parameters
|
||||
llm_params = {
|
||||
"model": model,
|
||||
"prompt": params["prompt"],
|
||||
}
|
||||
|
||||
# Add optional parameters
|
||||
if "system_prompt" in params:
|
||||
llm_params["system_prompt"] = params["system_prompt"]
|
||||
if "temperature" in params:
|
||||
llm_params["temperature"] = params["temperature"]
|
||||
if "max_tokens" in params:
|
||||
llm_params["max_tokens"] = params["max_tokens"]
|
||||
if "json_mode" in params:
|
||||
llm_params["json_mode"] = params["json_mode"]
|
||||
|
||||
# Call the LLM tool
|
||||
return await self.llm_tool.call(ctx, **llm_params)
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
|
||||
|
||||
@final
|
||||
class OpenAITool(BaseProviderTool):
|
||||
"""OpenAI-specific LLM tool."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
provider="openai",
|
||||
default_model="gpt-4o",
|
||||
model_variants={
|
||||
"4o": "gpt-4o",
|
||||
"4": "gpt-4",
|
||||
"3.5": "gpt-3.5-turbo",
|
||||
"o1": "o1-preview",
|
||||
"o1-mini": "o1-mini",
|
||||
"4-turbo": "gpt-4-turbo-preview",
|
||||
"4-vision": "gpt-4-vision-preview",
|
||||
},
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
return "openai"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
status = "✓ Available" if self.is_available else "✗ No API key"
|
||||
return f"""Query OpenAI models directly ({status}).
|
||||
|
||||
Models:
|
||||
- 4o (default): GPT-4o - Latest and most capable
|
||||
- 4: GPT-4 - Advanced reasoning
|
||||
- 3.5: GPT-3.5 Turbo - Fast and efficient
|
||||
- o1: O1 Preview - Chain of thought reasoning
|
||||
- o1-mini: O1 Mini - Smaller reasoning model
|
||||
|
||||
Examples:
|
||||
- openai --prompt "Explain quantum computing"
|
||||
- openai --model 4 --prompt "Write a Python function"
|
||||
- openai --model o1 --prompt "Solve this step by step"
|
||||
"""
|
||||
|
||||
|
||||
@final
|
||||
class AnthropicTool(BaseProviderTool):
|
||||
"""Anthropic-specific LLM tool."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
provider="anthropic",
|
||||
default_model="claude-3-sonnet-20240229",
|
||||
model_variants={
|
||||
"opus": "claude-3-opus-20240229",
|
||||
"sonnet": "claude-3-sonnet-20240229",
|
||||
"haiku": "claude-3-haiku-20240307",
|
||||
"2.1": "claude-2.1",
|
||||
"2": "claude-2",
|
||||
"instant": "claude-instant-1.2",
|
||||
},
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
return "anthropic"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
status = "✓ Available" if self.is_available else "✗ No API key"
|
||||
return f"""Query Anthropic Claude models directly ({status}).
|
||||
|
||||
Models:
|
||||
- sonnet (default): Claude 3 Sonnet - Balanced performance
|
||||
- opus: Claude 3 Opus - Most capable
|
||||
- haiku: Claude 3 Haiku - Fast and efficient
|
||||
- 2.1: Claude 2.1 - Previous generation
|
||||
- instant: Claude Instant - Very fast
|
||||
|
||||
Examples:
|
||||
- anthropic --prompt "Analyze this code"
|
||||
- anthropic --model opus --prompt "Write a detailed essay"
|
||||
- anthropic --model haiku --prompt "Quick question"
|
||||
"""
|
||||
|
||||
|
||||
@final
|
||||
class GeminiTool(BaseProviderTool):
|
||||
"""Google Gemini-specific LLM tool."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
provider="google",
|
||||
default_model="gemini/gemini-pro",
|
||||
model_variants={
|
||||
"pro": "gemini/gemini-pro",
|
||||
"pro-vision": "gemini/gemini-pro-vision",
|
||||
"1.5-pro": "gemini/gemini-1.5-pro-latest",
|
||||
"1.5-flash": "gemini/gemini-1.5-flash-latest",
|
||||
"ultra": "gemini/gemini-ultra",
|
||||
},
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
return "gemini"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
status = "✓ Available" if self.is_available else "✗ No API key"
|
||||
return f"""Query Google Gemini models directly ({status}).
|
||||
|
||||
Models:
|
||||
- pro (default): Gemini Pro - Balanced model
|
||||
- 1.5-pro: Gemini 1.5 Pro - Advanced with long context
|
||||
- 1.5-flash: Gemini 1.5 Flash - Fast and efficient
|
||||
- pro-vision: Gemini Pro Vision - Multimodal
|
||||
- ultra: Gemini Ultra - Most capable (if available)
|
||||
|
||||
Examples:
|
||||
- gemini --prompt "Explain this concept"
|
||||
- gemini --model 1.5-pro --prompt "Analyze this long document"
|
||||
- gemini --model 1.5-flash --prompt "Quick task"
|
||||
"""
|
||||
|
||||
|
||||
@final
|
||||
class GroqTool(BaseProviderTool):
|
||||
"""Groq-specific LLM tool."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
provider="groq",
|
||||
default_model="groq/mixtral-8x7b-32768",
|
||||
model_variants={
|
||||
"mixtral": "groq/mixtral-8x7b-32768",
|
||||
"llama3-70b": "groq/llama3-70b-8192",
|
||||
"llama3-8b": "groq/llama3-8b-8192",
|
||||
"llama2-70b": "groq/llama2-70b-4096",
|
||||
"gemma-7b": "groq/gemma-7b-it",
|
||||
},
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
return "groq"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
status = "✓ Available" if self.is_available else "✗ No API key"
|
||||
return f"""Query Groq LPU models - ultra-fast inference ({status}).
|
||||
|
||||
Models:
|
||||
- mixtral (default): Mixtral 8x7B - High quality
|
||||
- llama3-70b: Llama 3 70B - Very capable
|
||||
- llama3-8b: Llama 3 8B - Fast and efficient
|
||||
- llama2-70b: Llama 2 70B - Previous gen
|
||||
- gemma-7b: Google Gemma 7B - Efficient
|
||||
|
||||
Examples:
|
||||
- groq --prompt "Fast response needed"
|
||||
- groq --model llama3-70b --prompt "Complex reasoning"
|
||||
- groq --model gemma-7b --prompt "Quick task"
|
||||
"""
|
||||
|
||||
|
||||
@final
|
||||
class MistralTool(BaseProviderTool):
|
||||
"""Mistral-specific LLM tool."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
provider="mistral",
|
||||
default_model="mistral/mistral-medium",
|
||||
model_variants={
|
||||
"tiny": "mistral/mistral-tiny",
|
||||
"small": "mistral/mistral-small-latest",
|
||||
"medium": "mistral/mistral-medium-latest",
|
||||
"large": "mistral/mistral-large-latest",
|
||||
"embed": "mistral/mistral-embed",
|
||||
},
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
return "mistral"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
status = "✓ Available" if self.is_available else "✗ No API key"
|
||||
return f"""Query Mistral AI models directly ({status}).
|
||||
|
||||
Models:
|
||||
- medium (default): Mistral Medium - Balanced
|
||||
- large: Mistral Large - Most capable
|
||||
- small: Mistral Small - Efficient
|
||||
- tiny: Mistral Tiny - Very fast
|
||||
|
||||
Examples:
|
||||
- mistral --prompt "Explain this"
|
||||
- mistral --model large --prompt "Complex analysis"
|
||||
- mistral --model tiny --prompt "Quick response"
|
||||
"""
|
||||
|
||||
|
||||
@final
|
||||
class PerplexityTool(BaseProviderTool):
|
||||
"""Perplexity-specific LLM tool."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
provider="perplexity",
|
||||
default_model="perplexity/sonar-medium-online",
|
||||
model_variants={
|
||||
"sonar-small": "perplexity/sonar-small-online",
|
||||
"sonar-medium": "perplexity/sonar-medium-online",
|
||||
"sonar-small-chat": "perplexity/sonar-small-chat",
|
||||
"sonar-medium-chat": "perplexity/sonar-medium-chat",
|
||||
},
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
return "perplexity"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
status = "✓ Available" if self.is_available else "✗ No API key"
|
||||
return f"""Query Perplexity models with internet access ({status}).
|
||||
|
||||
Models:
|
||||
- sonar-medium (default): Online search + reasoning
|
||||
- sonar-small: Faster online search
|
||||
- sonar-medium-chat: Chat without search
|
||||
- sonar-small-chat: Fast chat without search
|
||||
|
||||
Examples:
|
||||
- perplexity --prompt "Latest news about AI"
|
||||
- perplexity --model sonar-small --prompt "Quick fact check"
|
||||
- perplexity --model sonar-medium-chat --prompt "Explain without search"
|
||||
"""
|
||||
|
||||
|
||||
# Export all provider tools
|
||||
PROVIDER_TOOLS = [
|
||||
OpenAITool,
|
||||
AnthropicTool,
|
||||
GeminiTool,
|
||||
GroqTool,
|
||||
MistralTool,
|
||||
PerplexityTool,
|
||||
]
|
||||
|
||||
|
||||
def create_provider_tools() -> list[BaseTool]:
|
||||
"""Create instances of all provider tools."""
|
||||
tools = []
|
||||
for tool_class in PROVIDER_TOOLS:
|
||||
tool = tool_class()
|
||||
# Only include tools with available API keys
|
||||
if tool.is_available:
|
||||
tools.append(tool)
|
||||
return tools
|
||||
@@ -4,24 +4,19 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hanzo-tools-llm"
|
||||
version = "0.2.0"
|
||||
description = "LLM tools for Hanzo AI - model management, consensus, providers"
|
||||
version = "0.2.1"
|
||||
description = "Models through Hanzo as MCP tools: llm (Enso routing, the catalog, feedback) and kai_decide (Kai decisions)"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
license = { text = "MIT" }
|
||||
authors = [{ name = "Hanzo Industries Inc", email = "[email protected]" }]
|
||||
keywords = ["hanzo", "tools", "llm", "mcp", "ai"]
|
||||
dependencies = [
|
||||
"hanzo-tools>=0.3.0",
|
||||
"hanzo-consensus>=0.1.0",
|
||||
"hanzo-tools>=0.3.8",
|
||||
"httpx>=0.27.0",
|
||||
"mcp>=1.25.0",
|
||||
"fastmcp>=2.14.1",
|
||||
"pydantic>=2.12.5",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
full = ["hanzo-llm>=1.0.0", "openai>=1.0.0", "anthropic>=0.40.0"]
|
||||
|
||||
[project.entry-points."hanzo.tools"]
|
||||
llm = "hanzo_tools.llm:TOOLS"
|
||||
|
||||
|
||||
@@ -1,17 +1,191 @@
|
||||
"""Tests for hanzo-tools-llm."""
|
||||
"""Tests for hanzo-tools-llm: llm (Enso, the catalog, feedback) and kai_decide.
|
||||
|
||||
The tools run against a real HanzoCloud whose httpx transport is a recorder
|
||||
answering as api.hanzo.ai does, so each test reads the exact request that went
|
||||
out (method, path, bearer, Enso's headers, body) and what the tool made of the
|
||||
answer. No socket is opened.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from hanzo_tools.core import HanzoCloud
|
||||
from hanzo_tools.llm import TOOLS, KaiDecideTool, UnifiedLLMTool
|
||||
|
||||
class TestImports:
|
||||
"""Test that all modules can be imported."""
|
||||
COMPLETION = {
|
||||
"id": "chatcmpl-ae7c311d-5c64-47da-886b-86ceb2fd185f",
|
||||
"object": "chat.completion",
|
||||
"model": "enso-free",
|
||||
"choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Hi there, friend."}}],
|
||||
"usage": {"prompt_tokens": 110, "completion_tokens": 6, "total_tokens": 116},
|
||||
}
|
||||
CATALOG = {
|
||||
"object": "list",
|
||||
"data": [
|
||||
{"id": "kai", "family": "kai", "class": "ours", "outputs": ["decision"], "pricing": {"input_per_million": 0.021, "output_per_million": 0}},
|
||||
{"id": "enso-auto", "family": "enso", "class": "free", "outputs": ["text"], "context_window": 1000000, "pricing": {"input_per_million": 0, "output_per_million": 0}},
|
||||
{"id": "zen5", "family": "zen", "class": "ours", "pricing": {"input_per_million": 0.3, "output_per_million": 1.2}},
|
||||
],
|
||||
}
|
||||
DECISION = {
|
||||
"id": "dec_ff9388f6b8b85328960ba301635ac108",
|
||||
"model": "kai",
|
||||
"provider": "Hanzo",
|
||||
"answers": {"team": {"type": "choice", "choice": "payments", "confidence": 0.9983, "probabilities": {"account": 0.0006, "payments": 0.9989}, "answer_confidence": 0.9989}},
|
||||
"usage": {"input_tokens": 30, "output_tokens": 0},
|
||||
}
|
||||
|
||||
def test_import_package(self):
|
||||
from hanzo_tools import llm
|
||||
|
||||
assert llm is not None
|
||||
class Gateway:
|
||||
"""Records each request and answers it the way api.hanzo.ai does."""
|
||||
|
||||
def test_import_tools(self):
|
||||
from hanzo_tools.llm import TOOLS
|
||||
def __init__(self):
|
||||
self.seen: list[httpx.Request] = []
|
||||
self.answer = None # (status, body) to force one reply
|
||||
|
||||
assert len(TOOLS) > 0
|
||||
def __call__(self, request: httpx.Request) -> httpx.Response:
|
||||
self.seen.append(request)
|
||||
if self.answer:
|
||||
status, body = self.answer
|
||||
return httpx.Response(status, json=body)
|
||||
path = request.url.path
|
||||
if path == "/v1/models":
|
||||
return httpx.Response(200, json=CATALOG)
|
||||
if path == "/v1/decisions":
|
||||
return httpx.Response(200, json=DECISION)
|
||||
if path == "/v1/ai/feedback":
|
||||
sent = json.loads(request.content)
|
||||
return httpx.Response(200, json={"status": "ok", "msg": "", "data": {"request_id": sent["request_id"].removeprefix("chatcmpl-"), "reward": 1, "recorded": True}})
|
||||
return httpx.Response(200, json=COMPLETION, headers={"X-Routed-Model": "enso-free"})
|
||||
|
||||
@property
|
||||
def last(self) -> httpx.Request:
|
||||
return self.seen[-1]
|
||||
|
||||
|
||||
def wired(tool_class):
|
||||
gw = Gateway()
|
||||
cloud = HanzoCloud(base_url="https://api.hanzo.ai", key="sk-test")
|
||||
cloud._client = httpx.AsyncClient(base_url="https://api.hanzo.ai", transport=httpx.MockTransport(gw))
|
||||
tool = tool_class()
|
||||
tool._cloud = cloud
|
||||
return tool, gw
|
||||
|
||||
|
||||
def run(tool, **kw):
|
||||
return asyncio.run(tool.call(None, **kw))
|
||||
|
||||
|
||||
def body(req: httpx.Request):
|
||||
return json.loads(req.content) if req.content else None
|
||||
|
||||
|
||||
def test_the_package_carries_llm_and_kai_decide():
|
||||
assert [t.name for t in TOOLS] == ["llm", "kai_decide"]
|
||||
assert UnifiedLLMTool.DEFAULT_ACTION == "query"
|
||||
assert KaiDecideTool.DEFAULT_ACTION == "decide"
|
||||
|
||||
|
||||
def test_query_defaults_to_enso_auto_and_sends_no_unasked_bound():
|
||||
tool, gw = wired(UnifiedLLMTool)
|
||||
env = run(tool, prompt="Say hi.")
|
||||
assert env["ok"], env
|
||||
assert (gw.last.method, gw.last.url.path) == ("POST", "/v1/chat/completions")
|
||||
assert gw.last.headers["authorization"] == "Bearer sk-test"
|
||||
assert "x-max-cost" not in gw.last.headers and "x-max-latency-ms" not in gw.last.headers
|
||||
assert body(gw.last) == {"model": "enso-auto", "messages": [{"role": "user", "content": "Say hi."}]}
|
||||
|
||||
|
||||
def test_query_carries_auto_and_the_bounds_and_returns_the_model_that_served():
|
||||
tool, gw = wired(UnifiedLLMTool)
|
||||
env = run(tool, action="query", model="auto", system="Be brief.", prompt="Say hi.", max_tokens=20, max_cost=0.01, max_latency_ms=800.4)
|
||||
assert gw.last.headers["x-max-cost"] == "0.01"
|
||||
assert gw.last.headers["x-max-latency-ms"] == "800"
|
||||
assert body(gw.last) == {
|
||||
"model": "auto",
|
||||
"messages": [{"role": "system", "content": "Be brief."}, {"role": "user", "content": "Say hi."}],
|
||||
"max_tokens": 20,
|
||||
}
|
||||
assert env["data"] == {"id": COMPLETION["id"], "model": "enso-free", "content": "Hi there, friend.", "finish_reason": "stop", "usage": COMPLETION["usage"]}
|
||||
|
||||
|
||||
def test_query_takes_a_conversation_and_refuses_a_bad_bound_before_sending():
|
||||
tool, gw = wired(UnifiedLLMTool)
|
||||
msgs = [{"role": "user", "content": "a"}, {"role": "assistant", "content": "b"}, {"role": "user", "content": "c"}]
|
||||
run(tool, messages=msgs, model="enso-pro")
|
||||
assert body(gw.last) == {"model": "enso-pro", "messages": msgs}
|
||||
env = run(tool, prompt="hi", max_cost=-1)
|
||||
assert not env["ok"] and "max_cost must be a positive number" in env["error"]["message"]
|
||||
assert len(gw.seen) == 1
|
||||
|
||||
|
||||
def test_a_refusal_carries_the_status_and_the_servers_sentence():
|
||||
tool, gw = wired(UnifiedLLMTool)
|
||||
gw.answer = (402, {"error": {"message": "Pick a plan", "type": "billing_error"}})
|
||||
env = run(tool, prompt="hi")
|
||||
assert not env["ok"]
|
||||
assert env["error"]["code"] == "UPSTREAM"
|
||||
assert "402" in env["error"]["message"] and "Pick a plan" in env["error"]["message"]
|
||||
|
||||
|
||||
def test_models_lists_the_catalogs_own_prices_by_family():
|
||||
tool, gw = wired(UnifiedLLMTool)
|
||||
env = run(tool, action="models", family="kai")
|
||||
assert (gw.last.method, gw.last.url.path) == ("GET", "/v1/models")
|
||||
assert env["data"] == {"count": 1, "models": [{"id": "kai", "family": "kai", "class": "ours", "outputs": ["decision"], "context_window": None, "input_per_million": 0.021, "output_per_million": 0}]}
|
||||
assert run(tool, action="models")["data"]["count"] == 3
|
||||
|
||||
|
||||
def test_feedback_posts_the_signal_and_returns_what_was_recorded():
|
||||
tool, gw = wired(UnifiedLLMTool)
|
||||
env = run(tool, action="feedback", request_id=COMPLETION["id"], signal="rating", rating=3)
|
||||
assert (gw.last.method, gw.last.url.path) == ("POST", "/v1/ai/feedback")
|
||||
assert body(gw.last) == {"request_id": COMPLETION["id"], "signal": "rating", "rating": 3}
|
||||
assert env["data"] == {"request_id": "ae7c311d-5c64-47da-886b-86ceb2fd185f", "reward": 1, "recorded": True}
|
||||
|
||||
|
||||
def test_feedback_refuses_what_the_server_would_and_reads_an_error_envelope():
|
||||
tool, gw = wired(UnifiedLLMTool)
|
||||
assert "signal must be one of" in run(tool, action="feedback", request_id="chatcmpl-1", signal="bogus")["error"]["message"]
|
||||
assert "rating must be 1, 2 or 3" in run(tool, action="feedback", request_id="chatcmpl-1", signal="rating", rating=4)["error"]["message"]
|
||||
assert gw.seen == []
|
||||
gw.answer = (200, {"status": "error", "msg": "request not found", "data": None})
|
||||
env = run(tool, action="feedback", request_id="chatcmpl-1", signal="down")
|
||||
assert not env["ok"] and env["error"]["message"] == "request not found"
|
||||
|
||||
|
||||
def test_kai_decide_posts_the_decision_and_returns_it():
|
||||
tool, gw = wired(KaiDecideTool)
|
||||
questions = {"team": {"type": "choice", "instructions": "Which team?", "criteria": {"account": "logins", "payments": "charges and refunds"}}}
|
||||
env = run(tool, state={"message": "charged twice"}, questions=questions)
|
||||
assert env["ok"], env
|
||||
assert env["data"] == DECISION
|
||||
assert (gw.last.method, gw.last.url.path) == ("POST", "/v1/decisions")
|
||||
assert body(gw.last) == {"model": "kai", "state": {"message": "charged twice"}, "questions": questions}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"question,said",
|
||||
[
|
||||
({"type": "vote"}, "type must be one of"),
|
||||
({"type": "choice", "criteria": ["a", "a"]}, "at least 2 labels, not 1"),
|
||||
({"type": "score", "criteria": []}, "at least 1 level"),
|
||||
({"type": "score", "criteria": ["a", None]}, "score level 1 is null"),
|
||||
({"type": "noul", "criteria": {"maybe": "x"}}, "not maybe"),
|
||||
],
|
||||
)
|
||||
def test_kai_decide_names_each_shape_the_contract_refuses(question, said):
|
||||
tool, gw = wired(KaiDecideTool)
|
||||
env = run(tool, state="s", questions={"q": question})
|
||||
assert not env["ok"] and said in env["error"]["message"]
|
||||
assert gw.seen == []
|
||||
|
||||
|
||||
def test_kai_decide_carries_the_plan_cap_refusal():
|
||||
tool, gw = wired(KaiDecideTool)
|
||||
gw.answer = (429, {"error": {"message": "Free plan: today's Kai requests are used.", "code": "free_plan_cap"}})
|
||||
env = run(tool, state="s", questions={"q": {"type": "noul"}})
|
||||
assert not env["ok"] and "429" in env["error"]["message"] and "free_plan_cap" in env["error"]["message"]
|
||||
|
||||
Reference in New Issue
Block a user