From 496f37a538cca76cff560bf1f96d1cbf9dc1ff66 Mon Sep 17 00:00:00 2001 From: Jeroen Schweitzer Date: Thu, 5 Feb 2026 07:19:29 +0100 Subject: [PATCH] feat: add Claude backend with automatic Ollama fallback (Claudification Phase 1) All agents now prefer Claude API when ANTHROPIC_API_KEY is configured, with automatic fallback to Ollama when offline or unconfigured. New src/anthropic/ module provides model selection via get_model() factory. Co-Authored-By: Claude Opus 4.5 --- .env.example | 11 +- CHANGELOG.md | 38 ++- PROJECT_CLAUDIFICATION.md | 401 ++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- requirements.txt | 5 +- src/agents/biographer/agent.py | 17 +- src/agents/housekeeper/agent.py | 17 +- src/agents/librarian/agent.py | 17 +- src/agents/steward/agent.py | 123 +++++++--- src/agents/tatlock.py | 84 ++----- src/anthropic/__init__.py | 17 ++ src/anthropic/model_selector.py | 151 ++++++++++++ src/core/config.py | 16 +- src/core/startup.py | 19 +- src/main.py | 6 +- tests/core/test_benchmarks.py | 352 ---------------------------- 16 files changed, 790 insertions(+), 486 deletions(-) create mode 100644 PROJECT_CLAUDIFICATION.md create mode 100644 src/anthropic/__init__.py create mode 100644 src/anthropic/model_selector.py delete mode 100644 tests/core/test_benchmarks.py diff --git a/.env.example b/.env.example index 26d9bbd..4a7719e 100644 --- a/.env.example +++ b/.env.example @@ -8,7 +8,14 @@ API_HOST=0.0.0.0 API_PORT=8000 API_PREFIX=/v1 -# Ollama Configuration +# Anthropic Configuration (Claude - preferred backend) +# Set ANTHROPIC_API_KEY to enable Claude as the default backend +# Without an API key, Tatlock uses Ollama exclusively +# ANTHROPIC_API_KEY=sk-ant-api03-your-key-here +ANTHROPIC_MODEL=claude-sonnet-4-20250514 +PREFER_CLOUD_BACKEND=true + +# Ollama Configuration (local fallback when Claude unavailable) OLLAMA_HOST=http://localhost:11434 OLLAMA_DEFAULT_MODEL=mistral-nemo:latest OLLAMA_TIMEOUT=120 @@ -21,7 +28,6 @@ SEARXNG_TIMEOUT=30 REDIS_HOST=localhost REDIS_PORT=6379 REDIS_MEMORY_DB=1 -REDIS_BENCHMARK_DB=6 REDIS_TIMEOUT=5 # Qdrant Configuration @@ -33,7 +39,6 @@ QDRANT_PORT=6333 # - development: DEBUG (maximum verbosity) # - production: WARNING (minimal noise) # Uncomment to override: LOG_LEVEL=INFO -ENABLE_BENCHMARKS=true # Note: Log format is auto-selected based on ENVIRONMENT (console for dev, json for production) # User Configuration diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d159ad..e3932e5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,40 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [2.0.0] - 2026-02-05 + +### Added + +- **Claude backend support (Claudification Phase 1)** - All agents now prefer Claude over Ollama + - New `src/anthropic/` module with model selector and health check + - `get_model()` factory returns Claude if available, Ollama as fallback + - Startup health check caches Claude API availability + - Configuration: `ANTHROPIC_API_KEY`, `ANTHROPIC_MODEL`, `PREFER_CLOUD_BACKEND` + - 200k token context when using Claude backend + +- **Steward dual-backend support** - Direct API calls to Claude or Ollama + - `_call_claude()`: Anthropic Messages API path + - `_call_ollama()`: Existing Ollama generate API path (preserved) + - Automatic fallback: if Claude call fails mid-request, retries with Ollama + +- **Claudification project tracking** - `PROJECT_CLAUDIFICATION.md` with Phase 1/2 roadmap + +### Changed + +- **All PydanticAI agents refactored to use `get_model()`**: + - Tatlock (6 instantiation locations) + - Librarian + - Biographer + - Housekeeper +- **`initialize_application()` is now async** - Supports async Claude health check at startup +- **Dependencies**: `pydantic-ai-slim[openai,anthropic]` replaces `pydantic-ai-slim[openai]` +- **Startup logging** now includes backend selection info (claude/ollama) +- **Agent creation logging** now includes backend and model info + +### Removed + +- Stale `tests/core/test_benchmarks.py` (benchmark system was removed in v1.10.0) + ## [1.11.0] - 2025-12-30 ### Added @@ -856,7 +890,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - CORS middleware - Exception handlers (OpenAI-compatible error format) -[Unreleased]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v1.10.0...main +[Unreleased]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v2.0.0...main +[2.0.0]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v1.11.0...v2.0.0 +[1.11.0]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v1.10.0...v1.11.0 [1.10.0]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v1.9.0...v1.10.0 [1.9.0]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v1.8.6...v1.9.0 [1.8.6]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v1.8.5...v1.8.6 diff --git a/PROJECT_CLAUDIFICATION.md b/PROJECT_CLAUDIFICATION.md new file mode 100644 index 0000000..91b5f39 --- /dev/null +++ b/PROJECT_CLAUDIFICATION.md @@ -0,0 +1,401 @@ +# Tatlock Enhancement Plan: Bidirectional Claude Integration + +## Executive Summary + +Implement a **bidirectional architecture** that: +1. **Superpowers Tatlock** by swapping Ollama→Claude backend (200k context, better reasoning, same butler personality) +2. **Exposes Tatlock as MCP server** for Claude instances on any device (phone, browser, desktop) + +This gives you the flexibility to use whichever AI is best/most accessible at any moment. + +## Key Insight: Blanket Backend Swap (Simpler Than Sidecar) + +Instead of adding a Claude "Analyst" sidecar agent, **swap the underlying model for ALL agents**: + +``` +CURRENT: TatlockAgent → OpenAIChatModel → OllamaProvider → Ollama (mistral-nemo) +PROPOSED: TatlockAgent → AnthropicModel → AnthropicProvider → Claude API + ↘ (fallback when offline) → OllamaProvider → Ollama +``` + +**Why this works:** +- PydanticAI natively supports Anthropic via `AnthropicModel` + `AnthropicProvider` +- The same `TATLOCK_SYSTEM_PROMPT` is passed to Claude - butler personality preserved +- Claude is **better** at following system prompts than mistral-nemo +- 200k context for ALL queries, not just "complex" ones +- Simpler architecture: no routing logic, no sidecar delegation + +--- + +## Research Findings + +### Industry Best Practices (2025-2026) + +**MCP Protocol Updates** ([MCP Spec Updates June 2025](https://auth0.com/blog/mcp-specs-update-all-about-auth/)): +- Streamable HTTP replaced SSE (March 2025) - better for cloud deployment +- OAuth 2.0 required for remote servers - MCP servers are OAuth Resource Servers +- Tool Output Schemas now available - better structured data handling +- MCP Registry launched (Sept 2025) - community server discovery + +**Community Patterns** ([Claude Code Router](https://github.com/musistudio/claude-code-router)): +- Task-based routing is becoming standard: route simple→local, complex→cloud +- Translation proxies bridge Anthropic Messages API ↔ OpenAI format +- Cost savings of up to 98% reported with smart routing + +**Home Automation MCP** ([ha-mcp](https://github.com/homeassistant-ai/ha-mcp)): +- Production-ready MCP servers exist for Home Assistant +- Support Claude Code, Gemini CLI, Open WebUI, VSCode, Cursor +- Pattern: expose local tools securely to remote AI clients + +**Remote MCP Access** ([mcp-remote](https://www.npmjs.com/package/mcp-remote)): +- Bridge local MCP servers to Claude Desktop/Browser via proxy +- Supports authentication headers for security +- Works with ngrok/Cloudflare Tunnel for HTTPS + +--- + +## Recommended Architecture + +``` +┌─────────────────────────────────────────────────────────────────────────────────────┐ +│ BIDIRECTIONAL TATLOCK-CLAUDE ARCHITECTURE │ +├─────────────────────────────────────────────────────────────────────────────────────┤ +│ │ +│ ╔═══════════════════════════════════════════════════════════════════════════════╗ │ +│ ║ SCENARIO A: Using Tatlock (Open WebUI, local apps) ║ │ +│ ║ ───────────────────────────────────────────────── ║ │ +│ ║ ║ │ +│ ║ Request → Steward → Tatlock → Tools + Expert Delegation ║ │ +│ ║ │ ║ │ +│ ║ ├─→ Librarian (Claude) → research, wiki, RAG ║ │ +│ ║ ├─→ Biographer (Claude) → memory, preferences ║ │ +│ ║ ├─→ Housekeeper (Claude) → home automation ║ │ +│ ║ └─→ All powered by Claude with Ollama fallback ║ │ +│ ║ ║ │ +│ ║ Butler personality preserved, 200k context for all queries ║ │ +│ ╚═══════════════════════════════════════════════════════════════════════════════╝ │ +│ │ +│ ╔═══════════════════════════════════════════════════════════════════════════════╗ │ +│ ║ SCENARIO B: Using Claude.ai / Claude Desktop / Phone ║ │ +│ ║ ──────────────────────────────────────────────────── ║ │ +│ ║ ║ │ +│ ║ Claude ──[MCP over HTTPS]──► Tatlock MCP Server → Household Tools ║ │ +│ ║ │ ║ │ +│ ║ ├─→ calculator, datetime ║ │ +│ ║ ├─→ web_search, wiki_search ║ │ +│ ║ ├─→ hybrid_search (RAG) ║ │ +│ ║ ├─→ memory_recall, store_insight ║ │ +│ ║ └─→ home_control (lights, climate) ║ │ +│ ║ ║ │ +│ ║ Full 200k context, your local tools accessible from anywhere ║ │ +│ ╚═══════════════════════════════════════════════════════════════════════════════╝ │ +│ │ +│ ╔═══════════════════════════════════════════════════════════════════════════════╗ │ +│ ║ SCENARIO C: Offline (internet down) ║ │ +│ ║ ────────────────────────────────── ║ │ +│ ║ ║ │ +│ ║ Tatlock operates fully locally with Ollama ║ │ +│ ║ • All tools work (except web search) ║ │ +│ ║ • Graceful degradation with same butler personality ║ │ +│ ╚═══════════════════════════════════════════════════════════════════════════════╝ │ +│ │ +└─────────────────────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## Implementation Plan + +### Phase 1: Blanket Backend Swap (Claude for All Agents) + +Replace Ollama with Claude as the default backend for all PydanticAI agents, with automatic offline fallback. + +**New Files:** +``` +src/anthropic/ +├── __init__.py +├── provider.py # Claude provider with health check +└── model_selector.py # Chooses Claude or Ollama based on availability +``` + +**Key Implementation (`src/anthropic/provider.py`):** +```python +from pydantic_ai.models.anthropic import AnthropicModel +from pydantic_ai.providers.anthropic import AnthropicProvider +from pydantic_ai.models.openai import OpenAIChatModel +from src.ollama.provider import get_ollama_provider +from src.core.config import config + +_anthropic_available: bool | None = None + +async def check_anthropic_health() -> bool: + """Check if Anthropic API is reachable.""" + global _anthropic_available + try: + from anthropic import AsyncAnthropic + client = AsyncAnthropic(api_key=config.ANTHROPIC_API_KEY) + await client.messages.create( + model=config.ANTHROPIC_MODEL, + max_tokens=1, + messages=[{"role": "user", "content": "hi"}] + ) + _anthropic_available = True + except Exception: + _anthropic_available = False + return _anthropic_available + +def get_model(prefer_cloud: bool = True): + """Get the best available model. Returns Claude if available, otherwise Ollama.""" + if prefer_cloud and config.ANTHROPIC_API_KEY and _anthropic_available: + provider = AnthropicProvider(api_key=config.ANTHROPIC_API_KEY) + return AnthropicModel( + model_name=config.ANTHROPIC_MODEL, + provider=provider, + ) + else: + return OpenAIChatModel( + model_name=config.OLLAMA_DEFAULT_MODEL, + provider=get_ollama_provider() + ) +``` + +**Modify TatlockAgent (`src/agents/tatlock.py`):** +```python +def _ensure_agent(self): + if self._agent is not None: + return + + from src.anthropic.model_selector import get_model + + model = get_model(prefer_cloud=True) + + self._agent = Agent( + model, + system_prompt=TATLOCK_SYSTEM_PROMPT, # Same butler personality! + ) + self._register_tools() +``` + +--- + +### Phase 2: MCP Server (Expose Tools to Claude) + +Create an MCP server that exposes Tatlock's household tools to external Claude instances. + +**New Files:** +``` +src/mcp/ +├── __init__.py +├── server.py # MCP server using mcp Python SDK +├── tool_adapters.py # Convert PydanticAI tools → MCP schemas +├── auth.py # API key authentication +└── transport.py # Streamable HTTP transport +``` + +**Docker Stack Addition (`stacks/agents.yml`):** +```yaml +tatlock-mcp: + image: git.schweitz.internal/jpmschweitzer/tatlock:latest + command: ["python", "-m", "src.mcp.server"] + ports: + - "8002:8002" + environment: + - MCP_AUTH_TOKEN=${MCP_AUTH_TOKEN} + networks: + - docker-dataplane +``` + +**Claude Desktop Configuration:** +```json +{ + "mcpServers": { + "tatlock": { + "command": "npx", + "args": ["mcp-remote", "https://mcp.schweitz.net/sse", "--header", "Authorization: Bearer ${MCP_AUTH_TOKEN}"] + } + } +} +``` + +--- + +## Files to Modify + +### Phase 1 - Backend Swap + +**New Files:** +| File | Purpose | +|------|---------| +| `src/anthropic/__init__.py` | Package init | +| `src/anthropic/provider.py` | Claude provider with health check | +| `src/anthropic/model_selector.py` | Choose Claude or Ollama based on availability | + +**Modified Files:** +| File | Changes | +|------|---------| +| `src/core/config.py` | Add `ANTHROPIC_API_KEY`, `ANTHROPIC_MODEL`, `PREFER_CLOUD_BACKEND` | +| `src/agents/tatlock.py` | Use `get_model()` instead of hardcoded Ollama | +| `src/agents/librarian/agent.py` | Use `get_model()` instead of hardcoded Ollama | +| `src/agents/biographer/agent.py` | Use `get_model()` instead of hardcoded Ollama | +| `src/agents/steward/agent.py` | Convert to PydanticAI or add Anthropic API support | +| `src/core/startup.py` | Add Anthropic health check on startup | +| `requirements.txt` | Add `anthropic>=0.40.0` | +| `.env.example` | Document new environment variables | + +### Phase 2 - MCP Server + +**New Files:** +| File | Purpose | +|------|---------| +| `src/mcp/__init__.py` | Package init | +| `src/mcp/server.py` | MCP server implementation | +| `src/mcp/tool_adapters.py` | PydanticAI → MCP schema conversion | +| `src/mcp/auth.py` | Token-based authentication | + +--- + +## Cost Analysis + +- **Claude API**: $5-30/month (10-50 calls/day, ~2k input + 1k output tokens/call) +- **MCP via Claude Pro**: Included in subscription +- **Total**: ~$10-80/month for full bidirectional integration + +--- + +## Verification Plan + +### Phase 1 Testing +```bash +# 1. Run with Claude backend +ANTHROPIC_API_KEY=your-key docker-compose up -d tatlock + +# 2. Verify Claude is being used +docker logs tatlock 2>&1 | grep -i "anthropic\|claude" + +# 3. Test butler personality +curl -X POST http://tatlock.schweitz.internal:8000/v1/responses \ + -H "Content-Type: application/json" \ + -d '{"model": "Tatlock", "input": "Hello, who are you?"}' + +# 4. Test offline fallback +ANTHROPIC_API_KEY="" docker-compose up -d tatlock +docker logs tatlock 2>&1 | grep -i "ollama\|fallback" +``` + +### Phase 2 Testing +```bash +# 1. Start MCP server +docker-compose up -d tatlock-mcp + +# 2. Test MCP endpoint +curl -X POST https://mcp.schweitz.net/tools/list \ + -H "Authorization: Bearer $MCP_AUTH_TOKEN" +``` + +--- + +## Implementation Priority + +1. **Phase 1: Backend Swap** (~1 week) + - Immediate value: 200k context for ALL queries + - Low risk: provider abstraction, graceful offline fallback + +2. **Phase 2: MCP Server** (~2-3 weeks) + - Enables cross-device access + - Bidirectional: Tatlock superpowered by Claude AND accessible to Claude + +--- + +## Future Phases (Optional) + +- **Phase 3: LiteLLM Gateway** - Unified endpoint for all models, config-driven routing +- **Phase 4: Multi-Provider** - Add OpenAI, Vertex AI, etc. +- **Phase 5: Smart Routing** - Context-aware model selection, cost ceiling enforcement + +--- + +## Offline Behavior + +| Scenario | Behavior | +|----------|----------| +| No API key | Use Ollama exclusively | +| API unreachable | Use Ollama, log warning | +| API rate limited | Fallback to Ollama | + +| Aspect | Claude | Ollama | +|--------|--------|--------| +| Context | 200k tokens | ~8k tokens | +| Latency | 1-3s (network) | 0.5-1s (local) | +| Personality | Preserved | Preserved | +| Tools | All work | All work | +| Cost | API charges | Free | + +--- + +## Implementation Status + +### Phase 1: Backend Swap - CODE COMPLETE (awaiting API access) + +- [x] Add Anthropic config settings to `src/core/config.py` +- [x] Add `pydantic-ai-slim[openai,anthropic]` to requirements.txt +- [x] Create `src/anthropic/` module (model_selector.py) +- [x] Add Claude health check to startup.py +- [x] Refactor all PydanticAI agents to use `get_model()` + - [x] Librarian + - [x] Biographer + - [x] Housekeeper + - [x] Tatlock (6 locations) +- [x] Add Claude API path to Steward agent (direct API calls) +- [x] Update `.env.example` with new variables +- [x] Test Ollama fallback (working) +- [ ] Test with Claude API key (blocked: no API access currently) + +**Note:** Implementation complete. Currently runs in Ollama-only mode. Will automatically use Claude when `ANTHROPIC_API_KEY` is configured. + +### Phase 2: MCP Server - NOT STARTED + +- [ ] Create `src/mcp/` module +- [ ] Tool adapters (PydanticAI → MCP schema) +- [ ] Authentication middleware +- [ ] Streamable HTTP transport +- [ ] Docker stack configuration + +--- + +## Related Repository Handovers + +Handover documents created in each repo: `PROJECT_CLAUDIFICATION_HANDOVER.md` + +### library-desk - HANDOVER CREATED + +- [x] Write handover document +- [ ] Review HybridRAG response size limits +- [ ] Review smart_create endpoint for Claude optimization +- [ ] Evaluate response formats for LLM consumption + +### core-api - HANDOVER CREATED + +- [x] Write handover document +- [ ] Review list_devices response format +- [ ] Review error messages for LLM consumption +- [ ] Evaluate rate limiting for faster Claude processing + +### portainer-core - HANDOVER CREATED (blocking for production) + +- [x] Write handover document +- [ ] Update stack with new environment variables +- [ ] Configure secrets management for API key +- [ ] Update CONTAINERS.md documentation + +### webber - HANDOVER CREATED + +- [x] Write handover document +- [ ] Review content truncation limits +- [ ] Evaluate extraction quality for LLM consumption + +### tatlock-ui - HANDOVER CREATED + +- [x] Write handover document +- [ ] Test streaming responses with Claude backend +- [ ] Test conversation history with larger context +- [ ] Verify tool call display and reasoning rendering diff --git a/pyproject.toml b/pyproject.toml index b50a1d8..90e2e75 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "tatlock" -version = "1.11.0" +version = "2.0.0" description = "OpenAI-compatible API with Ollama backend" requires-python = ">=3.12" dependencies = [] diff --git a/requirements.txt b/requirements.txt index a178b8a..6c96844 100644 --- a/requirements.txt +++ b/requirements.txt @@ -21,10 +21,9 @@ pydantic-settings>=2.12,<2.13 # AI/LLM integration # PydanticAI: Agent framework for using Pydantic with LLMs -# Using slim version with only openai extra (Ollama uses OpenAI-compatible API) -# This avoids installing SDKs for anthropic, cohere, google, groq, huggingface, etc. +# Using slim version with openai (Ollama) and anthropic (Claude) extras # See DEPENDENCY_SLIM.md for rollback instructions if this breaks -pydantic-ai-slim[openai]>=1.27,<1.28 +pydantic-ai-slim[openai,anthropic]>=1.27,<1.28 # HTTP client for Ollama communication # Latest: 0.28.1 - No known CVEs diff --git a/src/agents/biographer/agent.py b/src/agents/biographer/agent.py index e5003ed..71a4eb4 100644 --- a/src/agents/biographer/agent.py +++ b/src/agents/biographer/agent.py @@ -102,16 +102,10 @@ _biographer_agent: Optional[Agent[None, str]] = None def _create_biographer_agent() -> Agent[None, str]: """Create The Biographer PydanticAI agent.""" - from pydantic_ai.models.openai import OpenAIChatModel + from src.anthropic.model_selector import get_model - from src.ollama.provider import get_ollama_provider - - # Create Ollama model with sanitized provider - # (fixes 'content: null' issue with tool calls) - model = OpenAIChatModel( - model_name=config.OLLAMA_DEFAULT_MODEL, - provider=get_ollama_provider(), - ) + # Get best available model (Claude if available, else Ollama) + model = get_model() agent: Agent[None, str] = Agent( model=model, @@ -131,9 +125,12 @@ def _create_biographer_agent() -> Agent[None, str]: # Register management tools agent.tool_plain(forget_memory) + from src.anthropic.model_selector import get_model_info + model_info = get_model_info() logger.info( "biographer_agent_created", - model=config.OLLAMA_DEFAULT_MODEL, + backend=model_info["backend"], + model=model_info["model"], tool_count=6, ) diff --git a/src/agents/housekeeper/agent.py b/src/agents/housekeeper/agent.py index 4870a64..9bc9044 100644 --- a/src/agents/housekeeper/agent.py +++ b/src/agents/housekeeper/agent.py @@ -103,16 +103,10 @@ _housekeeper_agent: Optional[Agent[None, str]] = None def _create_housekeeper_agent() -> Agent[None, str]: """Create the Housekeeper PydanticAI agent.""" - from pydantic_ai.models.openai import OpenAIChatModel + from src.anthropic.model_selector import get_model - from src.ollama.provider import get_ollama_provider - - # Create Ollama model with sanitized provider - # (fixes 'content: null' issue with tool calls) - model = OpenAIChatModel( - model_name=config.OLLAMA_DEFAULT_MODEL, - provider=get_ollama_provider(), - ) + # Get best available model (Claude if available, else Ollama) + model = get_model() agent: Agent[None, str] = Agent( model=model, @@ -145,9 +139,12 @@ def _create_housekeeper_agent() -> Agent[None, str]: # Register history tools agent.tool_plain(get_history) + from src.anthropic.model_selector import get_model_info + model_info = get_model_info() logger.info( "housekeeper_agent_created", - model=config.OLLAMA_DEFAULT_MODEL, + backend=model_info["backend"], + model=model_info["model"], tool_count=13, ) diff --git a/src/agents/librarian/agent.py b/src/agents/librarian/agent.py index f51d9aa..063ab6d 100644 --- a/src/agents/librarian/agent.py +++ b/src/agents/librarian/agent.py @@ -143,16 +143,10 @@ _librarian_agent: Optional[Agent[None, str]] = None def _create_librarian_agent() -> Agent[None, str]: """Create the Librarian PydanticAI agent.""" - from pydantic_ai.models.openai import OpenAIChatModel + from src.anthropic.model_selector import get_model - from src.ollama.provider import get_ollama_provider - - # Create Ollama model with sanitized provider - # (fixes 'content: null' issue with tool calls) - model = OpenAIChatModel( - model_name=config.OLLAMA_DEFAULT_MODEL, - provider=get_ollama_provider(), - ) + # Get best available model (Claude if available, else Ollama) + model = get_model() agent: Agent[None, str] = Agent( model=model, @@ -182,9 +176,12 @@ def _create_librarian_agent() -> Agent[None, str]: agent.tool_plain(update_wiki_page) agent.tool_plain(smart_create_wiki_page) + from src.anthropic.model_selector import get_model_info + model_info = get_model_info() logger.info( "librarian_agent_created", - model=config.OLLAMA_DEFAULT_MODEL, + backend=model_info["backend"], + model=model_info["model"], tool_count=14, # 7 research + 3 web + 1 wiki read + 3 wiki write ) diff --git a/src/agents/steward/agent.py b/src/agents/steward/agent.py index 6a8715a..55a1fce 100644 --- a/src/agents/steward/agent.py +++ b/src/agents/steward/agent.py @@ -5,11 +5,13 @@ The Steward analyzes incoming requests, identifies relevant household capabilities, and provides focused recommendations to Tatlock (the Butler). This creates a two-tier architecture that prevents cognitive overload. -Uses plain text output (not JSON) for reliability with Ollama models. +Uses plain text output (not JSON) for reliability. Supports both Claude +(preferred) and Ollama (fallback) backends via direct API calls. """ import httpx from typing import Optional +from src.anthropic.model_selector import is_claude_available, get_model_info from src.core.config import config from src.core.household_registry import get_household_registry from src.core.logging_config import get_logger @@ -105,22 +107,74 @@ class StewardAgent: Analyzes requests with full conversation context and recommends which household capabilities the Butler should use. - Uses plain text output for reliability with Ollama models. + Uses plain text output for reliability. Supports both Claude + (preferred) and Ollama (fallback) backends via direct API calls. """ def __init__(self): - """Initialize Steward with Ollama model (same as Tatlock for VRAM efficiency).""" + """Initialize Steward with backend selection based on availability.""" + # Ollama config (fallback) self.ollama_host = str(config.OLLAMA_HOST).rstrip('/') - self.model_name = config.OLLAMA_DEFAULT_MODEL + self.ollama_model = config.OLLAMA_DEFAULT_MODEL + + # Claude config (preferred) + self.claude_model = config.ANTHROPIC_MODEL + self._anthropic_client = None + + # Determine which backend to use + self._use_claude = config.PREFER_CLOUD_BACKEND and is_claude_available() + self.timeout = 30.0 # 30 second timeout for analysis + model_info = get_model_info() logger.info( "steward_agent_created", - ollama_host=self.ollama_host, - model=self.model_name, + backend=model_info["backend"], + model=model_info["model"], timeout=self.timeout, ) + def _get_anthropic_client(self): + """Get or create Anthropic client (lazy initialization).""" + if self._anthropic_client is None: + from anthropic import AsyncAnthropic + self._anthropic_client = AsyncAnthropic(api_key=config.ANTHROPIC_API_KEY) + return self._anthropic_client + + async def _call_claude(self, system_prompt: str, user_message: str) -> str: + """Call Claude API directly for plain text generation.""" + client = self._get_anthropic_client() + + response = await client.messages.create( + model=self.claude_model, + max_tokens=1024, + system=system_prompt, + messages=[{"role": "user", "content": user_message}], + temperature=0.3, # Lower = more consistent + ) + + return response.content[0].text.strip() + + async def _call_ollama(self, prompt: str) -> str: + """Call Ollama API directly for plain text generation.""" + async with httpx.AsyncClient(timeout=self.timeout) as client: + response = await client.post( + f"{self.ollama_host}/api/generate", + json={ + "model": self.ollama_model, + "prompt": prompt, + "stream": False, + "options": { + "temperature": 0.3, # Lower = more consistent + "top_p": 0.9 + } + } + ) + + response.raise_for_status() + result = response.json() + return result["response"].strip() + async def analyze( self, query: str, @@ -129,6 +183,8 @@ class StewardAgent: """ Analyze query and return plain text recommendation. + Uses Claude if available, falls back to Ollama. + Args: query: User's query to analyze conversation_history: Previous conversation turns @@ -144,35 +200,48 @@ class StewardAgent: history = conversation_history or [] prompt = build_steward_prompt(query, history) - logger.debug("steward_calling_ollama", query_preview=query[:100]) + backend = "claude" if self._use_claude else "ollama" + logger.debug( + "steward_calling_llm", + backend=backend, + query_preview=query[:100], + ) - # Call Ollama API directly (more reliable than PydanticAI for plain text) - async with httpx.AsyncClient(timeout=self.timeout) as client: - response = await client.post( - f"{self.ollama_host}/api/generate", - json={ - "model": self.model_name, - "prompt": prompt, - "stream": False, - "options": { - "temperature": 0.3, # Lower = more consistent - "top_p": 0.9 - } - } - ) - - response.raise_for_status() - result = response.json() - - analysis_text = result["response"].strip() + try: + if self._use_claude: + # For Claude, split into system + user message + # The prompt contains both, but Claude prefers explicit system + analysis_text = await self._call_claude( + system_prompt="You are the Steward of the household, advising the Butler (Tatlock) on which capabilities to use. Be concise and specific.", + user_message=prompt, + ) + else: + analysis_text = await self._call_ollama(prompt) logger.debug( "steward_analysis_received", - text_preview=analysis_text[:150] + backend=backend, + text_preview=analysis_text[:150], ) return analysis_text + except Exception as e: + # If Claude fails, try Ollama as fallback + if self._use_claude: + logger.warning( + "steward_claude_fallback", + error=str(e), + ) + analysis_text = await self._call_ollama(prompt) + logger.debug( + "steward_analysis_received", + backend="ollama_fallback", + text_preview=analysis_text[:150], + ) + return analysis_text + raise + # Global Steward instance _steward_agent = None diff --git a/src/agents/tatlock.py b/src/agents/tatlock.py index 16cbfdd..0d6d3db 100644 --- a/src/agents/tatlock.py +++ b/src/agents/tatlock.py @@ -138,10 +138,7 @@ class TatlockAgent(AgentInterface): """ def __init__(self): - """Initialize Tatlock configuration (lazy agent creation).""" - # Store Ollama configuration - self.ollama_host = str(config.OLLAMA_HOST) - self.model_name = config.OLLAMA_DEFAULT_MODEL + """Initialize Tatlock (lazy agent creation).""" self._agent = None # Lazy initialization def _ensure_agent(self): @@ -149,30 +146,21 @@ class TatlockAgent(AgentInterface): if self._agent is not None: return + from src.anthropic.model_selector import get_model, get_model_info + + model_info = get_model_info() logger.info( "tatlock_agent_initializing", - ollama_host=self.ollama_host, - model=self.model_name, + backend=model_info["backend"], + model=model_info["model"], ) - # Import required classes for Ollama configuration - from pydantic_ai.models.openai import OpenAIChatModel - from src.ollama.provider import get_ollama_provider + # Get best available model (Claude if available, else Ollama) + model = get_model() - # PydanticAI expects Ollama base URL to end with /v1 - # Remove trailing slash from ollama_host if present - clean_host = self.ollama_host.rstrip('/') - base_url = f"{clean_host}/v1" - - # Create Ollama model with provider - ollama_model = OpenAIChatModel( - model_name=self.model_name, - provider=get_ollama_provider() - ) - - # Create PydanticAI agent with Ollama model + # Create PydanticAI agent self._agent = Agent( - ollama_model, + model, system_prompt=TATLOCK_SYSTEM_PROMPT, ) @@ -461,8 +449,7 @@ class TatlockAgent(AgentInterface): ... tool_tracker=tracker, ... ) """ - from pydantic_ai.models.openai import OpenAIChatModel - from src.ollama.provider import get_ollama_provider + from src.anthropic.model_selector import get_model logger.info( "tatlock_run_with_scoped_tools", @@ -473,18 +460,12 @@ class TatlockAgent(AgentInterface): # Create a fresh agent instance with scoped tools only # This ensures Tatlock can ONLY use tools recommended by the Steward - clean_host = self.ollama_host.rstrip('/') - base_url = f"{clean_host}/v1" - - ollama_model = OpenAIChatModel( - model_name=self.model_name, - provider=get_ollama_provider() - ) + model = get_model() # Create agent with scoped tools # Tools from household registry are already PydanticAI Tool objects scoped_agent = Agent( - ollama_model, + model, system_prompt=TATLOCK_SYSTEM_PROMPT, tools=scoped_tools, # Pass tools directly to Agent constructor ) @@ -555,8 +536,7 @@ class TatlockAgent(AgentInterface): Yields: Text chunks from the streaming response """ - from pydantic_ai.models.openai import OpenAIChatModel - from src.ollama.provider import get_ollama_provider + from src.anthropic.model_selector import get_model logger.info( "tatlock_run_with_scoped_tools_stream", @@ -566,17 +546,11 @@ class TatlockAgent(AgentInterface): ) # Create a fresh agent instance with scoped tools only - clean_host = self.ollama_host.rstrip('/') - base_url = f"{clean_host}/v1" - - ollama_model = OpenAIChatModel( - model_name=self.model_name, - provider=get_ollama_provider() - ) + model = get_model() # Create agent with scoped tools scoped_agent = Agent( - ollama_model, + model, system_prompt=TATLOCK_SYSTEM_PROMPT, tools=scoped_tools, ) @@ -650,8 +624,6 @@ class TatlockAgent(AgentInterface): - tool_outputs: Dict mapping tool names to their outputs - raw_output: The agent's raw text output """ - from pydantic_ai.models.openai import OpenAIChatModel - from src.ollama.provider import get_ollama_provider from pydantic_ai.settings import ModelSettings from pydantic_ai.messages import ( ModelRequest, @@ -661,6 +633,7 @@ class TatlockAgent(AgentInterface): ToolCallPart, ToolReturnPart, ) + from src.anthropic.model_selector import get_model logger.info( "tatlock_orchestrate_tool_calls", @@ -680,17 +653,11 @@ class TatlockAgent(AgentInterface): ) # Create a fresh agent instance with scoped tools only - clean_host = self.ollama_host.rstrip('/') - base_url = f"{clean_host}/v1" - - ollama_model = OpenAIChatModel( - model_name=self.model_name, - provider=get_ollama_provider() - ) + model = get_model() # Create agent with scoped tools scoped_agent = Agent( - ollama_model, + model, system_prompt=TATLOCK_SYSTEM_PROMPT, tools=scoped_tools, ) @@ -799,9 +766,8 @@ class TatlockAgent(AgentInterface): Returns: str: Butler-toned response synthesized from all results """ - from pydantic_ai.models.openai import OpenAIChatModel - from src.ollama.provider import get_ollama_provider from pydantic_ai.messages import ModelRequest, ModelResponse, UserPromptPart, TextPart + from src.anthropic.model_selector import get_model logger.info( "tatlock_synthesize_from_results", @@ -848,17 +814,11 @@ class TatlockAgent(AgentInterface): synthesis_prompt = "\n".join(synthesis_parts) # Create synthesis agent (no tools needed) - clean_host = self.ollama_host.rstrip('/') - base_url = f"{clean_host}/v1" - - ollama_model = OpenAIChatModel( - model_name=self.model_name, - provider=get_ollama_provider() - ) + model = get_model() # Synthesis agent uses butler prompt but no tools synthesis_agent = Agent( - ollama_model, + model, system_prompt=TATLOCK_SYSTEM_PROMPT, # No tools for synthesis phase ) diff --git a/src/anthropic/__init__.py b/src/anthropic/__init__.py new file mode 100644 index 0000000..e101c66 --- /dev/null +++ b/src/anthropic/__init__.py @@ -0,0 +1,17 @@ +""" +Anthropic/Claude integration module. + +Provides model selection with automatic fallback between Claude and Ollama. +""" + +from src.anthropic.model_selector import ( + check_claude_health, + get_model, + is_claude_available, +) + +__all__ = [ + "check_claude_health", + "get_model", + "is_claude_available", +] diff --git a/src/anthropic/model_selector.py b/src/anthropic/model_selector.py new file mode 100644 index 0000000..dc0aa12 --- /dev/null +++ b/src/anthropic/model_selector.py @@ -0,0 +1,151 @@ +""" +Model selector for Claude/Ollama backend switching. + +Provides automatic model selection with Claude as preferred backend +and Ollama as offline fallback. +""" + +from typing import Union + +from pydantic_ai.models.anthropic import AnthropicModel +from pydantic_ai.models.openai import OpenAIChatModel + +from src.core.config import config +from src.core.logging_config import get_logger + +logger = get_logger(__name__) + +# Cached health check result (set once at startup) +_claude_available: bool | None = None + + +async def check_claude_health() -> bool: + """ + Check if Claude API is reachable and working. + + This should be called once at application startup. + The result is cached in `_claude_available`. + + Returns: + True if Claude API is accessible, False otherwise. + """ + global _claude_available + + # No API key configured - Claude not available + if not config.ANTHROPIC_API_KEY: + logger.info( + "claude_health_check_skipped", + reason="no_api_key", + ) + _claude_available = False + return False + + try: + from anthropic import AsyncAnthropic + + client = AsyncAnthropic(api_key=config.ANTHROPIC_API_KEY) + + # Minimal API call to verify connectivity + # Using a tiny max_tokens to minimize cost + await client.messages.create( + model=config.ANTHROPIC_MODEL, + max_tokens=1, + messages=[{"role": "user", "content": "hi"}], + ) + + _claude_available = True + logger.info( + "claude_health_check_passed", + model=config.ANTHROPIC_MODEL, + ) + return True + + except Exception as e: + _claude_available = False + logger.warning( + "claude_health_check_failed", + error=str(e), + model=config.ANTHROPIC_MODEL, + ) + return False + + +def is_claude_available() -> bool: + """ + Check if Claude is available (from cached health check result). + + Returns: + True if Claude API was reachable at startup, False otherwise. + + Note: + Returns False if health check hasn't been run yet. + Call `check_claude_health()` at startup first. + """ + return _claude_available is True + + +def get_model(prefer_cloud: bool | None = None) -> Union[AnthropicModel, OpenAIChatModel]: + """ + Get the best available model. + + Returns Claude if available and preferred, otherwise Ollama. + + Args: + prefer_cloud: Override config.PREFER_CLOUD_BACKEND for this call. + If None, uses the config value. + + Returns: + PydanticAI model instance (AnthropicModel or OpenAIChatModel). + + Example: + >>> model = get_model() + >>> agent = Agent(model, system_prompt="...") + """ + # Determine preference + use_cloud = prefer_cloud if prefer_cloud is not None else config.PREFER_CLOUD_BACKEND + + # Use Claude if available and preferred + if use_cloud and is_claude_available(): + logger.debug( + "model_selected", + backend="claude", + model=config.ANTHROPIC_MODEL, + ) + return AnthropicModel( + model_name=config.ANTHROPIC_MODEL, + api_key=config.ANTHROPIC_API_KEY, + ) + + # Fall back to Ollama + from src.ollama.provider import get_ollama_provider + + logger.debug( + "model_selected", + backend="ollama", + model=config.OLLAMA_DEFAULT_MODEL, + reason="fallback" if use_cloud else "preferred_local", + ) + return OpenAIChatModel( + model_name=config.OLLAMA_DEFAULT_MODEL, + provider=get_ollama_provider(), + ) + + +def get_model_info() -> dict: + """ + Get information about the current model configuration. + + Useful for health checks and debugging. + + Returns: + Dict with backend, model name, and availability info. + """ + use_cloud = config.PREFER_CLOUD_BACKEND and is_claude_available() + + return { + "backend": "claude" if use_cloud else "ollama", + "model": config.ANTHROPIC_MODEL if use_cloud else config.OLLAMA_DEFAULT_MODEL, + "claude_available": is_claude_available(), + "claude_configured": bool(config.ANTHROPIC_API_KEY), + "prefer_cloud": config.PREFER_CLOUD_BACKEND, + } diff --git a/src/core/config.py b/src/core/config.py index 267fa87..cf51257 100644 --- a/src/core/config.py +++ b/src/core/config.py @@ -64,7 +64,21 @@ class Config(BaseSettings): API_PORT: int = Field(default=8000, description="API port") API_PREFIX: str = Field(default="/v1", description="API route prefix") - # Ollama Configuration + # Anthropic Configuration (Claude - preferred backend) + ANTHROPIC_API_KEY: str | None = Field( + default=None, + description="Anthropic API key for Claude access" + ) + ANTHROPIC_MODEL: str = Field( + default="claude-sonnet-4-20250514", + description="Claude model to use" + ) + PREFER_CLOUD_BACKEND: bool = Field( + default=True, + description="Prefer Claude over Ollama when available" + ) + + # Ollama Configuration (local fallback) OLLAMA_HOST: HttpUrl = Field( default="http://localhost:11434", description="Ollama server URL" diff --git a/src/core/startup.py b/src/core/startup.py index 377998a..185453a 100644 --- a/src/core/startup.py +++ b/src/core/startup.py @@ -9,6 +9,7 @@ from src.agents.biographer import register_biographer from src.agents.housekeeper import register_housekeeper from src.agents.librarian import register_librarian from src.agents.tatlock_core import TATLOCK_CORE_CAPABILITY, tatlock_core_tools +from src.anthropic.model_selector import check_claude_health, get_model_info from src.core.household_registry import get_household_registry from src.core.logging_config import get_logger @@ -81,19 +82,29 @@ def register_household_members(): ) -def initialize_application(): +async def initialize_application(): """ Initialize the application. Performs all startup tasks: - 1. Register household members - 2. (Future) Initialize connections - 3. (Future) Load configuration + 1. Check Claude API health (for backend selection) + 2. Register household members + 3. (Future) Initialize connections This should be called once during application startup. """ logger.info("application_initialization_starting") + # Check Claude API health for backend selection + await check_claude_health() + model_info = get_model_info() + logger.info( + "model_backend_configured", + backend=model_info["backend"], + model=model_info["model"], + claude_available=model_info["claude_available"], + ) + # Register household members register_household_members() diff --git a/src/main.py b/src/main.py index d3028af..1f4bf4f 100644 --- a/src/main.py +++ b/src/main.py @@ -44,14 +44,16 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: app_name=config.APP_NAME, version=config.APP_VERSION, environment=config.ENVIRONMENT.value, + prefer_cloud=config.PREFER_CLOUD_BACKEND, + anthropic_model=config.ANTHROPIC_MODEL, ollama_host=str(config.OLLAMA_HOST), ollama_model=config.OLLAMA_DEFAULT_MODEL, redis_url=config.redis_memory_url, log_format=config.log_format, ) - # Initialize application (register household members, etc.) - initialize_application() + # Initialize application (check Claude health, register household members, etc.) + await initialize_application() yield diff --git a/tests/core/test_benchmarks.py b/tests/core/test_benchmarks.py deleted file mode 100644 index def6946..0000000 --- a/tests/core/test_benchmarks.py +++ /dev/null @@ -1,352 +0,0 @@ -""" -Tests for benchmark storage. - -Tests performance tracking, Redis storage, and analytics features. -""" -import json -from datetime import datetime, timedelta, timezone -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from src.core.benchmarks import ( - BenchmarkStore, - PerformanceBenchmark, - get_benchmark_store, -) - - -class TestPerformanceBenchmark: - """Test PerformanceBenchmark model.""" - - def test_benchmark_creation(self): - """Test creating a performance benchmark.""" - benchmark = PerformanceBenchmark( - operation="steward_analysis", - duration_seconds=1.23, - success=True, - recommendation_count=3, - ) - - assert benchmark.operation == "steward_analysis" - assert benchmark.duration_seconds == 1.23 - assert benchmark.success is True - assert benchmark.recommendation_count == 3 - assert isinstance(benchmark.timestamp, datetime) - - def test_benchmark_with_tool_fields(self): - """Test benchmark with tool-specific fields.""" - benchmark = PerformanceBenchmark( - operation="tool_call", - duration_seconds=0.5, - success=True, - tool_name="calculate", - was_recommended=True, - was_actually_used=True, - ) - - assert benchmark.tool_name == "calculate" - assert benchmark.was_recommended is True - assert benchmark.was_actually_used is True - - def test_benchmark_to_redis_dict(self): - """Test conversion to Redis dict.""" - benchmark = PerformanceBenchmark( - operation="test_op", - duration_seconds=1.0, - success=True, - metadata={"key": "value"}, - ) - - redis_dict = benchmark.to_redis_dict() - assert redis_dict["operation"] == "test_op" - assert redis_dict["duration_seconds"] == 1.0 - assert redis_dict["success"] == "True" # Booleans stored as strings in Redis - assert isinstance(redis_dict["timestamp"], str) - assert isinstance(redis_dict["metadata"], str) - - def test_benchmark_from_redis_dict(self): - """Test reconstruction from Redis dict.""" - now = datetime.now(timezone.utc) - redis_dict = { - "timestamp": now.isoformat(), - "operation": "test_op", - "duration_seconds": 1.5, - "success": "True", # Booleans stored as strings in Redis - "metadata": json.dumps({"test": "data"}), - "recommendation_count": None, - "confidence": None, - "tool_name": None, - "was_recommended": None, - "was_actually_used": None, - "conversation_id": None, - } - - benchmark = PerformanceBenchmark.from_redis_dict(redis_dict) - assert benchmark.operation == "test_op" - assert benchmark.duration_seconds == 1.5 - assert benchmark.success is True # Converted back to bool - assert benchmark.metadata == {"test": "data"} - - -class TestBenchmarkStore: - """Test BenchmarkStore functionality.""" - - @pytest.fixture - def mock_redis(self): - """Create mock Redis client.""" - mock = AsyncMock() - mock.hset = AsyncMock() - mock.expire = AsyncMock() - mock.zadd = AsyncMock() - mock.zrevrangebyscore = AsyncMock(return_value=[]) - mock.hgetall = AsyncMock(return_value={}) - mock.aclose = AsyncMock() - return mock - - @pytest.fixture - def store(self, mock_redis): - """Create benchmark store with mock Redis.""" - return BenchmarkStore(redis_client=mock_redis) - - @pytest.mark.asyncio - async def test_record_benchmark(self, store, mock_redis): - """Test recording a benchmark.""" - benchmark = PerformanceBenchmark( - operation="test_op", - duration_seconds=1.0, - success=True, - ) - - await store.record(benchmark) - - # Verify Redis calls - mock_redis.hset.assert_called_once() - mock_redis.expire.assert_called() - mock_redis.zadd.assert_called_once() - - @pytest.mark.asyncio - async def test_record_benchmark_disabled(self, mock_redis): - """Test recording when benchmarks are disabled.""" - with patch("src.core.benchmarks.config.ENABLE_BENCHMARKS", False): - store = BenchmarkStore(redis_client=mock_redis) - benchmark = PerformanceBenchmark( - operation="test_op", - duration_seconds=1.0, - success=True, - ) - - await store.record(benchmark) - - # Should not call Redis - mock_redis.hset.assert_not_called() - - @pytest.mark.asyncio - async def test_record_benchmark_handles_errors(self, store, mock_redis): - """Test recording handles Redis errors gracefully.""" - mock_redis.hset.side_effect = Exception("Redis error") - - benchmark = PerformanceBenchmark( - operation="test_op", - duration_seconds=1.0, - success=True, - ) - - # Should not raise exception - await store.record(benchmark) - - @pytest.mark.asyncio - async def test_query_benchmarks(self, store, mock_redis): - """Test querying benchmarks.""" - # Setup mock data - now = datetime.now(timezone.utc) - mock_key = f"benchmark:test_op:{int(now.timestamp() * 1000)}" - mock_redis.zrevrangebyscore.return_value = [mock_key] - - # Mock hgetall to return proper data (booleans as strings, like Redis) - mock_redis.hgetall.return_value = { - "timestamp": now.isoformat(), - "operation": "test_op", - "duration_seconds": 1.5, # Numeric, not string - "success": "True", # Booleans stored as strings in Redis - "metadata": "{}", - "recommendation_count": None, - "confidence": None, - "tool_name": None, - "was_recommended": None, - "was_actually_used": None, - "conversation_id": None, - } - - results = await store.query("test_op", limit=10) - - assert len(results) == 1 - assert results[0].operation == "test_op" - mock_redis.zrevrangebyscore.assert_called_once() - - @pytest.mark.asyncio - async def test_query_with_time_range(self, store, mock_redis): - """Test querying with time range.""" - now = datetime.now(timezone.utc) - start_time = now - timedelta(hours=1) - end_time = now - - await store.query("test_op", start_time=start_time, end_time=end_time) - - # Verify time range was converted to timestamps - call_args = mock_redis.zrevrangebyscore.call_args - assert call_args is not None - - @pytest.mark.asyncio - async def test_query_disabled_benchmarks(self, mock_redis): - """Test querying when benchmarks are disabled.""" - with patch("src.core.benchmarks.config.ENABLE_BENCHMARKS", False): - store = BenchmarkStore(redis_client=mock_redis) - results = await store.query("test_op") - assert results == [] - - @pytest.mark.asyncio - async def test_query_handles_errors(self, store, mock_redis): - """Test query handles errors gracefully.""" - mock_redis.zrevrangebyscore.side_effect = Exception("Redis error") - - results = await store.query("test_op") - assert results == [] - - @pytest.mark.asyncio - async def test_get_statistics(self, store, mock_redis): - """Test getting statistics.""" - # Setup mock data with multiple benchmarks - now = datetime.now(timezone.utc) - mock_keys = [ - f"benchmark:test_op:{int((now - timedelta(seconds=i)).timestamp() * 1000)}" - for i in range(3) - ] - mock_redis.zrevrangebyscore.return_value = mock_keys - - # Return different durations and success values - benchmarks_data = [ - {"duration_seconds": "1.0", "success": "True"}, - {"duration_seconds": "2.0", "success": "True"}, - {"duration_seconds": "3.0", "success": "False"}, - ] - - async def mock_hgetall(key): - idx = mock_keys.index(key) - data = benchmarks_data[idx] - return { - "timestamp": now.isoformat(), - "operation": "test_op", - "duration_seconds": float(data["duration_seconds"]), - "success": data["success"], # Pass string through, from_redis_dict converts - "metadata": "{}", - "recommendation_count": None, - "confidence": None, - "tool_name": None, - "was_recommended": None, - "was_actually_used": None, - "conversation_id": None, - } - - mock_redis.hgetall.side_effect = mock_hgetall - - stats = await store.get_statistics("test_op") - - assert stats["count"] == 3 - assert stats["avg_duration"] == 2.0 # (1 + 2 + 3) / 3 - assert stats["min_duration"] == 1.0 - assert stats["max_duration"] == 3.0 - assert stats["success_rate"] == pytest.approx(66.67, rel=0.01) - assert stats["total_successes"] == 2 - assert stats["total_failures"] == 1 - - @pytest.mark.asyncio - async def test_get_statistics_empty(self, store, mock_redis): - """Test statistics with no data.""" - mock_redis.zrevrangebyscore.return_value = [] - - stats = await store.get_statistics("test_op") - - assert stats["count"] == 0 - assert stats["avg_duration"] == 0.0 - assert stats["success_rate"] == 0.0 - - @pytest.mark.asyncio - async def test_get_tool_accuracy(self, store, mock_redis): - """Test tool accuracy calculation.""" - # Setup mock data - now = datetime.now(timezone.utc) - mock_keys = [ - f"benchmark:tool_call:{int((now - timedelta(seconds=i)).timestamp() * 1000)}" - for i in range(4) - ] - mock_redis.zrevrangebyscore.return_value = mock_keys - - # Different combinations of recommended/used - tool_data = [ - {"was_recommended": "True", "was_actually_used": "True"}, # Good - {"was_recommended": "True", "was_actually_used": "True"}, # Good - {"was_recommended": "False", "was_actually_used": "True"}, # Missed - {"was_recommended": "True", "was_actually_used": "False"}, # Not used - ] - - async def mock_hgetall(key): - idx = mock_keys.index(key) - data = tool_data[idx] - return { - "timestamp": now.isoformat(), - "operation": "tool_call", - "duration_seconds": 1.0, - "success": "True", # Booleans stored as strings in Redis - "metadata": "{}", - "recommendation_count": None, - "confidence": None, - "tool_name": "test_tool", - "conversation_id": None, - "was_recommended": data["was_recommended"], # Already strings - "was_actually_used": data["was_actually_used"], # Already strings - } - - mock_redis.hgetall.side_effect = mock_hgetall - - accuracy = await store.get_tool_accuracy() - - assert accuracy["total_calls"] == 4 - assert accuracy["total_used"] == 3 - assert accuracy["recommended_and_used"] == 2 - assert accuracy["not_recommended_but_used"] == 1 - assert accuracy["precision"] == pytest.approx(66.67, rel=0.01) - - @pytest.mark.asyncio - async def test_get_tool_accuracy_empty(self, store, mock_redis): - """Test tool accuracy with no data.""" - mock_redis.zrevrangebyscore.return_value = [] - - accuracy = await store.get_tool_accuracy() - - assert accuracy["total_calls"] == 0 - assert accuracy["precision"] == 0.0 - - @pytest.mark.asyncio - async def test_close(self, store, mock_redis): - """Test closing the store.""" - await store.close() - mock_redis.aclose.assert_called_once() - - # Client should be None after close - assert store._client is None - - -class TestGlobalBenchmarkStore: - """Test global benchmark store instance.""" - - def test_get_benchmark_store(self): - """Test getting global store instance.""" - store = get_benchmark_store() - assert isinstance(store, BenchmarkStore) - - def test_get_benchmark_store_singleton(self): - """Test store is singleton.""" - store1 = get_benchmark_store() - store2 = get_benchmark_store() - assert store1 is store2