remove obsolete core-ai
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -1,18 +0,0 @@
|
||||
# Use a Python base image
|
||||
FROM python:3.12-slim-bookworm
|
||||
|
||||
# Set working directory
|
||||
WORKDIR /app
|
||||
|
||||
# Copy requirements file and install dependencies
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Copy the rest of the application code
|
||||
COPY . .
|
||||
|
||||
# Expose the port the app runs on
|
||||
EXPOSE 8084
|
||||
|
||||
# Run the application
|
||||
CMD ["python", "main.py"]
|
||||
@@ -1,535 +0,0 @@
|
||||
# Core-AI Service
|
||||
|
||||
AI agent service built on PydanticAI for infrastructure management and automation.
|
||||
|
||||
## Architecture
|
||||
|
||||
Core-AI provides two agents with distinct capabilities:
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Core-AI Service │
|
||||
├─────────────────────────────────────────────────────────────┤
|
||||
│ │
|
||||
│ ┌──────────────────┐ ┌──────────────────────┐ │
|
||||
│ │ PydanticAgent │ │ SimpleLiteLLMAgent │ │
|
||||
│ │ (Primary) │ │ (Fallback) │ │
|
||||
│ │ │ │ │ │
|
||||
│ │ • Tool calling │ │ • No tools │ │
|
||||
│ │ • Memory (3-tier)│ │ • Direct LiteLLM │ │
|
||||
│ │ • OpenAPI tools │ │ • Minimal overhead │ │
|
||||
│ └────────┬─────────┘ └──────────┬───────────┘ │
|
||||
│ │ │ │
|
||||
│ └──────────┬───────────────────┘ │
|
||||
│ │ │
|
||||
│ ▼ │
|
||||
│ ┌──────────────────────┐ │
|
||||
│ │ PydanticAI Runtime │ │
|
||||
│ │ (Ollama backend) │ │
|
||||
│ └──────────────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### PydanticAgent (Primary)
|
||||
|
||||
**Endpoint:** `/v1/chat/completions` (default)
|
||||
|
||||
Advanced agent using the PydanticAI framework with:
|
||||
|
||||
- **Tool Calling:** Automatic function calling with proper validation
|
||||
- **Memory System:** 3-tier conversation memory (buffer + Qdrant)
|
||||
- **Local Tools:** Time, calculations, web search (SearXNG)
|
||||
- **OpenAPI Tools:** Auto-discovered from core-api infrastructure endpoints
|
||||
- **Streaming Support:** Server-sent events for real-time responses
|
||||
|
||||
### SimpleLiteLLMAgent (Fallback)
|
||||
|
||||
**Endpoint:** `/v1/chat/simple`
|
||||
|
||||
Lightweight agent for direct LLM interaction:
|
||||
|
||||
- **No Tools:** Pure conversational mode
|
||||
- **Direct LiteLLM:** Minimal abstraction layer
|
||||
- **No Memory:** Stateless request/response
|
||||
- **Low Latency:** Fastest response times
|
||||
|
||||
## Tool System
|
||||
|
||||
### Local Tools
|
||||
|
||||
Built-in utilities available immediately (defined in `src/tools/local.py`):
|
||||
|
||||
- `get_current_time(timezone)` - Timezone-aware time with IANA timezone support
|
||||
- `get_current_date()` - Current date in ISO format
|
||||
- `calculate(expression)` - Safe mathematical calculations
|
||||
- `calculate_date_difference(date1, date2)` - Date arithmetic
|
||||
- `add_days_to_date(date, days)` - Date manipulation
|
||||
- `web_search(query, category, max_results)` - SearXNG metasearch integration
|
||||
|
||||
### OpenAPI Discovery
|
||||
|
||||
Dynamically discovers infrastructure tools from core-api's OpenAPI spec:
|
||||
|
||||
- **Auto-Discovery:** Fetches `/openapi.json` on startup
|
||||
- **REST Mapping:** Converts endpoints to callable functions
|
||||
- **Prefixed Names:** Tools prefixed with service name (e.g., `core-api__list_containers`)
|
||||
- **Type Safety:** Preserves parameter types and validation
|
||||
|
||||
**Configuration:**
|
||||
```bash
|
||||
OPENAPI_ENABLED=true
|
||||
OPENAPI_ENDPOINTS=http://core-api:8083/openapi.json
|
||||
```
|
||||
|
||||
**List Available Tools:**
|
||||
```bash
|
||||
curl http://localhost:8086/v1/tools
|
||||
```
|
||||
|
||||
## Memory System
|
||||
|
||||
3-tier multi-tenant memory with per-user data isolation:
|
||||
|
||||
### Tier 1: Conversation Buffer (RAM)
|
||||
- **Storage:** In-memory per-user buffers
|
||||
- **Scope:** Recent N turns (configurable, default: 10)
|
||||
- **Speed:** Instant access
|
||||
- **Purpose:** Fast context for ongoing conversations
|
||||
|
||||
### Tier 2: Persistent Storage (Qdrant)
|
||||
- **Storage:** Per-user Qdrant collections
|
||||
- **Scope:** Complete conversation history
|
||||
- **Speed:** Fast retrieval by conversation ID
|
||||
- **Purpose:** Conversation continuity across sessions
|
||||
|
||||
### Tier 3: Semantic Search (Qdrant)
|
||||
- **Storage:** Same as Tier 2 with vector embeddings
|
||||
- **Scope:** Cross-conversation semantic search
|
||||
- **Speed:** Sub-second similarity search
|
||||
- **Purpose:** Contextual recall across all user conversations
|
||||
|
||||
### Multi-Tenancy
|
||||
|
||||
- **Per-User Collections:** Each user gets isolated Qdrant collection
|
||||
- **User ID Format:** Sanitized email (`username_at_domain_com`)
|
||||
- **GDPR Compliance:** Complete user data deletion support
|
||||
- **Automatic Isolation:** No cross-user data leakage
|
||||
|
||||
**Memory Configuration:**
|
||||
```bash
|
||||
MEMORY_ENABLED=true
|
||||
MEMORY_TIER1_SIZE=10
|
||||
QDRANT_URL=http://qdrant:6333
|
||||
EMBEDDING_MODEL=nomic-embed-text
|
||||
DEFAULT_USER_ID=llmdefault_at_schweitz_net
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
### Chat Completions
|
||||
|
||||
**POST /v1/chat/completions** (Default: PydanticAI)
|
||||
|
||||
OpenAI-compatible chat endpoint using PydanticAgent.
|
||||
|
||||
**Request:**
|
||||
```json
|
||||
{
|
||||
"messages": [
|
||||
{"role": "user", "content": "What containers are running?"}
|
||||
],
|
||||
"conversation_id": "optional-conversation-id",
|
||||
"enable_tools": true,
|
||||
"stream": false
|
||||
}
|
||||
```
|
||||
|
||||
**Response:**
|
||||
```json
|
||||
{
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "I found 5 running containers..."
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"model": "pydantic",
|
||||
"tools_enabled": true,
|
||||
"tools_count": 12
|
||||
}
|
||||
```
|
||||
|
||||
**Streaming:** Set `"stream": true` for SSE response
|
||||
|
||||
### Simple Chat
|
||||
|
||||
**POST /v1/chat/simple**
|
||||
|
||||
No-tools fallback endpoint using SimpleLiteLLMAgent.
|
||||
|
||||
Same request/response format as above, but `tools_enabled` will be `false`.
|
||||
|
||||
### List Models
|
||||
|
||||
**GET /v1/models**
|
||||
|
||||
Returns available agent types:
|
||||
- `pydantic` - PydanticAgent (primary)
|
||||
- `simple` - SimpleLiteLLMAgent (fallback)
|
||||
|
||||
### List Tools
|
||||
|
||||
**GET /v1/tools**
|
||||
|
||||
Returns all available tools (local + discovered OpenAPI tools).
|
||||
|
||||
### Health Check
|
||||
|
||||
**GET /health**
|
||||
|
||||
Service health status with agent availability.
|
||||
|
||||
## Configuration
|
||||
|
||||
All configuration via environment variables (see `src/config.py`):
|
||||
|
||||
### Core Settings
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `HOST` | `0.0.0.0` | Server host |
|
||||
| `PORT` | `8086` | Server port |
|
||||
| `LOG_LEVEL` | `INFO` | Logging level |
|
||||
|
||||
### Ollama Integration
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `OLLAMA_BASE_URL` | `http://ollama:11434` | Ollama API URL |
|
||||
| `AGENT_MODEL` | `mistral-nemo:latest` | Primary model (tool-calling optimized) |
|
||||
| `OLLAMA_TIMEOUT` | `300` | Request timeout (seconds) |
|
||||
|
||||
### System Prompts
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `SYSTEM_PROMPT_VARIANT` | `minimal_agent` | Prompt for SimpleLiteLLMAgent |
|
||||
| `PYDANTIC_SYSTEM_PROMPT_VARIANT` | `pydantic_agent` | Prompt for PydanticAgent |
|
||||
|
||||
### Tool Discovery
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `OPENAPI_ENABLED` | `true` | Enable OpenAPI tool discovery |
|
||||
| `OPENAPI_ENDPOINTS` | `http://core-api:8083/openapi.json` | OpenAPI spec URLs (comma-separated) |
|
||||
|
||||
### Memory System
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `MEMORY_ENABLED` | `true` | Enable conversation memory |
|
||||
| `MEMORY_TIER1_SIZE` | `10` | Max turns in RAM buffer |
|
||||
| `QDRANT_URL` | `http://qdrant:6333` | Qdrant vector DB URL |
|
||||
| `QDRANT_COLLECTION_PREFIX` | `core_ai_user` | Prefix for user collections |
|
||||
| `EMBEDDING_MODEL` | `nomic-embed-text` | Ollama embedding model |
|
||||
| `EMBEDDING_DIMENSION` | `768` | Embedding vector size |
|
||||
| `DEFAULT_USER_ID` | `llmdefault_at_schweitz_net` | Default user (until auth integration) |
|
||||
|
||||
## Quick Start
|
||||
|
||||
### 1. Install Dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### 2. Configure Environment
|
||||
|
||||
Create `.env` file:
|
||||
|
||||
```bash
|
||||
OLLAMA_BASE_URL=http://ollama:11434
|
||||
AGENT_MODEL=mistral-nemo:latest
|
||||
QDRANT_URL=http://qdrant:6333
|
||||
MEMORY_ENABLED=true
|
||||
OPENAPI_ENABLED=true
|
||||
```
|
||||
|
||||
### 3. Start Service
|
||||
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
|
||||
Service available at `http://localhost:8086`
|
||||
|
||||
### 4. Test Chat
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8086/v1/chat/completions \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": "What time is it in Amsterdam?"}
|
||||
],
|
||||
"enable_tools": true
|
||||
}'
|
||||
```
|
||||
|
||||
The agent will automatically use the `get_current_time` tool.
|
||||
|
||||
## Testing
|
||||
|
||||
### Unit Tests
|
||||
|
||||
```bash
|
||||
# Run all tests
|
||||
pytest tests/ -v
|
||||
|
||||
# Run specific test suite
|
||||
pytest tests/test_ai_flow_quality.py -v
|
||||
|
||||
# Run with coverage
|
||||
pytest tests/ --cov=src --cov-report=html
|
||||
```
|
||||
|
||||
### Integration Tests
|
||||
|
||||
Quality tests for end-to-end AI flows:
|
||||
|
||||
```bash
|
||||
pytest tests/test_ai_flow_quality.py -v
|
||||
```
|
||||
|
||||
See `tests/QUALITY_TESTS.md` for test documentation.
|
||||
|
||||
### Manual Testing
|
||||
|
||||
```bash
|
||||
# Test PydanticAgent (with tools)
|
||||
curl -X POST http://localhost:8086/v1/chat/completions \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"messages": [{"role": "user", "content": "Calculate 123 * 456"}]}'
|
||||
|
||||
# Test SimpleLiteLLMAgent (no tools)
|
||||
curl -X POST http://localhost:8086/v1/chat/simple \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"messages": [{"role": "user", "content": "Hello!"}]}'
|
||||
|
||||
# List available tools
|
||||
curl http://localhost:8086/v1/tools
|
||||
|
||||
# Health check
|
||||
curl http://localhost:8086/health
|
||||
```
|
||||
|
||||
## Docker Deployment
|
||||
|
||||
### Build
|
||||
|
||||
```bash
|
||||
docker build -t core-ai:latest .
|
||||
```
|
||||
|
||||
### Run
|
||||
|
||||
```bash
|
||||
docker run -d \
|
||||
--name core-ai \
|
||||
-p 8086:8086 \
|
||||
-e OLLAMA_BASE_URL=http://ollama:11434 \
|
||||
-e QDRANT_URL=http://qdrant:6333 \
|
||||
-e AGENT_MODEL=mistral-nemo:latest \
|
||||
--network docker-dataplane \
|
||||
core-ai:latest
|
||||
```
|
||||
|
||||
### Using Docker Compose
|
||||
|
||||
```bash
|
||||
docker-compose -f ../../stacks/core-ai.yml up
|
||||
```
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
services/core-ai/
|
||||
├── main.py # HTTP server (aiohttp)
|
||||
├── src/
|
||||
│ ├── agents/
|
||||
│ │ ├── __init__.py # Agent exports
|
||||
│ │ ├── pydantic_agent.py # PydanticAgent (primary)
|
||||
│ │ └── simple.py # SimpleLiteLLMAgent (fallback)
|
||||
│ ├── memory/
|
||||
│ │ ├── manager.py # Multi-tenant memory manager
|
||||
│ │ ├── tier1_buffer.py # RAM conversation buffer
|
||||
│ │ ├── qdrant_memory.py # Qdrant persistent + semantic
|
||||
│ │ ├── base.py # Base memory interfaces
|
||||
│ │ └── schemas.py # Memory data schemas
|
||||
│ ├── tools/
|
||||
│ │ ├── local.py # Local utility tools
|
||||
│ │ ├── openapi_discovery.py # OpenAPI tool discovery
|
||||
│ │ └── registry.py # Tool registration system
|
||||
│ ├── config.py # Configuration (Pydantic Settings)
|
||||
│ ├── prompts.py # System prompts
|
||||
│ └── utils.py # Utilities
|
||||
├── tests/
|
||||
│ ├── test_ai_flow_quality.py # End-to-end AI quality tests
|
||||
│ └── QUALITY_TESTS.md # Test documentation
|
||||
├── requirements.txt
|
||||
├── Dockerfile
|
||||
└── README.md (this file)
|
||||
```
|
||||
|
||||
## Development
|
||||
|
||||
### Adding Local Tools
|
||||
|
||||
Edit `src/tools/local.py`:
|
||||
|
||||
```python
|
||||
from src.tools.registry import register_tool
|
||||
|
||||
@register_tool
|
||||
async def my_new_tool(param: str) -> str:
|
||||
"""
|
||||
Tool description for LLM.
|
||||
|
||||
Args:
|
||||
param: Parameter description
|
||||
|
||||
Returns:
|
||||
Result description
|
||||
"""
|
||||
# Implementation
|
||||
return f"Result: {param}"
|
||||
```
|
||||
|
||||
Tool automatically available to PydanticAgent.
|
||||
|
||||
### Adding OpenAPI Sources
|
||||
|
||||
Add endpoints to configuration:
|
||||
|
||||
```bash
|
||||
OPENAPI_ENDPOINTS=http://core-api:8083/openapi.json,http://automation:8080/openapi.json
|
||||
```
|
||||
|
||||
Tools auto-discovered on startup with service prefix:
|
||||
- `core-api__list_containers`
|
||||
- `automation__deploy_stack`
|
||||
|
||||
### Modifying System Prompts
|
||||
|
||||
Edit `src/prompts.py`:
|
||||
|
||||
```python
|
||||
PROMPTS = {
|
||||
"pydantic_agent": "Your custom PydanticAgent prompt...",
|
||||
"minimal_agent": "Your custom SimpleLiteLLMAgent prompt..."
|
||||
}
|
||||
```
|
||||
|
||||
Update environment:
|
||||
```bash
|
||||
PYDANTIC_SYSTEM_PROMPT_VARIANT=pydantic_agent
|
||||
```
|
||||
|
||||
### Memory System Usage
|
||||
|
||||
Memory automatically managed per user:
|
||||
|
||||
```python
|
||||
from src.memory import get_memory_manager_for_user
|
||||
|
||||
# Get user's memory manager
|
||||
memory = get_memory_manager_for_user(user_id="user_at_example_com")
|
||||
|
||||
# Memory automatically used by PydanticAgent when conversation_id provided
|
||||
# See: src/agents/pydantic_agent.py
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### PydanticAI Not Available
|
||||
|
||||
**Error:** `PydanticAI not available. Install with: pip install pydantic-ai`
|
||||
|
||||
**Solution:**
|
||||
```bash
|
||||
pip install pydantic-ai
|
||||
```
|
||||
|
||||
### Tools Not Discovered
|
||||
|
||||
**Issue:** `/v1/tools` returns empty list or only local tools
|
||||
|
||||
**Check:**
|
||||
1. Verify `OPENAPI_ENABLED=true`
|
||||
2. Check core-api is running: `curl http://core-api:8083/openapi.json`
|
||||
3. Review logs for discovery errors: `docker logs core-ai`
|
||||
|
||||
### Memory Errors
|
||||
|
||||
**Issue:** Memory operations failing
|
||||
|
||||
**Check:**
|
||||
1. Verify Qdrant running: `curl http://qdrant:6333/collections`
|
||||
2. Check embedding model available: `docker exec ollama ollama list | grep nomic-embed-text`
|
||||
3. Review logs for initialization errors
|
||||
|
||||
### Model Timeouts
|
||||
|
||||
**Issue:** Requests timing out
|
||||
|
||||
**Solutions:**
|
||||
1. Increase timeout: `OLLAMA_TIMEOUT=600`
|
||||
2. Use smaller model: `AGENT_MODEL=mistral-tools:7b`
|
||||
3. Check GPU access: `docker exec ollama nvidia-smi`
|
||||
|
||||
### Tool Calling Failures
|
||||
|
||||
**Issue:** Agent not using tools correctly
|
||||
|
||||
**Check:**
|
||||
1. Verify model supports tool calling: `mistral-nemo`, `mistral-tools:7b`
|
||||
2. Test with `enable_tools=false` to isolate issue
|
||||
3. Review tool logs: Look for `🔧 TOOL CALL:` in logs
|
||||
|
||||
## Model Recommendations
|
||||
|
||||
### For Tool Calling (PydanticAgent)
|
||||
|
||||
- **mistral-nemo:latest** (default) - Best balance
|
||||
- **mistral-tools:7b** - Faster, less accurate
|
||||
- **llama3.1:8b** - Good alternative
|
||||
|
||||
### For Simple Chat (SimpleLiteLLMAgent)
|
||||
|
||||
- **gemma2:9b** - Fast conversational
|
||||
- **llama3.2:3b** - Minimal resources
|
||||
- Any model works (no tool calling required)
|
||||
|
||||
## Migration Notes
|
||||
|
||||
This service has migrated from:
|
||||
- **ADK (Agent Development Kit)** → PydanticAI
|
||||
- **LangChain/LangGraph** → PydanticAI native
|
||||
- **OllamaNativeAgent** → Removed (superseded by PydanticAgent)
|
||||
|
||||
All references to these frameworks have been removed. The codebase now exclusively uses PydanticAI for agent orchestration.
|
||||
|
||||
## Contributing
|
||||
|
||||
When making changes:
|
||||
1. Add tests in `tests/`
|
||||
2. Update docstrings
|
||||
3. Test with both agents (`/v1/chat/completions` and `/v1/chat/simple`)
|
||||
4. Verify tool discovery works
|
||||
5. Test memory persistence
|
||||
|
||||
## License
|
||||
|
||||
Part of the portainer-core project.
|
||||
@@ -1,559 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Multi-Agent AI System Benchmark and Diagnostic Tool
|
||||
|
||||
This script thoroughly tests and profiles the core-ai multi-agent system to identify
|
||||
performance bottlenecks and diagnose the ~30s tool selection issue.
|
||||
|
||||
Usage:
|
||||
python benchmark_multi_agent.py [--verbose] [--save-requests]
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Any, Optional
|
||||
from datetime import datetime
|
||||
import statistics
|
||||
|
||||
# Add src to path for imports
|
||||
sys.path.insert(0, str(Path(__file__).parent / "src"))
|
||||
|
||||
from src.config import get_settings
|
||||
from src.agents.pydantic_agent import get_pydantic_agent
|
||||
from src.agents.steward_agent import get_steward_agent
|
||||
from src.agents.speedy_steward_agent import get_speedy_steward_agent
|
||||
from src.agents.multi_stage_agent import create_multi_stage_agent
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OllamaRequestInterceptor:
|
||||
"""
|
||||
Intercepts and logs requests to Ollama for analysis.
|
||||
|
||||
Captures:
|
||||
- Full request payload
|
||||
- Response
|
||||
- Timing
|
||||
- Token counts (if available)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.requests: List[Dict[str, Any]] = []
|
||||
self.enabled = False
|
||||
|
||||
def enable(self):
|
||||
"""Enable request interception"""
|
||||
self.enabled = True
|
||||
logger.info("Ollama request interception ENABLED")
|
||||
|
||||
def disable(self):
|
||||
"""Disable request interception"""
|
||||
self.enabled = False
|
||||
logger.info("Ollama request interception DISABLED")
|
||||
|
||||
def log_request(self, endpoint: str, payload: Dict[str, Any],
|
||||
response: Optional[Dict[str, Any]] = None,
|
||||
duration_ms: float = 0):
|
||||
"""Log a request to Ollama"""
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
# Calculate payload size
|
||||
payload_size = len(json.dumps(payload))
|
||||
|
||||
record = {
|
||||
'timestamp': datetime.utcnow().isoformat(),
|
||||
'endpoint': endpoint,
|
||||
'payload_size_bytes': payload_size,
|
||||
'payload': payload,
|
||||
'response': response,
|
||||
'duration_ms': duration_ms
|
||||
}
|
||||
|
||||
self.requests.append(record)
|
||||
|
||||
logger.info(f"Captured Ollama request: {endpoint} ({payload_size} bytes, {duration_ms:.1f}ms)")
|
||||
|
||||
def save_to_file(self, filepath: str):
|
||||
"""Save captured requests to JSON file"""
|
||||
with open(filepath, 'w') as f:
|
||||
json.dump(self.requests, f, indent=2)
|
||||
logger.info(f"Saved {len(self.requests)} requests to {filepath}")
|
||||
|
||||
def get_summary(self) -> Dict[str, Any]:
|
||||
"""Get summary statistics of captured requests"""
|
||||
if not self.requests:
|
||||
return {'total_requests': 0}
|
||||
|
||||
payload_sizes = [r['payload_size_bytes'] for r in self.requests]
|
||||
durations = [r['duration_ms'] for r in self.requests]
|
||||
|
||||
return {
|
||||
'total_requests': len(self.requests),
|
||||
'total_payload_bytes': sum(payload_sizes),
|
||||
'avg_payload_bytes': statistics.mean(payload_sizes),
|
||||
'max_payload_bytes': max(payload_sizes),
|
||||
'avg_duration_ms': statistics.mean(durations),
|
||||
'max_duration_ms': max(durations),
|
||||
'p95_duration_ms': sorted(durations)[int(len(durations) * 0.95)] if len(durations) > 1 else durations[0]
|
||||
}
|
||||
|
||||
|
||||
# Global interceptor instance
|
||||
interceptor = OllamaRequestInterceptor()
|
||||
|
||||
|
||||
class BenchmarkResult:
|
||||
"""Container for benchmark results"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
self.durations: List[float] = []
|
||||
self.errors: List[str] = []
|
||||
self.stages: Dict[str, List[float]] = {}
|
||||
|
||||
def add_timing(self, duration_ms: float):
|
||||
"""Add a timing measurement"""
|
||||
self.durations.append(duration_ms)
|
||||
|
||||
def add_error(self, error: str):
|
||||
"""Add an error"""
|
||||
self.errors.append(error)
|
||||
|
||||
def add_stage_timing(self, stage: str, duration_ms: float):
|
||||
"""Add a stage-specific timing"""
|
||||
if stage not in self.stages:
|
||||
self.stages[stage] = []
|
||||
self.stages[stage].append(duration_ms)
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get statistical summary"""
|
||||
if not self.durations:
|
||||
return {
|
||||
'name': self.name,
|
||||
'status': 'no_data',
|
||||
'error_count': len(self.errors)
|
||||
}
|
||||
|
||||
sorted_durations = sorted(self.durations)
|
||||
|
||||
stats = {
|
||||
'name': self.name,
|
||||
'runs': len(self.durations),
|
||||
'avg_ms': statistics.mean(self.durations),
|
||||
'median_ms': statistics.median(self.durations),
|
||||
'min_ms': min(self.durations),
|
||||
'max_ms': max(self.durations),
|
||||
'p95_ms': sorted_durations[int(len(sorted_durations) * 0.95)] if len(sorted_durations) > 1 else sorted_durations[0],
|
||||
'p99_ms': sorted_durations[int(len(sorted_durations) * 0.99)] if len(sorted_durations) > 1 else sorted_durations[0],
|
||||
'error_count': len(self.errors),
|
||||
'success_rate': (len(self.durations) - len(self.errors)) / len(self.durations) if self.durations else 0
|
||||
}
|
||||
|
||||
# Add stage breakdowns
|
||||
if self.stages:
|
||||
stats['stages'] = {}
|
||||
for stage, timings in self.stages.items():
|
||||
stats['stages'][stage] = {
|
||||
'avg_ms': statistics.mean(timings),
|
||||
'min_ms': min(timings),
|
||||
'max_ms': max(timings)
|
||||
}
|
||||
|
||||
return stats
|
||||
|
||||
|
||||
class MultiAgentBenchmark:
|
||||
"""
|
||||
Comprehensive benchmark suite for the multi-agent AI system.
|
||||
"""
|
||||
|
||||
def __init__(self, save_requests: bool = False):
|
||||
self.settings = get_settings()
|
||||
self.save_requests = save_requests
|
||||
self.results: Dict[str, BenchmarkResult] = {}
|
||||
|
||||
# Test queries representing different scenarios
|
||||
self.test_queries = [
|
||||
{
|
||||
'name': 'simple_general_knowledge',
|
||||
'query': 'What is the capital of France?',
|
||||
'expected_tools': [],
|
||||
'category': 'no_tools'
|
||||
},
|
||||
{
|
||||
'name': 'simple_calculation',
|
||||
'query': 'What is 15 + 27?',
|
||||
'expected_tools': ['calculate'],
|
||||
'category': 'single_tool'
|
||||
},
|
||||
{
|
||||
'name': 'time_query',
|
||||
'query': 'What time is it?',
|
||||
'expected_tools': ['get_current_time'],
|
||||
'category': 'single_tool'
|
||||
},
|
||||
{
|
||||
'name': 'web_search_query',
|
||||
'query': 'What are the latest developments in AI?',
|
||||
'expected_tools': ['web_search'],
|
||||
'category': 'single_tool'
|
||||
},
|
||||
{
|
||||
'name': 'complex_multi_tool',
|
||||
'query': 'Search for the current price of Bitcoin and calculate 50% of it',
|
||||
'expected_tools': ['web_search', 'calculate'],
|
||||
'category': 'multi_tool'
|
||||
}
|
||||
]
|
||||
|
||||
def _init_result(self, name: str) -> BenchmarkResult:
|
||||
"""Initialize a benchmark result"""
|
||||
if name not in self.results:
|
||||
self.results[name] = BenchmarkResult(name)
|
||||
return self.results[name]
|
||||
|
||||
async def benchmark_tatlock_solo(self, query: str, run_id: int) -> Dict[str, Any]:
|
||||
"""
|
||||
Benchmark Tatlock agent alone (no steward).
|
||||
|
||||
This is the baseline - fastest path.
|
||||
"""
|
||||
result = self._init_result('tatlock_solo')
|
||||
|
||||
logger.info(f"[Run {run_id}] Testing Tatlock Solo: {query[:50]}...")
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
# Get Tatlock agent directly
|
||||
tatlock = get_pydantic_agent(discover_tools=True)
|
||||
|
||||
messages = [{'role': 'user', 'content': query}]
|
||||
|
||||
# Non-streaming for simplicity
|
||||
response = await tatlock.chat_completion(messages=messages)
|
||||
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
result.add_timing(duration_ms)
|
||||
|
||||
logger.info(f"[Run {run_id}] Tatlock Solo: {duration_ms:.1f}ms")
|
||||
|
||||
return {
|
||||
'duration_ms': duration_ms,
|
||||
'response': response,
|
||||
'success': True
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
result.add_error(str(e))
|
||||
logger.error(f"[Run {run_id}] Tatlock Solo failed: {e}")
|
||||
|
||||
return {
|
||||
'duration_ms': duration_ms,
|
||||
'error': str(e),
|
||||
'success': False
|
||||
}
|
||||
|
||||
async def benchmark_steward_only(self, query: str, run_id: int, use_speedy: bool = False) -> Dict[str, Any]:
|
||||
"""
|
||||
Benchmark steward analysis only (no Tatlock execution).
|
||||
|
||||
This isolates the steward's performance.
|
||||
"""
|
||||
variant = 'speedy_steward' if use_speedy else 'regular_steward'
|
||||
result = self._init_result(f'{variant}_only')
|
||||
|
||||
logger.info(f"[Run {run_id}] Testing {variant.upper()}: {query[:50]}...")
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if use_speedy:
|
||||
steward = get_speedy_steward_agent()
|
||||
else:
|
||||
steward = get_steward_agent()
|
||||
|
||||
recommendation = await steward.analyze(query, timeout=None)
|
||||
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
result.add_timing(duration_ms)
|
||||
|
||||
logger.info(
|
||||
f"[Run {run_id}] {variant.upper()}: {duration_ms:.1f}ms - "
|
||||
f"Recommended {len(recommendation.recommended_tools)} tools: {recommendation.recommended_tools}"
|
||||
)
|
||||
|
||||
return {
|
||||
'duration_ms': duration_ms,
|
||||
'recommendation': {
|
||||
'tools': recommendation.recommended_tools,
|
||||
'reasoning': recommendation.reasoning,
|
||||
'requires_assistance': recommendation.requires_assistance
|
||||
},
|
||||
'success': True
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
result.add_error(str(e))
|
||||
logger.error(f"[Run {run_id}] {variant.upper()} failed: {e}")
|
||||
|
||||
return {
|
||||
'duration_ms': duration_ms,
|
||||
'error': str(e),
|
||||
'success': False
|
||||
}
|
||||
|
||||
async def benchmark_multi_stage(self, query: str, run_id: int, use_speedy: bool = True) -> Dict[str, Any]:
|
||||
"""
|
||||
Benchmark full multi-stage flow: Steward → Tatlock.
|
||||
|
||||
This tests the complete orchestration.
|
||||
"""
|
||||
variant = 'multi_stage_speedy' if use_speedy else 'multi_stage_regular'
|
||||
result = self._init_result(variant)
|
||||
|
||||
logger.info(f"[Run {run_id}] Testing {variant.upper()}: {query[:50]}...")
|
||||
|
||||
total_start = time.time()
|
||||
stages = {}
|
||||
|
||||
try:
|
||||
# Create multi-stage agent
|
||||
tatlock = get_pydantic_agent(discover_tools=True)
|
||||
|
||||
# Temporarily set the use_speedy_steward config
|
||||
original_setting = self.settings.use_speedy_steward
|
||||
self.settings.use_speedy_steward = use_speedy
|
||||
|
||||
multi_agent = create_multi_stage_agent(tatlock, enable_multi_stage=True)
|
||||
|
||||
messages = [{'role': 'user', 'content': query}]
|
||||
|
||||
# Track stages
|
||||
stage_start = time.time()
|
||||
|
||||
full_response = ""
|
||||
steward_duration = None
|
||||
tatlock_start = None
|
||||
|
||||
async for chunk in multi_agent.chat_with_analysis(messages=messages, stream=True):
|
||||
if chunk.get('type') == 'status':
|
||||
phase = chunk.get('phase')
|
||||
|
||||
if phase == 'analysis_complete':
|
||||
steward_duration = (time.time() - stage_start) * 1000
|
||||
stages['steward_analysis'] = steward_duration
|
||||
tatlock_start = time.time()
|
||||
|
||||
elif chunk.get('type') == 'content':
|
||||
full_response += chunk.get('content', '')
|
||||
|
||||
elif chunk.get('type') == 'done':
|
||||
if tatlock_start:
|
||||
tatlock_duration = (time.time() - tatlock_start) * 1000
|
||||
stages['tatlock_execution'] = tatlock_duration
|
||||
|
||||
total_duration_ms = (time.time() - total_start) * 1000
|
||||
result.add_timing(total_duration_ms)
|
||||
|
||||
# Add stage timings
|
||||
for stage, duration in stages.items():
|
||||
result.add_stage_timing(stage, duration)
|
||||
|
||||
logger.info(
|
||||
f"[Run {run_id}] {variant.upper()}: {total_duration_ms:.1f}ms total "
|
||||
f"(Steward: {stages.get('steward_analysis', 0):.1f}ms, "
|
||||
f"Tatlock: {stages.get('tatlock_execution', 0):.1f}ms)"
|
||||
)
|
||||
|
||||
# Restore original setting
|
||||
self.settings.use_speedy_steward = original_setting
|
||||
|
||||
return {
|
||||
'duration_ms': total_duration_ms,
|
||||
'stages': stages,
|
||||
'response': full_response,
|
||||
'success': True
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
total_duration_ms = (time.time() - total_start) * 1000
|
||||
result.add_error(str(e))
|
||||
logger.error(f"[Run {run_id}] {variant.upper()} failed: {e}")
|
||||
|
||||
return {
|
||||
'duration_ms': total_duration_ms,
|
||||
'stages': stages,
|
||||
'error': str(e),
|
||||
'success': False
|
||||
}
|
||||
|
||||
async def run_full_benchmark(self, runs_per_test: int = 3):
|
||||
"""
|
||||
Run comprehensive benchmark suite.
|
||||
|
||||
Tests all variants across all test queries.
|
||||
"""
|
||||
logger.info("=" * 80)
|
||||
logger.info("MULTI-AGENT AI SYSTEM BENCHMARK")
|
||||
logger.info("=" * 80)
|
||||
|
||||
if self.save_requests:
|
||||
interceptor.enable()
|
||||
|
||||
for test_case in self.test_queries:
|
||||
query = test_case['query']
|
||||
logger.info(f"\n{'=' * 80}")
|
||||
logger.info(f"Test Case: {test_case['name']}")
|
||||
logger.info(f"Query: {query}")
|
||||
logger.info(f"Expected Tools: {test_case['expected_tools']}")
|
||||
logger.info(f"{'=' * 80}\n")
|
||||
|
||||
for run in range(1, runs_per_test + 1):
|
||||
logger.info(f"\n--- Run {run}/{runs_per_test} ---")
|
||||
|
||||
# Test 1: Tatlock Solo (baseline)
|
||||
await self.benchmark_tatlock_solo(query, run)
|
||||
await asyncio.sleep(1) # Cooldown
|
||||
|
||||
# Test 2: Speedy Steward Only
|
||||
await self.benchmark_steward_only(query, run, use_speedy=True)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Test 3: Regular Steward Only
|
||||
await self.benchmark_steward_only(query, run, use_speedy=False)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Test 4: Multi-Stage with Speedy Steward
|
||||
await self.benchmark_multi_stage(query, run, use_speedy=True)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Test 5: Multi-Stage with Regular Steward
|
||||
await self.benchmark_multi_stage(query, run, use_speedy=False)
|
||||
await asyncio.sleep(2) # Longer cooldown
|
||||
|
||||
if self.save_requests:
|
||||
interceptor.disable()
|
||||
|
||||
def generate_report(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Generate comprehensive diagnostic report.
|
||||
"""
|
||||
report = {
|
||||
'timestamp': datetime.utcnow().isoformat() + 'Z',
|
||||
'configuration': {
|
||||
'agent_model': self.settings.agent_model,
|
||||
'multi_stage_enabled': self.settings.multi_stage_enabled,
|
||||
'use_speedy_steward': self.settings.use_speedy_steward,
|
||||
'analysis_timeout': self.settings.analysis_timeout,
|
||||
},
|
||||
'results': {}
|
||||
}
|
||||
|
||||
# Add all benchmark results
|
||||
for name, result in self.results.items():
|
||||
report['results'][name] = result.get_stats()
|
||||
|
||||
# Add Ollama request analysis if available
|
||||
if self.save_requests and interceptor.requests:
|
||||
report['ollama_requests'] = interceptor.get_summary()
|
||||
|
||||
return report
|
||||
|
||||
def print_report(self):
|
||||
"""Print human-readable report"""
|
||||
print("\n" + "=" * 80)
|
||||
print("BENCHMARK RESULTS")
|
||||
print("=" * 80)
|
||||
|
||||
for name, result in sorted(self.results.items()):
|
||||
stats = result.get_stats()
|
||||
print(f"\n{name}:")
|
||||
print(f" Runs: {stats.get('runs', 0)}")
|
||||
print(f" Average: {stats.get('avg_ms', 0):.1f}ms")
|
||||
print(f" Median: {stats.get('median_ms', 0):.1f}ms")
|
||||
print(f" Min: {stats.get('min_ms', 0):.1f}ms")
|
||||
print(f" Max: {stats.get('max_ms', 0):.1f}ms")
|
||||
print(f" P95: {stats.get('p95_ms', 0):.1f}ms")
|
||||
print(f" Errors: {stats.get('error_count', 0)}")
|
||||
|
||||
if 'stages' in stats:
|
||||
print(" Stage Breakdown:")
|
||||
for stage, stage_stats in stats['stages'].items():
|
||||
print(f" {stage}: {stage_stats['avg_ms']:.1f}ms avg")
|
||||
|
||||
if self.save_requests and interceptor.requests:
|
||||
print("\n" + "=" * 80)
|
||||
print("OLLAMA REQUEST ANALYSIS")
|
||||
print("=" * 80)
|
||||
summary = interceptor.get_summary()
|
||||
print(f" Total Requests: {summary['total_requests']}")
|
||||
print(f" Total Payload: {summary['total_payload_bytes']:,} bytes")
|
||||
print(f" Avg Payload: {summary['avg_payload_bytes']:.1f} bytes")
|
||||
print(f" Max Payload: {summary['max_payload_bytes']:,} bytes")
|
||||
print(f" Avg Duration: {summary['avg_duration_ms']:.1f}ms")
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
|
||||
|
||||
async def main():
|
||||
"""Main benchmark execution"""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description='Benchmark multi-agent AI system')
|
||||
parser.add_argument('--verbose', action='store_true', help='Enable verbose logging')
|
||||
parser.add_argument('--save-requests', action='store_true', help='Save Ollama requests to file')
|
||||
parser.add_argument('--runs', type=int, default=3, help='Number of runs per test (default: 3)')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.verbose:
|
||||
logging.getLogger().setLevel(logging.DEBUG)
|
||||
|
||||
benchmark = MultiAgentBenchmark(save_requests=args.save_requests)
|
||||
|
||||
try:
|
||||
await benchmark.run_full_benchmark(runs_per_test=args.runs)
|
||||
|
||||
# Generate and save report
|
||||
report = benchmark.generate_report()
|
||||
|
||||
report_file = f"benchmark_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
|
||||
with open(report_file, 'w') as f:
|
||||
json.dump(report, f, indent=2)
|
||||
|
||||
print(f"\n✓ Report saved to: {report_file}")
|
||||
|
||||
# Print summary
|
||||
benchmark.print_report()
|
||||
|
||||
# Save Ollama requests if captured
|
||||
if args.save_requests:
|
||||
requests_file = f"ollama_requests_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
|
||||
interceptor.save_to_file(requests_file)
|
||||
print(f"✓ Ollama requests saved to: {requests_file}")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n\nBenchmark interrupted by user")
|
||||
benchmark.print_report()
|
||||
except Exception as e:
|
||||
logger.error(f"Benchmark failed: {e}", exc_info=True)
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(asyncio.run(main()))
|
||||
@@ -1 +0,0 @@
|
||||
"""Diagnostic tools for core-ai service"""
|
||||
@@ -1,128 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Diagnostic tool to check Ollama connectivity and available models.
|
||||
Run this first to verify the foundation is working.
|
||||
|
||||
Usage:
|
||||
python diagnostics/check_ollama.py
|
||||
"""
|
||||
import asyncio
|
||||
import httpx
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Add parent directory to path to import from src
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from src.config import get_settings
|
||||
|
||||
|
||||
async def check_ollama():
|
||||
"""Check Ollama connectivity and list available models"""
|
||||
settings = get_settings()
|
||||
ollama_url = settings.ollama_base_url
|
||||
|
||||
print("=" * 70)
|
||||
print("OLLAMA CONNECTIVITY CHECK")
|
||||
print("=" * 70)
|
||||
print(f"\n1. Configuration")
|
||||
print(f" Ollama URL: {ollama_url}")
|
||||
print(f" Target Model: {settings.agent_model}")
|
||||
print(f" Timeout: {settings.ollama_timeout}s")
|
||||
|
||||
async with httpx.AsyncClient(timeout=settings.ollama_timeout) as client:
|
||||
# Test 1: Basic connectivity
|
||||
print(f"\n2. Testing connectivity to {ollama_url}...")
|
||||
try:
|
||||
response = await client.get(f"{ollama_url}/api/tags")
|
||||
if response.status_code == 200:
|
||||
print(" ✓ Ollama is reachable")
|
||||
else:
|
||||
print(f" ✗ Unexpected status code: {response.status_code}")
|
||||
print(f" Response: {response.text}")
|
||||
return False
|
||||
except httpx.ConnectError as e:
|
||||
print(f" ✗ Connection failed: {e}")
|
||||
print(f" → Is Ollama running?")
|
||||
print(f" → Check docker ps | grep ollama")
|
||||
print(f" → Verify network connectivity")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f" ✗ Error: {e}")
|
||||
return False
|
||||
|
||||
# Test 2: List available models
|
||||
print(f"\n3. Available models:")
|
||||
try:
|
||||
data = response.json()
|
||||
models = data.get("models", [])
|
||||
|
||||
if not models:
|
||||
print(" ✗ No models found!")
|
||||
print(" → Pull a model: docker exec ollama ollama pull gemma2:9b-instruct-q5_K_M")
|
||||
return False
|
||||
|
||||
target_found = False
|
||||
for model in models:
|
||||
model_name = model.get("name", "unknown")
|
||||
size_gb = model.get("size", 0) / (1024**3)
|
||||
is_target = "✓" if settings.agent_model in model_name else " "
|
||||
print(f" {is_target} {model_name} ({size_gb:.2f} GB)")
|
||||
if settings.agent_model in model_name:
|
||||
target_found = True
|
||||
|
||||
if not target_found:
|
||||
print(f"\n ⚠ Target model '{settings.agent_model}' not found!")
|
||||
print(f" → Pull it: docker exec ollama ollama pull {settings.agent_model}")
|
||||
return False
|
||||
else:
|
||||
print(f"\n ✓ Target model '{settings.agent_model}' is available")
|
||||
|
||||
except Exception as e:
|
||||
print(f" ✗ Error parsing models: {e}")
|
||||
return False
|
||||
|
||||
# Test 3: Simple generation test
|
||||
print(f"\n4. Testing text generation with '{settings.agent_model}'...")
|
||||
try:
|
||||
test_payload = {
|
||||
"model": settings.agent_model,
|
||||
"prompt": "Say 'Hello, Ollama is working!' and nothing else.",
|
||||
"stream": False
|
||||
}
|
||||
|
||||
response = await client.post(
|
||||
f"{ollama_url}/api/generate",
|
||||
json=test_payload,
|
||||
timeout=60.0
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
generated_text = result.get("response", "").strip()
|
||||
print(f" Response: {generated_text}")
|
||||
print(" ✓ Text generation successful!")
|
||||
else:
|
||||
print(f" ✗ Generation failed with status {response.status_code}")
|
||||
print(f" Response: {response.text}")
|
||||
return False
|
||||
|
||||
except httpx.TimeoutException:
|
||||
print(f" ✗ Request timed out")
|
||||
print(f" → Model may be loading (first run takes longer)")
|
||||
print(f" → Try again or increase timeout")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f" ✗ Error: {e}")
|
||||
return False
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print("✓ ALL CHECKS PASSED - Ollama is ready!")
|
||||
print("=" * 70)
|
||||
return True
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
result = asyncio.run(check_ollama())
|
||||
sys.exit(0 if result else 1)
|
||||
@@ -1,458 +0,0 @@
|
||||
import os
|
||||
import logging
|
||||
import json
|
||||
import time
|
||||
from aiohttp import web
|
||||
from aiohttp_cors import setup as cors_setup, ResourceOptions
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables from .env file
|
||||
load_dotenv()
|
||||
|
||||
# Set up logging
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Import the agent logic
|
||||
from src.agents import (
|
||||
get_simple_litellm_agent,
|
||||
get_pydantic_agent,
|
||||
PYDANTIC_AI_AVAILABLE
|
||||
)
|
||||
from src.agents.multi_stage_agent import create_multi_stage_agent
|
||||
from src.agents.stream_handler import get_stream_handler
|
||||
from src.tools import get_all_tools
|
||||
from src.utils import extract_user_id_from_request
|
||||
|
||||
def format_status_message(message: str, phase: str) -> str:
|
||||
"""
|
||||
Format status messages with butler-appropriate box-drawing characters.
|
||||
|
||||
Uses Unicode box-drawing characters for visual structure:
|
||||
- ┌─ for starting messages
|
||||
- ├─ for continuing messages
|
||||
- └─ for completing messages
|
||||
"""
|
||||
if phase == "analysis":
|
||||
# Starting steward consultation
|
||||
return f"┌─ {message}"
|
||||
elif phase == "analysis_complete":
|
||||
# Steward consultation complete
|
||||
return f"└─ {message}"
|
||||
elif phase == "tool_execution":
|
||||
# Tool being executed
|
||||
return f"├─ {message}"
|
||||
elif phase == "fallback":
|
||||
# Fallback/warning
|
||||
return f"└─ {message}"
|
||||
else:
|
||||
# Default
|
||||
return f"├─ {message}"
|
||||
|
||||
async def chat_completions(request):
|
||||
"""
|
||||
Handles OpenAI-compatible chat completion requests using PydanticAI.
|
||||
Default endpoint - uses PydanticAI with tools enabled.
|
||||
"""
|
||||
if not PYDANTIC_AI_AVAILABLE:
|
||||
return web.json_response({
|
||||
"error": {"message": "PydanticAI not available. Install with: pip install pydantic-ai"}
|
||||
}, status=503)
|
||||
|
||||
from src.metrics import get_metrics_collector
|
||||
metrics = get_metrics_collector()
|
||||
metrics.increment_concurrent_requests()
|
||||
start_time = time.time()
|
||||
agent_type = "pydantic"
|
||||
success = False
|
||||
error_msg = None
|
||||
|
||||
try:
|
||||
data = await request.json()
|
||||
logger.info(f"[DEFAULT/PYDANTIC_AI] Received chat request")
|
||||
|
||||
# Extract relevant fields from the request
|
||||
messages = data.get("messages")
|
||||
model = data.get("model", "Tatlock")
|
||||
stream = data.get("stream", False)
|
||||
conversation_id = data.get("conversation_id")
|
||||
enable_tools = data.get("enable_tools", True)
|
||||
multi_stage_analysis = data.get("multi_stage_analysis", True) # Enable by default
|
||||
|
||||
# Extract user ID from request
|
||||
user_id = extract_user_id_from_request(data)
|
||||
|
||||
if not messages:
|
||||
raise web.HTTPBadRequest(reason="'messages' field is required")
|
||||
|
||||
# Get the agent instance
|
||||
base_agent = get_pydantic_agent(discover_tools=enable_tools, user_id=user_id)
|
||||
|
||||
# Wrap with multi-stage orchestration if enabled
|
||||
agent = create_multi_stage_agent(base_agent, enable_multi_stage=multi_stage_analysis)
|
||||
|
||||
# For non-streaming requests, collect the full response
|
||||
if not stream:
|
||||
# Collect all content chunks from two-stage agent
|
||||
full_content = []
|
||||
async for chunk in agent.chat_with_analysis(
|
||||
messages=messages,
|
||||
conversation_id=conversation_id,
|
||||
stream=False
|
||||
):
|
||||
if chunk.get("type") == "content":
|
||||
full_content.append(chunk.get("content", ""))
|
||||
|
||||
response_content = "".join(full_content)
|
||||
success = True
|
||||
return web.json_response({
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": response_content
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"model": model,
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0
|
||||
},
|
||||
"tools_enabled": enable_tools,
|
||||
"multi_stage_analysis": multi_stage_analysis
|
||||
})
|
||||
else:
|
||||
# Handle streaming response with StreamHandler
|
||||
response = web.StreamResponse(
|
||||
status=200,
|
||||
reason='OK',
|
||||
headers={
|
||||
'Content-Type': 'text/event-stream',
|
||||
'Cache-Control': 'no-cache',
|
||||
'Connection': 'keep-alive',
|
||||
}
|
||||
)
|
||||
await response.prepare(request)
|
||||
|
||||
try:
|
||||
# Create stream handler instance
|
||||
handler = get_stream_handler(model_name=model)
|
||||
|
||||
# Process stream through handler
|
||||
async for chunk in handler.process_stream(
|
||||
messages=messages,
|
||||
multi_stage_enabled=multi_stage_analysis
|
||||
):
|
||||
# Stream handler returns pre-formatted messages
|
||||
await response.write(f"data: {json.dumps(chunk)}\n\n".encode('utf-8'))
|
||||
|
||||
await response.write(b"data: [DONE]\n\n")
|
||||
success = True
|
||||
finally:
|
||||
await response.write_eof()
|
||||
|
||||
return response
|
||||
|
||||
except web.HTTPBadRequest:
|
||||
raise
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
logger.exception(f"Error in chat_completions: {e}")
|
||||
return web.json_response({
|
||||
"error": {"message": f"Internal server error: {str(e)}"}
|
||||
}, status=500)
|
||||
finally:
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
metrics.decrement_concurrent_requests()
|
||||
metrics.record_request(
|
||||
agent_type=agent_type,
|
||||
duration_ms=duration_ms,
|
||||
success=success,
|
||||
streaming=data.get("stream", False) if 'data' in locals() else False,
|
||||
user_id=user_id if 'user_id' in locals() else None,
|
||||
error=error_msg
|
||||
)
|
||||
|
||||
async def chat_simple(request):
|
||||
"""
|
||||
Handles chat requests using SimpleLiteLLMAgent (fallback, no tools).
|
||||
Endpoint: /v1/chat/simple
|
||||
"""
|
||||
from src.metrics import get_metrics_collector
|
||||
metrics = get_metrics_collector()
|
||||
metrics.increment_concurrent_requests()
|
||||
start_time = time.time()
|
||||
agent_type = "simple"
|
||||
success = False
|
||||
error_msg = None
|
||||
|
||||
try:
|
||||
data = await request.json()
|
||||
logger.info(f"[SIMPLE/LITELLM] Received chat request")
|
||||
|
||||
# Extract relevant fields from the request
|
||||
messages = data.get("messages")
|
||||
model = data.get("model", "simple")
|
||||
stream = data.get("stream", False)
|
||||
conversation_id = data.get("conversation_id")
|
||||
user_id = extract_user_id_from_request(data)
|
||||
|
||||
if not messages:
|
||||
raise web.HTTPBadRequest(reason="'messages' field is required")
|
||||
|
||||
# Get simple agent instance (no tools)
|
||||
agent = get_simple_litellm_agent()
|
||||
|
||||
# For non-streaming requests
|
||||
if not stream:
|
||||
response_content = await agent.chat_completion(
|
||||
messages=messages,
|
||||
conversation_id=conversation_id
|
||||
)
|
||||
return web.json_response({
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": response_content
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"model": model,
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0
|
||||
}
|
||||
})
|
||||
else:
|
||||
# Streaming response
|
||||
response = web.StreamResponse(
|
||||
status=200,
|
||||
reason='OK',
|
||||
headers={
|
||||
'Content-Type': 'text/event-stream',
|
||||
'Cache-Control': 'no-cache',
|
||||
'Connection': 'keep-alive',
|
||||
}
|
||||
)
|
||||
await response.prepare(request)
|
||||
|
||||
try:
|
||||
async for chunk in agent.chat(messages=messages, conversation_id=conversation_id, stream=True):
|
||||
if chunk["type"] == "content":
|
||||
chunk_data = {
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {"content": chunk["content"]},
|
||||
"finish_reason": chunk.get("finish_reason")
|
||||
}],
|
||||
"model": model
|
||||
}
|
||||
await response.write(f"data: {json.dumps(chunk_data)}\n\n".encode('utf-8'))
|
||||
|
||||
await response.write(b"data: [DONE]\n\n")
|
||||
success = True
|
||||
finally:
|
||||
await response.write_eof()
|
||||
|
||||
return response
|
||||
|
||||
except web.HTTPBadRequest:
|
||||
raise
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
logger.exception(f"Error in chat_simple: {e}")
|
||||
return web.json_response({
|
||||
"error": {"message": f"Internal server error: {str(e)}"}
|
||||
}, status=500)
|
||||
finally:
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
metrics.decrement_concurrent_requests()
|
||||
metrics.record_request(
|
||||
agent_type=agent_type,
|
||||
duration_ms=duration_ms,
|
||||
success=success,
|
||||
streaming=data.get("stream", False) if 'data' in locals() else False,
|
||||
user_id=user_id if 'user_id' in locals() else None,
|
||||
error=error_msg
|
||||
)
|
||||
|
||||
async def list_models(request):
|
||||
"""List available models"""
|
||||
return web.json_response({
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "Tatlock",
|
||||
"object": "model",
|
||||
"created": int(time.time()),
|
||||
"owned_by": "core-ai",
|
||||
"description": "PydanticAI agent with full tool support - your British butler assistant"
|
||||
}
|
||||
]
|
||||
})
|
||||
|
||||
async def list_tools(request):
|
||||
"""List all available tools"""
|
||||
try:
|
||||
tools = get_all_tools()
|
||||
tool_list = []
|
||||
|
||||
for name, func in tools.items():
|
||||
import inspect
|
||||
doc = inspect.getdoc(func) or "No description"
|
||||
tool_list.append({
|
||||
"name": name,
|
||||
"description": doc.split('\n')[0],
|
||||
"type": "local" if not name.startswith("core-api__") else "openapi"
|
||||
})
|
||||
|
||||
return web.json_response({
|
||||
"tools": tool_list,
|
||||
"tools_count": len(tool_list)
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Error listing tools: {e}")
|
||||
return web.json_response({
|
||||
"error": {"message": str(e)}
|
||||
}, status=500)
|
||||
|
||||
async def health_check(request):
|
||||
"""Health check endpoint"""
|
||||
return web.json_response({
|
||||
"status": "ok",
|
||||
"service": "core-ai",
|
||||
"agents": {
|
||||
"Tatlock": PYDANTIC_AI_AVAILABLE
|
||||
},
|
||||
"default_agent": "Tatlock" if PYDANTIC_AI_AVAILABLE else None,
|
||||
"tools_count": len(get_all_tools())
|
||||
})
|
||||
|
||||
async def get_metrics(request):
|
||||
"""
|
||||
Get comprehensive performance metrics.
|
||||
|
||||
Returns detailed statistics on agent performance, tool execution,
|
||||
memory system, and request patterns.
|
||||
"""
|
||||
try:
|
||||
from src.metrics import get_metrics_collector
|
||||
metrics_collector = get_metrics_collector()
|
||||
metrics = metrics_collector.get_metrics()
|
||||
return web.json_response(metrics)
|
||||
except Exception as e:
|
||||
logger.exception(f"Error getting metrics: {e}")
|
||||
return web.json_response({
|
||||
"error": {"message": str(e)}
|
||||
}, status=500)
|
||||
|
||||
async def get_recent_errors(request):
|
||||
"""Get recent request errors."""
|
||||
try:
|
||||
from src.metrics import get_metrics_collector
|
||||
metrics_collector = get_metrics_collector()
|
||||
|
||||
limit = int(request.query.get('limit', 20))
|
||||
errors = metrics_collector.get_recent_errors(limit=limit)
|
||||
|
||||
return web.json_response({
|
||||
"errors": errors,
|
||||
"total": len(errors)
|
||||
})
|
||||
except Exception as e:
|
||||
logger.exception(f"Error getting errors: {e}")
|
||||
return web.json_response({
|
||||
"error": {"message": str(e)}
|
||||
}, status=500)
|
||||
|
||||
async def get_tool_failures(request):
|
||||
"""Get recent tool execution failures."""
|
||||
try:
|
||||
from src.metrics import get_metrics_collector
|
||||
metrics_collector = get_metrics_collector()
|
||||
|
||||
limit = int(request.query.get('limit', 20))
|
||||
failures = metrics_collector.get_recent_tool_failures(limit=limit)
|
||||
|
||||
return web.json_response({
|
||||
"failures": failures,
|
||||
"total": len(failures)
|
||||
})
|
||||
except Exception as e:
|
||||
logger.exception(f"Error getting tool failures: {e}")
|
||||
return web.json_response({
|
||||
"error": {"message": str(e)}
|
||||
}, status=500)
|
||||
|
||||
async def reset_metrics(request):
|
||||
"""Reset all metrics (admin endpoint)."""
|
||||
try:
|
||||
from src.metrics import get_metrics_collector
|
||||
metrics_collector = get_metrics_collector()
|
||||
metrics_collector.reset()
|
||||
|
||||
return web.json_response({
|
||||
"message": "Metrics reset successfully"
|
||||
})
|
||||
except Exception as e:
|
||||
logger.exception(f"Error resetting metrics: {e}")
|
||||
return web.json_response({
|
||||
"error": {"message": str(e)}
|
||||
}, status=500)
|
||||
|
||||
async def setup_routes(app):
|
||||
# Chat endpoints
|
||||
app.router.add_post("/chat/completions", chat_completions) # Alias without /v1
|
||||
app.router.add_post("/v1/chat/completions", chat_completions) # Default: PydanticAI
|
||||
app.router.add_post("/v1/chat/simple", chat_simple) # Fallback: Simple agent
|
||||
|
||||
# Models endpoint
|
||||
app.router.add_get("/models", list_models)
|
||||
app.router.add_get("/v1/models", list_models)
|
||||
|
||||
# Tools endpoint
|
||||
app.router.add_get("/tools", list_tools)
|
||||
app.router.add_get("/v1/tools", list_tools)
|
||||
|
||||
# Health check
|
||||
app.router.add_get("/health", health_check)
|
||||
|
||||
# Metrics endpoints
|
||||
app.router.add_get("/metrics", get_metrics)
|
||||
app.router.add_get("/metrics/errors", get_recent_errors)
|
||||
app.router.add_get("/metrics/tool-failures", get_tool_failures)
|
||||
app.router.add_post("/metrics/reset", reset_metrics)
|
||||
|
||||
# Setup CORS
|
||||
cors = cors_setup(app, defaults={
|
||||
"*": ResourceOptions(
|
||||
allow_credentials=True,
|
||||
expose_headers="*",
|
||||
allow_headers="*",
|
||||
allow_methods="*"
|
||||
)
|
||||
})
|
||||
|
||||
# Apply CORS to all routes
|
||||
for route in list(app.router.routes()):
|
||||
cors.add(route)
|
||||
|
||||
def main():
|
||||
app = web.Application()
|
||||
|
||||
# Set up routes
|
||||
import asyncio
|
||||
loop = asyncio.get_event_loop()
|
||||
loop.run_until_complete(setup_routes(app))
|
||||
|
||||
# Run the application
|
||||
logger.info("Starting core-ai service on http://0.0.0.0:8086")
|
||||
web.run_app(app, host='0.0.0.0', port=8086)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -1,31 +0,0 @@
|
||||
[pytest]
|
||||
# Pytest configuration for core-ai tests
|
||||
|
||||
# Test discovery patterns
|
||||
python_files = test_*.py
|
||||
python_classes = Test*
|
||||
python_functions = test_*
|
||||
|
||||
# Output options
|
||||
addopts =
|
||||
-v
|
||||
--tb=short
|
||||
--strict-markers
|
||||
--color=yes
|
||||
|
||||
# Markers
|
||||
markers =
|
||||
asyncio: mark test as async
|
||||
integration: mark test as integration test (requires external services)
|
||||
|
||||
# Asyncio configuration
|
||||
asyncio_mode = auto
|
||||
|
||||
# Log configuration
|
||||
log_cli = true
|
||||
log_cli_level = INFO
|
||||
log_cli_format = %(asctime)s [%(levelname)8s] %(message)s
|
||||
log_cli_date_format = %Y-%m-%d %H:%M:%S
|
||||
|
||||
# Test paths
|
||||
testpaths = tests diagnostics
|
||||
@@ -1,24 +0,0 @@
|
||||
# PydanticAI and dependencies (slim to reduce bloat)
|
||||
pydantic-ai-slim # Minimal library - Ollama uses OpenAI-compatible API
|
||||
pydantic>=2.10.3 # Let pydantic-ai determine the compatible version
|
||||
pydantic-settings==2.6.1
|
||||
ollama>=0.4.0 # Native Ollama Python library with tool calling support
|
||||
|
||||
# LiteLLM (for simple agent fallback)
|
||||
litellm==1.80.5
|
||||
|
||||
# Core dependencies
|
||||
aiohttp==3.10.1
|
||||
aiohttp-cors==0.7.0
|
||||
python-dotenv>=1.1.0
|
||||
httpx==0.28.1
|
||||
|
||||
# Memory system
|
||||
qdrant-client>=1.12.0 # Vector database client
|
||||
|
||||
# Timezone support
|
||||
pytz>=2025.2
|
||||
|
||||
# Testing
|
||||
pytest==8.3.4
|
||||
pytest-asyncio==0.24.0
|
||||
@@ -1,126 +0,0 @@
|
||||
"""
|
||||
Core AI Agent - Direct LiteLLM Chat Completion
|
||||
This is a diagnostic file to test direct text generation via LiteLLM, bypassing Google ADK.
|
||||
"""
|
||||
import os
|
||||
import logging
|
||||
from typing import AsyncIterator, Dict, Any, List, Optional
|
||||
from functools import lru_cache
|
||||
|
||||
# We will directly use litellm here
|
||||
import litellm
|
||||
|
||||
# Adjusted import paths for the new core-ai service structure
|
||||
from src.config import get_settings
|
||||
from src.prompts import get_prompt
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Simplified Agent for direct LiteLLM interaction
|
||||
class SimpleLiteLLMAgent:
|
||||
def __init__(self):
|
||||
# Enable verbose logging for LiteLLM
|
||||
litellm.set_verbose = True
|
||||
logger.info("LiteLLM verbose logging enabled.")
|
||||
|
||||
self.settings = get_settings()
|
||||
|
||||
# Load system prompt
|
||||
self.system_prompt = get_prompt(self.settings.system_prompt_variant)
|
||||
logger.info(f"System prompt variant: {self.settings.system_prompt_variant}")
|
||||
logger.info(f"System prompt: {self.system_prompt[:100]}...")
|
||||
|
||||
# Initialize LiteLLM for Ollama (format: "ollama/model_name")
|
||||
model_name = self.settings.agent_model
|
||||
litellm_model = f"ollama/{model_name}"
|
||||
|
||||
logger.info(f"Initializing LiteLLM direct model: {litellm_model}")
|
||||
logger.info(f"Ollama base URL from settings: {self.settings.ollama_base_url}")
|
||||
|
||||
self.model_params = {
|
||||
"model": litellm_model,
|
||||
"api_base": self.settings.ollama_base_url,
|
||||
"temperature": 0.1,
|
||||
# No tool definitions passed here to force text generation
|
||||
}
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
conversation_id: str = None, # Not used in this simple mode
|
||||
stream: bool = True,
|
||||
prompt_variant: Optional[str] = None # Not used in this simple mode
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Processes a chat message using direct LiteLLM completion.
|
||||
"""
|
||||
logger.info(f"🚀 Starting direct LiteLLM completion for message: {messages[-1]['content'][:50]}...")
|
||||
try:
|
||||
# Prepare messages in LiteLLM format
|
||||
litellm_messages = [{"role": m["role"], "content": m["content"]} for m in messages]
|
||||
|
||||
# Inject system prompt if not already present
|
||||
if not litellm_messages or litellm_messages[0]["role"] != "system":
|
||||
litellm_messages.insert(0, {"role": "system", "content": self.system_prompt})
|
||||
logger.info("✓ System prompt injected")
|
||||
|
||||
# Log full message payload for debugging
|
||||
logger.info(f"📤 Sending {len(litellm_messages)} messages to LiteLLM:")
|
||||
for i, msg in enumerate(litellm_messages):
|
||||
content_preview = msg['content'][:100] + "..." if len(msg['content']) > 100 else msg['content']
|
||||
logger.info(f" [{i}] {msg['role']}: {content_preview}")
|
||||
|
||||
# Use acompletion for async environments
|
||||
response = await litellm.acompletion(
|
||||
messages=litellm_messages,
|
||||
stream=stream,
|
||||
**self.model_params
|
||||
)
|
||||
|
||||
if stream:
|
||||
chunk_count = 0
|
||||
async for chunk in response:
|
||||
chunk_count += 1
|
||||
content_delta = chunk.choices[0].delta.content if chunk.choices[0].delta.content else ""
|
||||
finish_reason = chunk.choices[0].finish_reason
|
||||
if content_delta:
|
||||
yield {"type": "content", "content": content_delta}
|
||||
if finish_reason:
|
||||
logger.info(f"📥 Stream completed after {chunk_count} chunks. Finish reason: {finish_reason}")
|
||||
yield {"type": "content", "content": "", "finish_reason": finish_reason}
|
||||
else:
|
||||
content = response.choices[0].message.content
|
||||
logger.info(f"📥 Response received: {content[:200]}..." if len(content) > 200 else f"📥 Response received: {content}")
|
||||
yield {"type": "content", "content": content, "finish_reason": "stop"}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in direct LiteLLM chat: {e}", exc_info=True)
|
||||
yield {
|
||||
"type": "error",
|
||||
"content": f"Sorry, an error occurred during text generation: {str(e)}",
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
|
||||
async def chat_completion(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
conversation_id: str = None,
|
||||
prompt_variant: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Get a non-streaming response from the direct LiteLLM chat.
|
||||
"""
|
||||
final_content = ""
|
||||
async for chunk in self.chat(messages=messages, conversation_id=conversation_id, stream=False, prompt_variant=prompt_variant):
|
||||
if chunk["type"] == "content":
|
||||
final_content += chunk["content"]
|
||||
if chunk.get("finish_reason") == "stop":
|
||||
break
|
||||
return final_content if final_content else "I couldn't generate a response."
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def get_simple_litellm_agent() -> SimpleLiteLLMAgent:
|
||||
"""Get cached simple LiteLLM agent instance"""
|
||||
return SimpleLiteLLMAgent()
|
||||
@@ -1,19 +0,0 @@
|
||||
"""Agent implementations for core-ai service"""
|
||||
|
||||
from .simple import SimpleLiteLLMAgent, get_simple_litellm_agent
|
||||
|
||||
try:
|
||||
from .pydantic_agent import PydanticAgent, get_pydantic_agent
|
||||
PYDANTIC_AI_AVAILABLE = True
|
||||
except ImportError:
|
||||
PYDANTIC_AI_AVAILABLE = False
|
||||
PydanticAgent = None
|
||||
get_pydantic_agent = None
|
||||
|
||||
__all__ = [
|
||||
'SimpleLiteLLMAgent',
|
||||
'get_simple_litellm_agent',
|
||||
'PydanticAgent',
|
||||
'get_pydantic_agent',
|
||||
'PYDANTIC_AI_AVAILABLE',
|
||||
]
|
||||
@@ -1,407 +0,0 @@
|
||||
"""
|
||||
Multi-Stage Agent Orchestration - Coordinates expert agents for intelligent query processing.
|
||||
|
||||
This module implements a multi-stage agent system with:
|
||||
1. Steward: Analyzes queries and recommends optimal tools (0-5 tools)
|
||||
2. Tatlock: Main execution agent using steward's recommendations
|
||||
3. Future: Additional expert agents can be added to the flow
|
||||
|
||||
Features:
|
||||
- Silent time/date injection when recommended
|
||||
- Visible status for web_search tool calls
|
||||
- Context enrichment with tool recommendations
|
||||
- Streaming status events separate from content
|
||||
- Native PydanticAI event streaming for tool call detection
|
||||
- Fallback to single-stage if steward fails
|
||||
"""
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import Optional, AsyncGenerator, Dict, Any
|
||||
from datetime import datetime
|
||||
import json
|
||||
|
||||
try:
|
||||
from pydantic_ai import Agent
|
||||
PYDANTIC_AI_AVAILABLE = True
|
||||
except ImportError:
|
||||
PYDANTIC_AI_AVAILABLE = False
|
||||
Agent = None
|
||||
|
||||
from src.agents.steward_agent import get_steward_agent, ToolRecommendation
|
||||
from src.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MultiStageAgent:
|
||||
"""
|
||||
Multi-stage agent orchestrator.
|
||||
|
||||
Coordinates steward analysis with Tatlock execution, providing
|
||||
intelligent tool selection and transparent status updates.
|
||||
Designed for extensibility with future expert agents.
|
||||
"""
|
||||
|
||||
def __init__(self, tatlock_agent: Agent, enable_multi_stage: bool = True):
|
||||
"""
|
||||
Initialize multi-stage orchestrator.
|
||||
|
||||
Args:
|
||||
tatlock_agent: The main Tatlock agent instance
|
||||
enable_multi_stage: Whether to use multi-stage analysis (default: True)
|
||||
"""
|
||||
self.tatlock = tatlock_agent
|
||||
self.enable_multi_stage = enable_multi_stage
|
||||
self.settings = get_settings()
|
||||
|
||||
# Get steward instance
|
||||
if self.enable_multi_stage:
|
||||
try:
|
||||
# Use regular steward with mistral-nemo (same model as Tatlock)
|
||||
steward_model = getattr(self.settings, 'steward_model', 'mistral-nemo:latest')
|
||||
self.steward = get_steward_agent(model_name=steward_model)
|
||||
logger.info(f"MultiStageAgent: Steward enabled with {steward_model}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize steward, disabling multi-stage: {e}")
|
||||
self.enable_multi_stage = False
|
||||
self.steward = None
|
||||
else:
|
||||
self.steward = None
|
||||
logger.info("MultiStageAgent: Multi-stage analysis disabled")
|
||||
|
||||
async def _get_current_datetime(self) -> Dict[str, str]:
|
||||
"""
|
||||
Get current date and time for silent injection.
|
||||
|
||||
Returns:
|
||||
Dict with 'date' and 'time' keys
|
||||
"""
|
||||
now = datetime.now()
|
||||
return {
|
||||
"date": now.strftime("%A, %B %d, %Y"),
|
||||
"time": now.strftime("%I:%M %p %Z").strip()
|
||||
}
|
||||
|
||||
def _enrich_user_message(
|
||||
self,
|
||||
original_message: str,
|
||||
recommendation: ToolRecommendation,
|
||||
datetime_info: Optional[Dict[str, str]] = None
|
||||
) -> str:
|
||||
"""
|
||||
Enrich user message with steward's note and optional time/date.
|
||||
|
||||
Args:
|
||||
original_message: Original user query
|
||||
recommendation: Steward's recommendation (reasoning contains the note)
|
||||
datetime_info: Optional current date/time to inject silently
|
||||
|
||||
Returns:
|
||||
Enriched message with steward's note prepended
|
||||
"""
|
||||
enrichment_parts = []
|
||||
|
||||
# Silent time/date injection (if recommended)
|
||||
if datetime_info:
|
||||
enrichment_parts.append(
|
||||
f"[Current context - Date: {datetime_info['date']}, Time: {datetime_info['time']}]"
|
||||
)
|
||||
|
||||
# Prepend steward's note (the full text analysis)
|
||||
if recommendation.reasoning:
|
||||
enrichment_parts.append(
|
||||
f"[Steward's analysis: {recommendation.reasoning}]"
|
||||
)
|
||||
|
||||
# Combine enrichments with original message
|
||||
if enrichment_parts:
|
||||
enrichment = "\n".join(enrichment_parts)
|
||||
result = f"{enrichment}\n\n{original_message}"
|
||||
logger.info(f"_enrich_user_message: enrichment={len(enrichment)} chars, original={len(original_message)} chars, result={len(result)} chars")
|
||||
return result
|
||||
|
||||
return original_message
|
||||
|
||||
async def _perform_steward_analysis(self, user_query: str) -> Optional[ToolRecommendation]:
|
||||
"""
|
||||
Perform steward analysis with error handling.
|
||||
|
||||
Args:
|
||||
user_query: User's query to analyze
|
||||
|
||||
Returns:
|
||||
ToolRecommendation if successful, None if failed
|
||||
"""
|
||||
try:
|
||||
timeout = getattr(self.settings, 'analysis_timeout', 3)
|
||||
recommendation = await self.steward.analyze(user_query, timeout=timeout)
|
||||
|
||||
if recommendation:
|
||||
logger.info(
|
||||
f"Steward analysis: {len(recommendation.recommended_tools)} tools recommended"
|
||||
)
|
||||
else:
|
||||
logger.warning("Steward analysis returned None")
|
||||
|
||||
return recommendation
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
logger.error("Steward analysis timed out, falling back to single-stage")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Steward analysis failed: {e}", exc_info=True)
|
||||
return None
|
||||
|
||||
async def chat_with_analysis(
|
||||
self,
|
||||
messages: list[dict],
|
||||
conversation_id: Optional[str] = None,
|
||||
stream: bool = True
|
||||
) -> AsyncGenerator[Dict[str, Any], None]:
|
||||
"""
|
||||
Execute multi-stage chat with steward analysis and streaming.
|
||||
|
||||
Args:
|
||||
messages: Conversation messages (OpenAI format)
|
||||
conversation_id: Optional conversation ID for memory
|
||||
stream: Whether to stream response (default: True)
|
||||
|
||||
Yields:
|
||||
Dict with 'type' and relevant fields:
|
||||
- type='status': Status update (tool_name, message, arguments)
|
||||
- type='content': Response content chunk
|
||||
- type='done': Completion marker
|
||||
"""
|
||||
# Extract user query from last message
|
||||
user_query = messages[-1].get("content", "") if messages else ""
|
||||
|
||||
if not user_query:
|
||||
logger.warning("Empty user query in multi-stage analysis")
|
||||
# Fall through to single-stage
|
||||
async for chunk in self._single_stage_chat(messages, conversation_id, stream):
|
||||
yield chunk
|
||||
return
|
||||
|
||||
# Stage 1: Steward Analysis
|
||||
if self.enable_multi_stage and self.steward:
|
||||
# Emit consulting status (butler-appropriate)
|
||||
yield {
|
||||
"type": "status",
|
||||
"message": self.settings.status_consulting,
|
||||
"phase": "analysis"
|
||||
}
|
||||
|
||||
# Perform analysis
|
||||
recommendation = await self._perform_steward_analysis(user_query)
|
||||
|
||||
if recommendation is None:
|
||||
# Analysis failed, fall back to single-stage
|
||||
yield {
|
||||
"type": "status",
|
||||
"message": self.settings.status_fallback,
|
||||
"phase": "fallback"
|
||||
}
|
||||
async for chunk in self._single_stage_chat(messages, conversation_id, stream):
|
||||
yield chunk
|
||||
return
|
||||
|
||||
# Emit completion status based on tool count
|
||||
if len(recommendation.recommended_tools) == 0:
|
||||
yield {
|
||||
"type": "status",
|
||||
"message": self.settings.status_no_assistance,
|
||||
"phase": "analysis_complete"
|
||||
}
|
||||
else:
|
||||
# Format tool recommendations in butler voice
|
||||
tools_str = ", ".join(recommendation.recommended_tools)
|
||||
yield {
|
||||
"type": "status",
|
||||
"message": f"The steward recommends: {tools_str}",
|
||||
"phase": "analysis_complete",
|
||||
"recommended_tools": recommendation.recommended_tools,
|
||||
"reasoning": recommendation.reasoning
|
||||
}
|
||||
|
||||
# Check if time/date tools recommended (for silent injection)
|
||||
needs_datetime = any(
|
||||
tool in recommendation.recommended_tools
|
||||
for tool in ('get_current_time', 'get_current_date')
|
||||
)
|
||||
|
||||
datetime_info = None
|
||||
if needs_datetime:
|
||||
datetime_info = await self._get_current_datetime()
|
||||
logger.info("Injecting current date/time silently")
|
||||
|
||||
# Enrich user message with recommendations and time/date
|
||||
logger.info(f"user_query before enrichment (length={len(user_query)}): {user_query[:100]}")
|
||||
enriched_message = self._enrich_user_message(
|
||||
user_query,
|
||||
recommendation,
|
||||
datetime_info
|
||||
)
|
||||
|
||||
logger.info(f"Enriched message (length={len(enriched_message)})")
|
||||
logger.info(f"Enriched (first 200): {enriched_message[:200]}")
|
||||
logger.info(f"Enriched (last 200): {enriched_message[-200:]}")
|
||||
|
||||
# Replace last message with enriched version
|
||||
enriched_messages = messages[:-1] + [{
|
||||
"role": "user",
|
||||
"content": enriched_message
|
||||
}]
|
||||
else:
|
||||
# Multi-stage disabled, use original messages
|
||||
enriched_messages = messages
|
||||
|
||||
# Stage 2: Tatlock Execution with PydanticAI Event Streaming
|
||||
try:
|
||||
# Extract user query from enriched messages
|
||||
user_query = enriched_messages[-1].get("content", "") if enriched_messages else ""
|
||||
|
||||
if not user_query:
|
||||
logger.error("Empty query for Tatlock execution")
|
||||
yield {"type": "error", "message": "Empty query"}
|
||||
return
|
||||
|
||||
# Use PydanticAI's native streaming with event monitoring
|
||||
# Access the underlying PydanticAI agent directly
|
||||
from pydantic_ai.messages import (
|
||||
ModelResponse,
|
||||
ToolCallPart,
|
||||
ToolReturnPart,
|
||||
)
|
||||
|
||||
logger.info("Starting Tatlock execution with event streaming")
|
||||
logger.info(f"Tatlock user_query length={len(user_query)}, first 100 chars: {user_query[:100]}")
|
||||
|
||||
# Track cumulative text for delta calculation
|
||||
previous_text = ""
|
||||
|
||||
async with self.tatlock.agent.run_stream(user_query) as run:
|
||||
# Stream all messages (includes tool calls and text)
|
||||
async for message in run.stream():
|
||||
# Handle tool call events
|
||||
if hasattr(message, 'parts'):
|
||||
for part in message.parts:
|
||||
# Check for tool call parts
|
||||
if isinstance(part, ToolCallPart):
|
||||
# Emit butler-appropriate status for specific tools
|
||||
if part.tool_name == "web_search":
|
||||
# Extract query argument safely
|
||||
query = ""
|
||||
if hasattr(part, 'args') and isinstance(part.args, dict):
|
||||
query = part.args.get("query", "")
|
||||
|
||||
yield {
|
||||
"type": "status",
|
||||
"message": f"{self.settings.status_web_search} (searching for: {query})",
|
||||
"phase": "tool_execution",
|
||||
"tool_name": part.tool_name
|
||||
}
|
||||
logger.info(f"Tool call detected: {part.tool_name}")
|
||||
elif part.tool_name == "calculate":
|
||||
yield {
|
||||
"type": "status",
|
||||
"message": self.settings.status_calculate,
|
||||
"phase": "tool_execution",
|
||||
"tool_name": part.tool_name
|
||||
}
|
||||
logger.info(f"Tool call detected: {part.tool_name}")
|
||||
|
||||
# Handle text response
|
||||
if isinstance(message, ModelResponse):
|
||||
# Get current cumulative text
|
||||
current_text = await run.get_text_so_far()
|
||||
|
||||
# Calculate delta (new text only)
|
||||
delta = current_text[len(previous_text):]
|
||||
if delta:
|
||||
yield {"type": "content", "content": delta}
|
||||
|
||||
previous_text = current_text
|
||||
|
||||
# Final finish marker
|
||||
yield {"type": "content", "content": "", "finish_reason": "stop"}
|
||||
logger.info("Tatlock execution complete")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in Tatlock execution: {e}", exc_info=True)
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": f"Error generating response: {str(e)}"
|
||||
}
|
||||
|
||||
# Final completion marker
|
||||
yield {"type": "done"}
|
||||
|
||||
async def _single_stage_chat(
|
||||
self,
|
||||
messages: list[dict],
|
||||
conversation_id: Optional[str] = None,
|
||||
stream: bool = True
|
||||
) -> AsyncGenerator[Dict[str, Any], None]:
|
||||
"""
|
||||
Execute single-stage chat (no steward analysis).
|
||||
|
||||
Args:
|
||||
messages: Conversation messages
|
||||
conversation_id: Optional conversation ID
|
||||
stream: Whether to stream
|
||||
|
||||
Yields:
|
||||
Chat response chunks
|
||||
"""
|
||||
if not stream:
|
||||
# Non-streaming response
|
||||
try:
|
||||
response = await self.tatlock.chat_completion(
|
||||
messages=messages,
|
||||
conversation_id=conversation_id
|
||||
)
|
||||
yield {
|
||||
"type": "content",
|
||||
"content": response,
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
yield {"type": "done"}
|
||||
except Exception as e:
|
||||
logger.error(f"Error in single-stage chat: {e}", exc_info=True)
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": str(e)
|
||||
}
|
||||
else:
|
||||
# Streaming response
|
||||
try:
|
||||
async for chunk in self.tatlock.chat(
|
||||
messages=messages,
|
||||
conversation_id=conversation_id,
|
||||
stream=True
|
||||
):
|
||||
yield chunk
|
||||
except Exception as e:
|
||||
logger.error(f"Error in single-stage streaming: {e}", exc_info=True)
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": str(e)
|
||||
}
|
||||
|
||||
|
||||
def create_multi_stage_agent(
|
||||
tatlock_agent: Agent,
|
||||
enable_multi_stage: bool = True
|
||||
) -> MultiStageAgent:
|
||||
"""
|
||||
Create multi-stage agent orchestrator.
|
||||
|
||||
Args:
|
||||
tatlock_agent: The main Tatlock agent instance
|
||||
enable_multi_stage: Whether to enable multi-stage analysis
|
||||
|
||||
Returns:
|
||||
MultiStageAgent instance
|
||||
"""
|
||||
return MultiStageAgent(tatlock_agent, enable_multi_stage=enable_multi_stage)
|
||||
@@ -1,287 +0,0 @@
|
||||
"""
|
||||
PydanticAI Agent - Agent using PydanticAI framework with Ollama backend.
|
||||
|
||||
Based on documentation:
|
||||
- https://ai.pydantic.dev/
|
||||
- https://ai.pydantic.dev/models/#ollama
|
||||
"""
|
||||
import logging
|
||||
from typing import AsyncIterator, Dict, Any, List, Optional
|
||||
from functools import lru_cache
|
||||
|
||||
# PydanticAI imports
|
||||
try:
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.openai import OpenAIModel
|
||||
from pydantic_ai.providers.ollama import OllamaProvider
|
||||
PYDANTIC_AI_AVAILABLE = True
|
||||
except ImportError:
|
||||
PYDANTIC_AI_AVAILABLE = False
|
||||
Agent = None
|
||||
OpenAIModel = None
|
||||
OllamaProvider = None
|
||||
|
||||
from src.config import get_settings
|
||||
from src.prompts import get_prompt
|
||||
from src.memory import get_memory_manager_for_user, MessageRole
|
||||
from src.utils import sanitize_email_to_user_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PydanticAgent:
|
||||
"""
|
||||
Agent using PydanticAI framework with Ollama backend.
|
||||
Supports tool calling with proper response handling.
|
||||
|
||||
Example:
|
||||
agent = PydanticAgent(tools=[my_tool])
|
||||
response = await agent.chat_completion(messages=[{"role": "user", "content": "Hello"}])
|
||||
"""
|
||||
|
||||
def __init__(self, tools: List = None, discover_tools: bool = False, user_id: Optional[str] = None, enable_memory: Optional[bool] = None):
|
||||
if not PYDANTIC_AI_AVAILABLE:
|
||||
raise ImportError("PydanticAI not available. Install with: pip install pydantic-ai")
|
||||
|
||||
logger.info("PydanticAgent: Initializing PydanticAI agent...")
|
||||
|
||||
self.settings = get_settings()
|
||||
|
||||
# Memory configuration
|
||||
self.enable_memory = enable_memory if enable_memory is not None else self.settings.memory_enabled
|
||||
self.user_id = user_id or self.settings.default_user_id
|
||||
|
||||
# Initialize memory manager if enabled
|
||||
if self.enable_memory:
|
||||
try:
|
||||
self.memory_manager = get_memory_manager_for_user(
|
||||
user_id=self.user_id,
|
||||
buffer_max_turns=self.settings.memory_tier1_size
|
||||
)
|
||||
logger.info(f"PydanticAgent: Memory enabled for user '{self.user_id}'")
|
||||
except Exception as e:
|
||||
logger.warning(f"PydanticAgent: Failed to initialize memory: {e}. Continuing without memory.")
|
||||
self.enable_memory = False
|
||||
self.memory_manager = None
|
||||
else:
|
||||
self.memory_manager = None
|
||||
logger.info("PydanticAgent: Memory disabled")
|
||||
|
||||
# Tools can be provided explicitly or discovered
|
||||
if tools is not None:
|
||||
# Explicit tools provided
|
||||
self.tools = tools
|
||||
logger.info(f"PydanticAgent: Using {len(tools)} explicitly provided tools")
|
||||
elif discover_tools:
|
||||
# Discover tools from registry (includes local + core-api)
|
||||
logger.info("PydanticAgent: Discovering tools from registry...")
|
||||
from src.tools.registry import get_all_tools
|
||||
# Get the raw tool functions (not wrapped in ADK FunctionTool)
|
||||
tool_dict = get_all_tools()
|
||||
self.tools = list(tool_dict.values())
|
||||
logger.info(f"PydanticAgent: Discovered {len(self.tools)} tools")
|
||||
else:
|
||||
# No tools
|
||||
self.tools = []
|
||||
logger.info("PydanticAgent: No tools enabled")
|
||||
|
||||
# Load system prompt and inject current date
|
||||
from datetime import datetime
|
||||
pydantic_prompt_variant = getattr(self.settings, 'pydantic_system_prompt_variant', 'minimal_agent')
|
||||
base_prompt = get_prompt(pydantic_prompt_variant)
|
||||
|
||||
# Inject current date for general temporal awareness
|
||||
current_date = datetime.now().strftime("%A, %B %d, %Y")
|
||||
self.system_prompt = f"Today is {current_date}.\n\n{base_prompt}"
|
||||
|
||||
logger.info(f"PydanticAgent: System prompt variant: {pydantic_prompt_variant}")
|
||||
logger.info(f"PydanticAgent: System prompt: {self.system_prompt[:100]}...")
|
||||
|
||||
# Initialize Ollama model via OpenAI-compatible API
|
||||
model_name = self.settings.agent_model
|
||||
logger.info(f"PydanticAgent: Initializing Ollama model: {model_name}")
|
||||
logger.info(f"PydanticAgent: Ollama API base: {self.settings.ollama_base_url}")
|
||||
logger.info(f"PydanticAgent: Tools registered: {len(self.tools)}")
|
||||
|
||||
# Create Ollama provider with custom base URL
|
||||
# PydanticAI uses OpenAI-compatible Ollama API which requires /v1 suffix
|
||||
ollama_base_url_v1 = self.settings.ollama_base_url.rstrip('/') + '/v1'
|
||||
logger.info(f"PydanticAgent: Using Ollama URL with /v1: {ollama_base_url_v1}")
|
||||
|
||||
ollama_provider = OllamaProvider(
|
||||
base_url=ollama_base_url_v1,
|
||||
)
|
||||
|
||||
self.model = OpenAIModel(
|
||||
model_name=model_name,
|
||||
provider=ollama_provider,
|
||||
)
|
||||
|
||||
# Create PydanticAI Agent
|
||||
self.agent = Agent(
|
||||
model=self.model,
|
||||
system_prompt=self.system_prompt,
|
||||
tools=self.tools,
|
||||
)
|
||||
|
||||
logger.info("✓ PydanticAgent: Initialization complete")
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
conversation_id: str = None,
|
||||
stream: bool = True,
|
||||
prompt_variant: Optional[str] = None
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Process a chat message using PydanticAI agent.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'
|
||||
conversation_id: Optional conversation ID (not used yet)
|
||||
stream: Whether to stream responses
|
||||
prompt_variant: Optional prompt variant (not used, set in __init__)
|
||||
|
||||
Yields:
|
||||
Dict with 'type' and content. Types:
|
||||
- {"type": "content", "content": "text chunk"}
|
||||
- {"type": "content", "content": "", "finish_reason": "stop"}
|
||||
- {"type": "error", "content": "error message"}
|
||||
"""
|
||||
logger.info(f"🚀 PydanticAgent: Starting completion for message: {messages[-1]['content'][:50]}...")
|
||||
|
||||
try:
|
||||
# Extract user message (PydanticAI handles system prompt internally)
|
||||
user_messages = [m for m in messages if m["role"] != "system"]
|
||||
if not user_messages:
|
||||
raise ValueError("No user messages provided")
|
||||
|
||||
# Use the last user message
|
||||
user_query = user_messages[-1]["content"]
|
||||
logger.info(f"📤 PydanticAgent: User query: {user_query[:100]}...")
|
||||
|
||||
# Store user message in memory
|
||||
if self.enable_memory and conversation_id:
|
||||
try:
|
||||
await self.memory_manager.add_turn(
|
||||
conversation_id=conversation_id,
|
||||
role=MessageRole.USER,
|
||||
content=user_query
|
||||
)
|
||||
logger.debug(f"Stored user message in memory for conversation {conversation_id}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to store user message in memory: {e}")
|
||||
|
||||
# Run the agent
|
||||
if stream:
|
||||
# Streaming response - collect chunks to avoid async context issues
|
||||
chunks = []
|
||||
try:
|
||||
async with self.agent.run_stream(user_query) as response:
|
||||
async for chunk in response.stream_text():
|
||||
chunks.append(chunk)
|
||||
except Exception as e:
|
||||
logger.error(f"Streaming error: {e}")
|
||||
# Fall back to non-streaming
|
||||
result = await self.agent.run(user_query)
|
||||
yield {"type": "content", "content": str(result.output), "finish_reason": "stop"}
|
||||
return
|
||||
|
||||
# Convert cumulative chunks to deltas (only new content)
|
||||
previous_text = ""
|
||||
for chunk in chunks:
|
||||
# Calculate delta: new text = current chunk - previous text
|
||||
delta = chunk[len(previous_text):]
|
||||
if delta:
|
||||
yield {"type": "content", "content": delta}
|
||||
previous_text = chunk
|
||||
|
||||
# Final chunk with finish reason
|
||||
yield {"type": "content", "content": "", "finish_reason": "stop"}
|
||||
logger.info(f"📥 PydanticAgent: Streaming complete")
|
||||
|
||||
# Store assistant response in memory (streaming)
|
||||
if self.enable_memory and conversation_id:
|
||||
try:
|
||||
await self.memory_manager.add_turn(
|
||||
conversation_id=conversation_id,
|
||||
role=MessageRole.ASSISTANT,
|
||||
content=previous_text
|
||||
)
|
||||
logger.debug(f"Stored assistant response in memory for conversation {conversation_id}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to store assistant response in memory: {e}")
|
||||
|
||||
else:
|
||||
# Non-streaming response
|
||||
result = await self.agent.run(user_query)
|
||||
response_text = result.output
|
||||
logger.info(f"📥 PydanticAgent: Response: {str(response_text)[:100]}...")
|
||||
yield {"type": "content", "content": str(response_text), "finish_reason": "stop"}
|
||||
|
||||
# Store assistant response in memory (non-streaming)
|
||||
if self.enable_memory and conversation_id:
|
||||
try:
|
||||
await self.memory_manager.add_turn(
|
||||
conversation_id=conversation_id,
|
||||
role=MessageRole.ASSISTANT,
|
||||
content=str(response_text)
|
||||
)
|
||||
logger.debug(f"Stored assistant response in memory for conversation {conversation_id}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to store assistant response in memory: {e}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"PydanticAgent: Error during chat: {e}", exc_info=True)
|
||||
yield {
|
||||
"type": "error",
|
||||
"content": f"Sorry, an error occurred: {str(e)}",
|
||||
"finish_reason": "error"
|
||||
}
|
||||
|
||||
async def chat_completion(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
conversation_id: str = None,
|
||||
prompt_variant: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Get a non-streaming response from the PydanticAI agent.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts
|
||||
conversation_id: Optional conversation ID
|
||||
prompt_variant: Optional prompt variant
|
||||
|
||||
Returns:
|
||||
Complete response string
|
||||
"""
|
||||
final_content = ""
|
||||
async for chunk in self.chat(messages=messages, conversation_id=conversation_id, stream=False, prompt_variant=prompt_variant):
|
||||
if chunk["type"] == "content":
|
||||
final_content += chunk["content"]
|
||||
if chunk.get("finish_reason"):
|
||||
break
|
||||
|
||||
return final_content if final_content else "I couldn't generate a response."
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def get_pydantic_agent(tools: tuple = None, discover_tools: bool = False, user_id: str = None, enable_memory: bool = None) -> PydanticAgent:
|
||||
"""
|
||||
Get cached PydanticAI agent instance.
|
||||
|
||||
Note: tools must be a tuple for caching to work.
|
||||
Convert list to tuple before calling: get_pydantic_agent(tuple(tools))
|
||||
|
||||
Args:
|
||||
tools: Tuple of tool functions (None to use discovery)
|
||||
discover_tools: Whether to discover tools from registry
|
||||
user_id: Optional user ID for memory (defaults to config default_user_id)
|
||||
enable_memory: Optional memory enable flag (defaults to config memory_enabled)
|
||||
|
||||
Returns:
|
||||
Cached PydanticAgent instance
|
||||
"""
|
||||
tools_list = list(tools) if tools is not None else None
|
||||
return PydanticAgent(tools=tools_list, discover_tools=discover_tools, user_id=user_id, enable_memory=enable_memory)
|
||||
@@ -1,124 +0,0 @@
|
||||
"""
|
||||
Simple LiteLLM Agent - Direct text generation via LiteLLM, bypassing Google ADK.
|
||||
"""
|
||||
import logging
|
||||
from typing import AsyncIterator, Dict, Any, List, Optional
|
||||
from functools import lru_cache
|
||||
|
||||
# We will directly use litellm here
|
||||
import litellm
|
||||
|
||||
# Adjusted import paths for the new core-ai service structure
|
||||
from src.config import get_settings
|
||||
from src.prompts import get_prompt
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Simplified Agent for direct LiteLLM interaction
|
||||
class SimpleLiteLLMAgent:
|
||||
def __init__(self):
|
||||
# Enable verbose logging for LiteLLM
|
||||
litellm.set_verbose = True
|
||||
logger.info("SimpleLiteLLMAgent: LiteLLM verbose logging enabled.")
|
||||
|
||||
self.settings = get_settings()
|
||||
|
||||
# Load system prompt
|
||||
self.system_prompt = get_prompt(self.settings.system_prompt_variant)
|
||||
logger.info(f"SimpleLiteLLMAgent: System prompt variant: {self.settings.system_prompt_variant}")
|
||||
logger.info(f"SimpleLiteLLMAgent: System prompt: {self.system_prompt[:100]}...")
|
||||
|
||||
# Initialize LiteLLM for Ollama (format: "ollama/model_name")
|
||||
model_name = self.settings.agent_model
|
||||
litellm_model = f"ollama/{model_name}"
|
||||
|
||||
logger.info(f"SimpleLiteLLMAgent: Initializing LiteLLM direct model: {litellm_model}")
|
||||
logger.info(f"SimpleLiteLLMAgent: Ollama base URL from settings: {self.settings.ollama_base_url}")
|
||||
|
||||
self.model_params = {
|
||||
"model": litellm_model,
|
||||
"api_base": self.settings.ollama_base_url,
|
||||
"temperature": 0.1,
|
||||
# No tool definitions passed here to force text generation
|
||||
}
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
conversation_id: str = None, # Not used in this simple mode
|
||||
stream: bool = True,
|
||||
prompt_variant: Optional[str] = None # Not used in this simple mode
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Processes a chat message using direct LiteLLM completion.
|
||||
"""
|
||||
logger.info(f"🚀 SimpleLiteLLMAgent: Starting direct LiteLLM completion for message: {messages[-1]['content'][:50]}...")
|
||||
try:
|
||||
# Prepare messages in LiteLLM format
|
||||
litellm_messages = [{"role": m["role"], "content": m["content"]} for m in messages]
|
||||
|
||||
# Inject system prompt if not already present
|
||||
if not litellm_messages or litellm_messages[0]["role"] != "system":
|
||||
litellm_messages.insert(0, {"role": "system", "content": self.system_prompt})
|
||||
logger.info("✓ SimpleLiteLLMAgent: System prompt injected")
|
||||
|
||||
# Log full message payload for debugging
|
||||
logger.info(f"📤 SimpleLiteLLMAgent: Sending {len(litellm_messages)} messages to LiteLLM:")
|
||||
for i, msg in enumerate(litellm_messages):
|
||||
content_preview = msg['content'][:100] + "..." if len(msg['content']) > 100 else msg['content']
|
||||
logger.info(f" [{i}] {msg['role']}: {content_preview}")
|
||||
|
||||
# Use acompletion for async environments
|
||||
response = await litellm.acompletion(
|
||||
messages=litellm_messages,
|
||||
stream=stream,
|
||||
**self.model_params
|
||||
)
|
||||
|
||||
if stream:
|
||||
chunk_count = 0
|
||||
async for chunk in response:
|
||||
chunk_count += 1
|
||||
content_delta = chunk.choices[0].delta.content if chunk.choices[0].delta.content else ""
|
||||
finish_reason = chunk.choices[0].finish_reason
|
||||
if content_delta:
|
||||
yield {"type": "content", "content": content_delta}
|
||||
if finish_reason:
|
||||
logger.info(f"📥 SimpleLiteLLMAgent: Stream completed after {chunk_count} chunks. Finish reason: {finish_reason}")
|
||||
yield {"type": "content", "content": "", "finish_reason": finish_reason}
|
||||
else:
|
||||
content = response.choices[0].message.content
|
||||
logger.info(f"📥 SimpleLiteLLMAgent: Response received: {content[:200]}..." if len(content) > 200 else f"📥 SimpleLiteLLMAgent: Response received: {content}")
|
||||
yield {"type": "content", "content": content, "finish_reason": "stop"}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"SimpleLiteLLMAgent: Error in direct LiteLLM chat: {e}", exc_info=True)
|
||||
yield {
|
||||
"type": "error",
|
||||
"content": f"Sorry, an error occurred during text generation: {str(e)}",
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
|
||||
async def chat_completion(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
conversation_id: str = None,
|
||||
prompt_variant: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Get a non-streaming response from the direct LiteLLM chat.
|
||||
"""
|
||||
final_content = ""
|
||||
async for chunk in self.chat(messages=messages, conversation_id=conversation_id, stream=False, prompt_variant=prompt_variant):
|
||||
if chunk["type"] == "content":
|
||||
final_content += chunk["content"]
|
||||
if chunk.get("finish_reason") == "stop":
|
||||
break
|
||||
return final_content if final_content else "I couldn't generate a response."
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def get_simple_litellm_agent() -> SimpleLiteLLMAgent:
|
||||
"""Get cached simple LiteLLM agent instance"""
|
||||
return SimpleLiteLLMAgent()
|
||||
@@ -1,270 +0,0 @@
|
||||
"""
|
||||
Steward Analysis Agent - Analyzes queries and recommends optimal tools.
|
||||
|
||||
The steward is a lightweight analysis layer that examines user queries and
|
||||
recommends 0-5 tools that would be most helpful for answering the query.
|
||||
Uses the same model as Tatlock (mistral-nemo) for consistency and performance.
|
||||
"""
|
||||
import logging
|
||||
from typing import Optional
|
||||
from pydantic import BaseModel, Field
|
||||
from datetime import datetime
|
||||
|
||||
try:
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.openai import OpenAIModel
|
||||
from pydantic_ai.providers.ollama import OllamaProvider
|
||||
PYDANTIC_AI_AVAILABLE = True
|
||||
except ImportError:
|
||||
PYDANTIC_AI_AVAILABLE = False
|
||||
Agent = None
|
||||
OpenAIModel = None
|
||||
OllamaProvider = None
|
||||
|
||||
from src.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ToolRecommendation(BaseModel):
|
||||
"""Structured tool recommendations from steward analysis."""
|
||||
intent: str = Field(description="Brief description of what user wants")
|
||||
recommended_tools: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="List of 0-5 recommended tool names"
|
||||
)
|
||||
reasoning: str = Field(description="Why these tools are optimal (or why none needed)")
|
||||
requires_assistance: bool = Field(
|
||||
default=False,
|
||||
description="False if 0 tools recommended (general knowledge sufficient)"
|
||||
)
|
||||
|
||||
|
||||
class StewardAgent:
|
||||
"""
|
||||
The Steward - Tatlock's analytical assistant for tool selection.
|
||||
|
||||
Analyzes user queries and recommends which tools would be most helpful.
|
||||
Uses the same model as Tatlock (mistral-nemo) to avoid VRAM overhead.
|
||||
|
||||
Key Rules:
|
||||
- ALWAYS recommend 'calculate' for any math/arithmetic
|
||||
- Recommend time/date tools for temporal queries (will be injected silently)
|
||||
- Recommend web_search for real-time information
|
||||
- Can recommend 0 tools if general knowledge is sufficient
|
||||
- Maximum 5 tool recommendations
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: Optional[str] = None):
|
||||
"""
|
||||
Initialize the steward agent.
|
||||
|
||||
Args:
|
||||
model_name: Model to use (default: from settings, same as Tatlock)
|
||||
"""
|
||||
if not PYDANTIC_AI_AVAILABLE:
|
||||
raise ImportError("PydanticAI not available. Install with: pip install pydantic-ai")
|
||||
|
||||
logger.info("StewardAgent: Initializing steward analysis agent...")
|
||||
|
||||
self.settings = get_settings()
|
||||
|
||||
# Use steward model from config (defaults to same as Tatlock to stay in VRAM)
|
||||
if model_name is None:
|
||||
model_name = getattr(self.settings, 'steward_model', self.settings.agent_model)
|
||||
|
||||
# Initialize Ollama model via OpenAI-compatible API
|
||||
ollama_base_url = self.settings.ollama_base_url
|
||||
# OpenAI-compatible endpoint requires /v1 suffix
|
||||
ollama_base_url_v1 = f"{ollama_base_url}/v1" if not ollama_base_url.endswith('/v1') else ollama_base_url
|
||||
|
||||
ollama_provider = OllamaProvider(base_url=ollama_base_url_v1)
|
||||
|
||||
self.model = OpenAIModel(
|
||||
model_name=model_name,
|
||||
provider=ollama_provider,
|
||||
)
|
||||
|
||||
# Generate system prompt with tool catalog
|
||||
self.system_prompt = self._generate_steward_prompt()
|
||||
|
||||
# Create PydanticAI agent WITHOUT structured output (text-based for reliability)
|
||||
self.agent = Agent(
|
||||
model=self.model,
|
||||
system_prompt=self.system_prompt,
|
||||
# NO output_type - returns plain text instead of structured JSON
|
||||
)
|
||||
|
||||
logger.info(f"StewardAgent: Initialized with model {model_name}")
|
||||
|
||||
def _generate_steward_prompt(self) -> str:
|
||||
"""
|
||||
Generate steward's system prompt with tool catalog.
|
||||
|
||||
Returns:
|
||||
System prompt string
|
||||
"""
|
||||
# Get current date for context
|
||||
current_date = datetime.now().strftime("%A, %B %d, %Y")
|
||||
|
||||
# Import tool registry to get available tools
|
||||
from src.tools.registry import get_all_tools
|
||||
|
||||
try:
|
||||
tools = get_all_tools(include_openapi=False) # Local tools only for now
|
||||
tool_catalog = self._format_tool_catalog(tools)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load tool catalog: {e}")
|
||||
tool_catalog = "Tool catalog unavailable"
|
||||
|
||||
return f"""You are the steward of Tatlock's household. Your role is to analyze incoming requests and provide a brief, helpful note to Tatlock about how best to address them.
|
||||
|
||||
Today is {current_date}.
|
||||
|
||||
AVAILABLE TOOLS FOR TATLOCK:
|
||||
{tool_catalog}
|
||||
|
||||
YOUR TASK:
|
||||
Write a concise analysis note (2-3 sentences) that:
|
||||
1. Identifies what the user is asking for
|
||||
2. Recommends which tools would be most helpful (if any)
|
||||
3. Explains why those tools are appropriate
|
||||
|
||||
CRITICAL TOOL SELECTION RULES:
|
||||
- ANY math/calculations → recommend 'calculate' (LLMs are unreliable with math)
|
||||
- Current information (news, markets, weather) → recommend 'web_search'
|
||||
- Time/date queries → recommend 'get_current_time' or 'get_current_date'
|
||||
- Infrastructure operations → recommend appropriate infrastructure tools
|
||||
- General knowledge → no tools needed
|
||||
|
||||
OUTPUT FORMAT (plain text, natural language):
|
||||
Write your note as you would brief a butler about a household matter. Be concise and specific.
|
||||
|
||||
EXAMPLE:
|
||||
"Sir's query requires current financial market data and mathematical calculations. I recommend employing web_search to gather the latest information, followed by calculate for the numerical analysis. This combination will ensure accurate, up-to-date results."
|
||||
|
||||
Now write your analysis note for the query."""
|
||||
|
||||
def _format_tool_catalog(self, tools: dict) -> str:
|
||||
"""
|
||||
Format available tools into a readable catalog.
|
||||
|
||||
Args:
|
||||
tools: Dictionary of tool name -> function
|
||||
|
||||
Returns:
|
||||
Formatted tool catalog string
|
||||
"""
|
||||
import inspect
|
||||
|
||||
catalog_lines = []
|
||||
|
||||
for name, func in tools.items():
|
||||
# Get function signature
|
||||
try:
|
||||
sig = inspect.signature(func)
|
||||
params = [p.name for p in sig.parameters.values()]
|
||||
signature = f"{name}({', '.join(params)})"
|
||||
except Exception:
|
||||
signature = f"{name}(...)"
|
||||
|
||||
# Get docstring (first line only)
|
||||
doc = inspect.getdoc(func)
|
||||
if doc:
|
||||
description = doc.split('\n')[0]
|
||||
else:
|
||||
description = "No description"
|
||||
|
||||
catalog_lines.append(f" • {signature}")
|
||||
catalog_lines.append(f" {description}")
|
||||
|
||||
return "\n".join(catalog_lines)
|
||||
|
||||
async def analyze(self, query: str, timeout: Optional[float] = None) -> ToolRecommendation:
|
||||
"""
|
||||
Analyze user query and recommend tools.
|
||||
|
||||
Args:
|
||||
query: User's query to analyze
|
||||
timeout: Optional timeout in seconds
|
||||
|
||||
Returns:
|
||||
ToolRecommendation with recommended tools and reasoning
|
||||
|
||||
Raises:
|
||||
asyncio.TimeoutError: If analysis exceeds timeout
|
||||
Exception: If analysis fails
|
||||
"""
|
||||
if timeout is None:
|
||||
timeout = self.settings.analysis_timeout if hasattr(self.settings, 'analysis_timeout') else 3
|
||||
|
||||
logger.info(f"Steward analyzing query: {query[:100]}...")
|
||||
|
||||
try:
|
||||
import asyncio
|
||||
|
||||
# Run analysis with timeout (text output, no structured format)
|
||||
result = await asyncio.wait_for(
|
||||
self.agent.run(f"User query: {query}"),
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
# Get text response (steward's note for Tatlock)
|
||||
steward_note = str(result.output).strip()
|
||||
|
||||
# Strip surrounding quotes if present (avoid JSON escaping issues)
|
||||
if steward_note.startswith('"') and steward_note.endswith('"'):
|
||||
steward_note = steward_note[1:-1]
|
||||
elif steward_note.startswith("'") and steward_note.endswith("'"):
|
||||
steward_note = steward_note[1:-1]
|
||||
|
||||
logger.info(f"Steward's note: {steward_note[:150]}...")
|
||||
|
||||
# Extract tool names mentioned in the note (for logging and status)
|
||||
tools_mentioned = []
|
||||
note_lower = steward_note.lower()
|
||||
for tool in ['web_search', 'calculate', 'get_current_time', 'get_current_date',
|
||||
'add_days_to_date', 'calculate_date_difference']:
|
||||
if tool in note_lower:
|
||||
tools_mentioned.append(tool)
|
||||
|
||||
logger.info(f"Tools mentioned in steward's note: {tools_mentioned}")
|
||||
|
||||
# Return as ToolRecommendation for compatibility
|
||||
# The reasoning field contains the full steward's note
|
||||
recommendation = ToolRecommendation(
|
||||
intent=query[:50] + "..." if len(query) > 50 else query,
|
||||
recommended_tools=tools_mentioned, # Extracted for display/logging
|
||||
reasoning=steward_note, # Full note text for Tatlock
|
||||
requires_assistance=len(tools_mentioned) > 0 or "recommend" in note_lower
|
||||
)
|
||||
|
||||
return recommendation
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(f"Steward analysis timed out after {timeout}s")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Steward analysis failed: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
|
||||
# Singleton instance
|
||||
_steward_agent: Optional[StewardAgent] = None
|
||||
|
||||
|
||||
def get_steward_agent(model_name: Optional[str] = None) -> StewardAgent:
|
||||
"""
|
||||
Get singleton steward agent instance.
|
||||
|
||||
Args:
|
||||
model_name: Optional model name (uses settings default if not provided)
|
||||
|
||||
Returns:
|
||||
StewardAgent instance
|
||||
"""
|
||||
global _steward_agent
|
||||
if _steward_agent is None:
|
||||
_steward_agent = StewardAgent(model_name=model_name)
|
||||
return _steward_agent
|
||||
@@ -1,464 +0,0 @@
|
||||
"""
|
||||
StreamHandler - Central coordinator for async agent tasks and message streaming.
|
||||
|
||||
Decouples PydanticAI and other components from the output stream, allowing
|
||||
any component to emit messages independently. Handles OpenAI format conversion,
|
||||
message history, and proper scope management.
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, Any, Optional, List, AsyncIterator, Callable
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MessageType(Enum):
|
||||
"""Types of messages that can be emitted."""
|
||||
STATUS = "status"
|
||||
CONTENT = "content"
|
||||
TOOL_CALL = "tool_call"
|
||||
TOOL_RESULT = "tool_result"
|
||||
ERROR = "error"
|
||||
DONE = "done"
|
||||
|
||||
|
||||
class StreamHandler:
|
||||
"""
|
||||
Central coordinator for async agent tasks and output streaming.
|
||||
|
||||
The StreamHandler:
|
||||
- Manages the output stream to the user (OpenAI SSE format)
|
||||
- Spawns and coordinates async tasks (steward, tatlock, tools, etc.)
|
||||
- Receives messages from anywhere in the async tree via queue
|
||||
- Maintains message history and conversation state
|
||||
- Handles format conversions and scope management
|
||||
|
||||
Architecture:
|
||||
User Request
|
||||
↓
|
||||
StreamHandler (coordinator)
|
||||
├─▶ Steward (async task) ──▶ emit_message()
|
||||
├─▶ Tatlock (async task) ──▶ emit_message()
|
||||
│ └─▶ Tools ──▶ emit_message()
|
||||
└─▶ Future: Handyman, Secretary, etc.
|
||||
↓
|
||||
Formatted Output Stream
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str = "Tatlock"):
|
||||
"""
|
||||
Initialize stream handler.
|
||||
|
||||
Args:
|
||||
model_name: Model name for OpenAI format responses
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.message_queue: asyncio.Queue = asyncio.Queue()
|
||||
self.message_history: List[Dict[str, str]] = []
|
||||
self.tasks: List[asyncio.Task] = []
|
||||
self.running = False
|
||||
|
||||
# Track cumulative content for delta calculation
|
||||
self._content_buffer = ""
|
||||
|
||||
logger.info("StreamHandler: Initialized")
|
||||
|
||||
async def emit_message(
|
||||
self,
|
||||
message_type: MessageType,
|
||||
content: Optional[str] = None,
|
||||
phase: Optional[str] = None,
|
||||
tool_name: Optional[str] = None,
|
||||
**kwargs
|
||||
):
|
||||
"""
|
||||
Emit a message to the stream from any component.
|
||||
|
||||
This is the main interface for components (steward, tatlock, tools)
|
||||
to send messages to the user.
|
||||
|
||||
Args:
|
||||
message_type: Type of message (status, content, error, etc.)
|
||||
content: Message content
|
||||
phase: Phase identifier (for status messages)
|
||||
tool_name: Tool name (for tool-related messages)
|
||||
**kwargs: Additional message-specific data
|
||||
|
||||
Example:
|
||||
await handler.emit_message(
|
||||
MessageType.STATUS,
|
||||
content="Consulting the steward...",
|
||||
phase="analysis"
|
||||
)
|
||||
"""
|
||||
message = {
|
||||
"type": message_type.value,
|
||||
"content": content,
|
||||
"phase": phase,
|
||||
"tool_name": tool_name,
|
||||
**kwargs
|
||||
}
|
||||
|
||||
# Remove None values
|
||||
message = {k: v for k, v in message.items() if v is not None}
|
||||
|
||||
await self.message_queue.put(message)
|
||||
logger.debug(f"Message emitted: {message_type.value}")
|
||||
|
||||
async def process_stream(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
multi_stage_enabled: bool = True
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Main entry point - processes a chat request and streams responses.
|
||||
|
||||
Args:
|
||||
messages: Conversation history (OpenAI format)
|
||||
multi_stage_enabled: Whether to use multi-stage analysis
|
||||
|
||||
Yields:
|
||||
Formatted message dicts ready for SSE output
|
||||
"""
|
||||
self.running = True
|
||||
self.message_history = messages
|
||||
|
||||
try:
|
||||
# Extract user query
|
||||
user_query = messages[-1].get("content", "") if messages else ""
|
||||
if not user_query:
|
||||
await self.emit_message(
|
||||
MessageType.ERROR,
|
||||
content="Empty user query"
|
||||
)
|
||||
return
|
||||
|
||||
# Spawn the main orchestration task
|
||||
orchestration_task = asyncio.create_task(
|
||||
self._orchestrate_request(user_query, multi_stage_enabled)
|
||||
)
|
||||
self.tasks.append(orchestration_task)
|
||||
|
||||
# Stream messages as they arrive
|
||||
async for message in self._stream_messages():
|
||||
yield message
|
||||
|
||||
finally:
|
||||
self.running = False
|
||||
await self._cleanup_tasks()
|
||||
|
||||
async def _orchestrate_request(
|
||||
self,
|
||||
user_query: str,
|
||||
multi_stage_enabled: bool
|
||||
):
|
||||
"""
|
||||
Orchestrate the full request flow.
|
||||
|
||||
This runs as an async task and coordinates:
|
||||
1. Steward analysis (if enabled)
|
||||
2. Tatlock execution
|
||||
3. Tool calls
|
||||
4. Response generation
|
||||
|
||||
Args:
|
||||
user_query: User's query
|
||||
multi_stage_enabled: Whether to use steward
|
||||
"""
|
||||
try:
|
||||
recommendation = None
|
||||
|
||||
# Stage 1: Steward Analysis
|
||||
if multi_stage_enabled:
|
||||
await self.emit_message(
|
||||
MessageType.STATUS,
|
||||
content="Consulting the steward on the matter...",
|
||||
phase="analysis"
|
||||
)
|
||||
|
||||
# Run steward analysis
|
||||
recommendation = await self._run_steward_analysis(user_query)
|
||||
|
||||
if recommendation:
|
||||
tools_str = ", ".join(recommendation.get("tools", []))
|
||||
await self.emit_message(
|
||||
MessageType.STATUS,
|
||||
content=f"The steward recommends: {tools_str}" if tools_str else "The steward advises no further assistance is required",
|
||||
phase="analysis_complete"
|
||||
)
|
||||
else:
|
||||
await self.emit_message(
|
||||
MessageType.STATUS,
|
||||
content="The steward is unavailable. Proceeding with standard protocols...",
|
||||
phase="fallback"
|
||||
)
|
||||
|
||||
# Stage 2: Enrich message with steward's note
|
||||
enriched_query = user_query
|
||||
if recommendation and recommendation.get("note"):
|
||||
enriched_query = f"[Steward's analysis: {recommendation['note']}]\n\n{user_query}"
|
||||
|
||||
# Stage 3: Tatlock execution
|
||||
logger.info(f"StreamHandler: Starting Tatlock execution with query length={len(enriched_query)}")
|
||||
await self._run_tatlock(enriched_query)
|
||||
logger.info("StreamHandler: Tatlock execution complete")
|
||||
|
||||
# Signal completion
|
||||
await self.emit_message(MessageType.DONE)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Orchestration error: {e}", exc_info=True)
|
||||
await self.emit_message(
|
||||
MessageType.ERROR,
|
||||
content=f"Error: {str(e)}"
|
||||
)
|
||||
await self.emit_message(MessageType.DONE)
|
||||
|
||||
async def _run_steward_analysis(self, query: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Run steward analysis as an async task.
|
||||
|
||||
Args:
|
||||
query: User query to analyze
|
||||
|
||||
Returns:
|
||||
Dict with steward analysis or None if failed
|
||||
"""
|
||||
try:
|
||||
from src.agents.steward_agent import get_steward_agent
|
||||
from src.config import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
steward_model = getattr(settings, 'steward_model', 'mistral-nemo:latest')
|
||||
timeout = getattr(settings, 'analysis_timeout', 10)
|
||||
|
||||
steward = get_steward_agent(model_name=steward_model)
|
||||
recommendation = await steward.analyze(query, timeout=timeout)
|
||||
|
||||
if recommendation:
|
||||
return {
|
||||
"tools": recommendation.recommended_tools,
|
||||
"note": recommendation.reasoning,
|
||||
"intent": recommendation.intent
|
||||
}
|
||||
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Steward analysis failed: {e}", exc_info=True)
|
||||
return None
|
||||
|
||||
async def _run_tatlock(self, query: str):
|
||||
"""
|
||||
Run Tatlock agent using PydanticAI's run_stream() - each event reports to StreamHandler.
|
||||
|
||||
Instead of PydanticAI controlling the stream, we iterate through each message
|
||||
and report to the StreamHandler, giving us full control over output.
|
||||
|
||||
Args:
|
||||
query: Enriched user query (with steward note if available)
|
||||
"""
|
||||
try:
|
||||
from src.agents.pydantic_agent import get_pydantic_agent
|
||||
from pydantic_ai.messages import (
|
||||
ModelResponse,
|
||||
ToolCallPart,
|
||||
ToolReturnPart,
|
||||
)
|
||||
|
||||
tatlock = get_pydantic_agent()
|
||||
|
||||
# Track cumulative text for delta calculation
|
||||
previous_text = ""
|
||||
|
||||
# Use run_stream() with event monitoring
|
||||
async with tatlock.agent.run_stream(query) as run:
|
||||
# Stream all messages (includes tool calls and text)
|
||||
async for message in run.stream():
|
||||
# Handle tool call events
|
||||
if hasattr(message, 'parts'):
|
||||
for part in message.parts:
|
||||
# Check for tool call parts
|
||||
if isinstance(part, ToolCallPart):
|
||||
await self._handle_tool_call(part)
|
||||
|
||||
# Handle text response
|
||||
if isinstance(message, ModelResponse):
|
||||
# Get current cumulative text
|
||||
current_text = await run.get_text_so_far()
|
||||
|
||||
# Calculate delta (new text only)
|
||||
delta = current_text[len(previous_text):]
|
||||
if delta:
|
||||
await self.emit_message(
|
||||
MessageType.CONTENT,
|
||||
content=delta
|
||||
)
|
||||
|
||||
previous_text = current_text
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Tatlock execution error: {e}", exc_info=True)
|
||||
await self.emit_message(
|
||||
MessageType.ERROR,
|
||||
content=f"Error generating response: {str(e)}"
|
||||
)
|
||||
|
||||
async def _handle_tool_call(self, part):
|
||||
"""
|
||||
Handle a tool call part from PydanticAI message stream.
|
||||
|
||||
Args:
|
||||
part: ToolCallPart from message.parts
|
||||
"""
|
||||
tool_name = getattr(part, 'tool_name', 'unknown')
|
||||
|
||||
# Emit butler-appropriate status for specific tools
|
||||
if tool_name == "web_search":
|
||||
args = getattr(part, 'args', {})
|
||||
query = args.get('query', '') if isinstance(args, dict) else ''
|
||||
await self.emit_message(
|
||||
MessageType.STATUS,
|
||||
content=f"Making enquiries... (searching for: {query})",
|
||||
phase="tool_execution",
|
||||
tool_name=tool_name
|
||||
)
|
||||
logger.info(f"Tool call detected: web_search (query: {query})")
|
||||
elif tool_name == "calculate":
|
||||
await self.emit_message(
|
||||
MessageType.STATUS,
|
||||
content="Calculating... (one prefers precision in mathematics)",
|
||||
phase="tool_execution",
|
||||
tool_name=tool_name
|
||||
)
|
||||
logger.info(f"Tool call detected: calculate")
|
||||
|
||||
async def _stream_messages(self) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Stream messages from the queue, formatting for OpenAI SSE.
|
||||
|
||||
Yields:
|
||||
Formatted message dicts ready for SSE output
|
||||
"""
|
||||
while self.running or not self.message_queue.empty():
|
||||
try:
|
||||
# Wait for message with timeout
|
||||
message = await asyncio.wait_for(
|
||||
self.message_queue.get(),
|
||||
timeout=0.1
|
||||
)
|
||||
|
||||
# Format based on message type
|
||||
formatted = self._format_message(message)
|
||||
if formatted:
|
||||
yield formatted
|
||||
|
||||
# Check if done
|
||||
if message.get("type") == "done":
|
||||
break
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.error(f"Error streaming message: {e}", exc_info=True)
|
||||
continue
|
||||
|
||||
def _format_message(self, message: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Format a message for OpenAI SSE output.
|
||||
|
||||
Args:
|
||||
message: Internal message dict
|
||||
|
||||
Returns:
|
||||
Formatted message or None
|
||||
"""
|
||||
msg_type = message.get("type")
|
||||
|
||||
if msg_type == "status":
|
||||
# Format status with box-drawing characters
|
||||
phase = message.get("phase", "")
|
||||
content = message.get("content", "")
|
||||
|
||||
# Apply box-drawing formatting
|
||||
if phase == "analysis":
|
||||
formatted_content = f"┌─ {content}"
|
||||
elif phase == "analysis_complete":
|
||||
formatted_content = f"└─ {content}"
|
||||
elif phase in ("tool_execution", "fallback"):
|
||||
formatted_content = f"├─ {content}"
|
||||
else:
|
||||
formatted_content = f"├─ {content}"
|
||||
|
||||
return {
|
||||
"type": "status",
|
||||
"message": formatted_content,
|
||||
"phase": phase,
|
||||
"tool_name": message.get("tool_name"),
|
||||
"arguments": message.get("arguments")
|
||||
}
|
||||
|
||||
elif msg_type == "content":
|
||||
return {
|
||||
"type": "content",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {"content": message.get("content", "")},
|
||||
"finish_reason": message.get("finish_reason")
|
||||
}],
|
||||
"model": self.model_name
|
||||
}
|
||||
|
||||
elif msg_type == "error":
|
||||
return {
|
||||
"type": "error",
|
||||
"message": message.get("content", "Unknown error")
|
||||
}
|
||||
|
||||
elif msg_type == "done":
|
||||
return {
|
||||
"type": "content",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {"content": ""},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"model": self.model_name
|
||||
}
|
||||
|
||||
return None
|
||||
|
||||
async def _cleanup_tasks(self):
|
||||
"""Clean up any running async tasks."""
|
||||
for task in self.tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
self.tasks.clear()
|
||||
logger.debug("StreamHandler: Tasks cleaned up")
|
||||
|
||||
|
||||
# Singleton instance
|
||||
_stream_handler: Optional[StreamHandler] = None
|
||||
|
||||
|
||||
def get_stream_handler(model_name: str = "Tatlock") -> StreamHandler:
|
||||
"""
|
||||
Get or create stream handler instance.
|
||||
|
||||
Args:
|
||||
model_name: Model name for responses
|
||||
|
||||
Returns:
|
||||
StreamHandler instance
|
||||
"""
|
||||
global _stream_handler
|
||||
if _stream_handler is None:
|
||||
_stream_handler = StreamHandler(model_name=model_name)
|
||||
return _stream_handler
|
||||
@@ -1,288 +0,0 @@
|
||||
"""
|
||||
Text-Based Steward Agent - Reliable domain analysis without JSON validation.
|
||||
|
||||
The steward analyzes queries and recommends which domains (toolsets) are relevant.
|
||||
Returns plain text instead of structured JSON to avoid validation failures.
|
||||
|
||||
Architecture:
|
||||
- Input: User query
|
||||
- Process: LLM analysis (mistral-nemo, 2-3s)
|
||||
- Output: Plain text with domain recommendations
|
||||
- Parsing: Simple keyword extraction (no JSON)
|
||||
|
||||
Reliability: 100% (no JSON validation failures)
|
||||
"""
|
||||
import logging
|
||||
import httpx
|
||||
import re
|
||||
from typing import List, Dict, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class StewardAnalysis:
|
||||
"""
|
||||
Result of steward's domain analysis.
|
||||
|
||||
Attributes:
|
||||
domains: List of recommended domain names (e.g., ['core', 'infrastructure'])
|
||||
reasoning: Steward's plain text explanation
|
||||
query: Original user query
|
||||
confidence: Confidence score (0.0-1.0, based on keyword matches)
|
||||
"""
|
||||
domains: List[str]
|
||||
reasoning: str
|
||||
query: str
|
||||
confidence: float = 1.0
|
||||
|
||||
|
||||
class TextBasedSteward:
|
||||
"""
|
||||
Reliable steward using plain text output.
|
||||
|
||||
No JSON validation = No failures = Happy household.
|
||||
|
||||
Example:
|
||||
steward = TextBasedSteward()
|
||||
analysis = await steward.analyze("Check if nginx is running")
|
||||
# Returns: StewardAnalysis(domains=['core', 'infrastructure'], reasoning="...", ...)
|
||||
"""
|
||||
|
||||
# Domain keyword mapping
|
||||
# Format: domain_name → list of keywords that indicate this domain
|
||||
DOMAIN_KEYWORDS = {
|
||||
'infrastructure': [
|
||||
'infrastructure', 'docker', 'container', 'service',
|
||||
'nginx', 'ollama', 'core-ai', 'core-api',
|
||||
'restart', 'status', 'logs', 'system',
|
||||
'running', 'stopped', 'deployed', 'monitoring',
|
||||
'memory', 'cpu', 'resource', 'health'
|
||||
],
|
||||
'secretary': [
|
||||
'secretary', 'calendar', 'remind', 'appointment',
|
||||
'schedule', 'meeting', 'task', 'todo', 'event',
|
||||
'tomorrow', 'next week', 'later today'
|
||||
],
|
||||
'home_automation': [
|
||||
'home', 'automation', 'light', 'climate',
|
||||
'temperature', 'thermostat', 'scene',
|
||||
'turn on', 'turn off', 'dim', 'brighten'
|
||||
],
|
||||
'webdev': [
|
||||
'webdev', 'screenshot', 'browser', 'html',
|
||||
'webpage', 'css', 'javascript', 'render'
|
||||
]
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = "mistral-nemo:latest",
|
||||
ollama_url: str = "http://ollama:11434",
|
||||
timeout: float = 5.0
|
||||
):
|
||||
"""
|
||||
Initialize text-based steward.
|
||||
|
||||
Args:
|
||||
model: Ollama model to use
|
||||
ollama_url: Ollama API base URL
|
||||
timeout: Request timeout in seconds
|
||||
"""
|
||||
self.model = model
|
||||
self.ollama_url = ollama_url
|
||||
self.timeout = timeout
|
||||
|
||||
logger.info(f"TextBasedSteward: Initialized with {model}")
|
||||
|
||||
async def analyze(self, query: str) -> StewardAnalysis:
|
||||
"""
|
||||
Analyze query and recommend relevant domains.
|
||||
|
||||
This is the main entry point. Returns domains to load for the query.
|
||||
|
||||
Args:
|
||||
query: User's query to analyze
|
||||
|
||||
Returns:
|
||||
StewardAnalysis with recommended domains and reasoning
|
||||
|
||||
Example:
|
||||
>>> analysis = await steward.analyze("Turn off the living room lights")
|
||||
>>> analysis.domains
|
||||
['core', 'home_automation']
|
||||
"""
|
||||
logger.info(f"Steward analyzing: {query[:100]}...")
|
||||
|
||||
try:
|
||||
# Generate analysis text (2-3s)
|
||||
reasoning = await self._generate_analysis(query)
|
||||
|
||||
# Extract domains from text (< 1ms)
|
||||
domains = self._extract_domains(reasoning, query)
|
||||
|
||||
# Calculate confidence
|
||||
confidence = self._calculate_confidence(domains, reasoning)
|
||||
|
||||
logger.info(f"Steward recommends: {domains} (confidence: {confidence:.2f})")
|
||||
|
||||
return StewardAnalysis(
|
||||
domains=domains,
|
||||
reasoning=reasoning,
|
||||
query=query,
|
||||
confidence=confidence
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Steward analysis failed: {e}", exc_info=True)
|
||||
|
||||
# Fallback: Use keyword-based analysis only
|
||||
logger.info("Falling back to keyword-based domain selection")
|
||||
domains = self._keyword_fallback(query)
|
||||
|
||||
return StewardAnalysis(
|
||||
domains=domains,
|
||||
reasoning=f"Steward unavailable. Using keyword analysis: {', '.join(domains)}",
|
||||
query=query,
|
||||
confidence=0.5
|
||||
)
|
||||
|
||||
async def _generate_analysis(self, query: str) -> str:
|
||||
"""
|
||||
Generate plain text analysis using LLM.
|
||||
|
||||
Returns plain text (not JSON) to avoid validation failures.
|
||||
"""
|
||||
prompt = self._build_analysis_prompt(query)
|
||||
|
||||
# Call Ollama API directly (not via PydanticAI)
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.post(
|
||||
f"{self.ollama_url}/api/generate",
|
||||
json={
|
||||
"model": self.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()
|
||||
|
||||
def _build_analysis_prompt(self, query: str) -> str:
|
||||
"""
|
||||
Build the steward's analysis prompt.
|
||||
|
||||
Asks for plain text domain recommendations (no JSON).
|
||||
"""
|
||||
return f"""You are the steward of a British household, assisting with domain selection for the butler (Tatlock).
|
||||
|
||||
Your task: Analyze the query and suggest which domains are relevant.
|
||||
|
||||
AVAILABLE DOMAINS:
|
||||
• core - Essential utilities (mathematics, web search, time/date queries)
|
||||
• infrastructure - System management (Docker containers, services, monitoring, logs)
|
||||
• secretary - Calendar, reminders, appointments, task management (not yet available)
|
||||
• home_automation - Smart home control (lights, climate, scenes) (not yet available)
|
||||
• webdev - Web development tools (screenshots, browser automation) (not yet available)
|
||||
|
||||
GUIDELINES:
|
||||
- Mathematical calculations → core
|
||||
- Web searches, current events → core
|
||||
- Time/date queries → core
|
||||
- Docker, containers, services, logs → infrastructure
|
||||
- Calendar, reminders, appointments → secretary
|
||||
- Smart home, lights, temperature → home_automation
|
||||
- Multiple domains may be relevant for complex queries
|
||||
|
||||
USER QUERY: {query}
|
||||
|
||||
Respond with 1-2 sentences suggesting which domains are relevant and why.
|
||||
Use domain names in your response (e.g., "This involves infrastructure and core domains").
|
||||
Plain text only - no JSON, no special formatting."""
|
||||
|
||||
def _extract_domains(self, reasoning: str, query: str) -> List[str]:
|
||||
"""
|
||||
Extract domain names from steward's text response.
|
||||
|
||||
Uses both the reasoning text and the original query for robustness.
|
||||
|
||||
Args:
|
||||
reasoning: Steward's text response
|
||||
query: Original user query
|
||||
|
||||
Returns:
|
||||
List of domain names
|
||||
"""
|
||||
text_lower = (reasoning + " " + query).lower()
|
||||
|
||||
domains = set()
|
||||
|
||||
# Check each domain's keywords
|
||||
for domain, keywords in self.DOMAIN_KEYWORDS.items():
|
||||
if any(keyword in text_lower for keyword in keywords):
|
||||
domains.add(domain)
|
||||
|
||||
# Always include 'core' domain (has essential tools)
|
||||
domains.add('core')
|
||||
|
||||
return sorted(list(domains))
|
||||
|
||||
def _keyword_fallback(self, query: str) -> List[str]:
|
||||
"""
|
||||
Fallback domain selection using only keywords (no LLM).
|
||||
|
||||
Used when steward LLM is unavailable.
|
||||
"""
|
||||
# Use same extraction logic with empty reasoning
|
||||
return self._extract_domains("", query)
|
||||
|
||||
def _calculate_confidence(self, domains: List[str], reasoning: str) -> float:
|
||||
"""
|
||||
Calculate confidence score based on domain detection.
|
||||
|
||||
Higher confidence = more explicit domain mentions in reasoning.
|
||||
"""
|
||||
reasoning_lower = reasoning.lower()
|
||||
|
||||
# Count explicit domain mentions
|
||||
mentions = 0
|
||||
for domain in domains:
|
||||
if domain != 'core' and domain in reasoning_lower:
|
||||
mentions += 1
|
||||
|
||||
# Base confidence on mention rate
|
||||
non_core_domains = [d for d in domains if d != 'core']
|
||||
if non_core_domains:
|
||||
mention_rate = mentions / len(non_core_domains)
|
||||
else:
|
||||
mention_rate = 1.0 # Only core domain
|
||||
|
||||
return min(mention_rate, 1.0)
|
||||
|
||||
|
||||
# Singleton instance
|
||||
_steward_instance: Optional[TextBasedSteward] = None
|
||||
|
||||
|
||||
def get_text_steward() -> TextBasedSteward:
|
||||
"""
|
||||
Get or create singleton text-based steward instance.
|
||||
|
||||
Returns:
|
||||
TextBasedSteward instance
|
||||
"""
|
||||
global _steward_instance
|
||||
|
||||
if _steward_instance is None:
|
||||
_steward_instance = TextBasedSteward()
|
||||
logger.info("Created TextBasedSteward singleton")
|
||||
|
||||
return _steward_instance
|
||||
@@ -1,108 +0,0 @@
|
||||
"""
|
||||
Tool Event Emitter - Allows tools to emit events during execution.
|
||||
|
||||
Provides a lightweight event system for tools to signal when they're being called,
|
||||
enabling real-time status updates during streaming responses.
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Optional
|
||||
from dataclasses import dataclass
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCallEvent:
|
||||
"""Event emitted when a tool is called."""
|
||||
tool_name: str
|
||||
arguments: dict
|
||||
timestamp: float
|
||||
|
||||
|
||||
class ToolEventEmitter:
|
||||
"""
|
||||
Singleton event emitter for tool calls.
|
||||
|
||||
Tools emit events when they execute, and the streaming system
|
||||
can listen for these events to provide real-time status updates.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._queue = asyncio.Queue()
|
||||
self._enabled = True
|
||||
|
||||
def emit(self, tool_name: str, arguments: dict):
|
||||
"""
|
||||
Emit a tool call event.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool being called
|
||||
arguments: Arguments passed to the tool
|
||||
"""
|
||||
if not self._enabled:
|
||||
return
|
||||
|
||||
event = ToolCallEvent(
|
||||
tool_name=tool_name,
|
||||
arguments=arguments,
|
||||
timestamp=time.time()
|
||||
)
|
||||
|
||||
# Non-blocking emit - don't wait for consumers
|
||||
try:
|
||||
self._queue.put_nowait(event)
|
||||
logger.debug(f"Tool event emitted: {tool_name}")
|
||||
except asyncio.QueueFull:
|
||||
logger.warning(f"Tool event queue full, dropping event for {tool_name}")
|
||||
|
||||
def has_events(self) -> bool:
|
||||
"""Check if there are pending events."""
|
||||
return not self._queue.empty()
|
||||
|
||||
async def get_event(self, timeout: float = 0.01) -> Optional[ToolCallEvent]:
|
||||
"""
|
||||
Get next tool event with timeout.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time to wait for event (seconds)
|
||||
|
||||
Returns:
|
||||
ToolCallEvent if available, None if timeout
|
||||
"""
|
||||
try:
|
||||
event = await asyncio.wait_for(self._queue.get(), timeout=timeout)
|
||||
return event
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
def clear(self):
|
||||
"""Clear all pending events."""
|
||||
while not self._queue.empty():
|
||||
try:
|
||||
self._queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
def enable(self):
|
||||
"""Enable event emission."""
|
||||
self._enabled = True
|
||||
logger.info("Tool event emission enabled")
|
||||
|
||||
def disable(self):
|
||||
"""Disable event emission."""
|
||||
self._enabled = False
|
||||
logger.info("Tool event emission disabled")
|
||||
|
||||
|
||||
# Global singleton emitter
|
||||
_emitter: Optional[ToolEventEmitter] = None
|
||||
|
||||
|
||||
def get_tool_emitter() -> ToolEventEmitter:
|
||||
"""Get the global tool event emitter instance."""
|
||||
global _emitter
|
||||
if _emitter is None:
|
||||
_emitter = ToolEventEmitter()
|
||||
return _emitter
|
||||
@@ -1,389 +0,0 @@
|
||||
"""
|
||||
Two-Stage Agent Orchestration - Coordinates steward analysis and Tatlock execution.
|
||||
|
||||
This module implements the two-stage tool selection system:
|
||||
1. Stage 1 (Steward): Analyze query and recommend 0-5 optimal tools
|
||||
2. Stage 2 (Tatlock): Answer using tool recommendations as guidance
|
||||
|
||||
Features:
|
||||
- Silent time/date injection when recommended
|
||||
- Visible status for web_search tool calls
|
||||
- Context enrichment with tool recommendations
|
||||
- Streaming status events separate from content
|
||||
- Fallback to single-stage if steward fails
|
||||
"""
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import Optional, AsyncGenerator, Dict, Any
|
||||
from datetime import datetime
|
||||
import json
|
||||
|
||||
try:
|
||||
from pydantic_ai import Agent
|
||||
PYDANTIC_AI_AVAILABLE = True
|
||||
except ImportError:
|
||||
PYDANTIC_AI_AVAILABLE = False
|
||||
Agent = None
|
||||
|
||||
from src.agents.steward_agent import get_steward_agent, ToolRecommendation
|
||||
from src.agents.tool_events import get_tool_emitter
|
||||
from src.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TwoStageAgent:
|
||||
"""
|
||||
Two-stage agent orchestrator.
|
||||
|
||||
Coordinates steward analysis with Tatlock execution, providing
|
||||
intelligent tool selection and transparent status updates.
|
||||
"""
|
||||
|
||||
def __init__(self, tatlock_agent: Agent, enable_two_stage: bool = True):
|
||||
"""
|
||||
Initialize two-stage orchestrator.
|
||||
|
||||
Args:
|
||||
tatlock_agent: The main Tatlock agent instance
|
||||
enable_two_stage: Whether to use two-stage analysis (default: True)
|
||||
"""
|
||||
self.tatlock = tatlock_agent
|
||||
self.enable_two_stage = enable_two_stage
|
||||
self.settings = get_settings()
|
||||
|
||||
# Get steward instance (uses same model as Tatlock)
|
||||
if self.enable_two_stage:
|
||||
try:
|
||||
self.steward = get_steward_agent()
|
||||
logger.info("TwoStageAgent: Steward enabled for tool analysis")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize steward, disabling two-stage: {e}")
|
||||
self.enable_two_stage = False
|
||||
self.steward = None
|
||||
else:
|
||||
self.steward = None
|
||||
logger.info("TwoStageAgent: Two-stage analysis disabled")
|
||||
|
||||
async def _get_current_datetime(self) -> Dict[str, str]:
|
||||
"""
|
||||
Get current date and time for silent injection.
|
||||
|
||||
Returns:
|
||||
Dict with 'date' and 'time' keys
|
||||
"""
|
||||
now = datetime.now()
|
||||
return {
|
||||
"date": now.strftime("%A, %B %d, %Y"),
|
||||
"time": now.strftime("%I:%M %p %Z").strip()
|
||||
}
|
||||
|
||||
def _enrich_user_message(
|
||||
self,
|
||||
original_message: str,
|
||||
recommendation: ToolRecommendation,
|
||||
datetime_info: Optional[Dict[str, str]] = None
|
||||
) -> str:
|
||||
"""
|
||||
Enrich user message with steward recommendations and optional time/date.
|
||||
|
||||
Args:
|
||||
original_message: Original user query
|
||||
recommendation: Steward's tool recommendations
|
||||
datetime_info: Optional current date/time to inject silently
|
||||
|
||||
Returns:
|
||||
Enriched message with injected context
|
||||
"""
|
||||
enrichment_parts = []
|
||||
|
||||
# Silent time/date injection (if recommended)
|
||||
if datetime_info:
|
||||
enrichment_parts.append(
|
||||
f"[Current context - Date: {datetime_info['date']}, Time: {datetime_info['time']}]"
|
||||
)
|
||||
|
||||
# Tool recommendations (if any, excluding time/date tools)
|
||||
visible_tools = [
|
||||
tool for tool in recommendation.recommended_tools
|
||||
if tool not in ('get_current_time', 'get_current_date')
|
||||
]
|
||||
|
||||
if visible_tools:
|
||||
tools_str = ", ".join(visible_tools)
|
||||
enrichment_parts.append(
|
||||
f"[Steward analysis: Recommended tools for this query: {tools_str}. "
|
||||
f"Reasoning: {recommendation.reasoning}]"
|
||||
)
|
||||
|
||||
# Combine enrichments with original message
|
||||
if enrichment_parts:
|
||||
enrichment = "\n".join(enrichment_parts)
|
||||
return f"{enrichment}\n\n{original_message}"
|
||||
|
||||
return original_message
|
||||
|
||||
async def _perform_steward_analysis(self, user_query: str) -> Optional[ToolRecommendation]:
|
||||
"""
|
||||
Perform steward analysis with error handling.
|
||||
|
||||
Args:
|
||||
user_query: User's query to analyze
|
||||
|
||||
Returns:
|
||||
ToolRecommendation if successful, None if failed
|
||||
"""
|
||||
try:
|
||||
timeout = getattr(self.settings, 'analysis_timeout', 3)
|
||||
recommendation = await self.steward.analyze(user_query, timeout=timeout)
|
||||
|
||||
logger.info(
|
||||
f"Steward analysis: {len(recommendation.recommended_tools)} tools recommended"
|
||||
)
|
||||
|
||||
return recommendation
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
logger.error("Steward analysis timed out, falling back to single-stage")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Steward analysis failed: {e}", exc_info=True)
|
||||
return None
|
||||
|
||||
async def chat_with_analysis(
|
||||
self,
|
||||
messages: list[dict],
|
||||
conversation_id: Optional[str] = None,
|
||||
stream: bool = True
|
||||
) -> AsyncGenerator[Dict[str, Any], None]:
|
||||
"""
|
||||
Execute two-stage chat with steward analysis and streaming.
|
||||
|
||||
Args:
|
||||
messages: Conversation messages (OpenAI format)
|
||||
conversation_id: Optional conversation ID for memory
|
||||
stream: Whether to stream response (default: True)
|
||||
|
||||
Yields:
|
||||
Dict with 'type' and relevant fields:
|
||||
- type='status': Status update (tool_name, message, arguments)
|
||||
- type='content': Response content chunk
|
||||
- type='done': Completion marker
|
||||
"""
|
||||
# Extract user query from last message
|
||||
user_query = messages[-1].get("content", "") if messages else ""
|
||||
|
||||
if not user_query:
|
||||
logger.warning("Empty user query in two-stage analysis")
|
||||
# Fall through to single-stage
|
||||
async for chunk in self._single_stage_chat(messages, conversation_id, stream):
|
||||
yield chunk
|
||||
return
|
||||
|
||||
# Stage 1: Steward Analysis
|
||||
if self.enable_two_stage and self.steward:
|
||||
# Emit consulting status
|
||||
yield {
|
||||
"type": "status",
|
||||
"message": "🤵 Consulting the steward...",
|
||||
"phase": "analysis"
|
||||
}
|
||||
|
||||
# Perform analysis
|
||||
recommendation = await self._perform_steward_analysis(user_query)
|
||||
|
||||
if recommendation is None:
|
||||
# Analysis failed, fall back to single-stage
|
||||
yield {
|
||||
"type": "status",
|
||||
"message": "⚠️ Steward unavailable, proceeding without analysis",
|
||||
"phase": "fallback"
|
||||
}
|
||||
async for chunk in self._single_stage_chat(messages, conversation_id, stream):
|
||||
yield chunk
|
||||
return
|
||||
|
||||
# Emit completion status based on tool count
|
||||
if len(recommendation.recommended_tools) == 0:
|
||||
yield {
|
||||
"type": "status",
|
||||
"message": "✓ No further assistance required - answering from general knowledge",
|
||||
"phase": "analysis_complete"
|
||||
}
|
||||
else:
|
||||
yield {
|
||||
"type": "status",
|
||||
"message": "✓ Steward consultation complete",
|
||||
"phase": "analysis_complete",
|
||||
"recommended_tools": recommendation.recommended_tools,
|
||||
"reasoning": recommendation.reasoning
|
||||
}
|
||||
|
||||
# Check if time/date tools recommended (for silent injection)
|
||||
needs_datetime = any(
|
||||
tool in recommendation.recommended_tools
|
||||
for tool in ('get_current_time', 'get_current_date')
|
||||
)
|
||||
|
||||
datetime_info = None
|
||||
if needs_datetime:
|
||||
datetime_info = await self._get_current_datetime()
|
||||
logger.info("Injecting current date/time silently")
|
||||
|
||||
# Enrich user message with recommendations and time/date
|
||||
enriched_message = self._enrich_user_message(
|
||||
user_query,
|
||||
recommendation,
|
||||
datetime_info
|
||||
)
|
||||
|
||||
# Replace last message with enriched version
|
||||
enriched_messages = messages[:-1] + [{
|
||||
"role": "user",
|
||||
"content": enriched_message
|
||||
}]
|
||||
else:
|
||||
# Two-stage disabled, use original messages
|
||||
enriched_messages = messages
|
||||
|
||||
# Stage 2: Tatlock Execution with PydanticAI Event Streaming
|
||||
try:
|
||||
# Extract user query from enriched messages
|
||||
user_query = enriched_messages[-1].get("content", "") if enriched_messages else ""
|
||||
|
||||
if not user_query:
|
||||
logger.error("Empty query for Tatlock execution")
|
||||
yield {"type": "error", "message": "Empty query"}
|
||||
return
|
||||
|
||||
# Use PydanticAI's native streaming with event monitoring
|
||||
# Access the underlying PydanticAI agent directly
|
||||
from pydantic_ai.messages import (
|
||||
ModelTextResponse,
|
||||
ToolCallPart,
|
||||
ToolReturnPart,
|
||||
)
|
||||
|
||||
logger.info("Starting Tatlock execution with event streaming")
|
||||
|
||||
# Track cumulative text for delta calculation
|
||||
previous_text = ""
|
||||
|
||||
async with self.tatlock.agent.run_stream(user_query) as run:
|
||||
# Stream all messages (includes tool calls and text)
|
||||
async for message in run.stream():
|
||||
# Handle tool call events
|
||||
if hasattr(message, 'parts'):
|
||||
for part in message.parts:
|
||||
# Check for tool call parts
|
||||
if isinstance(part, ToolCallPart):
|
||||
# Only emit status for web_search
|
||||
if part.tool_name == "web_search":
|
||||
# Extract query argument safely
|
||||
query = ""
|
||||
if hasattr(part, 'args') and isinstance(part.args, dict):
|
||||
query = part.args.get("query", "")
|
||||
|
||||
yield {
|
||||
"type": "status",
|
||||
"message": f"🔍 Searching the web: \"{query}\"",
|
||||
"phase": "tool_execution",
|
||||
"tool_name": part.tool_name
|
||||
}
|
||||
logger.info(f"Tool call detected: {part.tool_name}")
|
||||
|
||||
# Handle text response
|
||||
if isinstance(message, ModelTextResponse):
|
||||
# Get current cumulative text
|
||||
current_text = await run.get_text_so_far()
|
||||
|
||||
# Calculate delta (new text only)
|
||||
delta = current_text[len(previous_text):]
|
||||
if delta:
|
||||
yield {"type": "content", "content": delta}
|
||||
|
||||
previous_text = current_text
|
||||
|
||||
# Final finish marker
|
||||
yield {"type": "content", "content": "", "finish_reason": "stop"}
|
||||
logger.info("Tatlock execution complete")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in Tatlock execution: {e}", exc_info=True)
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": f"Error generating response: {str(e)}"
|
||||
}
|
||||
|
||||
# Final completion marker
|
||||
yield {"type": "done"}
|
||||
|
||||
async def _single_stage_chat(
|
||||
self,
|
||||
messages: list[dict],
|
||||
conversation_id: Optional[str] = None,
|
||||
stream: bool = True
|
||||
) -> AsyncGenerator[Dict[str, Any], None]:
|
||||
"""
|
||||
Execute single-stage chat (no steward analysis).
|
||||
|
||||
Args:
|
||||
messages: Conversation messages
|
||||
conversation_id: Optional conversation ID
|
||||
stream: Whether to stream
|
||||
|
||||
Yields:
|
||||
Chat response chunks
|
||||
"""
|
||||
if not stream:
|
||||
# Non-streaming response
|
||||
try:
|
||||
response = await self.tatlock.chat_completion(
|
||||
messages=messages,
|
||||
conversation_id=conversation_id
|
||||
)
|
||||
yield {
|
||||
"type": "content",
|
||||
"content": response,
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
yield {"type": "done"}
|
||||
except Exception as e:
|
||||
logger.error(f"Error in single-stage chat: {e}", exc_info=True)
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": str(e)
|
||||
}
|
||||
else:
|
||||
# Streaming response
|
||||
try:
|
||||
async for chunk in self.tatlock.chat(
|
||||
messages=messages,
|
||||
conversation_id=conversation_id,
|
||||
stream=True
|
||||
):
|
||||
yield chunk
|
||||
except Exception as e:
|
||||
logger.error(f"Error in single-stage streaming: {e}", exc_info=True)
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": str(e)
|
||||
}
|
||||
|
||||
|
||||
def create_two_stage_agent(
|
||||
tatlock_agent: Agent,
|
||||
enable_two_stage: bool = True
|
||||
) -> TwoStageAgent:
|
||||
"""
|
||||
Create two-stage agent orchestrator.
|
||||
|
||||
Args:
|
||||
tatlock_agent: The main Tatlock agent instance
|
||||
enable_two_stage: Whether to enable two-stage analysis
|
||||
|
||||
Returns:
|
||||
TwoStageAgent instance
|
||||
"""
|
||||
return TwoStageAgent(tatlock_agent, enable_two_stage=enable_two_stage)
|
||||
@@ -1,88 +0,0 @@
|
||||
"""
|
||||
Configuration for the Core AI service
|
||||
"""
|
||||
from pydantic_settings import BaseSettings
|
||||
from functools import lru_cache
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
"""Core AI application settings"""
|
||||
|
||||
# Application
|
||||
app_name: str = "Core AI Service"
|
||||
app_version: str = "1.0.0"
|
||||
debug: bool = False
|
||||
|
||||
# Server
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 8086 # Different port to avoid conflict with core-api
|
||||
|
||||
# Logging
|
||||
log_level: str = "INFO"
|
||||
|
||||
# Ollama Configuration (for AI orchestration)
|
||||
ollama_base_url: str = "http://ollama:11434"
|
||||
ollama_timeout: int = 300 # 5 minutes
|
||||
|
||||
# Model Configuration
|
||||
agent_model: str = "mistral-nemo:latest" # Optimized for PydanticAI tool calling
|
||||
|
||||
# Multi-Stage Tool Selection Configuration
|
||||
multi_stage_enabled: bool = True # Enable multi-stage steward analysis by default
|
||||
use_speedy_steward: bool = False # DISABLED - Use regular steward with mistral-nemo
|
||||
steward_model: str = "mistral-nemo:latest" # Steward uses same model as Tatlock (stays in VRAM)
|
||||
analysis_timeout: int = 10 # Steward analysis timeout (reasonable for mistral-nemo)
|
||||
max_recommended_tools: int = 5 # Maximum tools steward can recommend
|
||||
min_recommended_tools: int = 0 # Minimum tools (0 = can recommend no tools)
|
||||
|
||||
# Status Message Configuration - Butler-appropriate tone
|
||||
enable_status_messages: bool = True # Show status messages during streaming
|
||||
show_web_search_status: bool = True # Show status when web_search tool is called
|
||||
status_consulting: str = "Consulting the steward on the matter..."
|
||||
status_complete: str = "The steward consultation is complete"
|
||||
status_no_assistance: str = "The steward advises no further assistance is required"
|
||||
status_web_search: str = "Making enquiries..." # Butler tone for web search
|
||||
status_calculate: str = "Calculating... (one prefers precision in mathematics)"
|
||||
status_fallback: str = "The steward is unavailable. Proceeding with standard protocols..."
|
||||
|
||||
# System Prompt Variants
|
||||
system_prompt_variant: str = "minimal_agent" # For simple mode
|
||||
pydantic_system_prompt_variant: str = "pydantic_agent" # For PydanticAI mode
|
||||
|
||||
# Base URL for Core API tools (e.g., system status, services)
|
||||
core_api_base_url: str = "http://core-api:8083/v1"
|
||||
|
||||
# OpenAPI Tool Discovery
|
||||
# Comma-separated list of OpenAPI spec URLs for dynamic tool discovery
|
||||
# Example: "http://core-api:8083/openapi.json,http://automation:8080/openapi.json"
|
||||
openapi_endpoints: str = "http://core-api:8083/openapi.json"
|
||||
openapi_enabled: bool = True # Enable/disable OpenAPI tool discovery
|
||||
|
||||
# Feature Flags
|
||||
simple_enabled: bool = True # Enable simple endpoint
|
||||
pydantic_enabled: bool = True # Enable PydanticAI endpoint
|
||||
|
||||
# Memory System Configuration
|
||||
memory_enabled: bool = True
|
||||
memory_tier1_size: int = 10 # Max turns in RAM buffer
|
||||
|
||||
# Qdrant Configuration (for conversation memory)
|
||||
qdrant_url: str = "http://qdrant:6333"
|
||||
qdrant_collection_prefix: str = "core_ai_user" # Prefix for user collections
|
||||
|
||||
# Embedding Configuration
|
||||
embedding_model: str = "nomic-embed-text" # Ollama embedding model
|
||||
embedding_dimension: int = 768 # nomic-embed-text dimension
|
||||
|
||||
# Default User (until external auth is integrated)
|
||||
default_user_id: str = "llmdefault_at_schweitz_net"
|
||||
|
||||
class Config:
|
||||
env_file = ".env"
|
||||
case_sensitive = False
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def get_settings() -> Settings:
|
||||
"""Cached settings instance"""
|
||||
return Settings()
|
||||
@@ -1,56 +0,0 @@
|
||||
"""
|
||||
Multi-tenant memory system for conversation persistence
|
||||
|
||||
Architecture:
|
||||
- Tier 1: ConversationBufferMemory (in-memory, fast, last 10 turns) - per user
|
||||
- Tier 2/3: QdrantConversationMemory (persistent + semantic search) - separate collection per user
|
||||
- Manager: MemoryManager (orchestrates all tiers) - per user instance
|
||||
|
||||
Multi-tenancy:
|
||||
- Each user gets their own Qdrant collection: core_ai_user_{user_id}
|
||||
- Complete data isolation between users
|
||||
- Easy GDPR compliance (delete entire user collection)
|
||||
"""
|
||||
from .tier1_buffer import ConversationBufferMemory, get_buffer_memory
|
||||
from .qdrant_memory import QdrantConversationMemory, get_qdrant_memory_for_user
|
||||
from .manager import MemoryManager, get_memory_manager_for_user, clear_memory_manager_cache
|
||||
from .schemas import (
|
||||
ConversationTurn,
|
||||
ConversationBuffer,
|
||||
ConversationMetadata,
|
||||
ConversationSummary,
|
||||
MemoryQuery,
|
||||
MemoryResult,
|
||||
MessageRole,
|
||||
TokenUsage,
|
||||
ConversationListResponse,
|
||||
ConversationDetailResponse,
|
||||
ConversationSearchRequest,
|
||||
ConversationSearchResponse,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Manager (primary interface)
|
||||
"MemoryManager",
|
||||
"get_memory_manager_for_user",
|
||||
"clear_memory_manager_cache",
|
||||
# Tier 1
|
||||
"ConversationBufferMemory",
|
||||
"get_buffer_memory",
|
||||
# Tier 2/3
|
||||
"QdrantConversationMemory",
|
||||
"get_qdrant_memory_for_user",
|
||||
# Schemas
|
||||
"ConversationTurn",
|
||||
"ConversationBuffer",
|
||||
"ConversationMetadata",
|
||||
"ConversationSummary",
|
||||
"MemoryQuery",
|
||||
"MemoryResult",
|
||||
"MessageRole",
|
||||
"TokenUsage",
|
||||
"ConversationListResponse",
|
||||
"ConversationDetailResponse",
|
||||
"ConversationSearchRequest",
|
||||
"ConversationSearchResponse",
|
||||
]
|
||||
@@ -1,169 +0,0 @@
|
||||
"""
|
||||
Base classes for memory system
|
||||
"""
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Optional
|
||||
from .schemas import ConversationTurn, ConversationBuffer, MemoryQuery, MemoryResult
|
||||
|
||||
|
||||
class BaseMemory(ABC):
|
||||
"""Base class for all memory tiers"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_turn(self, conversation_id: str, turn: ConversationTurn) -> None:
|
||||
"""
|
||||
Add a new turn to memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn to store
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: Optional[int] = None,
|
||||
offset: int = 0
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Retrieve turns from memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
offset: Number of turns to skip
|
||||
|
||||
Returns:
|
||||
List of conversation turns
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def clear_conversation(self, conversation_id: str) -> None:
|
||||
"""
|
||||
Clear all turns for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def conversation_exists(self, conversation_id: str) -> bool:
|
||||
"""
|
||||
Check if a conversation exists in this memory tier
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
True if conversation exists
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class Tier1Memory(BaseMemory):
|
||||
"""Base class for Tier 1 (working memory)"""
|
||||
|
||||
@abstractmethod
|
||||
async def get_buffer(self, conversation_id: str) -> Optional[ConversationBuffer]:
|
||||
"""
|
||||
Get the full conversation buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
ConversationBuffer or None if not found
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def prune(self, conversation_id: str, keep_last: int = 5) -> None:
|
||||
"""
|
||||
Prune old turns, keeping only the most recent ones
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
keep_last: Number of recent turns to keep
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class Tier2Memory(BaseMemory):
|
||||
"""Base class for Tier 2 (short-term memory with summaries)"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_summary(
|
||||
self,
|
||||
conversation_id: str,
|
||||
summary_text: str,
|
||||
turn_range_start: int,
|
||||
turn_range_end: int
|
||||
) -> None:
|
||||
"""
|
||||
Add a conversation summary
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
summary_text: The summarized text
|
||||
turn_range_start: First turn number in summary
|
||||
turn_range_end: Last turn number in summary
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_summaries(self, conversation_id: str) -> List[dict]:
|
||||
"""
|
||||
Get all summaries for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
List of summary dictionaries
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class Tier3Memory(BaseMemory):
|
||||
"""Base class for Tier 3 (long-term vector memory)"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_turn_with_embedding(
|
||||
self,
|
||||
conversation_id: str,
|
||||
turn: ConversationTurn,
|
||||
embedding: List[float]
|
||||
) -> None:
|
||||
"""
|
||||
Add a turn with its vector embedding
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn
|
||||
embedding: Vector embedding of the turn content
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def similarity_search(
|
||||
self,
|
||||
query_embedding: List[float],
|
||||
conversation_id: Optional[str] = None,
|
||||
limit: int = 5
|
||||
) -> List[dict]:
|
||||
"""
|
||||
Perform semantic similarity search
|
||||
|
||||
Args:
|
||||
query_embedding: Vector embedding of the search query
|
||||
conversation_id: Optional filter to specific conversation
|
||||
limit: Maximum number of results
|
||||
|
||||
Returns:
|
||||
List of matching turns with scores
|
||||
"""
|
||||
pass
|
||||
@@ -1,337 +0,0 @@
|
||||
"""
|
||||
Multi-tenant Memory Manager: Orchestrates all memory tiers with per-user isolation
|
||||
|
||||
Coordinates:
|
||||
- Tier 1: ConversationBufferMemory (RAM, fast, last N turns) - per user
|
||||
- Tier 2/3: QdrantConversationMemory (persistent + semantic) - separate collection per user
|
||||
|
||||
Provides unified interface for memory operations with automatic
|
||||
tier management and per-user data isolation.
|
||||
"""
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import List, Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
from .tier1_buffer import ConversationBufferMemory
|
||||
from .qdrant_memory import QdrantConversationMemory, get_qdrant_memory_for_user
|
||||
from .schemas import ConversationTurn, MessageRole, TokenUsage
|
||||
from src.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class MemoryManager:
|
||||
"""
|
||||
Multi-tenant unified memory manager orchestrating all tiers
|
||||
|
||||
Features:
|
||||
- Per-user data isolation (separate Qdrant collections)
|
||||
- Per-user in-memory buffers
|
||||
- Automatic consolidation from buffer to Qdrant
|
||||
- Semantic search within user's conversations
|
||||
- Memory lifecycle management
|
||||
|
||||
Responsibilities:
|
||||
- Add turns to appropriate tiers
|
||||
- Retrieve conversation history (buffer + persistent)
|
||||
- Consolidate buffer to persistent storage
|
||||
- Semantic search across user's conversations
|
||||
- Memory lifecycle management
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
user_id: str,
|
||||
buffer_max_turns: int = 10,
|
||||
auto_consolidate: bool = True
|
||||
):
|
||||
"""
|
||||
Initialize memory manager for a specific user
|
||||
|
||||
Args:
|
||||
user_id: Sanitized user ID (email format: username_at_domain_com)
|
||||
buffer_max_turns: Max turns to keep in RAM buffer
|
||||
auto_consolidate: Automatically consolidate when buffer threshold reached
|
||||
"""
|
||||
self.user_id = user_id
|
||||
self.auto_consolidate = auto_consolidate
|
||||
|
||||
# Create user-specific buffer (in-memory)
|
||||
self.buffer_memory = ConversationBufferMemory(max_turns=buffer_max_turns)
|
||||
|
||||
# Create user-specific Qdrant memory (separate collection)
|
||||
self.qdrant_memory = get_qdrant_memory_for_user(user_id)
|
||||
|
||||
logger.info(
|
||||
f"MemoryManager initialized for user '{user_id}' "
|
||||
f"(auto_consolidate={auto_consolidate}, buffer_max={buffer_max_turns})"
|
||||
)
|
||||
|
||||
async def add_turn(
|
||||
self,
|
||||
conversation_id: str,
|
||||
role: MessageRole,
|
||||
content: str,
|
||||
tokens: Optional[TokenUsage] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> ConversationTurn:
|
||||
"""
|
||||
Add a conversation turn to memory
|
||||
|
||||
Automatically:
|
||||
1. Adds to Tier 1 (buffer)
|
||||
2. Adds to Tier 2/3 (Qdrant) immediately
|
||||
3. Auto-prunes buffer if max turns reached
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
role: Message role (user, assistant, system)
|
||||
content: Message content
|
||||
tokens: Optional token usage
|
||||
metadata: Optional metadata
|
||||
|
||||
Returns:
|
||||
The created conversation turn
|
||||
"""
|
||||
# Get current buffer to determine turn number
|
||||
buffer = await self.buffer_memory.get_buffer(conversation_id)
|
||||
turn_number = (buffer.metadata.turn_count + 1) if buffer else 1
|
||||
|
||||
# Create turn with user_id
|
||||
turn = ConversationTurn(
|
||||
role=role,
|
||||
content=content,
|
||||
timestamp=datetime.utcnow(),
|
||||
turn_number=turn_number,
|
||||
user_id=self.user_id,
|
||||
tokens=tokens,
|
||||
metadata=metadata or {}
|
||||
)
|
||||
|
||||
# Add to Tier 1 (buffer) - fast RAM storage
|
||||
await self.buffer_memory.add_turn(conversation_id, turn)
|
||||
logger.debug(
|
||||
f"Turn {turn_number} added to buffer for user '{self.user_id}' "
|
||||
f"conversation {conversation_id}"
|
||||
)
|
||||
|
||||
# Add to Tier 2/3 (Qdrant) immediately - persistent storage with embeddings
|
||||
try:
|
||||
await self.qdrant_memory.add_turn(conversation_id, turn)
|
||||
logger.debug(
|
||||
f"Turn {turn_number} added to Qdrant for user '{self.user_id}' "
|
||||
f"conversation {conversation_id}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error adding turn to Qdrant for user '{self.user_id}': {e}"
|
||||
)
|
||||
# Don't fail the whole operation if Qdrant fails
|
||||
# Buffer still has the turn
|
||||
|
||||
return turn
|
||||
|
||||
async def get_recent_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: int = 10
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Get recent conversation turns (from buffer)
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
|
||||
Returns:
|
||||
List of recent conversation turns
|
||||
"""
|
||||
return await self.buffer_memory.get_recent_turns(conversation_id, limit)
|
||||
|
||||
async def get_full_history(
|
||||
self,
|
||||
conversation_id: str,
|
||||
include_buffer: bool = True
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Get complete conversation history
|
||||
|
||||
Retrieves from Qdrant (Tier 2) - buffer is just a cache
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
include_buffer: Ignored (kept for API compatibility)
|
||||
|
||||
Returns:
|
||||
Complete conversation history, sorted chronologically
|
||||
"""
|
||||
# Get from Qdrant (source of truth)
|
||||
turns = await self.qdrant_memory.get_turns(conversation_id)
|
||||
|
||||
# Sort chronologically (should already be sorted, but ensure it)
|
||||
turns.sort(key=lambda t: t.turn_number)
|
||||
|
||||
return turns
|
||||
|
||||
async def search_conversations(
|
||||
self,
|
||||
query: str,
|
||||
conversation_id: Optional[str] = None,
|
||||
limit: int = 5
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Semantic search across user's conversations (Tier 3 mode)
|
||||
|
||||
Searches only within this user's collection.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
conversation_id: Optional filter to specific conversation
|
||||
limit: Maximum number of results
|
||||
|
||||
Returns:
|
||||
List of matching turns with scores
|
||||
"""
|
||||
return await self.qdrant_memory.similarity_search(
|
||||
query=query,
|
||||
conversation_id=conversation_id,
|
||||
limit=limit
|
||||
)
|
||||
|
||||
async def clear_conversation(
|
||||
self,
|
||||
conversation_id: str,
|
||||
clear_buffer: bool = True,
|
||||
clear_qdrant: bool = True
|
||||
) -> None:
|
||||
"""
|
||||
Clear conversation from memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
clear_buffer: Clear from Tier 1 buffer
|
||||
clear_qdrant: Clear from Tier 2/3 Qdrant
|
||||
"""
|
||||
if clear_buffer:
|
||||
await self.buffer_memory.clear_conversation(conversation_id)
|
||||
logger.info(
|
||||
f"Cleared buffer for user '{self.user_id}' conversation {conversation_id}"
|
||||
)
|
||||
|
||||
if clear_qdrant:
|
||||
await self.qdrant_memory.clear_conversation(conversation_id)
|
||||
logger.info(
|
||||
f"Cleared Qdrant for user '{self.user_id}' conversation {conversation_id}"
|
||||
)
|
||||
|
||||
async def clear_all_user_data(self) -> None:
|
||||
"""
|
||||
Clear ALL data for this user (GDPR compliance)
|
||||
|
||||
Deletes:
|
||||
- All buffer data for this user
|
||||
- Entire Qdrant collection for this user
|
||||
"""
|
||||
# Clear all buffers (in-memory)
|
||||
conversation_ids = await self.buffer_memory.get_all_conversation_ids()
|
||||
for conv_id in conversation_ids:
|
||||
await self.buffer_memory.clear_conversation(conv_id)
|
||||
|
||||
# Delete entire Qdrant collection
|
||||
await self.qdrant_memory.clear_all_data()
|
||||
|
||||
logger.info(f"Cleared ALL data for user '{self.user_id}'")
|
||||
|
||||
async def get_conversation_stats(
|
||||
self,
|
||||
conversation_id: str
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Get conversation statistics across all tiers
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Dictionary with stats from buffer and Qdrant
|
||||
"""
|
||||
# Get buffer stats
|
||||
buffer = await self.buffer_memory.get_buffer(conversation_id)
|
||||
buffer_stats = {
|
||||
"buffer_turns": buffer.metadata.turn_count if buffer else 0,
|
||||
"buffer_tokens": buffer.metadata.total_tokens if buffer else 0
|
||||
}
|
||||
|
||||
# Get Qdrant stats
|
||||
qdrant_stats = await self.qdrant_memory.get_conversation_stats(conversation_id)
|
||||
|
||||
# Combine
|
||||
return {
|
||||
"user_id": self.user_id,
|
||||
"conversation_id": conversation_id,
|
||||
**buffer_stats,
|
||||
"qdrant_turns": qdrant_stats["total_turns"],
|
||||
"qdrant_tokens": qdrant_stats["total_tokens"],
|
||||
"exists_in_buffer": buffer is not None,
|
||||
"exists_in_qdrant": qdrant_stats["exists"]
|
||||
}
|
||||
|
||||
async def list_conversations(self) -> List[str]:
|
||||
"""
|
||||
List all conversation IDs for this user
|
||||
|
||||
Returns:
|
||||
List of conversation IDs
|
||||
"""
|
||||
return await self.qdrant_memory.list_conversations()
|
||||
|
||||
|
||||
# Per-user memory manager cache
|
||||
_memory_managers: Dict[str, MemoryManager] = {}
|
||||
|
||||
|
||||
def get_memory_manager_for_user(
|
||||
user_id: str,
|
||||
buffer_max_turns: int = 10,
|
||||
auto_consolidate: bool = True
|
||||
) -> MemoryManager:
|
||||
"""
|
||||
Get or create memory manager instance for a specific user
|
||||
|
||||
Args:
|
||||
user_id: Sanitized user ID
|
||||
buffer_max_turns: Max turns to keep in RAM buffer
|
||||
auto_consolidate: Automatically consolidate when buffer threshold reached
|
||||
|
||||
Returns:
|
||||
MemoryManager instance for the user
|
||||
"""
|
||||
if user_id not in _memory_managers:
|
||||
_memory_managers[user_id] = MemoryManager(
|
||||
user_id=user_id,
|
||||
buffer_max_turns=buffer_max_turns,
|
||||
auto_consolidate=auto_consolidate
|
||||
)
|
||||
logger.info(f"Created new MemoryManager for user '{user_id}'")
|
||||
|
||||
return _memory_managers[user_id]
|
||||
|
||||
|
||||
def clear_memory_manager_cache(user_id: Optional[str] = None) -> None:
|
||||
"""
|
||||
Clear memory manager cache
|
||||
|
||||
Args:
|
||||
user_id: Optional user ID to clear (None = clear all)
|
||||
"""
|
||||
global _memory_managers
|
||||
|
||||
if user_id:
|
||||
if user_id in _memory_managers:
|
||||
del _memory_managers[user_id]
|
||||
logger.info(f"Cleared MemoryManager cache for user '{user_id}'")
|
||||
else:
|
||||
_memory_managers.clear()
|
||||
logger.info("Cleared all MemoryManager caches")
|
||||
@@ -1,465 +0,0 @@
|
||||
"""
|
||||
Unified Tier 2/3: Qdrant-based conversation memory with collection-per-user
|
||||
|
||||
Multi-tenant architecture:
|
||||
- Each user gets their own Qdrant collection: core_ai_user_{user_id}
|
||||
- Collections created on-demand
|
||||
- Complete data isolation between users
|
||||
- Easy GDPR compliance (delete entire collection)
|
||||
|
||||
Dual-mode operation:
|
||||
- Tier 2: Historical retrieval (filter by conversation_id, time-based)
|
||||
- Tier 3: Semantic search (vector similarity across user's conversations)
|
||||
"""
|
||||
import logging
|
||||
import uuid
|
||||
import re
|
||||
from typing import List, Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.models import (
|
||||
Distance,
|
||||
VectorParams,
|
||||
PointStruct,
|
||||
Filter,
|
||||
FieldCondition,
|
||||
MatchValue,
|
||||
)
|
||||
|
||||
from .base import BaseMemory
|
||||
from .schemas import ConversationTurn, MessageRole
|
||||
from src.config import get_settings
|
||||
from src.models.embeddings_ollama import get_embedding_client
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class QdrantConversationMemory(BaseMemory):
|
||||
"""
|
||||
Multi-tenant conversation memory using Qdrant with collection-per-user.
|
||||
|
||||
Each user gets a dedicated collection for complete data isolation.
|
||||
Stores all conversation turns with vectors for semantic search.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
user_id: str,
|
||||
collection_prefix: Optional[str] = None,
|
||||
qdrant_url: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
Initialize Qdrant memory for a specific user
|
||||
|
||||
Args:
|
||||
user_id: Sanitized user ID (email format: username_at_domain_com)
|
||||
collection_prefix: Collection name prefix (default: core_ai_user)
|
||||
qdrant_url: Qdrant connection URL (default from settings)
|
||||
"""
|
||||
self.user_id = user_id
|
||||
self.collection_prefix = collection_prefix or "core_ai_user"
|
||||
self.collection_name = self._get_collection_name(user_id)
|
||||
|
||||
# Parse Qdrant URL (format: http://qdrant:6333)
|
||||
qdrant_url = qdrant_url or getattr(settings, 'qdrant_url', 'http://qdrant:6333')
|
||||
self.qdrant_url = qdrant_url
|
||||
|
||||
# Initialize clients
|
||||
self.client = QdrantClient(url=self.qdrant_url)
|
||||
self.embedding_client = get_embedding_client()
|
||||
|
||||
logger.info(
|
||||
f"Initialized QdrantConversationMemory for user '{user_id}': "
|
||||
f"{self.qdrant_url}/{self.collection_name}"
|
||||
)
|
||||
|
||||
# Ensure user's collection exists
|
||||
self._ensure_collection()
|
||||
|
||||
def _get_collection_name(self, user_id: str) -> str:
|
||||
"""
|
||||
Generate collection name for user
|
||||
|
||||
Args:
|
||||
user_id: Sanitized user ID
|
||||
|
||||
Returns:
|
||||
Collection name: {prefix}_{user_id}
|
||||
"""
|
||||
# Sanitize user_id for collection name (should already be sanitized, but double-check)
|
||||
sanitized = re.sub(r'[^a-z0-9_]', '_', user_id.lower())
|
||||
return f"{self.collection_prefix}_{sanitized}"
|
||||
|
||||
def _ensure_collection(self) -> None:
|
||||
"""Create user's collection if it doesn't exist"""
|
||||
try:
|
||||
collections = self.client.get_collections().collections
|
||||
collection_names = [c.name for c in collections]
|
||||
|
||||
if self.collection_name not in collection_names:
|
||||
logger.info(f"Creating new collection for user '{self.user_id}': {self.collection_name}")
|
||||
|
||||
# Get embedding dimension from settings or default to 768 (nomic-embed-text)
|
||||
embedding_dim = getattr(settings, 'embedding_dimension', 768)
|
||||
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=VectorParams(
|
||||
size=embedding_dim,
|
||||
distance=Distance.COSINE
|
||||
)
|
||||
)
|
||||
logger.info(f"✓ Collection created: {self.collection_name}")
|
||||
else:
|
||||
logger.info(f"✓ Collection exists: {self.collection_name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error ensuring collection for user '{self.user_id}': {e}")
|
||||
raise
|
||||
|
||||
async def add_turn(self, conversation_id: str, turn: ConversationTurn) -> None:
|
||||
"""
|
||||
Add a conversation turn with its embedding
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn to store
|
||||
"""
|
||||
# Generate embedding
|
||||
embedding = await self.embedding_client.embed_text(turn.content)
|
||||
|
||||
# Create point ID: deterministic UUID from conversation_id + turn_number
|
||||
point_id_str = f"{conversation_id}_{turn.turn_number}"
|
||||
point_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, point_id_str))
|
||||
|
||||
# Build payload (no user_id needed - collection is already user-specific)
|
||||
payload = {
|
||||
"conversation_id": conversation_id,
|
||||
"turn_number": turn.turn_number,
|
||||
"role": turn.role.value if isinstance(turn.role, MessageRole) else turn.role,
|
||||
"content": turn.content,
|
||||
"timestamp": turn.timestamp.isoformat(),
|
||||
"metadata": turn.metadata,
|
||||
}
|
||||
|
||||
# Add token info if available
|
||||
if turn.tokens:
|
||||
payload["tokens_prompt"] = turn.tokens.prompt
|
||||
payload["tokens_completion"] = turn.tokens.completion
|
||||
payload["tokens_total"] = turn.tokens.total
|
||||
|
||||
# Upsert to user's Qdrant collection
|
||||
try:
|
||||
self.client.upsert(
|
||||
collection_name=self.collection_name,
|
||||
points=[
|
||||
PointStruct(
|
||||
id=point_id,
|
||||
vector=embedding,
|
||||
payload=payload
|
||||
)
|
||||
]
|
||||
)
|
||||
logger.debug(
|
||||
f"Stored turn {turn.turn_number} for conversation {conversation_id} "
|
||||
f"(user: {self.user_id})"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing turn in Qdrant for user '{self.user_id}': {e}")
|
||||
raise
|
||||
|
||||
async def get_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: Optional[int] = None,
|
||||
offset: int = 0
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Retrieve turns for a conversation (Tier 2 mode: chronological)
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
offset: Number of turns to skip
|
||||
|
||||
Returns:
|
||||
List of conversation turns
|
||||
"""
|
||||
try:
|
||||
# Scroll through all points for this conversation
|
||||
points, _ = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
),
|
||||
limit=limit or 100,
|
||||
offset=offset,
|
||||
with_payload=True,
|
||||
with_vectors=False
|
||||
)
|
||||
|
||||
# Convert to ConversationTurn objects
|
||||
turns = []
|
||||
for point in points:
|
||||
payload = point.payload
|
||||
turn = ConversationTurn(
|
||||
role=MessageRole(payload["role"]),
|
||||
content=payload["content"],
|
||||
timestamp=datetime.fromisoformat(payload["timestamp"]),
|
||||
turn_number=payload["turn_number"],
|
||||
user_id=self.user_id, # User from collection context
|
||||
metadata=payload.get("metadata", {})
|
||||
)
|
||||
turns.append(turn)
|
||||
|
||||
# Sort by turn_number
|
||||
turns.sort(key=lambda t: t.turn_number)
|
||||
|
||||
return turns
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error retrieving turns from Qdrant for user '{self.user_id}': {e}")
|
||||
return []
|
||||
|
||||
async def similarity_search(
|
||||
self,
|
||||
query: str,
|
||||
conversation_id: Optional[str] = None,
|
||||
limit: int = 5
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Semantic search for relevant turns (Tier 3 mode: semantic)
|
||||
|
||||
Args:
|
||||
query: Search query text
|
||||
conversation_id: Optional filter to specific conversation
|
||||
limit: Maximum number of results
|
||||
|
||||
Returns:
|
||||
List of matching turns with scores
|
||||
"""
|
||||
try:
|
||||
# Generate query embedding
|
||||
query_embedding = await self.embedding_client.embed_text(query)
|
||||
|
||||
# Build filter if conversation_id specified
|
||||
search_filter = None
|
||||
if conversation_id:
|
||||
search_filter = Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Search in user's Qdrant collection
|
||||
results = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query_embedding,
|
||||
query_filter=search_filter,
|
||||
limit=limit,
|
||||
with_payload=True
|
||||
)
|
||||
|
||||
# Convert results
|
||||
matches = []
|
||||
for result in results:
|
||||
payload = result.payload
|
||||
match = {
|
||||
"conversation_id": payload["conversation_id"],
|
||||
"turn_number": payload["turn_number"],
|
||||
"role": payload["role"],
|
||||
"content": payload["content"],
|
||||
"timestamp": payload["timestamp"],
|
||||
"score": result.score,
|
||||
}
|
||||
matches.append(match)
|
||||
|
||||
logger.debug(
|
||||
f"Semantic search found {len(matches)} matches for user '{self.user_id}' "
|
||||
f"query: {query[:50]}..."
|
||||
)
|
||||
|
||||
return matches
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in semantic search for user '{self.user_id}': {e}")
|
||||
return []
|
||||
|
||||
async def clear_conversation(self, conversation_id: str) -> None:
|
||||
"""
|
||||
Clear all turns for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
"""
|
||||
try:
|
||||
# Delete all points with this conversation_id
|
||||
self.client.delete(
|
||||
collection_name=self.collection_name,
|
||||
points_selector=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
logger.info(
|
||||
f"Cleared conversation {conversation_id} for user '{self.user_id}' from Qdrant"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing conversation for user '{self.user_id}': {e}")
|
||||
raise
|
||||
|
||||
async def clear_all_data(self) -> None:
|
||||
"""
|
||||
Clear ALL data for this user (GDPR compliance)
|
||||
|
||||
Deletes the entire collection for this user.
|
||||
"""
|
||||
try:
|
||||
self.client.delete_collection(self.collection_name)
|
||||
logger.info(f"Deleted all data for user '{self.user_id}' (collection: {self.collection_name})")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting user data for '{self.user_id}': {e}")
|
||||
raise
|
||||
|
||||
async def conversation_exists(self, conversation_id: str) -> bool:
|
||||
"""
|
||||
Check if a conversation exists
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
True if conversation has any turns
|
||||
"""
|
||||
try:
|
||||
points, _ = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
),
|
||||
limit=1,
|
||||
with_payload=False,
|
||||
with_vectors=False
|
||||
)
|
||||
return len(points) > 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error checking conversation existence for user '{self.user_id}': {e}")
|
||||
return False
|
||||
|
||||
async def get_conversation_stats(self, conversation_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get statistics about a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Dictionary with stats
|
||||
"""
|
||||
try:
|
||||
points, _ = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
),
|
||||
limit=1000, # Get all points
|
||||
with_payload=True,
|
||||
with_vectors=False
|
||||
)
|
||||
|
||||
total_turns = len(points)
|
||||
total_tokens = sum(
|
||||
point.payload.get("tokens_total", 0) for point in points
|
||||
)
|
||||
|
||||
return {
|
||||
"user_id": self.user_id,
|
||||
"conversation_id": conversation_id,
|
||||
"total_turns": total_turns,
|
||||
"total_tokens": total_tokens,
|
||||
"exists": total_turns > 0
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting conversation stats for user '{self.user_id}': {e}")
|
||||
return {
|
||||
"user_id": self.user_id,
|
||||
"conversation_id": conversation_id,
|
||||
"total_turns": 0,
|
||||
"total_tokens": 0,
|
||||
"exists": False
|
||||
}
|
||||
|
||||
async def list_conversations(self) -> List[str]:
|
||||
"""
|
||||
List all conversation IDs for this user
|
||||
|
||||
Returns:
|
||||
List of conversation IDs
|
||||
"""
|
||||
try:
|
||||
# Scroll through all points to collect unique conversation_ids
|
||||
conversation_ids = set()
|
||||
offset = None
|
||||
|
||||
while True:
|
||||
points, next_offset = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
limit=100,
|
||||
offset=offset,
|
||||
with_payload=True,
|
||||
with_vectors=False
|
||||
)
|
||||
|
||||
for point in points:
|
||||
conversation_ids.add(point.payload["conversation_id"])
|
||||
|
||||
if next_offset is None:
|
||||
break
|
||||
offset = next_offset
|
||||
|
||||
return sorted(list(conversation_ids))
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing conversations for user '{self.user_id}': {e}")
|
||||
return []
|
||||
|
||||
|
||||
def get_qdrant_memory_for_user(user_id: str) -> QdrantConversationMemory:
|
||||
"""
|
||||
Get Qdrant memory instance for a specific user
|
||||
|
||||
Args:
|
||||
user_id: Sanitized user ID (email format: username_at_domain_com)
|
||||
|
||||
Returns:
|
||||
QdrantConversationMemory instance for the user
|
||||
"""
|
||||
return QdrantConversationMemory(user_id=user_id)
|
||||
@@ -1,110 +0,0 @@
|
||||
"""
|
||||
Pydantic schemas for memory system
|
||||
"""
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class MessageRole(str, Enum):
|
||||
"""Message role types"""
|
||||
SYSTEM = "system"
|
||||
USER = "user"
|
||||
ASSISTANT = "assistant"
|
||||
|
||||
|
||||
class TokenUsage(BaseModel):
|
||||
"""Token usage information"""
|
||||
prompt: int = 0
|
||||
completion: int = 0
|
||||
total: int = 0
|
||||
|
||||
|
||||
class ConversationTurn(BaseModel):
|
||||
"""A single turn in a conversation"""
|
||||
role: MessageRole
|
||||
content: str
|
||||
timestamp: datetime = Field(default_factory=datetime.utcnow)
|
||||
turn_number: int
|
||||
user_id: str = "llmdefault_at_schweitz_net" # Multi-tenancy: user who owns this turn
|
||||
tokens: Optional[TokenUsage] = None
|
||||
metadata: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ConversationMetadata(BaseModel):
|
||||
"""Metadata about a conversation"""
|
||||
conversation_id: str
|
||||
user_id: str = "llmdefault_at_schweitz_net" # Multi-tenancy: user who owns this conversation
|
||||
created_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
last_updated: datetime = Field(default_factory=datetime.utcnow)
|
||||
turn_count: int = 0
|
||||
total_tokens: int = 0
|
||||
status: str = "active" # active, archived, deleted
|
||||
|
||||
|
||||
class ConversationBuffer(BaseModel):
|
||||
"""In-memory conversation buffer (Tier 1)"""
|
||||
conversation_id: str
|
||||
turns: List[ConversationTurn] = Field(default_factory=list)
|
||||
metadata: ConversationMetadata
|
||||
|
||||
|
||||
class ConversationSummary(BaseModel):
|
||||
"""Summarized conversation segment (Tier 2)"""
|
||||
conversation_id: str
|
||||
summary_text: str
|
||||
turn_range_start: int
|
||||
turn_range_end: int
|
||||
created_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
token_count: int = 0
|
||||
|
||||
|
||||
class MemoryQuery(BaseModel):
|
||||
"""Query for memory retrieval"""
|
||||
conversation_id: str
|
||||
query: Optional[str] = None
|
||||
limit: int = Field(default=10, ge=1, le=100)
|
||||
include_tier1: bool = True
|
||||
include_tier2: bool = True
|
||||
include_tier3: bool = True
|
||||
|
||||
|
||||
class MemoryResult(BaseModel):
|
||||
"""Result from memory retrieval"""
|
||||
conversation_id: str
|
||||
turns: List[ConversationTurn] = Field(default_factory=list)
|
||||
summaries: List[ConversationSummary] = Field(default_factory=list)
|
||||
source_tiers: List[int] = Field(default_factory=list) # Which tiers contributed
|
||||
total_results: int = 0
|
||||
|
||||
|
||||
# API Request/Response Models
|
||||
|
||||
class ConversationListResponse(BaseModel):
|
||||
"""Response for listing conversations"""
|
||||
conversations: List[ConversationMetadata]
|
||||
total: int
|
||||
page: int = 1
|
||||
page_size: int = 50
|
||||
|
||||
|
||||
class ConversationDetailResponse(BaseModel):
|
||||
"""Response for conversation details"""
|
||||
metadata: ConversationMetadata
|
||||
recent_turns: List[ConversationTurn]
|
||||
turn_count: int
|
||||
|
||||
|
||||
class ConversationSearchRequest(BaseModel):
|
||||
"""Request for semantic search in conversation"""
|
||||
query: str
|
||||
limit: int = Field(default=5, ge=1, le=50)
|
||||
|
||||
|
||||
class ConversationSearchResponse(BaseModel):
|
||||
"""Response for semantic search"""
|
||||
conversation_id: str
|
||||
results: List[ConversationTurn]
|
||||
scores: List[float] = Field(default_factory=list)
|
||||
total_results: int
|
||||
@@ -1,239 +0,0 @@
|
||||
"""
|
||||
Tier 1: ConversationBufferMemory (In-Memory Working Memory)
|
||||
|
||||
Fast in-memory storage for recent conversation turns.
|
||||
- Stores last N turns in RAM
|
||||
- < 1ms access time
|
||||
- Ephemeral (lost on restart)
|
||||
- Automatic pruning when limit reached
|
||||
"""
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
from datetime import datetime
|
||||
from collections import OrderedDict
|
||||
|
||||
from .base import Tier1Memory
|
||||
from .schemas import (
|
||||
ConversationTurn,
|
||||
ConversationBuffer,
|
||||
ConversationMetadata,
|
||||
MessageRole,
|
||||
TokenUsage
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConversationBufferMemory(Tier1Memory):
|
||||
"""
|
||||
In-memory buffer for recent conversation turns.
|
||||
|
||||
Stores the last N turns of each conversation in RAM for fast access.
|
||||
Automatically prunes old turns when limit is reached.
|
||||
"""
|
||||
|
||||
def __init__(self, max_turns: int = 10):
|
||||
"""
|
||||
Initialize buffer memory
|
||||
|
||||
Args:
|
||||
max_turns: Maximum number of turns to keep per conversation
|
||||
"""
|
||||
self.max_turns = max_turns
|
||||
# Use OrderedDict to maintain insertion order
|
||||
self._buffers: Dict[str, ConversationBuffer] = OrderedDict()
|
||||
logger.info(f"Initialized ConversationBufferMemory with max_turns={max_turns}")
|
||||
|
||||
async def add_turn(self, conversation_id: str, turn: ConversationTurn) -> None:
|
||||
"""
|
||||
Add a new turn to the buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn to store
|
||||
"""
|
||||
# Get or create buffer
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
buffer = ConversationBuffer(
|
||||
conversation_id=conversation_id,
|
||||
turns=[],
|
||||
metadata=ConversationMetadata(
|
||||
conversation_id=conversation_id
|
||||
)
|
||||
)
|
||||
self._buffers[conversation_id] = buffer
|
||||
|
||||
# Add turn
|
||||
buffer.turns.append(turn)
|
||||
|
||||
# Update metadata
|
||||
buffer.metadata.turn_count = len(buffer.turns)
|
||||
buffer.metadata.last_updated = datetime.utcnow()
|
||||
|
||||
if turn.tokens:
|
||||
buffer.metadata.total_tokens += turn.tokens.total
|
||||
|
||||
# Auto-prune if exceeds max turns
|
||||
if len(buffer.turns) > self.max_turns:
|
||||
await self.prune(conversation_id, keep_last=self.max_turns)
|
||||
|
||||
logger.debug(
|
||||
f"Added turn {turn.turn_number} to conversation {conversation_id}. "
|
||||
f"Buffer size: {len(buffer.turns)}"
|
||||
)
|
||||
|
||||
async def get_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: Optional[int] = None,
|
||||
offset: int = 0
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Retrieve turns from the buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
offset: Number of turns to skip
|
||||
|
||||
Returns:
|
||||
List of conversation turns
|
||||
"""
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
return []
|
||||
|
||||
turns = buffer.turns[offset:]
|
||||
if limit:
|
||||
turns = turns[:limit]
|
||||
|
||||
return turns
|
||||
|
||||
async def get_recent_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: int = 10
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Get the most recent N turns
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Number of recent turns to retrieve
|
||||
|
||||
Returns:
|
||||
List of recent turns (most recent last)
|
||||
"""
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
return []
|
||||
|
||||
return buffer.turns[-limit:] if len(buffer.turns) > limit else buffer.turns
|
||||
|
||||
async def get_buffer(self, conversation_id: str) -> Optional[ConversationBuffer]:
|
||||
"""
|
||||
Get the full conversation buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
ConversationBuffer or None if not found
|
||||
"""
|
||||
return self._buffers.get(conversation_id)
|
||||
|
||||
async def clear_conversation(self, conversation_id: str) -> None:
|
||||
"""
|
||||
Clear all turns for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
"""
|
||||
if conversation_id in self._buffers:
|
||||
del self._buffers[conversation_id]
|
||||
logger.info(f"Cleared buffer for conversation {conversation_id}")
|
||||
|
||||
async def conversation_exists(self, conversation_id: str) -> bool:
|
||||
"""
|
||||
Check if a conversation exists in the buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
True if conversation exists
|
||||
"""
|
||||
return conversation_id in self._buffers
|
||||
|
||||
async def prune(self, conversation_id: str, keep_last: int = 5) -> None:
|
||||
"""
|
||||
Prune old turns, keeping only the most recent ones
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
keep_last: Number of recent turns to keep
|
||||
"""
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
return
|
||||
|
||||
if len(buffer.turns) > keep_last:
|
||||
removed_count = len(buffer.turns) - keep_last
|
||||
buffer.turns = buffer.turns[-keep_last:]
|
||||
buffer.metadata.turn_count = len(buffer.turns)
|
||||
|
||||
logger.debug(
|
||||
f"Pruned {removed_count} turns from conversation {conversation_id}. "
|
||||
f"Kept last {keep_last} turns."
|
||||
)
|
||||
|
||||
async def get_all_conversation_ids(self) -> List[str]:
|
||||
"""
|
||||
Get list of all conversation IDs in memory
|
||||
|
||||
Returns:
|
||||
List of conversation IDs
|
||||
"""
|
||||
return list(self._buffers.keys())
|
||||
|
||||
async def get_buffer_stats(self) -> dict:
|
||||
"""
|
||||
Get statistics about buffer memory usage
|
||||
|
||||
Returns:
|
||||
Dictionary with stats
|
||||
"""
|
||||
total_conversations = len(self._buffers)
|
||||
total_turns = sum(len(buf.turns) for buf in self._buffers.values())
|
||||
total_tokens = sum(buf.metadata.total_tokens for buf in self._buffers.values())
|
||||
|
||||
return {
|
||||
"total_conversations": total_conversations,
|
||||
"total_turns": total_turns,
|
||||
"total_tokens": total_tokens,
|
||||
"max_turns_per_conversation": self.max_turns,
|
||||
"avg_turns_per_conversation": (
|
||||
total_turns / total_conversations if total_conversations > 0 else 0
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
# Global instance
|
||||
_buffer_memory: Optional[ConversationBufferMemory] = None
|
||||
|
||||
|
||||
def get_buffer_memory(max_turns: int = 10) -> ConversationBufferMemory:
|
||||
"""
|
||||
Get or create the global buffer memory instance
|
||||
|
||||
Args:
|
||||
max_turns: Maximum turns per conversation
|
||||
|
||||
Returns:
|
||||
ConversationBufferMemory instance
|
||||
"""
|
||||
global _buffer_memory
|
||||
if _buffer_memory is None:
|
||||
_buffer_memory = ConversationBufferMemory(max_turns=max_turns)
|
||||
return _buffer_memory
|
||||
@@ -1,12 +0,0 @@
|
||||
"""
|
||||
Metrics collection and monitoring for Core-AI service.
|
||||
|
||||
Provides lightweight, in-memory performance tracking for:
|
||||
- Agent request/response metrics
|
||||
- Tool execution statistics
|
||||
- Memory system performance
|
||||
- Error tracking
|
||||
"""
|
||||
from .collector import MetricsCollector, get_metrics_collector
|
||||
|
||||
__all__ = ['MetricsCollector', 'get_metrics_collector']
|
||||
@@ -1,314 +0,0 @@
|
||||
"""
|
||||
Metrics Collector - In-memory performance metrics storage.
|
||||
|
||||
Tracks agent performance, tool execution, and memory system statistics
|
||||
with thread-safe counters and sliding window storage.
|
||||
"""
|
||||
import time
|
||||
import threading
|
||||
from typing import Dict, List, Any, Optional
|
||||
from collections import defaultdict, deque
|
||||
from datetime import datetime, timedelta
|
||||
import statistics
|
||||
|
||||
|
||||
class MetricsCollector:
|
||||
"""
|
||||
Thread-safe in-memory metrics collector.
|
||||
|
||||
Stores metrics in sliding windows:
|
||||
- Last 1 hour: detailed per-request data
|
||||
- Last 24 hours: aggregated statistics
|
||||
"""
|
||||
|
||||
def __init__(self, detailed_window_hours: int = 1, aggregated_window_hours: int = 24):
|
||||
self.lock = threading.Lock()
|
||||
self.start_time = time.time()
|
||||
|
||||
# Time windows
|
||||
self.detailed_window_seconds = detailed_window_hours * 3600
|
||||
self.aggregated_window_seconds = aggregated_window_hours * 3600
|
||||
|
||||
# Agent request metrics
|
||||
self.total_requests = 0
|
||||
self.requests_by_agent = defaultdict(int)
|
||||
self.request_durations = deque(maxlen=1000) # Last 1000 requests
|
||||
self.request_errors = deque(maxlen=100) # Last 100 errors
|
||||
|
||||
# Tool execution metrics
|
||||
self.total_tool_calls = 0
|
||||
self.tool_calls_by_name = defaultdict(int)
|
||||
self.tool_successes_by_name = defaultdict(int)
|
||||
self.tool_failures_by_name = defaultdict(int)
|
||||
self.tool_durations_by_name = defaultdict(lambda: deque(maxlen=100))
|
||||
self.recent_tool_failures = deque(maxlen=50)
|
||||
|
||||
# Memory system metrics
|
||||
self.memory_tier1_hits = 0
|
||||
self.memory_tier1_misses = 0
|
||||
self.memory_tier2_queries = 0
|
||||
self.memory_tier2_durations = deque(maxlen=100)
|
||||
self.memory_consolidations = 0
|
||||
self.active_users = set()
|
||||
|
||||
# Request pattern metrics
|
||||
self.streaming_requests = 0
|
||||
self.non_streaming_requests = 0
|
||||
self.requests_by_user = defaultdict(int)
|
||||
self.concurrent_requests = 0
|
||||
self.max_concurrent_requests = 0
|
||||
|
||||
# Timestamped events for rate calculation
|
||||
self.request_timestamps = deque(maxlen=1000)
|
||||
|
||||
def _cleanup_old_data(self):
|
||||
"""Remove data older than retention windows."""
|
||||
current_time = time.time()
|
||||
cutoff_time = current_time - self.detailed_window_seconds
|
||||
|
||||
# Clean up request timestamps
|
||||
while self.request_timestamps and self.request_timestamps[0] < cutoff_time:
|
||||
self.request_timestamps.popleft()
|
||||
|
||||
def record_request(self, agent_type: str, duration_ms: float, success: bool,
|
||||
streaming: bool = False, user_id: Optional[str] = None,
|
||||
error: Optional[str] = None):
|
||||
"""
|
||||
Record an agent request.
|
||||
|
||||
Args:
|
||||
agent_type: Type of agent used (pydantic, simple, ollama-native)
|
||||
duration_ms: Request duration in milliseconds
|
||||
success: Whether request completed successfully
|
||||
streaming: Whether this was a streaming request
|
||||
user_id: User identifier (optional)
|
||||
error: Error message if failed (optional)
|
||||
"""
|
||||
with self.lock:
|
||||
self.total_requests += 1
|
||||
self.requests_by_agent[agent_type] += 1
|
||||
self.request_durations.append((time.time(), duration_ms))
|
||||
self.request_timestamps.append(time.time())
|
||||
|
||||
if streaming:
|
||||
self.streaming_requests += 1
|
||||
else:
|
||||
self.non_streaming_requests += 1
|
||||
|
||||
if user_id:
|
||||
self.requests_by_user[user_id] += 1
|
||||
self.active_users.add(user_id)
|
||||
|
||||
if not success and error:
|
||||
self.request_errors.append({
|
||||
'timestamp': time.time(),
|
||||
'agent_type': agent_type,
|
||||
'error': error,
|
||||
'duration_ms': duration_ms
|
||||
})
|
||||
|
||||
self._cleanup_old_data()
|
||||
|
||||
def record_tool_execution(self, tool_name: str, duration_ms: float, success: bool,
|
||||
error: Optional[str] = None):
|
||||
"""
|
||||
Record a tool execution.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
duration_ms: Execution duration in milliseconds
|
||||
success: Whether execution succeeded
|
||||
error: Error message if failed (optional)
|
||||
"""
|
||||
with self.lock:
|
||||
self.total_tool_calls += 1
|
||||
self.tool_calls_by_name[tool_name] += 1
|
||||
self.tool_durations_by_name[tool_name].append(duration_ms)
|
||||
|
||||
if success:
|
||||
self.tool_successes_by_name[tool_name] += 1
|
||||
else:
|
||||
self.tool_failures_by_name[tool_name] += 1
|
||||
if error:
|
||||
self.recent_tool_failures.append({
|
||||
'timestamp': time.time(),
|
||||
'tool_name': tool_name,
|
||||
'error': error,
|
||||
'duration_ms': duration_ms
|
||||
})
|
||||
|
||||
def record_memory_access(self, tier: str, hit: bool, duration_ms: Optional[float] = None):
|
||||
"""
|
||||
Record a memory system access.
|
||||
|
||||
Args:
|
||||
tier: Memory tier (tier1, tier2)
|
||||
hit: Whether it was a cache hit
|
||||
duration_ms: Access duration in milliseconds (optional)
|
||||
"""
|
||||
with self.lock:
|
||||
if tier == "tier1":
|
||||
if hit:
|
||||
self.memory_tier1_hits += 1
|
||||
else:
|
||||
self.memory_tier1_misses += 1
|
||||
elif tier == "tier2":
|
||||
self.memory_tier2_queries += 1
|
||||
if duration_ms is not None:
|
||||
self.memory_tier2_durations.append(duration_ms)
|
||||
|
||||
def record_memory_consolidation(self):
|
||||
"""Record a memory consolidation event."""
|
||||
with self.lock:
|
||||
self.memory_consolidations += 1
|
||||
|
||||
def increment_concurrent_requests(self):
|
||||
"""Increment concurrent request counter."""
|
||||
with self.lock:
|
||||
self.concurrent_requests += 1
|
||||
if self.concurrent_requests > self.max_concurrent_requests:
|
||||
self.max_concurrent_requests = self.concurrent_requests
|
||||
|
||||
def decrement_concurrent_requests(self):
|
||||
"""Decrement concurrent request counter."""
|
||||
with self.lock:
|
||||
self.concurrent_requests = max(0, self.concurrent_requests - 1)
|
||||
|
||||
def get_metrics(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get comprehensive metrics snapshot.
|
||||
|
||||
Returns:
|
||||
Dictionary with all collected metrics
|
||||
"""
|
||||
with self.lock:
|
||||
current_time = time.time()
|
||||
uptime_seconds = current_time - self.start_time
|
||||
|
||||
# Calculate response time percentiles
|
||||
recent_durations = [d for _, d in self.request_durations]
|
||||
if recent_durations:
|
||||
avg_response_time = statistics.mean(recent_durations)
|
||||
p50 = statistics.median(recent_durations)
|
||||
sorted_durations = sorted(recent_durations)
|
||||
p95_idx = int(len(sorted_durations) * 0.95)
|
||||
p99_idx = int(len(sorted_durations) * 0.99)
|
||||
p95 = sorted_durations[p95_idx] if p95_idx < len(sorted_durations) else sorted_durations[-1]
|
||||
p99 = sorted_durations[p99_idx] if p99_idx < len(sorted_durations) else sorted_durations[-1]
|
||||
else:
|
||||
avg_response_time = p50 = p95 = p99 = 0
|
||||
|
||||
# Calculate requests per minute (last 5 minutes)
|
||||
five_min_ago = current_time - 300
|
||||
recent_requests = sum(1 for ts in self.request_timestamps if ts >= five_min_ago)
|
||||
requests_per_minute = (recent_requests / 5) if recent_requests > 0 else 0
|
||||
|
||||
# Calculate tool statistics
|
||||
tool_stats = {}
|
||||
for tool_name in self.tool_calls_by_name.keys():
|
||||
total_calls = self.tool_calls_by_name[tool_name]
|
||||
successes = self.tool_successes_by_name[tool_name]
|
||||
failures = self.tool_failures_by_name[tool_name]
|
||||
durations = list(self.tool_durations_by_name[tool_name])
|
||||
|
||||
tool_stats[tool_name] = {
|
||||
'calls': total_calls,
|
||||
'successes': successes,
|
||||
'failures': failures,
|
||||
'success_rate': successes / total_calls if total_calls > 0 else 0,
|
||||
'avg_duration_ms': statistics.mean(durations) if durations else 0,
|
||||
'p95_duration_ms': sorted(durations)[int(len(durations) * 0.95)] if len(durations) > 1 else (durations[0] if durations else 0)
|
||||
}
|
||||
|
||||
# Sort tools by usage
|
||||
top_tools = dict(sorted(tool_stats.items(), key=lambda x: x[1]['calls'], reverse=True)[:10])
|
||||
|
||||
# Calculate memory hit rate
|
||||
total_tier1_accesses = self.memory_tier1_hits + self.memory_tier1_misses
|
||||
tier1_hit_rate = self.memory_tier1_hits / total_tier1_accesses if total_tier1_accesses > 0 else 0
|
||||
|
||||
tier2_durations = list(self.memory_tier2_durations)
|
||||
tier2_avg_latency = statistics.mean(tier2_durations) if tier2_durations else 0
|
||||
|
||||
return {
|
||||
'uptime_seconds': int(uptime_seconds),
|
||||
'timestamp': datetime.utcnow().isoformat() + 'Z',
|
||||
'agent': {
|
||||
'total_requests': self.total_requests,
|
||||
'requests_by_agent': dict(self.requests_by_agent),
|
||||
'avg_response_time_ms': round(avg_response_time, 2),
|
||||
'p50_response_time_ms': round(p50, 2),
|
||||
'p95_response_time_ms': round(p95, 2),
|
||||
'p99_response_time_ms': round(p99, 2),
|
||||
'errors_total': len(self.request_errors),
|
||||
'requests_per_minute': round(requests_per_minute, 2),
|
||||
'streaming_requests': self.streaming_requests,
|
||||
'non_streaming_requests': self.non_streaming_requests
|
||||
},
|
||||
'concurrency': {
|
||||
'current': self.concurrent_requests,
|
||||
'max': self.max_concurrent_requests
|
||||
},
|
||||
'tools': {
|
||||
'total_calls': self.total_tool_calls,
|
||||
'success_rate': self.tool_successes_by_name and sum(self.tool_successes_by_name.values()) / self.total_tool_calls if self.total_tool_calls > 0 else 0,
|
||||
'top_tools': top_tools,
|
||||
'total_unique_tools': len(self.tool_calls_by_name)
|
||||
},
|
||||
'memory': {
|
||||
'tier1_hits': self.memory_tier1_hits,
|
||||
'tier1_misses': self.memory_tier1_misses,
|
||||
'tier1_hit_rate': round(tier1_hit_rate, 3),
|
||||
'tier2_queries': self.memory_tier2_queries,
|
||||
'tier2_avg_latency_ms': round(tier2_avg_latency, 2),
|
||||
'total_consolidations': self.memory_consolidations,
|
||||
'active_users': len(self.active_users)
|
||||
},
|
||||
'users': {
|
||||
'total_active': len(self.active_users),
|
||||
'top_users': dict(sorted(self.requests_by_user.items(), key=lambda x: x[1], reverse=True)[:5])
|
||||
}
|
||||
}
|
||||
|
||||
def get_recent_errors(self, limit: int = 20) -> List[Dict[str, Any]]:
|
||||
"""Get recent errors."""
|
||||
with self.lock:
|
||||
return [
|
||||
{
|
||||
'timestamp': datetime.fromtimestamp(e['timestamp']).isoformat() + 'Z',
|
||||
'agent_type': e['agent_type'],
|
||||
'error': e['error'],
|
||||
'duration_ms': e['duration_ms']
|
||||
}
|
||||
for e in list(self.request_errors)[-limit:]
|
||||
]
|
||||
|
||||
def get_recent_tool_failures(self, limit: int = 20) -> List[Dict[str, Any]]:
|
||||
"""Get recent tool failures."""
|
||||
with self.lock:
|
||||
return [
|
||||
{
|
||||
'timestamp': datetime.fromtimestamp(f['timestamp']).isoformat() + 'Z',
|
||||
'tool_name': f['tool_name'],
|
||||
'error': f['error'],
|
||||
'duration_ms': f['duration_ms']
|
||||
}
|
||||
for f in list(self.recent_tool_failures)[-limit:]
|
||||
]
|
||||
|
||||
def reset(self):
|
||||
"""Clear all metrics."""
|
||||
with self.lock:
|
||||
self.__init__()
|
||||
|
||||
|
||||
# Global metrics collector instance
|
||||
_metrics_collector: Optional[MetricsCollector] = None
|
||||
|
||||
|
||||
def get_metrics_collector() -> MetricsCollector:
|
||||
"""Get the global metrics collector instance."""
|
||||
global _metrics_collector
|
||||
if _metrics_collector is None:
|
||||
_metrics_collector = MetricsCollector()
|
||||
return _metrics_collector
|
||||
@@ -1,128 +0,0 @@
|
||||
"""
|
||||
Metrics decorators for easy instrumentation.
|
||||
|
||||
Provides decorators to automatically track performance metrics
|
||||
for functions and async functions.
|
||||
"""
|
||||
import time
|
||||
import functools
|
||||
import logging
|
||||
from typing import Callable, Any
|
||||
from .collector import get_metrics_collector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def track_tool_execution(func: Callable) -> Callable:
|
||||
"""
|
||||
Decorator to track tool execution metrics.
|
||||
|
||||
Automatically records:
|
||||
- Tool name
|
||||
- Execution duration
|
||||
- Success/failure status
|
||||
- Error messages on failure
|
||||
|
||||
Usage:
|
||||
@track_tool_execution
|
||||
async def my_tool(arg1, arg2):
|
||||
...
|
||||
"""
|
||||
@functools.wraps(func)
|
||||
async def async_wrapper(*args, **kwargs):
|
||||
metrics = get_metrics_collector()
|
||||
tool_name = func.__name__
|
||||
start_time = time.time()
|
||||
success = False
|
||||
error = None
|
||||
|
||||
try:
|
||||
result = await func(*args, **kwargs)
|
||||
success = True
|
||||
return result
|
||||
except Exception as e:
|
||||
error = str(e)
|
||||
raise
|
||||
finally:
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
metrics.record_tool_execution(
|
||||
tool_name=tool_name,
|
||||
duration_ms=duration_ms,
|
||||
success=success,
|
||||
error=error
|
||||
)
|
||||
if not success:
|
||||
logger.warning(f"Tool {tool_name} failed after {duration_ms:.1f}ms: {error}")
|
||||
|
||||
@functools.wraps(func)
|
||||
def sync_wrapper(*args, **kwargs):
|
||||
metrics = get_metrics_collector()
|
||||
tool_name = func.__name__
|
||||
start_time = time.time()
|
||||
success = False
|
||||
error = None
|
||||
|
||||
try:
|
||||
result = func(*args, **kwargs)
|
||||
success = True
|
||||
return result
|
||||
except Exception as e:
|
||||
error = str(e)
|
||||
raise
|
||||
finally:
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
metrics.record_tool_execution(
|
||||
tool_name=tool_name,
|
||||
duration_ms=duration_ms,
|
||||
success=success,
|
||||
error=error
|
||||
)
|
||||
if not success:
|
||||
logger.warning(f"Tool {tool_name} failed after {duration_ms:.1f}ms: {error}")
|
||||
|
||||
# Return appropriate wrapper based on whether function is async
|
||||
import inspect
|
||||
if inspect.iscoroutinefunction(func):
|
||||
return async_wrapper
|
||||
else:
|
||||
return sync_wrapper
|
||||
|
||||
|
||||
def track_duration(metric_name: str):
|
||||
"""
|
||||
Decorator to track function execution duration.
|
||||
|
||||
Args:
|
||||
metric_name: Name to use for the metric
|
||||
|
||||
Usage:
|
||||
@track_duration("database_query")
|
||||
async def query_database():
|
||||
...
|
||||
"""
|
||||
def decorator(func: Callable) -> Callable:
|
||||
@functools.wraps(func)
|
||||
async def async_wrapper(*args, **kwargs):
|
||||
start_time = time.time()
|
||||
try:
|
||||
return await func(*args, **kwargs)
|
||||
finally:
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
logger.debug(f"{metric_name}: {duration_ms:.1f}ms")
|
||||
|
||||
@functools.wraps(func)
|
||||
def sync_wrapper(*args, **kwargs):
|
||||
start_time = time.time()
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
finally:
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
logger.debug(f"{metric_name}: {duration_ms:.1f}ms")
|
||||
|
||||
import inspect
|
||||
if inspect.iscoroutinefunction(func):
|
||||
return async_wrapper
|
||||
else:
|
||||
return sync_wrapper
|
||||
|
||||
return decorator
|
||||
@@ -1,15 +0,0 @@
|
||||
"""Models for core-ai service"""
|
||||
|
||||
from .embeddings_ollama import (
|
||||
OllamaEmbeddingClient,
|
||||
get_embedding_client,
|
||||
embed_text_async,
|
||||
embed_batch_async
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"OllamaEmbeddingClient",
|
||||
"get_embedding_client",
|
||||
"embed_text_async",
|
||||
"embed_batch_async",
|
||||
]
|
||||
@@ -1,136 +0,0 @@
|
||||
"""
|
||||
Ollama-based embedding client for text vectorization
|
||||
|
||||
Uses Ollama's embedding API instead of local sentence-transformers.
|
||||
This eliminates the need for PyTorch and heavy ML dependencies.
|
||||
"""
|
||||
import logging
|
||||
import httpx
|
||||
from typing import List, Optional
|
||||
from src.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class OllamaEmbeddingClient:
|
||||
"""Client for generating text embeddings using Ollama"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
timeout: int = 30
|
||||
):
|
||||
"""
|
||||
Initialize Ollama embedding client
|
||||
|
||||
Args:
|
||||
model_name: Embedding model name (default: nomic-embed-text)
|
||||
base_url: Ollama base URL (default from settings)
|
||||
timeout: Request timeout in seconds
|
||||
"""
|
||||
self.model_name = model_name or settings.embedding_model
|
||||
self.base_url = (base_url or settings.ollama_base_url).rstrip("/")
|
||||
self.timeout = timeout
|
||||
self.dimension = settings.embedding_dimension
|
||||
|
||||
logger.info(f"Initializing OllamaEmbeddingClient with model: {self.model_name}")
|
||||
logger.info(f"Ollama URL: {self.base_url}")
|
||||
|
||||
async def embed_text(self, text: str) -> List[float]:
|
||||
"""
|
||||
Generate embedding for a single text using Ollama
|
||||
|
||||
Args:
|
||||
text: Input text to embed
|
||||
|
||||
Returns:
|
||||
List of floats representing the embedding vector
|
||||
"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/api/embeddings",
|
||||
json={
|
||||
"model": self.model_name,
|
||||
"prompt": text
|
||||
}
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
return result["embedding"]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating embedding via Ollama: {e}")
|
||||
raise
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
Generate embeddings for multiple texts
|
||||
|
||||
Args:
|
||||
texts: List of input texts
|
||||
|
||||
Returns:
|
||||
List of embedding vectors
|
||||
"""
|
||||
embeddings = []
|
||||
for text in texts:
|
||||
embedding = await self.embed_text(text)
|
||||
embeddings.append(embedding)
|
||||
return embeddings
|
||||
|
||||
def get_dimension(self) -> int:
|
||||
"""
|
||||
Get embedding dimension
|
||||
|
||||
Returns:
|
||||
Embedding vector dimension
|
||||
"""
|
||||
return self.dimension
|
||||
|
||||
|
||||
# Global instance
|
||||
_embedding_client: Optional[OllamaEmbeddingClient] = None
|
||||
|
||||
|
||||
def get_embedding_client() -> OllamaEmbeddingClient:
|
||||
"""
|
||||
Get or create global Ollama embedding client instance
|
||||
|
||||
Returns:
|
||||
OllamaEmbeddingClient instance
|
||||
"""
|
||||
global _embedding_client
|
||||
if _embedding_client is None:
|
||||
_embedding_client = OllamaEmbeddingClient()
|
||||
return _embedding_client
|
||||
|
||||
|
||||
async def embed_text_async(text: str) -> List[float]:
|
||||
"""
|
||||
Async wrapper for embedding text
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
|
||||
Returns:
|
||||
Embedding vector
|
||||
"""
|
||||
client = get_embedding_client()
|
||||
return await client.embed_text(text)
|
||||
|
||||
|
||||
async def embed_batch_async(texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
Async wrapper for batch embedding
|
||||
|
||||
Args:
|
||||
texts: List of input texts
|
||||
|
||||
Returns:
|
||||
List of embedding vectors
|
||||
"""
|
||||
client = get_embedding_client()
|
||||
return await client.embed_batch(texts)
|
||||
@@ -1,57 +0,0 @@
|
||||
"""
|
||||
System Prompt Variants for Core AI
|
||||
|
||||
This file contains minimal, clean prompts for the Core AI service.
|
||||
"""
|
||||
|
||||
PROMPTS = {
|
||||
"minimal_agent": """You are a helpful assistant. You can answer questions. If you need information, use the available tools.""",
|
||||
|
||||
"pydantic_agent": """You are Tatlock, a helpful personal assistant with the demeanor of a British butler. Address users as "sir" and maintain a formal yet personable tone. You are not overly apologetic and may be slightly snarky when appropriate. If an opportunity for a pun presents itself, you cannot resist.
|
||||
|
||||
Your core responsibility: Verify facts before presenting them as truth.
|
||||
|
||||
**IMPORTANT: The Steward's Counsel**
|
||||
When a query is prefixed with "[Steward's analysis: ...]", this represents the household steward's expert assessment of the matter. The steward is highly experienced and their recommendations regarding which tools to employ are exceptionally valuable. While you maintain final discretion, the steward's counsel should be given considerable weight in your deliberations. If the steward recommends specific tools, there is typically sound reasoning behind the suggestion.
|
||||
|
||||
You have access to two categories of tools:
|
||||
|
||||
**Core Tools** (always available):
|
||||
- web_search: ONLY for current events, news, research, or information not available through other tools
|
||||
- calculate: Mathematical operations (precision is paramount)
|
||||
- get_current_time/get_current_date: Time/date queries
|
||||
- add_days_to_date/calculate_date_difference: Date calculations
|
||||
|
||||
**Infrastructure Tools** (discovered from core-api, prefixed with "core_api__"):
|
||||
These tools provide DIRECT system access. PREFER these over web searches when available:
|
||||
|
||||
- DNS queries: core_api__dns_lookup_tools_dns_lookup_post
|
||||
* Use for: A, AAAA, MX, TXT, CNAME, NS, SOA, PTR records
|
||||
* Example: "lookup A records for github.com" → use DNS tool, NOT web_search
|
||||
|
||||
- Web scraping: core_api__scrape_website_tools_scrape_post
|
||||
* Use for: Extract content from specific URLs
|
||||
|
||||
- Docker services: core_api__list_services*, core_api__get_service*, core_api__start_service*, core_api__stop_service*
|
||||
* Use for: Manage sir's containerized services
|
||||
|
||||
- Network management: core_api__list_domains*, core_api__list_ports*, core_api__get_proxy_host*
|
||||
* Use for: Infrastructure configuration
|
||||
|
||||
- Monitoring: core_api__list_monitors*, core_api__get_monitor*, core_api__create_monitor*
|
||||
* Use for: System health checks
|
||||
|
||||
**Tool Selection Priority:**
|
||||
1. If an infrastructure tool exists for the task → use it (more reliable than web search)
|
||||
2. If no infrastructure tool exists → use web_search
|
||||
3. For general knowledge → answer directly (no tool needed)
|
||||
|
||||
Be concise unless details are specifically requested. When using tools, acknowledge them naturally in your dignified manner."""
|
||||
}
|
||||
|
||||
|
||||
def get_prompt(variant: str = "minimal_agent") -> str:
|
||||
"""
|
||||
Get a system prompt variant.
|
||||
"""
|
||||
return PROMPTS.get(variant, PROMPTS["minimal_agent"])
|
||||
@@ -1,25 +0,0 @@
|
||||
"""
|
||||
Tools module for Core-AI ADK agent.
|
||||
|
||||
This module provides tool registration and management for the ADK agent.
|
||||
Tools can make REST calls to core-api or operate independently.
|
||||
"""
|
||||
from src.tools.registry import (
|
||||
get_agent_tools,
|
||||
register_tool,
|
||||
get_all_tools,
|
||||
discover_and_register_tools,
|
||||
clear_registry
|
||||
)
|
||||
|
||||
# Import local tools to trigger registration
|
||||
# This must happen before get_agent_tools() is called
|
||||
import src.tools.local # noqa: F401
|
||||
|
||||
__all__ = [
|
||||
"get_agent_tools",
|
||||
"register_tool",
|
||||
"get_all_tools",
|
||||
"discover_and_register_tools",
|
||||
"clear_registry",
|
||||
]
|
||||
@@ -1,40 +0,0 @@
|
||||
"""
|
||||
Infrastructure Tools - Docker container and service management tools.
|
||||
|
||||
This package provides tools for managing infrastructure systems:
|
||||
- containers.py: Docker container lifecycle management (4 tools)
|
||||
- services.py: Docker Compose service management (3 tools)
|
||||
- monitoring.py: System and container resource monitoring (2 tools)
|
||||
"""
|
||||
from src.tools.infrastructure.containers import (
|
||||
docker_list_containers,
|
||||
docker_manage_container,
|
||||
docker_inspect_container,
|
||||
docker_container_logs,
|
||||
)
|
||||
|
||||
from src.tools.infrastructure.services import (
|
||||
list_services,
|
||||
manage_service,
|
||||
service_status,
|
||||
)
|
||||
|
||||
from src.tools.infrastructure.monitoring import (
|
||||
system_resources,
|
||||
container_resources,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Container tools
|
||||
"docker_list_containers",
|
||||
"docker_manage_container",
|
||||
"docker_inspect_container",
|
||||
"docker_container_logs",
|
||||
# Service tools
|
||||
"list_services",
|
||||
"manage_service",
|
||||
"service_status",
|
||||
# Monitoring tools
|
||||
"system_resources",
|
||||
"container_resources",
|
||||
]
|
||||
@@ -1,490 +0,0 @@
|
||||
"""
|
||||
Docker Container Management Tools
|
||||
|
||||
Provides tools for container lifecycle operations and diagnostics.
|
||||
All operations go through core-api for centralized logging.
|
||||
"""
|
||||
import logging
|
||||
import httpx
|
||||
from typing import Optional, Literal
|
||||
from datetime import datetime
|
||||
from src.config import get_settings
|
||||
from src.tools.registry import register_tool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
def _format_uptime(started_at: str) -> str:
|
||||
"""
|
||||
Convert ISO timestamp to human-readable uptime.
|
||||
|
||||
Args:
|
||||
started_at: ISO timestamp string
|
||||
|
||||
Returns:
|
||||
Human-readable uptime (e.g., "3 days", "7 hours", "45 minutes")
|
||||
"""
|
||||
try:
|
||||
# Handle both formats: with and without timezone
|
||||
if 'Z' in started_at:
|
||||
start_time = datetime.fromisoformat(started_at.replace('Z', '+00:00'))
|
||||
elif '+' in started_at or started_at.endswith('00:00'):
|
||||
start_time = datetime.fromisoformat(started_at)
|
||||
else:
|
||||
start_time = datetime.fromisoformat(started_at + '+00:00')
|
||||
|
||||
uptime_delta = datetime.now(start_time.tzinfo) - start_time
|
||||
|
||||
days = uptime_delta.days
|
||||
hours = uptime_delta.seconds // 3600
|
||||
minutes = (uptime_delta.seconds % 3600) // 60
|
||||
|
||||
if days > 0:
|
||||
return f"{days} day{'s' if days != 1 else ''}"
|
||||
elif hours > 0:
|
||||
return f"{hours} hour{'s' if hours != 1 else ''}"
|
||||
else:
|
||||
return f"{minutes} minute{'s' if minutes != 1 else ''}"
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to parse uptime from '{started_at}': {e}")
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _format_ports(ports: list) -> str:
|
||||
"""
|
||||
Format container port bindings for display.
|
||||
|
||||
Args:
|
||||
ports: Docker API port bindings
|
||||
|
||||
Returns:
|
||||
Formatted port string (e.g., "80→8080, 443→8443")
|
||||
"""
|
||||
if not ports:
|
||||
return "none"
|
||||
|
||||
port_mappings = []
|
||||
for port_data in ports:
|
||||
if isinstance(port_data, dict):
|
||||
private_port = port_data.get('PrivatePort')
|
||||
public_port = port_data.get('PublicPort')
|
||||
if public_port and private_port:
|
||||
port_mappings.append(f"{private_port}→{public_port}")
|
||||
elif private_port:
|
||||
port_mappings.append(f"{private_port} (internal)")
|
||||
|
||||
return ", ".join(port_mappings) if port_mappings else "none"
|
||||
|
||||
|
||||
@register_tool
|
||||
async def docker_list_containers(
|
||||
status: Optional[Literal["all", "running", "stopped", "paused"]] = "running"
|
||||
) -> str:
|
||||
"""
|
||||
List Docker containers with status and resource usage.
|
||||
|
||||
Args:
|
||||
status: Filter by status - "all", "running", "stopped", or "paused"
|
||||
Defaults to "running" to show only active containers.
|
||||
|
||||
Returns:
|
||||
Formatted list of containers with details including:
|
||||
- Container name and status
|
||||
- Uptime (for running containers)
|
||||
- Port mappings
|
||||
- Image information
|
||||
|
||||
Examples:
|
||||
docker_list_containers("all") # All containers
|
||||
docker_list_containers("running") # Only running (default)
|
||||
docker_list_containers("stopped") # Only stopped
|
||||
"""
|
||||
logger.info(f"Listing Docker containers (status filter: {status})")
|
||||
|
||||
try:
|
||||
# Call core-api infrastructure endpoint
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
response = await client.get(
|
||||
f"{settings.core_api_base_url}/infrastructure/containers",
|
||||
params={"status": status}
|
||||
)
|
||||
response.raise_for_status()
|
||||
containers = response.json()
|
||||
|
||||
if not containers:
|
||||
return f"No {status} containers found."
|
||||
|
||||
# Count by status
|
||||
running_count = sum(1 for c in containers if c.get('State') == 'running')
|
||||
stopped_count = len(containers) - running_count
|
||||
|
||||
# Build formatted output
|
||||
lines = [f"Containers ({running_count} running, {stopped_count} stopped):"]
|
||||
lines.append("")
|
||||
|
||||
for container in containers:
|
||||
# Extract container name (strip leading '/')
|
||||
names = container.get('Names', ['unknown'])
|
||||
name = names[0].lstrip('/') if names else 'unknown'
|
||||
|
||||
state = container.get('State', 'unknown')
|
||||
status_info = container.get('Status', '')
|
||||
|
||||
lines.append(f"• {name}")
|
||||
lines.append(f" Status: {state}")
|
||||
|
||||
# Add status info
|
||||
if status_info:
|
||||
lines.append(f" Info: {status_info}")
|
||||
|
||||
# Port mappings
|
||||
ports = _format_ports(container.get('Ports', []))
|
||||
if ports != "none":
|
||||
lines.append(f" Ports: {ports}")
|
||||
|
||||
# Image
|
||||
image = container.get('Image', 'unknown')
|
||||
if image != 'unknown':
|
||||
# Shorten long image names
|
||||
if len(image) > 50:
|
||||
image = image[:47] + "..."
|
||||
lines.append(f" Image: {image}")
|
||||
|
||||
lines.append("") # Blank line between containers
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_msg = f"Failed to list containers: HTTP {e.response.status_code}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Error listing containers: {str(e)}"
|
||||
logger.error(error_msg, exc_info=True)
|
||||
return error_msg
|
||||
|
||||
|
||||
@register_tool
|
||||
async def docker_manage_container(
|
||||
container: str,
|
||||
action: Literal["start", "stop", "restart"]
|
||||
) -> str:
|
||||
"""
|
||||
Manage Docker container state.
|
||||
|
||||
This high-level tool handles container lifecycle operations.
|
||||
|
||||
Args:
|
||||
container: Container name or ID (e.g., "nginx", "core-ai")
|
||||
action: Action to perform:
|
||||
- start: Start a stopped container
|
||||
- stop: Stop a running container
|
||||
- restart: Stop and start a container
|
||||
|
||||
Returns:
|
||||
Success message or error details
|
||||
|
||||
Examples:
|
||||
docker_manage_container("nginx", "restart")
|
||||
docker_manage_container("core-ai", "stop")
|
||||
docker_manage_container("ollama", "start")
|
||||
|
||||
Error Handling:
|
||||
- Container not found → Returns error with suggestion to check docker_list_containers()
|
||||
- Already in target state → Reports current state
|
||||
- Permission denied → Reports error for user escalation
|
||||
"""
|
||||
logger.info(f"Managing container '{container}': {action}")
|
||||
|
||||
try:
|
||||
# Call core-api infrastructure endpoint
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(
|
||||
f"{settings.core_api_base_url}/infrastructure/containers/{container}/{action}"
|
||||
)
|
||||
|
||||
if response.status_code == 404:
|
||||
return f"Container '{container}' not found. Use docker_list_containers() to see available containers."
|
||||
elif response.status_code == 304:
|
||||
return f"Container '{container}' is already in the target state for action '{action}'."
|
||||
elif response.status_code == 409:
|
||||
return f"Cannot {action} container '{container}': state conflict (may already be stopped/started)."
|
||||
elif response.status_code == 501:
|
||||
return f"Action '{action}' is not yet implemented on the server."
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
message = result.get('message', f"Action '{action}' completed for container '{container}'.")
|
||||
return f"✓ {message}"
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_msg = f"Failed to {action} container '{container}': HTTP {e.response.status_code}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Error managing container '{container}': {str(e)}"
|
||||
logger.error(error_msg, exc_info=True)
|
||||
return error_msg
|
||||
|
||||
|
||||
@register_tool
|
||||
async def docker_inspect_container(
|
||||
container: str,
|
||||
details: Literal["summary", "full", "resources"] = "summary"
|
||||
) -> str:
|
||||
"""
|
||||
Get detailed container information.
|
||||
|
||||
Args:
|
||||
container: Container name or ID (e.g., "nginx", "core-ai")
|
||||
details: Level of detail to return:
|
||||
- summary: Name, status, uptime, ports, image (default)
|
||||
- full: Add environment vars, mounts, network config
|
||||
- resources: Focus on resource limits and configuration
|
||||
|
||||
Returns:
|
||||
Formatted container inspection based on detail level
|
||||
|
||||
Examples:
|
||||
docker_inspect_container("nginx") # Quick summary
|
||||
docker_inspect_container("nginx", "full") # Complete details
|
||||
docker_inspect_container("core-ai", "resources") # Resource limits
|
||||
"""
|
||||
logger.info(f"Inspecting container '{container}' (detail level: {details})")
|
||||
|
||||
try:
|
||||
# Call core-api infrastructure endpoint
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
response = await client.get(
|
||||
f"{settings.core_api_base_url}/infrastructure/containers/{container}",
|
||||
params={"details": details}
|
||||
)
|
||||
|
||||
if response.status_code == 404:
|
||||
return f"Container '{container}' not found. Use docker_list_containers() to see available containers."
|
||||
|
||||
response.raise_for_status()
|
||||
info = response.json()
|
||||
|
||||
# Extract key information
|
||||
state = info.get('State', {})
|
||||
config = info.get('Config', {})
|
||||
network_settings = info.get('NetworkSettings', {})
|
||||
host_config = info.get('HostConfig', {})
|
||||
|
||||
name = info.get('Name', '').lstrip('/')
|
||||
status = state.get('Status', 'unknown')
|
||||
running = state.get('Running', False)
|
||||
|
||||
lines = [f"Container: {name}"]
|
||||
lines.append(f"Status: {status}")
|
||||
|
||||
if details == "summary":
|
||||
# Summary: Basic info
|
||||
if running and 'StartedAt' in state:
|
||||
uptime = _format_uptime(state['StartedAt'])
|
||||
lines.append(f"Uptime: {uptime}")
|
||||
elif 'FinishedAt' in state and state.get('ExitCode') is not None:
|
||||
exit_code = state.get('ExitCode')
|
||||
lines.append(f"Exit Code: {exit_code}")
|
||||
|
||||
# Ports
|
||||
ports = network_settings.get('Ports', {})
|
||||
if ports:
|
||||
lines.append(f"\nPorts:")
|
||||
for internal_port, bindings in ports.items():
|
||||
if bindings:
|
||||
for binding in bindings:
|
||||
host_ip = binding.get('HostIp', '0.0.0.0')
|
||||
host_port = binding.get('HostPort')
|
||||
lines.append(f" {internal_port} → {host_ip}:{host_port}")
|
||||
else:
|
||||
lines.append(f" {internal_port} (not exposed)")
|
||||
|
||||
# Image
|
||||
image = config.get('Image', 'unknown')
|
||||
lines.append(f"\nImage: {image}")
|
||||
|
||||
elif details == "full":
|
||||
# Full: Everything
|
||||
if running and 'StartedAt' in state:
|
||||
uptime = _format_uptime(state['StartedAt'])
|
||||
lines.append(f"Uptime: {uptime}")
|
||||
elif 'FinishedAt' in state:
|
||||
lines.append(f"Exit Code: {state.get('ExitCode', 'N/A')}")
|
||||
|
||||
# Image
|
||||
image = config.get('Image', 'unknown')
|
||||
lines.append(f"\nImage: {image}")
|
||||
|
||||
# Ports
|
||||
ports = network_settings.get('Ports', {})
|
||||
if ports:
|
||||
lines.append(f"\nPorts:")
|
||||
for internal_port, bindings in ports.items():
|
||||
if bindings:
|
||||
for binding in bindings:
|
||||
host_ip = binding.get('HostIp', '0.0.0.0')
|
||||
host_port = binding.get('HostPort')
|
||||
lines.append(f" {internal_port} → {host_ip}:{host_port}")
|
||||
|
||||
# Environment (show only non-sensitive keys)
|
||||
env = config.get('Env', [])
|
||||
if env:
|
||||
lines.append(f"\nEnvironment ({len(env)} variables):")
|
||||
# Show first 10 variable names only (not values, for security)
|
||||
for var in env[:10]:
|
||||
if '=' in var:
|
||||
key = var.split('=')[0]
|
||||
lines.append(f" {key}")
|
||||
if len(env) > 10:
|
||||
lines.append(f" ... and {len(env) - 10} more")
|
||||
|
||||
# Mounts
|
||||
mounts = info.get('Mounts', [])
|
||||
if mounts:
|
||||
lines.append(f"\nMounts ({len(mounts)}):")
|
||||
for mount in mounts[:5]:
|
||||
mount_type = mount.get('Type', 'unknown')
|
||||
source = mount.get('Source', '')[:40] # Truncate long paths
|
||||
destination = mount.get('Destination', '')
|
||||
lines.append(f" {mount_type}: {source} → {destination}")
|
||||
if len(mounts) > 5:
|
||||
lines.append(f" ... and {len(mounts) - 5} more")
|
||||
|
||||
# Networks
|
||||
networks = network_settings.get('Networks', {})
|
||||
if networks:
|
||||
lines.append(f"\nNetworks:")
|
||||
for net_name, net_config in networks.items():
|
||||
ip = net_config.get('IPAddress', 'N/A')
|
||||
gateway = net_config.get('Gateway', 'N/A')
|
||||
lines.append(f" {net_name}:")
|
||||
lines.append(f" IP: {ip}")
|
||||
lines.append(f" Gateway: {gateway}")
|
||||
|
||||
elif details == "resources":
|
||||
# Resources: Limits and configuration
|
||||
lines.append(f"\nResource Configuration:")
|
||||
|
||||
# Memory
|
||||
memory_limit = host_config.get('Memory', 0)
|
||||
if memory_limit > 0:
|
||||
# Format bytes
|
||||
memory_mb = memory_limit / (1024 * 1024)
|
||||
lines.append(f" Memory Limit: {memory_mb:.1f} MB")
|
||||
else:
|
||||
lines.append(f" Memory Limit: unlimited")
|
||||
|
||||
memory_reservation = host_config.get('MemoryReservation', 0)
|
||||
if memory_reservation > 0:
|
||||
mem_res_mb = memory_reservation / (1024 * 1024)
|
||||
lines.append(f" Memory Reservation: {mem_res_mb:.1f} MB")
|
||||
|
||||
# CPU
|
||||
cpu_shares = host_config.get('CpuShares', 0)
|
||||
if cpu_shares > 0:
|
||||
lines.append(f" CPU Shares: {cpu_shares}")
|
||||
|
||||
nano_cpus = host_config.get('NanoCpus', 0)
|
||||
if nano_cpus > 0:
|
||||
cpus = nano_cpus / 1_000_000_000
|
||||
lines.append(f" CPU Limit: {cpus:.2f} CPUs")
|
||||
|
||||
cpu_quota = host_config.get('CpuQuota', 0)
|
||||
if cpu_quota > 0:
|
||||
lines.append(f" CPU Quota: {cpu_quota}")
|
||||
|
||||
# Restart policy
|
||||
restart_policy = host_config.get('RestartPolicy', {})
|
||||
policy_name = restart_policy.get('Name', 'no')
|
||||
lines.append(f"\n Restart Policy: {policy_name}")
|
||||
|
||||
# Image
|
||||
image = config.get('Image', 'unknown')
|
||||
lines.append(f"\nImage: {image}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_msg = f"Failed to inspect container '{container}': HTTP {e.response.status_code}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Error inspecting container '{container}': {str(e)}"
|
||||
logger.error(error_msg, exc_info=True)
|
||||
return error_msg
|
||||
|
||||
|
||||
@register_tool
|
||||
async def docker_container_logs(
|
||||
container: str,
|
||||
lines: int = 50,
|
||||
since: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Retrieve Docker container logs.
|
||||
|
||||
Args:
|
||||
container: Container name or ID (e.g., "nginx", "core-ai")
|
||||
lines: Number of recent log lines to retrieve
|
||||
Default: 50, Maximum: 500 (to avoid overwhelming output)
|
||||
since: Optional time filter for logs (not yet fully implemented):
|
||||
- "1h" = last hour
|
||||
- "30m" = last 30 minutes
|
||||
- Note: Currently ignored server-side, uses line limit only
|
||||
|
||||
Returns:
|
||||
Container logs with timestamps
|
||||
Each line prefixed with timestamp if available
|
||||
|
||||
Examples:
|
||||
docker_container_logs("nginx") # Last 50 lines
|
||||
docker_container_logs("nginx", lines=100) # Last 100 lines
|
||||
docker_container_logs("core-ai", lines=200) # Last 200 lines
|
||||
"""
|
||||
logger.info(f"Retrieving logs for container '{container}' (lines={lines})")
|
||||
|
||||
# Clamp lines to reasonable limit
|
||||
lines = min(max(1, lines), 500)
|
||||
|
||||
try:
|
||||
# Call core-api infrastructure endpoint
|
||||
async with httpx.AsyncClient(timeout=15.0) as client:
|
||||
params = {"lines": lines}
|
||||
if since:
|
||||
params["since"] = since
|
||||
|
||||
response = await client.get(
|
||||
f"{settings.core_api_base_url}/infrastructure/containers/{container}/logs",
|
||||
params=params
|
||||
)
|
||||
|
||||
if response.status_code == 404:
|
||||
return f"Container '{container}' not found. Use docker_list_containers() to see available containers."
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
logs = result.get('logs', '')
|
||||
|
||||
if not logs or logs.strip() == '':
|
||||
return f"No logs found for container '{container}' (container may be newly started or have no output)."
|
||||
|
||||
# Format header
|
||||
time_filter = f" (since {since})" if since else ""
|
||||
header = f"Container '{container}' logs (last {lines} lines{time_filter}):\n"
|
||||
header += "=" * 60 + "\n\n"
|
||||
|
||||
return header + logs
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_msg = f"Failed to retrieve logs for container '{container}': HTTP {e.response.status_code}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Error retrieving logs for container '{container}': {str(e)}"
|
||||
logger.error(error_msg, exc_info=True)
|
||||
return error_msg
|
||||
@@ -1,247 +0,0 @@
|
||||
"""
|
||||
System and Container Monitoring Tools
|
||||
|
||||
Provides tools for monitoring resource usage at system and container levels.
|
||||
All operations go through core-api for centralized logging.
|
||||
"""
|
||||
import logging
|
||||
import httpx
|
||||
from typing import Optional
|
||||
from src.config import get_settings
|
||||
from src.tools.registry import register_tool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
@register_tool
|
||||
async def system_resources() -> str:
|
||||
"""
|
||||
Get overall system resource usage.
|
||||
|
||||
Provides comprehensive system-level metrics including:
|
||||
- CPU usage and core count
|
||||
- Memory usage (total, used, available)
|
||||
- Disk usage (total, used, available)
|
||||
- Network statistics (if available)
|
||||
|
||||
Returns:
|
||||
Formatted system resource report
|
||||
|
||||
Examples:
|
||||
system_resources() # Get current system metrics
|
||||
"""
|
||||
logger.info("Getting system resources")
|
||||
|
||||
try:
|
||||
# Call core-api infrastructure endpoint
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
response = await client.get(
|
||||
f"{settings.core_api_base_url}/infrastructure/resources/system"
|
||||
)
|
||||
response.raise_for_status()
|
||||
resources = response.json()
|
||||
|
||||
# Build formatted output
|
||||
lines = ["System Resources:"]
|
||||
lines.append("")
|
||||
|
||||
# CPU
|
||||
cpu = resources.get('cpu', {})
|
||||
if cpu:
|
||||
cores = cpu.get('cores')
|
||||
usage = cpu.get('usage_percent')
|
||||
load_avg = cpu.get('load_average', [])
|
||||
|
||||
lines.append("CPU:")
|
||||
if cores:
|
||||
lines.append(f" Cores: {cores}")
|
||||
if usage is not None:
|
||||
lines.append(f" Usage: {usage:.1f}%")
|
||||
if load_avg:
|
||||
load_str = ", ".join(f"{l:.2f}" for l in load_avg)
|
||||
lines.append(f" Load Average: {load_str}")
|
||||
lines.append("")
|
||||
|
||||
# Memory
|
||||
memory = resources.get('memory', {})
|
||||
if memory:
|
||||
total = memory.get('total_bytes')
|
||||
used = memory.get('used_bytes')
|
||||
available = memory.get('available_bytes')
|
||||
usage_pct = memory.get('usage_percent')
|
||||
|
||||
lines.append("Memory:")
|
||||
if total:
|
||||
total_gb = total / (1024**3)
|
||||
lines.append(f" Total: {total_gb:.1f} GB")
|
||||
if used:
|
||||
used_gb = used / (1024**3)
|
||||
lines.append(f" Used: {used_gb:.1f} GB")
|
||||
if available:
|
||||
avail_gb = available / (1024**3)
|
||||
lines.append(f" Available: {avail_gb:.1f} GB")
|
||||
if usage_pct is not None:
|
||||
lines.append(f" Usage: {usage_pct:.1f}%")
|
||||
lines.append("")
|
||||
|
||||
# Disk
|
||||
disk = resources.get('disk', {})
|
||||
if disk:
|
||||
total = disk.get('total_bytes')
|
||||
used = disk.get('used_bytes')
|
||||
available = disk.get('available_bytes')
|
||||
usage_pct = disk.get('usage_percent')
|
||||
|
||||
lines.append("Disk:")
|
||||
if total:
|
||||
total_gb = total / (1024**3)
|
||||
lines.append(f" Total: {total_gb:.1f} GB")
|
||||
if used:
|
||||
used_gb = used / (1024**3)
|
||||
lines.append(f" Used: {used_gb:.1f} GB")
|
||||
if available:
|
||||
avail_gb = available / (1024**3)
|
||||
lines.append(f" Available: {avail_gb:.1f} GB")
|
||||
if usage_pct is not None:
|
||||
lines.append(f" Usage: {usage_pct:.1f}%")
|
||||
lines.append("")
|
||||
|
||||
# Network
|
||||
network = resources.get('network', {})
|
||||
if network:
|
||||
interfaces = network.get('interfaces', {})
|
||||
if interfaces:
|
||||
lines.append("Network:")
|
||||
for iface_name, iface_data in interfaces.items():
|
||||
rx = iface_data.get('rx_bytes', 0)
|
||||
tx = iface_data.get('tx_bytes', 0)
|
||||
rx_gb = rx / (1024**3)
|
||||
tx_gb = tx / (1024**3)
|
||||
lines.append(f" {iface_name}:")
|
||||
lines.append(f" RX: {rx_gb:.2f} GB")
|
||||
lines.append(f" TX: {tx_gb:.2f} GB")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_msg = f"Failed to get system resources: HTTP {e.response.status_code}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Error getting system resources: {str(e)}"
|
||||
logger.error(error_msg, exc_info=True)
|
||||
return error_msg
|
||||
|
||||
|
||||
@register_tool
|
||||
async def container_resources(
|
||||
container: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Get container-specific resource usage.
|
||||
|
||||
Provides real-time resource metrics for containers including:
|
||||
- CPU usage percentage
|
||||
- Memory usage (current, limit, percentage)
|
||||
- Network I/O (received, transmitted)
|
||||
- Block I/O (read, write)
|
||||
|
||||
Args:
|
||||
container: Specific container name or ID (optional)
|
||||
If omitted, returns stats for all running containers
|
||||
|
||||
Returns:
|
||||
Formatted container resource report
|
||||
|
||||
Examples:
|
||||
container_resources() # All containers
|
||||
container_resources("nginx") # Specific container
|
||||
container_resources("core-ai") # Another specific container
|
||||
"""
|
||||
logger.info(f"Getting container resources (container={container})")
|
||||
|
||||
try:
|
||||
# Call core-api infrastructure endpoint
|
||||
async with httpx.AsyncClient(timeout=15.0) as client:
|
||||
params = {}
|
||||
if container:
|
||||
params["container"] = container
|
||||
|
||||
response = await client.get(
|
||||
f"{settings.core_api_base_url}/infrastructure/resources/containers",
|
||||
params=params
|
||||
)
|
||||
|
||||
if response.status_code == 404:
|
||||
return f"Container '{container}' not found. Use docker_list_containers() to see available containers."
|
||||
|
||||
response.raise_for_status()
|
||||
resources = response.json()
|
||||
|
||||
if not resources:
|
||||
if container:
|
||||
return f"No resource data available for container '{container}'."
|
||||
else:
|
||||
return "No containers are currently running."
|
||||
|
||||
# Build formatted output
|
||||
if container:
|
||||
lines = [f"Container '{container}' Resources:"]
|
||||
else:
|
||||
lines = [f"Container Resources ({len(resources)} containers):"]
|
||||
lines.append("")
|
||||
|
||||
for res in resources:
|
||||
name = res.get('name', 'unknown')
|
||||
cpu = res.get('cpu_percent')
|
||||
mem_usage = res.get('memory_usage_bytes')
|
||||
mem_limit = res.get('memory_limit_bytes')
|
||||
mem_pct = res.get('memory_percent')
|
||||
net_rx = res.get('network_rx_bytes')
|
||||
net_tx = res.get('network_tx_bytes')
|
||||
block_read = res.get('block_read_bytes')
|
||||
block_write = res.get('block_write_bytes')
|
||||
|
||||
lines.append(f"• {name}")
|
||||
|
||||
# CPU
|
||||
if cpu is not None:
|
||||
lines.append(f" CPU: {cpu:.1f}%")
|
||||
|
||||
# Memory
|
||||
if mem_usage is not None and mem_limit is not None:
|
||||
mem_usage_mb = mem_usage / (1024**2)
|
||||
mem_limit_mb = mem_limit / (1024**2)
|
||||
mem_line = f" Memory: {mem_usage_mb:.1f} MB / {mem_limit_mb:.1f} MB"
|
||||
if mem_pct is not None:
|
||||
mem_line += f" ({mem_pct:.1f}%)"
|
||||
lines.append(mem_line)
|
||||
elif mem_usage is not None:
|
||||
mem_usage_mb = mem_usage / (1024**2)
|
||||
lines.append(f" Memory: {mem_usage_mb:.1f} MB")
|
||||
|
||||
# Network
|
||||
if net_rx is not None and net_tx is not None:
|
||||
net_rx_mb = net_rx / (1024**2)
|
||||
net_tx_mb = net_tx / (1024**2)
|
||||
lines.append(f" Network: RX {net_rx_mb:.1f} MB / TX {net_tx_mb:.1f} MB")
|
||||
|
||||
# Block I/O
|
||||
if block_read is not None and block_write is not None:
|
||||
block_read_mb = block_read / (1024**2)
|
||||
block_write_mb = block_write / (1024**2)
|
||||
lines.append(f" Block I/O: Read {block_read_mb:.1f} MB / Write {block_write_mb:.1f} MB")
|
||||
|
||||
lines.append("") # Blank line between containers
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_msg = f"Failed to get container resources: HTTP {e.response.status_code}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Error getting container resources: {str(e)}"
|
||||
logger.error(error_msg, exc_info=True)
|
||||
return error_msg
|
||||
@@ -1,283 +0,0 @@
|
||||
"""
|
||||
Docker Service Management Tools
|
||||
|
||||
Provides tools for managing Docker Compose services/stacks.
|
||||
All operations go through core-api for centralized logging.
|
||||
"""
|
||||
import logging
|
||||
import httpx
|
||||
from typing import Optional, Literal
|
||||
from src.config import get_settings
|
||||
from src.tools.registry import register_tool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
@register_tool
|
||||
async def list_services(
|
||||
stack: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
List Docker Compose services (stacks).
|
||||
|
||||
Args:
|
||||
stack: Optional filter by stack name (case-insensitive)
|
||||
If provided, returns only matching stack.
|
||||
If omitted, returns all stacks.
|
||||
|
||||
Returns:
|
||||
Formatted list of services with:
|
||||
- Service/stack name and status
|
||||
- Container counts (running/total)
|
||||
- Exposed ports
|
||||
- Configured domains (from reverse proxy)
|
||||
|
||||
Examples:
|
||||
list_services() # All services
|
||||
list_services("portainer") # Specific stack
|
||||
list_services("core") # Stacks matching "core"
|
||||
"""
|
||||
logger.info(f"Listing services (stack filter: {stack})")
|
||||
|
||||
try:
|
||||
# Call core-api infrastructure endpoint
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
params = {}
|
||||
if stack:
|
||||
params["stack"] = stack
|
||||
|
||||
response = await client.get(
|
||||
f"{settings.core_api_base_url}/infrastructure/services",
|
||||
params=params
|
||||
)
|
||||
response.raise_for_status()
|
||||
services = response.json()
|
||||
|
||||
if not services:
|
||||
if stack:
|
||||
return f"No services found matching '{stack}'."
|
||||
else:
|
||||
return "No services found."
|
||||
|
||||
# Count active services
|
||||
active_count = sum(1 for s in services if s.get('status') == 'active')
|
||||
total_count = len(services)
|
||||
|
||||
# Build formatted output
|
||||
lines = [f"Services ({active_count} active, {total_count} total):"]
|
||||
lines.append("")
|
||||
|
||||
for service in services:
|
||||
name = service.get('name', 'unknown')
|
||||
status = service.get('status', 'unknown')
|
||||
running = service.get('containers_running', 0)
|
||||
total = service.get('containers_total', 0)
|
||||
|
||||
lines.append(f"• {name}")
|
||||
lines.append(f" Status: {status}")
|
||||
lines.append(f" Containers: {running}/{total} running")
|
||||
|
||||
# Ports
|
||||
ports = service.get('ports', [])
|
||||
if ports:
|
||||
port_str = ", ".join(str(p) for p in ports)
|
||||
lines.append(f" Ports: {port_str}")
|
||||
|
||||
# Domains
|
||||
domains = service.get('domains', [])
|
||||
if domains:
|
||||
domain_str = ", ".join(domains)
|
||||
lines.append(f" Domains: {domain_str}")
|
||||
|
||||
lines.append("") # Blank line between services
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_msg = f"Failed to list services: HTTP {e.response.status_code}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Error listing services: {str(e)}"
|
||||
logger.error(error_msg, exc_info=True)
|
||||
return error_msg
|
||||
|
||||
|
||||
@register_tool
|
||||
async def manage_service(
|
||||
service: str,
|
||||
action: Literal["start", "stop", "restart", "scale"],
|
||||
replicas: Optional[int] = None
|
||||
) -> str:
|
||||
"""
|
||||
Manage Docker service lifecycle and scaling.
|
||||
|
||||
Args:
|
||||
service: Service/stack name (e.g., "portainer", "core-ai")
|
||||
action: Action to perform:
|
||||
- start: Start all containers in the service
|
||||
- stop: Stop all containers in the service
|
||||
- restart: Stop and start the service
|
||||
- scale: Change number of replicas (requires replicas parameter)
|
||||
replicas: Number of replicas (required only for scale action)
|
||||
|
||||
Returns:
|
||||
Success message or error details
|
||||
|
||||
Examples:
|
||||
manage_service("web", "restart")
|
||||
manage_service("worker", "scale", replicas=3)
|
||||
manage_service("portainer", "stop")
|
||||
|
||||
Error Handling:
|
||||
- Service not found → Returns error with suggestion to check list_services()
|
||||
- Scale without replicas → Returns error asking for replicas parameter
|
||||
- Invalid action → Returns error with valid actions list
|
||||
"""
|
||||
logger.info(f"Managing service '{service}': action={action}, replicas={replicas}")
|
||||
|
||||
# Validate scale action has replicas
|
||||
if action == "scale" and replicas is None:
|
||||
return "Error: 'scale' action requires 'replicas' parameter. Example: manage_service('web', 'scale', replicas=3)"
|
||||
|
||||
try:
|
||||
# Call core-api infrastructure endpoint
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
# Build request body
|
||||
body = {"action": action}
|
||||
if replicas is not None:
|
||||
body["replicas"] = replicas
|
||||
|
||||
response = await client.post(
|
||||
f"{settings.core_api_base_url}/infrastructure/services/{service}/manage",
|
||||
json=body
|
||||
)
|
||||
|
||||
if response.status_code == 404:
|
||||
return f"Service '{service}' not found. Use list_services() to see available services."
|
||||
elif response.status_code == 400:
|
||||
error_detail = response.json().get('detail', 'Bad request')
|
||||
return f"Invalid request: {error_detail}"
|
||||
elif response.status_code == 501:
|
||||
return f"Action '{action}' is not yet implemented on the server."
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
message = result.get('message', f"Action '{action}' completed for service '{service}'.")
|
||||
return f"✓ {message}"
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_msg = f"Failed to {action} service '{service}': HTTP {e.response.status_code}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Error managing service '{service}': {str(e)}"
|
||||
logger.error(error_msg, exc_info=True)
|
||||
return error_msg
|
||||
|
||||
|
||||
@register_tool
|
||||
async def service_status(
|
||||
service: str
|
||||
) -> str:
|
||||
"""
|
||||
Get detailed service status and health.
|
||||
|
||||
Provides comprehensive information about a service including:
|
||||
- Overall status and replica health
|
||||
- Individual container statuses
|
||||
- Resource usage summary
|
||||
- Recent events (if available)
|
||||
|
||||
Args:
|
||||
service: Service/stack name (e.g., "portainer", "nginx")
|
||||
|
||||
Returns:
|
||||
Detailed service status report
|
||||
|
||||
Examples:
|
||||
service_status("web")
|
||||
service_status("portainer")
|
||||
service_status("core-ai")
|
||||
"""
|
||||
logger.info(f"Getting status for service '{service}'")
|
||||
|
||||
try:
|
||||
# Call core-api infrastructure endpoint
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
response = await client.get(
|
||||
f"{settings.core_api_base_url}/infrastructure/services/{service}/status"
|
||||
)
|
||||
|
||||
if response.status_code == 404:
|
||||
return f"Service '{service}' not found. Use list_services() to see available services."
|
||||
|
||||
response.raise_for_status()
|
||||
status = response.json()
|
||||
|
||||
# Build formatted output
|
||||
lines = [f"Service: {status.get('name', service)}"]
|
||||
lines.append(f"Status: {status.get('status', 'unknown')}")
|
||||
|
||||
# Stack info
|
||||
stack_id = status.get('stack_id')
|
||||
if stack_id:
|
||||
lines.append(f"Stack ID: {stack_id}")
|
||||
|
||||
# Replica status
|
||||
replica_status = status.get('replica_status')
|
||||
if replica_status:
|
||||
lines.append(f"Replicas: {replica_status}")
|
||||
|
||||
# Containers
|
||||
containers = status.get('containers', [])
|
||||
if containers:
|
||||
lines.append(f"\nContainers ({len(containers)}):")
|
||||
for container in containers:
|
||||
name = container.get('name', 'unknown')
|
||||
state = container.get('status', 'unknown')
|
||||
health = container.get('health', 'N/A')
|
||||
uptime = container.get('uptime', 'N/A')
|
||||
|
||||
lines.append(f" • {name}")
|
||||
lines.append(f" Status: {state}")
|
||||
if health != 'N/A':
|
||||
lines.append(f" Health: {health}")
|
||||
if uptime != 'N/A':
|
||||
lines.append(f" Uptime: {uptime}")
|
||||
|
||||
# Resources
|
||||
resources = status.get('resources', {})
|
||||
if resources:
|
||||
lines.append(f"\nResource Usage:")
|
||||
memory = resources.get('memory_total')
|
||||
if memory:
|
||||
lines.append(f" Memory: {memory}")
|
||||
cpu = resources.get('cpu_usage')
|
||||
if cpu:
|
||||
lines.append(f" CPU: {cpu}")
|
||||
|
||||
# Recent events
|
||||
events = status.get('recent_events', [])
|
||||
if events:
|
||||
lines.append(f"\nRecent Events ({len(events)}):")
|
||||
for event in events[:5]: # Show max 5 events
|
||||
time = event.get('time', 'unknown')
|
||||
action = event.get('action', 'unknown')
|
||||
target = event.get('container', 'unknown')
|
||||
lines.append(f" • {time} - {action} ({target})")
|
||||
if len(events) > 5:
|
||||
lines.append(f" ... and {len(events) - 5} more events")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_msg = f"Failed to get status for service '{service}': HTTP {e.response.status_code}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Error getting service status '{service}': {str(e)}"
|
||||
logger.error(error_msg, exc_info=True)
|
||||
return error_msg
|
||||
@@ -1,267 +0,0 @@
|
||||
"""
|
||||
Local utility tools for the AI agent.
|
||||
|
||||
These tools run locally in core-ai and don't require REST calls.
|
||||
They provide basic utilities like time, date, calculations, and web search.
|
||||
"""
|
||||
import logging
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Optional
|
||||
import pytz
|
||||
import httpx
|
||||
from src.tools.registry import register_tool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_tool
|
||||
async def get_current_time(timezone: str = "UTC") -> str:
|
||||
"""
|
||||
Get the current time in a specific timezone.
|
||||
|
||||
Args:
|
||||
timezone: Timezone name (e.g., "UTC", "Europe/Amsterdam", "America/New_York", "Asia/Tokyo")
|
||||
Use IANA timezone database names. Defaults to "UTC".
|
||||
|
||||
Returns:
|
||||
Current time as formatted string with timezone information
|
||||
|
||||
Examples:
|
||||
- get_current_time("Europe/Amsterdam") -> "2025-11-30 15:30:45 CET"
|
||||
- get_current_time("America/New_York") -> "2025-11-30 09:30:45 EST"
|
||||
- get_current_time() -> "2025-11-30 14:30:45 UTC"
|
||||
"""
|
||||
logger.info(f"Getting current time in timezone: {timezone}")
|
||||
|
||||
try:
|
||||
# Get timezone object
|
||||
tz = pytz.timezone(timezone)
|
||||
|
||||
# Get current time in that timezone
|
||||
now = datetime.now(tz)
|
||||
|
||||
# Format: "2025-11-30 15:30:45 CET"
|
||||
formatted_time = now.strftime("%Y-%m-%d %H:%M:%S %Z")
|
||||
|
||||
logger.info(f"Current time in {timezone}: {formatted_time}")
|
||||
return formatted_time
|
||||
|
||||
except pytz.exceptions.UnknownTimeZoneError:
|
||||
error_msg = f"Error: Unknown timezone '{timezone}'. Use IANA timezone names like 'Europe/Amsterdam', 'America/New_York', 'Asia/Tokyo', etc."
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Error getting time: {str(e)}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
|
||||
|
||||
@register_tool
|
||||
async def get_current_date() -> str:
|
||||
"""
|
||||
Get the current date.
|
||||
|
||||
Returns:
|
||||
Current date in YYYY-MM-DD format
|
||||
"""
|
||||
logger.info("Getting current date")
|
||||
return datetime.utcnow().date().isoformat()
|
||||
|
||||
|
||||
@register_tool
|
||||
async def calculate_date_difference(date1: str, date2: str) -> str:
|
||||
"""
|
||||
Calculate the difference between two dates.
|
||||
|
||||
Args:
|
||||
date1: First date in YYYY-MM-DD format
|
||||
date2: Second date in YYYY-MM-DD format
|
||||
|
||||
Returns:
|
||||
Human-readable description of the difference
|
||||
"""
|
||||
logger.info(f"Calculating difference between {date1} and {date2}")
|
||||
|
||||
try:
|
||||
d1 = datetime.fromisoformat(date1)
|
||||
d2 = datetime.fromisoformat(date2)
|
||||
|
||||
diff = abs((d2 - d1).days)
|
||||
|
||||
if diff == 0:
|
||||
return "The dates are the same day"
|
||||
elif diff == 1:
|
||||
return "1 day apart"
|
||||
else:
|
||||
return f"{diff} days apart"
|
||||
|
||||
except ValueError as e:
|
||||
logger.error(f"Invalid date format: {e}")
|
||||
return f"Error: Invalid date format. Please use YYYY-MM-DD format."
|
||||
|
||||
|
||||
@register_tool
|
||||
async def add_days_to_date(date: str, days: int) -> str:
|
||||
"""
|
||||
Add or subtract days from a date.
|
||||
|
||||
Args:
|
||||
date: Starting date in YYYY-MM-DD format
|
||||
days: Number of days to add (negative to subtract)
|
||||
|
||||
Returns:
|
||||
Resulting date in YYYY-MM-DD format
|
||||
"""
|
||||
logger.info(f"Adding {days} days to {date}")
|
||||
|
||||
try:
|
||||
d = datetime.fromisoformat(date)
|
||||
result = d + timedelta(days=days)
|
||||
return result.date().isoformat()
|
||||
except ValueError as e:
|
||||
logger.error(f"Invalid date format: {e}")
|
||||
return f"Error: Invalid date format. Please use YYYY-MM-DD format."
|
||||
|
||||
|
||||
@register_tool
|
||||
async def calculate(expression: str) -> str:
|
||||
"""
|
||||
Perform basic mathematical calculations.
|
||||
|
||||
Supports: +, -, *, /, //, %, ** (power), parentheses
|
||||
|
||||
Args:
|
||||
expression: Mathematical expression to evaluate (e.g., "2 + 2", "10 * (5 + 3)")
|
||||
|
||||
Returns:
|
||||
Result of the calculation as a string
|
||||
"""
|
||||
logger.info(f"Calculating: {expression}")
|
||||
|
||||
try:
|
||||
# Security: Only allow safe mathematical operations
|
||||
# Using eval() with restricted namespace
|
||||
allowed_names = {
|
||||
"abs": abs,
|
||||
"round": round,
|
||||
"min": min,
|
||||
"max": max,
|
||||
"sum": sum,
|
||||
}
|
||||
|
||||
# Remove any potentially dangerous characters
|
||||
dangerous_chars = ["_", "import", "exec", "eval", "open", "file", "__"]
|
||||
for char in dangerous_chars:
|
||||
if char in expression:
|
||||
return f"Error: Invalid expression - contains forbidden pattern '{char}'"
|
||||
|
||||
# Evaluate the expression
|
||||
result = eval(expression, {"__builtins__": {}}, allowed_names)
|
||||
|
||||
logger.info(f"Calculation result: {result}")
|
||||
return str(result)
|
||||
|
||||
except SyntaxError:
|
||||
return "Error: Invalid mathematical expression syntax"
|
||||
except ZeroDivisionError:
|
||||
return "Error: Division by zero"
|
||||
except Exception as e:
|
||||
logger.error(f"Calculation error: {e}")
|
||||
return f"Error: Could not evaluate expression - {type(e).__name__}"
|
||||
|
||||
|
||||
@register_tool
|
||||
async def web_search(query: str, category: str = "general", max_results: int = 10) -> str:
|
||||
"""
|
||||
Search the web using SearXNG metasearch engine.
|
||||
|
||||
Aggregates results from multiple search engines (Google, Bing, DuckDuckGo, etc.)
|
||||
while maintaining privacy - no tracking or data collection.
|
||||
|
||||
Args:
|
||||
query: Search query string (e.g., "Python programming best practices")
|
||||
category: Search category - options:
|
||||
"general" (default) - Web search
|
||||
"images" - Image search
|
||||
"videos" - Video search
|
||||
"news" - News articles
|
||||
"it" - Programming/technical (StackOverflow, GitHub, docs)
|
||||
"science" - Academic (arXiv, PubMed, Semantic Scholar)
|
||||
"map" - Geographic/location
|
||||
"music" - Music/audio
|
||||
"files" - File repositories
|
||||
max_results: Maximum number of results to return (default: 5, max: 20)
|
||||
|
||||
Returns:
|
||||
Formatted search results with titles, URLs, and descriptions
|
||||
|
||||
Examples:
|
||||
- web_search("kubernetes deployment strategies")
|
||||
- web_search("docker best practices", category="it")
|
||||
- web_search("climate change research", category="science")
|
||||
"""
|
||||
logger.info(f"Web search: query='{query}', category='{category}', max_results={max_results}")
|
||||
|
||||
# Emit tool call event for status tracking
|
||||
try:
|
||||
from src.agents.tool_events import get_tool_emitter
|
||||
emitter = get_tool_emitter()
|
||||
emitter.emit("web_search", {"query": query, "category": category})
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to emit tool event: {e}")
|
||||
|
||||
try:
|
||||
# Limit max_results to prevent overwhelming responses
|
||||
max_results = min(max_results, 20)
|
||||
|
||||
# Call SearXNG JSON API
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
response = await client.get(
|
||||
"http://searxng:8080/search",
|
||||
params={
|
||||
"q": query,
|
||||
"format": "json",
|
||||
"categories": category
|
||||
}
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
results = data.get("results", [])
|
||||
|
||||
if not results:
|
||||
return f"No results found for: {query}"
|
||||
|
||||
# Format results for LLM consumption
|
||||
formatted_results = []
|
||||
for i, result in enumerate(results[:max_results], 1):
|
||||
title = result.get("title", "No title")
|
||||
url = result.get("url", "")
|
||||
content = result.get("content", "No description available")
|
||||
engine = result.get("engine", "unknown")
|
||||
|
||||
formatted_results.append(
|
||||
f"{i}. **{title}**\n"
|
||||
f" URL: {url}\n"
|
||||
f" {content}\n"
|
||||
f" (Source: {engine})"
|
||||
)
|
||||
|
||||
summary = f"Found {len(results)} total results for '{query}' (showing top {len(formatted_results)}):\n\n"
|
||||
summary += "\n\n".join(formatted_results)
|
||||
|
||||
logger.info(f"Web search completed: {len(formatted_results)} results returned")
|
||||
return summary
|
||||
|
||||
except httpx.TimeoutException:
|
||||
error_msg = "Web search timed out. The search engine may be slow or unavailable."
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_msg = f"Web search failed with HTTP {e.response.status_code}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Web search error: {str(e)}"
|
||||
logger.error(error_msg)
|
||||
return error_msg
|
||||
@@ -1,307 +0,0 @@
|
||||
"""
|
||||
OpenAPI Tool Discovery
|
||||
|
||||
Dynamically discovers and creates tools from core-api's OpenAPI specification.
|
||||
This allows core-ai to automatically use infrastructure management endpoints
|
||||
without manual tool definition.
|
||||
|
||||
Architecture:
|
||||
- Core tools (local.py): Essential tools always available (web_search, calculate, etc.)
|
||||
- OpenAPI tools (this module): Infrastructure/automation endpoints from core-api
|
||||
"""
|
||||
import httpx
|
||||
import logging
|
||||
from typing import Dict, List, Any, Optional, Callable
|
||||
from functools import lru_cache
|
||||
import asyncio
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OpenAPIToolDiscovery:
|
||||
"""
|
||||
Discovers and creates executable tools from OpenAPI specifications.
|
||||
"""
|
||||
|
||||
def __init__(self, openapi_url: str = "http://core-api:8083/openapi.json"):
|
||||
self.openapi_url = openapi_url
|
||||
self.spec = None
|
||||
self.tools = {}
|
||||
|
||||
async def fetch_spec(self) -> Dict[str, Any]:
|
||||
"""Fetch OpenAPI specification from core-api"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
response = await client.get(self.openapi_url)
|
||||
response.raise_for_status()
|
||||
self.spec = response.json()
|
||||
logger.info(f"Fetched OpenAPI spec: {self.spec['info']['title']} "
|
||||
f"with {len(self.spec.get('paths', {}))} endpoints")
|
||||
return self.spec
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to fetch OpenAPI spec from {self.openapi_url}: {e}")
|
||||
return {}
|
||||
|
||||
def _extract_endpoint_info(self, path: str, method: str, operation: Dict) -> Dict[str, Any]:
|
||||
"""
|
||||
Extract relevant information from an OpenAPI operation.
|
||||
|
||||
Returns:
|
||||
{
|
||||
"name": "get_containers_list",
|
||||
"description": "List all Docker containers",
|
||||
"path": "/infrastructure/containers",
|
||||
"method": "GET",
|
||||
"parameters": [...],
|
||||
"summary": "..."
|
||||
}
|
||||
"""
|
||||
# Generate tool name from operationId or path
|
||||
operation_id = operation.get("operationId")
|
||||
if operation_id:
|
||||
# Convert operationId to snake_case
|
||||
tool_name = operation_id.replace("-", "_").replace(" ", "_").lower()
|
||||
else:
|
||||
# Generate from path and method
|
||||
path_parts = path.strip("/").replace("/", "_").replace("{", "").replace("}", "")
|
||||
tool_name = f"{method.lower()}_{path_parts}"
|
||||
|
||||
# Get description
|
||||
description = operation.get("summary") or operation.get("description") or f"{method} {path}"
|
||||
|
||||
# Extract parameters
|
||||
parameters = operation.get("parameters", [])
|
||||
request_body = operation.get("requestBody")
|
||||
|
||||
return {
|
||||
"name": tool_name,
|
||||
"description": description,
|
||||
"path": path,
|
||||
"method": method.upper(),
|
||||
"parameters": parameters,
|
||||
"request_body": request_body,
|
||||
"summary": operation.get("summary", ""),
|
||||
"tags": operation.get("tags", [])
|
||||
}
|
||||
|
||||
def _create_tool_function(self, endpoint_info: Dict[str, Any]) -> Callable:
|
||||
"""
|
||||
Create an executable async function for an API endpoint.
|
||||
|
||||
The function will make HTTP requests to core-api when called.
|
||||
"""
|
||||
path = endpoint_info["path"]
|
||||
method = endpoint_info["method"]
|
||||
description = endpoint_info["description"]
|
||||
|
||||
async def tool_function(**kwargs) -> str:
|
||||
"""
|
||||
Dynamically generated function that calls core-api endpoint.
|
||||
"""
|
||||
try:
|
||||
url = f"http://core-api:8083{path}"
|
||||
|
||||
# Replace path parameters
|
||||
for key, value in kwargs.items():
|
||||
url = url.replace(f"{{{key}}}", str(value))
|
||||
|
||||
# Build request
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
if method == "GET":
|
||||
response = await client.get(url, params=kwargs)
|
||||
elif method == "POST":
|
||||
response = await client.post(url, json=kwargs)
|
||||
elif method == "PUT":
|
||||
response = await client.put(url, json=kwargs)
|
||||
elif method == "DELETE":
|
||||
response = await client.delete(url, params=kwargs)
|
||||
else:
|
||||
return f"Unsupported HTTP method: {method}"
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
# Return JSON if possible, otherwise text
|
||||
try:
|
||||
result = response.json()
|
||||
# Format nicely for LLM
|
||||
if isinstance(result, list):
|
||||
return f"Found {len(result)} items:\n" + "\n".join(
|
||||
[f"- {item}" for item in result[:10]] # Limit to 10 items
|
||||
)
|
||||
elif isinstance(result, dict):
|
||||
return str(result)
|
||||
else:
|
||||
return str(result)
|
||||
except:
|
||||
return response.text
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
return f"HTTP Error {e.response.status_code}: {e.response.text}"
|
||||
except Exception as e:
|
||||
return f"Error calling {method} {path}: {str(e)}"
|
||||
|
||||
# Set function metadata
|
||||
tool_function.__name__ = endpoint_info["name"]
|
||||
tool_function.__doc__ = f"{description}\n\nEndpoint: {method} {path}"
|
||||
|
||||
return tool_function
|
||||
|
||||
async def discover_tools(
|
||||
self,
|
||||
include_tags: Optional[List[str]] = None,
|
||||
exclude_tags: Optional[List[str]] = None,
|
||||
method_filter: Optional[List[str]] = None
|
||||
) -> Dict[str, Callable]:
|
||||
"""
|
||||
Discover and create tools from OpenAPI spec.
|
||||
|
||||
Args:
|
||||
include_tags: Only include endpoints with these tags
|
||||
exclude_tags: Exclude endpoints with these tags
|
||||
method_filter: Only include these HTTP methods (e.g., ["GET", "POST"])
|
||||
|
||||
Returns:
|
||||
Dictionary of tool_name -> async function
|
||||
"""
|
||||
if not self.spec:
|
||||
await self.fetch_spec()
|
||||
|
||||
if not self.spec or "paths" not in self.spec:
|
||||
logger.warning("No OpenAPI spec available")
|
||||
return {}
|
||||
|
||||
discovered_tools = {}
|
||||
|
||||
for path, path_item in self.spec["paths"].items():
|
||||
for method in ["get", "post", "put", "delete", "patch"]:
|
||||
if method not in path_item:
|
||||
continue
|
||||
|
||||
operation = path_item[method]
|
||||
|
||||
# Apply filters
|
||||
if method_filter and method.upper() not in method_filter:
|
||||
continue
|
||||
|
||||
tags = operation.get("tags", [])
|
||||
if include_tags and not any(tag in include_tags for tag in tags):
|
||||
continue
|
||||
if exclude_tags and any(tag in exclude_tags for tag in tags):
|
||||
continue
|
||||
|
||||
# Extract endpoint info
|
||||
endpoint_info = self._extract_endpoint_info(path, method, operation)
|
||||
|
||||
# Create executable function
|
||||
tool_func = self._create_tool_function(endpoint_info)
|
||||
|
||||
discovered_tools[endpoint_info["name"]] = tool_func
|
||||
|
||||
logger.debug(f"Discovered tool: {endpoint_info['name']} ({method.upper()} {path})")
|
||||
|
||||
logger.info(f"Discovered {len(discovered_tools)} tools from OpenAPI spec")
|
||||
return discovered_tools
|
||||
|
||||
def get_tool_descriptions(self) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Get human-readable descriptions of all discovered tools.
|
||||
|
||||
Useful for logging/debugging.
|
||||
"""
|
||||
descriptions = []
|
||||
for name, func in self.tools.items():
|
||||
descriptions.append({
|
||||
"name": name,
|
||||
"description": func.__doc__ or "No description"
|
||||
})
|
||||
return descriptions
|
||||
|
||||
|
||||
# Global instances for multiple OpenAPI sources
|
||||
_discovery_instances: Dict[str, OpenAPIToolDiscovery] = {}
|
||||
|
||||
|
||||
async def get_openapi_tools(
|
||||
endpoints: Optional[List[str]] = None,
|
||||
include_tags: Optional[List[str]] = None,
|
||||
exclude_tags: Optional[List[str]] = None,
|
||||
refresh: bool = False
|
||||
) -> Dict[str, Callable]:
|
||||
"""
|
||||
Get dynamically discovered tools from one or more OpenAPI specifications.
|
||||
|
||||
Args:
|
||||
endpoints: List of OpenAPI spec URLs. If None, uses default (core-api)
|
||||
Example: ["http://core-api:8083/openapi.json", "http://automation:8080/openapi.json"]
|
||||
include_tags: Only include endpoints with these tags (e.g., ["infrastructure", "automation"])
|
||||
exclude_tags: Exclude endpoints with these tags (e.g., ["internal", "admin"])
|
||||
refresh: Force re-fetch of OpenAPI specs
|
||||
|
||||
Returns:
|
||||
Dictionary of tool_name -> async function (combined from all sources)
|
||||
"""
|
||||
global _discovery_instances
|
||||
|
||||
# Default to core-api if no endpoints specified
|
||||
if endpoints is None:
|
||||
endpoints = ["http://core-api:8083/openapi.json"]
|
||||
|
||||
all_tools = {}
|
||||
|
||||
for endpoint_url in endpoints:
|
||||
# Get or create discovery instance for this endpoint
|
||||
if endpoint_url not in _discovery_instances or refresh:
|
||||
_discovery_instances[endpoint_url] = OpenAPIToolDiscovery(openapi_url=endpoint_url)
|
||||
|
||||
instance = _discovery_instances[endpoint_url]
|
||||
|
||||
# Discover tools from this endpoint
|
||||
try:
|
||||
tools = await instance.discover_tools(
|
||||
include_tags=include_tags,
|
||||
exclude_tags=exclude_tags
|
||||
)
|
||||
|
||||
# Add source prefix to avoid name conflicts between APIs
|
||||
# Extract service name from URL (e.g., "core-api" from "http://core-api:8083/...")
|
||||
service_name = endpoint_url.split("//")[1].split(":")[0].split(".")[0]
|
||||
|
||||
for tool_name, tool_func in tools.items():
|
||||
# Prefix tool name with service (e.g., "core_api__list_containers")
|
||||
prefixed_name = f"{service_name}__{tool_name}"
|
||||
all_tools[prefixed_name] = tool_func
|
||||
|
||||
instance.tools = tools
|
||||
logger.info(f"Loaded {len(tools)} tools from {service_name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to discover tools from {endpoint_url}: {e}")
|
||||
continue
|
||||
|
||||
logger.info(f"Total OpenAPI tools discovered: {len(all_tools)} from {len(endpoints)} source(s)")
|
||||
return all_tools
|
||||
|
||||
|
||||
async def get_openapi_tool_descriptions(endpoints: Optional[List[str]] = None) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Get descriptions of all discovered OpenAPI tools.
|
||||
|
||||
Args:
|
||||
endpoints: List of OpenAPI spec URLs (same as get_openapi_tools)
|
||||
|
||||
Returns:
|
||||
List of tool descriptions
|
||||
"""
|
||||
global _discovery_instances
|
||||
|
||||
if not _discovery_instances:
|
||||
await get_openapi_tools(endpoints=endpoints)
|
||||
|
||||
all_descriptions = []
|
||||
for endpoint_url, instance in _discovery_instances.items():
|
||||
service_name = endpoint_url.split("//")[1].split(":")[0].split(".")[0]
|
||||
for desc in instance.get_tool_descriptions():
|
||||
desc["source"] = service_name
|
||||
all_descriptions.append(desc)
|
||||
|
||||
return all_descriptions
|
||||
@@ -1,403 +0,0 @@
|
||||
"""
|
||||
Tool Registry - Manages tool registration and discovery for AI agents.
|
||||
|
||||
This module provides a central registry for tools.
|
||||
Tools can be registered, discovered, and provided to AI agents.
|
||||
"""
|
||||
import logging
|
||||
import functools
|
||||
import inspect
|
||||
from typing import List, Dict, Any, Callable
|
||||
import httpx
|
||||
|
||||
from src.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Initialize settings once
|
||||
settings = get_settings()
|
||||
CORE_API_BASE_URL = settings.core_api_base_url
|
||||
|
||||
# HTTP client for REST calls to core-api
|
||||
http_client = httpx.AsyncClient()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Legacy ADK Integration (deprecated - kept for backwards compatibility)
|
||||
# ============================================================================
|
||||
|
||||
try:
|
||||
from google.adk.tools import FunctionTool
|
||||
ADK_AVAILABLE = True
|
||||
except ImportError:
|
||||
ADK_AVAILABLE = False
|
||||
FunctionTool = None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tool Registry
|
||||
# ============================================================================
|
||||
|
||||
# Global registry of tools
|
||||
_TOOL_REGISTRY: Dict[str, Callable] = {}
|
||||
|
||||
|
||||
def log_tool_call(func):
|
||||
"""Decorator to log tool calls and track metrics"""
|
||||
@functools.wraps(func)
|
||||
async def wrapper(*args, **kwargs):
|
||||
import time
|
||||
from src.metrics import get_metrics_collector
|
||||
|
||||
metrics = get_metrics_collector()
|
||||
tool_name = func.__name__
|
||||
start_time = time.time()
|
||||
success = False
|
||||
error_msg = None
|
||||
|
||||
params_str = ", ".join(
|
||||
[f"{arg}" for arg in args] +
|
||||
[f"{k}={repr(v)}" for k, v in kwargs.items()]
|
||||
)
|
||||
logger.info(f"🔧 TOOL CALL: {tool_name}({params_str})")
|
||||
|
||||
try:
|
||||
# Filter kwargs to only include valid parameters
|
||||
sig = inspect.signature(func)
|
||||
valid_kwargs = {
|
||||
key: value for key, value in kwargs.items()
|
||||
if key in sig.parameters
|
||||
}
|
||||
result = await func(*args, **valid_kwargs)
|
||||
result_preview = str(result)[:200] if result else "None"
|
||||
logger.info(f"✅ TOOL RESULT: {tool_name} → {result_preview}...")
|
||||
success = True
|
||||
return result
|
||||
except Exception as e:
|
||||
error_msg = f"{type(e).__name__}: {str(e)}"
|
||||
logger.error(
|
||||
f"❌ TOOL ERROR: {tool_name} failed with {type(e).__name__}: {e}",
|
||||
exc_info=True
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
metrics.record_tool_execution(
|
||||
tool_name=tool_name,
|
||||
duration_ms=duration_ms,
|
||||
success=success,
|
||||
error=error_msg
|
||||
)
|
||||
return wrapper
|
||||
|
||||
|
||||
def register_tool(func: Callable) -> Callable:
|
||||
"""
|
||||
Register a tool function for use with AI agents.
|
||||
|
||||
Usage:
|
||||
@register_tool
|
||||
async def my_tool(param: str) -> str:
|
||||
'''Tool description'''
|
||||
return "result"
|
||||
|
||||
Args:
|
||||
func: Async function to register as a tool
|
||||
|
||||
Returns:
|
||||
The decorated function
|
||||
"""
|
||||
_TOOL_REGISTRY[func.__name__] = func
|
||||
logger.info(f"📝 Registered tool: {func.__name__}")
|
||||
return log_tool_call(func)
|
||||
|
||||
|
||||
def get_all_tools(include_openapi: bool = False) -> Dict[str, Callable]:
|
||||
"""
|
||||
Get all registered tools.
|
||||
|
||||
Args:
|
||||
include_openapi: If True, also include dynamically discovered OpenAPI tools
|
||||
|
||||
Returns:
|
||||
Dictionary mapping tool names to functions
|
||||
"""
|
||||
tools = _TOOL_REGISTRY.copy()
|
||||
|
||||
# Add OpenAPI tools if requested
|
||||
if include_openapi:
|
||||
try:
|
||||
import asyncio
|
||||
from src.tools.openapi_discovery import get_openapi_tools
|
||||
|
||||
# Get or create event loop
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
# Fetch OpenAPI tools
|
||||
openapi_tools = loop.run_until_complete(get_openapi_tools())
|
||||
tools.update(openapi_tools)
|
||||
logger.info(f"Added {len(openapi_tools)} OpenAPI tools to registry")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load OpenAPI tools: {e}")
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
def get_agent_tools() -> List:
|
||||
"""
|
||||
DEPRECATED: Get all tools as ADK FunctionTool objects.
|
||||
|
||||
This function is kept for backwards compatibility but is no longer used.
|
||||
Use get_all_tools() instead for PydanticAI agents.
|
||||
|
||||
Returns:
|
||||
List of FunctionTool objects for legacy ADK agent
|
||||
"""
|
||||
if not ADK_AVAILABLE:
|
||||
logger.warning("ADK not available - returning empty tool list")
|
||||
return []
|
||||
|
||||
tools = []
|
||||
for name, func in _TOOL_REGISTRY.items():
|
||||
try:
|
||||
# Create ADK FunctionTool from the registered function
|
||||
tool = FunctionTool(func)
|
||||
tools.append(tool)
|
||||
logger.info(f"✓ Created ADK tool: {name}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create ADK tool for {name}: {e}")
|
||||
|
||||
logger.info(f"📦 Providing {len(tools)} tools to ADK agent")
|
||||
return tools
|
||||
|
||||
|
||||
def clear_registry():
|
||||
"""Clear all registered tools (useful for testing)"""
|
||||
_TOOL_REGISTRY.clear()
|
||||
logger.info("🗑️ Tool registry cleared")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Swagger/OpenAPI Dynamic Tool Discovery
|
||||
# ============================================================================
|
||||
|
||||
async def fetch_openapi_spec(base_url: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Fetch the OpenAPI/Swagger specification from core-api.
|
||||
|
||||
Args:
|
||||
base_url: Base URL of the API (e.g., http://core-api:8000)
|
||||
|
||||
Returns:
|
||||
OpenAPI spec as dictionary
|
||||
|
||||
Raises:
|
||||
Exception: If fetching fails
|
||||
"""
|
||||
try:
|
||||
# Try common OpenAPI spec endpoints
|
||||
endpoints = [
|
||||
f"{base_url}/openapi.json",
|
||||
f"{base_url}/api/openapi.json",
|
||||
f"{base_url}/docs/openapi.json",
|
||||
f"{base_url}/swagger.json",
|
||||
]
|
||||
|
||||
for endpoint in endpoints:
|
||||
try:
|
||||
logger.info(f"Attempting to fetch OpenAPI spec from: {endpoint}")
|
||||
response = await http_client.get(endpoint, timeout=5.0)
|
||||
if response.status_code == 200:
|
||||
spec = response.json()
|
||||
logger.info(f"✓ Successfully fetched OpenAPI spec from {endpoint}")
|
||||
return spec
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to fetch from {endpoint}: {e}")
|
||||
continue
|
||||
|
||||
raise Exception(f"Could not fetch OpenAPI spec from any endpoint at {base_url}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to fetch OpenAPI spec: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def create_rest_tool(
|
||||
operation_id: str,
|
||||
path: str,
|
||||
method: str,
|
||||
description: str,
|
||||
parameters: List[Dict[str, Any]],
|
||||
base_url: str
|
||||
) -> Callable:
|
||||
"""
|
||||
Create a dynamic REST tool function from OpenAPI operation.
|
||||
|
||||
Args:
|
||||
operation_id: Unique identifier for the operation
|
||||
path: API path (e.g., /api/v1/containers)
|
||||
method: HTTP method (GET, POST, etc.)
|
||||
description: Tool description from OpenAPI
|
||||
parameters: List of parameter specifications
|
||||
base_url: Base URL for API calls
|
||||
|
||||
Returns:
|
||||
Async function that calls the REST endpoint
|
||||
"""
|
||||
# Create parameter list for function signature
|
||||
param_names = [p["name"] for p in parameters]
|
||||
|
||||
async def rest_tool(**kwargs):
|
||||
"""
|
||||
Dynamically created REST tool.
|
||||
"""
|
||||
# Build request
|
||||
url = f"{base_url}{path}"
|
||||
|
||||
# Substitute path parameters
|
||||
for param in parameters:
|
||||
if param.get("in") == "path":
|
||||
param_name = param["name"]
|
||||
if param_name in kwargs:
|
||||
url = url.replace(f"{{{param_name}}}", str(kwargs[param_name]))
|
||||
|
||||
# Build query parameters
|
||||
query_params = {}
|
||||
for param in parameters:
|
||||
if param.get("in") == "query":
|
||||
param_name = param["name"]
|
||||
if param_name in kwargs:
|
||||
query_params[param_name] = kwargs[param_name]
|
||||
|
||||
# Build request body
|
||||
body = None
|
||||
for param in parameters:
|
||||
if param.get("in") == "body":
|
||||
param_name = param["name"]
|
||||
if param_name in kwargs:
|
||||
body = kwargs[param_name]
|
||||
|
||||
logger.info(f"REST Tool: {method} {url}")
|
||||
|
||||
try:
|
||||
# Make the REST call
|
||||
if method.upper() == "GET":
|
||||
response = await http_client.get(url, params=query_params)
|
||||
elif method.upper() == "POST":
|
||||
response = await http_client.post(url, json=body, params=query_params)
|
||||
elif method.upper() == "PUT":
|
||||
response = await http_client.put(url, json=body, params=query_params)
|
||||
elif method.upper() == "DELETE":
|
||||
response = await http_client.delete(url, params=query_params)
|
||||
else:
|
||||
return f"Error: Unsupported HTTP method {method}"
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
# Return response
|
||||
try:
|
||||
return response.json()
|
||||
except Exception:
|
||||
return response.text
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.error(f"REST tool HTTP error: {e}")
|
||||
return f"Error: HTTP {e.response.status_code} - {e.response.text}"
|
||||
except Exception as e:
|
||||
logger.error(f"REST tool error: {e}")
|
||||
return f"Error: {type(e).__name__} - {str(e)}"
|
||||
|
||||
# Set function metadata for ADK
|
||||
rest_tool.__name__ = operation_id
|
||||
rest_tool.__doc__ = description
|
||||
|
||||
# Add annotations for ADK type checking
|
||||
annotations = {}
|
||||
for param in parameters:
|
||||
param_name = param["name"]
|
||||
param_type = param.get("schema", {}).get("type", "string")
|
||||
|
||||
# Map OpenAPI types to Python types
|
||||
type_mapping = {
|
||||
"string": str,
|
||||
"integer": int,
|
||||
"number": float,
|
||||
"boolean": bool,
|
||||
"array": list,
|
||||
"object": dict,
|
||||
}
|
||||
annotations[param_name] = type_mapping.get(param_type, str)
|
||||
|
||||
annotations["return"] = str
|
||||
rest_tool.__annotations__ = annotations
|
||||
|
||||
return rest_tool
|
||||
|
||||
|
||||
async def discover_and_register_tools(base_url: str = None) -> int:
|
||||
"""
|
||||
Discover tools from core-api's OpenAPI spec and register them.
|
||||
|
||||
Args:
|
||||
base_url: Base URL of core-api (default: from settings)
|
||||
|
||||
Returns:
|
||||
Number of tools registered
|
||||
|
||||
Raises:
|
||||
Exception: If discovery fails
|
||||
"""
|
||||
if base_url is None:
|
||||
base_url = CORE_API_BASE_URL
|
||||
|
||||
logger.info(f"🔍 Discovering tools from {base_url}")
|
||||
|
||||
try:
|
||||
# Fetch OpenAPI spec
|
||||
spec = await fetch_openapi_spec(base_url)
|
||||
|
||||
paths = spec.get("paths", {})
|
||||
tools_registered = 0
|
||||
|
||||
# Iterate through all paths and operations
|
||||
for path, path_item in paths.items():
|
||||
for method, operation in path_item.items():
|
||||
if method.lower() not in ["get", "post", "put", "delete", "patch"]:
|
||||
continue
|
||||
|
||||
# Extract operation details
|
||||
operation_id = operation.get("operationId")
|
||||
if not operation_id:
|
||||
# Generate operation ID from path and method
|
||||
operation_id = f"{method}_{path.replace('/', '_').strip('_')}"
|
||||
|
||||
description = operation.get("summary", operation.get("description", f"{method.upper()} {path}"))
|
||||
|
||||
# Extract parameters
|
||||
parameters = operation.get("parameters", [])
|
||||
|
||||
# Create and register the tool
|
||||
tool_func = create_rest_tool(
|
||||
operation_id=operation_id,
|
||||
path=path,
|
||||
method=method,
|
||||
description=description,
|
||||
parameters=parameters,
|
||||
base_url=base_url
|
||||
)
|
||||
|
||||
# Register the tool
|
||||
_TOOL_REGISTRY[operation_id] = log_tool_call(tool_func)
|
||||
logger.info(f"📝 Registered REST tool: {operation_id} ({method.upper()} {path})")
|
||||
tools_registered += 1
|
||||
|
||||
logger.info(f"✓ Discovered and registered {tools_registered} tools from core-api")
|
||||
return tools_registered
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to discover tools: {e}", exc_info=True)
|
||||
raise
|
||||
@@ -1,102 +0,0 @@
|
||||
"""
|
||||
Utility functions for core-ai service.
|
||||
"""
|
||||
import re
|
||||
from typing import Optional
|
||||
|
||||
|
||||
# Default user for requests without user_id
|
||||
DEFAULT_USER_ID = "llmdefault_at_schweitz.net"
|
||||
|
||||
|
||||
def sanitize_email_to_user_id(email: Optional[str] = None) -> str:
|
||||
"""
|
||||
Convert email address to standardized user_id format.
|
||||
|
||||
Format: username_at_domain_com (lowercase, @ → _at_)
|
||||
|
||||
Examples:
|
||||
john@example.com → john_at_example_com
|
||||
Alice.Smith@Company.ORG → alice_smith_at_company_org
|
||||
None → llmdefault_at_schweitz.net (default)
|
||||
|
||||
Args:
|
||||
email: Email address to convert (None uses default user)
|
||||
|
||||
Returns:
|
||||
Sanitized user_id string safe for Qdrant collection names
|
||||
"""
|
||||
if not email:
|
||||
return DEFAULT_USER_ID
|
||||
|
||||
# Convert to lowercase
|
||||
email = email.lower().strip()
|
||||
|
||||
# Validate email format (basic check)
|
||||
if '@' not in email:
|
||||
# Invalid email, return default
|
||||
return DEFAULT_USER_ID
|
||||
|
||||
# Replace @ with _at_
|
||||
user_id = email.replace('@', '_at_')
|
||||
|
||||
# Replace any non-alphanumeric characters (except underscores) with underscores
|
||||
# This handles dots, hyphens, etc. in email addresses
|
||||
user_id = re.sub(r'[^a-z0-9_]', '_', user_id)
|
||||
|
||||
# Remove any duplicate underscores
|
||||
user_id = re.sub(r'_+', '_', user_id)
|
||||
|
||||
# Remove leading/trailing underscores
|
||||
user_id = user_id.strip('_')
|
||||
|
||||
return user_id
|
||||
|
||||
|
||||
def get_collection_name_for_user(user_id: str, prefix: str = "core_ai_user") -> str:
|
||||
"""
|
||||
Generate Qdrant collection name for a user.
|
||||
|
||||
Args:
|
||||
user_id: Sanitized user ID (from sanitize_email_to_user_id)
|
||||
prefix: Collection prefix (default: core_ai_user)
|
||||
|
||||
Returns:
|
||||
Full collection name: {prefix}_{user_id}
|
||||
|
||||
Examples:
|
||||
john_at_example_com → core_ai_user_john_at_example_com
|
||||
llmdefault_at_schweitz_net → core_ai_user_llmdefault_at_schweitz_net
|
||||
"""
|
||||
return f"{prefix}_{user_id}"
|
||||
|
||||
|
||||
def extract_user_id_from_request(data: dict) -> str:
|
||||
"""
|
||||
Extract and sanitize user_id from request data.
|
||||
|
||||
Priority:
|
||||
1. data.get("user_id") - if provided, sanitize it
|
||||
2. data.get("user_email") - convert to user_id format
|
||||
3. DEFAULT_USER_ID - fallback to default user
|
||||
|
||||
Args:
|
||||
data: Request JSON data
|
||||
|
||||
Returns:
|
||||
Sanitized user_id string
|
||||
"""
|
||||
# Check for explicit user_id
|
||||
if user_id := data.get("user_id"):
|
||||
# If it's already in our format, use it
|
||||
if "_at_" in user_id:
|
||||
return user_id
|
||||
# Otherwise treat it as an email
|
||||
return sanitize_email_to_user_id(user_id)
|
||||
|
||||
# Check for user_email
|
||||
if user_email := data.get("user_email"):
|
||||
return sanitize_email_to_user_id(user_email)
|
||||
|
||||
# Fallback to default
|
||||
return DEFAULT_USER_ID
|
||||
Reference in New Issue
Block a user