Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4efa717796 | ||
|
|
ac2ada89fe | ||
|
|
a53fd67f4f | ||
|
|
27375cd6d2 | ||
|
|
09e468e7f8 | ||
|
|
ebac19ba6e | ||
|
|
22d44b3071 | ||
|
|
27b46a9fe7 | ||
|
|
7ec6e03c65 | ||
|
|
f6f37b341b | ||
|
|
92c0d5d770 |
@@ -3,542 +3,48 @@
|
||||
This document contains instructions and documentation references for AI assistants working with this codebase.
|
||||
|
||||
> **📖 Important**: Before working on this project, read [PHILOSOPHY.md](PHILOSOPHY.md) to understand the system vision, architectural patterns, and design goals. All development should work towards realizing those patterns.
|
||||
# AGENTS.md
|
||||
|
||||
## Project Overview
|
||||
> **Start every session by reading this file.**
|
||||
> This file outlines the operational protocols, coding standards, and architectural decisions for this FastAPI project.
|
||||
|
||||
This project implements an OpenAI-compatible API with FastAPI, featuring a hybrid architecture that provides both the OpenAI Responses API and Chat Completions compatibility layer.
|
||||
## 1. Agent Operational Protocols
|
||||
|
||||
### Architecture Pattern
|
||||
### 🧠 Work Patterns (Plan-Act-Reflect)
|
||||
* **Plan:** Before writing code, briefly outline your plan. Identify which files you will touch and what the side effects might be.
|
||||
* **Act:** Execute the changes in small, atomic steps.
|
||||
* **Reflect:** After coding, verify your work. Did you break existing tests? Did you add new tests?
|
||||
|
||||
The **Orchestrator** infrastructure layer with hybrid API architecture:
|
||||
### 🛡️ Git Discipline
|
||||
* **NEVER commit to `main` or `master` directly.** Always create a feature branch: `feature/your-feature-name` or `fix/issue-description`.
|
||||
* **Commit Messages:** Use the [Conventional Commits](https://www.conventionalcommits.org/) format.
|
||||
* `feat: add user login endpoint`
|
||||
* `fix: resolve database connection timeout`
|
||||
* `refactor: split monolith dependency file`
|
||||
* **Atomic Commits:** Keep commits small. One logical change = one commit.
|
||||
|
||||
```
|
||||
Client (Open WebUI)
|
||||
↓
|
||||
Chat Completions (/v1/chat/completions) → Wrapper
|
||||
↓
|
||||
Responses API (/v1/responses) → Primary
|
||||
↓
|
||||
Agent Interface (lorem-tester, Tatlock)
|
||||
↓
|
||||
Mock Agents (lorem-tester) / Future: PydanticAI Agents (Tatlock, Steward, etc.)
|
||||
```
|
||||
### 📝 Changelog Maintenance
|
||||
* **Update `CHANGELOG.md`** with every user-facing change.
|
||||
* Format: `## [Unreleased] - YYYY-MM-DD` followed by `### Added`, `### Changed`, or `### Fixed`.
|
||||
|
||||
**Architectural Layers:**
|
||||
---
|
||||
|
||||
1. **The Orchestrator** (Current Implementation)
|
||||
- FastAPI application providing the infrastructure
|
||||
- HTTP/SSE endpoints, streaming coordination
|
||||
- Conversation history and context management
|
||||
- OpenAI-compatible API surface
|
||||
## 2. FastAPI Architecture & Best Practices
|
||||
*Reference: [FastAPI Best Practices](https://github.com/zhanymkanov/fastapi-best-practices)*
|
||||
|
||||
2. **Future: The Household** (Phases 1-4)
|
||||
- **Steward**: First-tier LLM for request analysis (PydanticAI agent)
|
||||
- **Tatlock**: Second-tier LLM with butler personality (PydanticAI agent)
|
||||
- **Expert Agents**: Domain specialists (Librarian, Developer, Handyman, etc.)
|
||||
### 📂 Project Structure (Directory-based, NOT File-type based)
|
||||
Do **not** group files by type (e.g., one huge `routers` folder). Group by **domain/module** inside a `src/` directory.
|
||||
|
||||
**Key Architectural Decisions:**
|
||||
|
||||
1. **Single Source of Truth**: All response generation happens in the Responses API
|
||||
- Structured output with reasoning, function_call, and message items
|
||||
- Real-time stop sequence and max tokens enforcement
|
||||
- Conversation history tracking
|
||||
- Context window management
|
||||
|
||||
2. **Chat Completions Wrapper**: Provides compatibility without duplicating logic
|
||||
- Calls Responses API internally
|
||||
- Automatically enables reasoning generation
|
||||
- Converts reasoning items to `<think>` tags for Open WebUI
|
||||
- Maintains OpenAI-compatible format
|
||||
|
||||
3. **Agent Interface**: Clean abstraction for multiple models
|
||||
- **lorem-tester**: Full-featured mock agent with realistic behavior
|
||||
- Reasoning summaries (adjustable effort levels)
|
||||
- Random tool/function calls
|
||||
- Error triggers for testing
|
||||
- Temperature variation
|
||||
- **Tatlock**: Advertised model name (currently mock, future: PydanticAI Butler agent)
|
||||
|
||||
4. **Hybrid Conversation History**:
|
||||
- Client MUST send full context in `input` array (OpenAI compatible)
|
||||
- Server optionally tracks via `metadata.conversation_id`
|
||||
- Auto-generates deterministic IDs from first message
|
||||
- Supports future vector memory integration (Qdrant)
|
||||
|
||||
**Why This Architecture?**
|
||||
|
||||
- **Open WebUI Compatibility**: Native Responses API support not yet in stable release
|
||||
- **Future-Proof**: Easy migration when Open WebUI adds native support
|
||||
- **Testability**: Full-featured mock agent (lorem-tester) for integration testing
|
||||
- **Clean Separation**: Responses API as stable core, wrappers can change
|
||||
|
||||
### Components
|
||||
|
||||
- **FastAPI**: Web framework for the API layer
|
||||
- **SSE-Starlette**: Server-Sent Events for streaming responses
|
||||
- **Pydantic**: Request/response validation with field validators
|
||||
- **Agent Interface**: Abstract base class for model implementations
|
||||
- **Conversation History**: Server-side tracking with configurable max turns
|
||||
- **Context Window**: Token counting and management
|
||||
- **PydanticAI**: Integrated with Tatlock agent (Ollama backend)
|
||||
- **Agent Tools**: Permanent tools module (`src/agents/tools.py`)
|
||||
- Calculator: Safe mathematical expression evaluation
|
||||
- Date/Time toolkit: Current time, relative dates, time differences
|
||||
- Web Search: SearXNG integration for privacy-preserving search
|
||||
|
||||
## Documentation References
|
||||
|
||||
### Core Framework Documentation
|
||||
|
||||
#### FastAPI
|
||||
- **Official Documentation**: https://fastapi.tiangolo.com/
|
||||
- **Version**: 0.123.9 (Dec 2025)
|
||||
- **Key Topics**:
|
||||
- Path operations and routing
|
||||
- Request/response models with Pydantic
|
||||
- Dependency injection
|
||||
- Background tasks
|
||||
- WebSocket and streaming support
|
||||
- **PyPI**: https://pypi.org/project/fastapi/
|
||||
|
||||
#### Uvicorn
|
||||
- **Official Documentation**: https://www.uvicorn.org/
|
||||
- **Version**: 0.38.0 (Oct 2025)
|
||||
- **Key Topics**:
|
||||
- ASGI server configuration
|
||||
- Deployment settings
|
||||
- Logging and monitoring
|
||||
- SSL/TLS configuration
|
||||
|
||||
### AI/LLM Integration
|
||||
|
||||
#### PydanticAI
|
||||
- **Official Documentation**: https://ai.pydantic.dev/
|
||||
- **Version**: 1.27.0 (Dec 2025)
|
||||
- **Status**: Dependency installed, ready for future integration
|
||||
- **Key Topics** (for future implementation):
|
||||
- Agent creation and configuration
|
||||
- LLM provider integration (Ollama support)
|
||||
- Structured outputs with Pydantic
|
||||
- Streaming responses
|
||||
- Tool/function calling
|
||||
- RunContext and dynamic configuration
|
||||
- MCP server integration
|
||||
- **GitHub**: https://github.com/pydantic/pydantic-ai
|
||||
- **PyPI**: https://pypi.org/project/pydantic-ai/
|
||||
|
||||
#### Pydantic
|
||||
- **Official Documentation**: https://docs.pydantic.dev/latest/
|
||||
- **Version**: 2.11+ (Required for PydanticAI, currently using >=2.11,<2.13)
|
||||
- **Key Topics**:
|
||||
- Data validation and serialization
|
||||
- Field types and validators
|
||||
- Model configuration
|
||||
- JSON schema generation
|
||||
|
||||
### HTTP and Streaming
|
||||
|
||||
#### HTTPX
|
||||
- **Official Documentation**: https://www.python-httpx.org/
|
||||
- **Version**: 0.28.1
|
||||
- **Key Topics**:
|
||||
- Async HTTP client for Ollama communication
|
||||
- Streaming responses
|
||||
- Timeout configuration
|
||||
- Connection pooling
|
||||
|
||||
#### SSE-Starlette
|
||||
- **GitHub**: https://github.com/sysid/sse-starlette
|
||||
- **Version**: 3.0.2 (Oct 2025)
|
||||
- **Key Topics**:
|
||||
- Server-Sent Events implementation
|
||||
- Streaming event responses
|
||||
- Integration with FastAPI/Starlette
|
||||
|
||||
### Ollama Integration
|
||||
|
||||
#### Ollama API
|
||||
- **Official Documentation**: https://github.com/ollama/ollama/blob/main/docs/api.md
|
||||
- **Status**: Async client implemented in `src/ollama/client.py`, ready for future integration
|
||||
- **Key Topics** (for future implementation):
|
||||
- REST API endpoints
|
||||
- Streaming responses
|
||||
- Model management
|
||||
- Generate and chat endpoints
|
||||
- Model configuration
|
||||
- **Current Model Target**: mistral-nemo:latest
|
||||
|
||||
### OpenAI API Compatibility
|
||||
|
||||
#### OpenAI API Reference
|
||||
- **Official Documentation**: https://platform.openai.com/docs/api-reference
|
||||
- **Key API Endpoints**:
|
||||
- `/v1/responses` - Responses API (PRIMARY) with structured output
|
||||
- `/v1/chat/completions` - OpenAI Chat Completions compatibility wrapper
|
||||
- `/v1/models` - List available models
|
||||
|
||||
- **Key Features for Development**:
|
||||
- **Responses API Format**: Structured output with reasoning, function_call, and message items
|
||||
- **Parameter Validation**: Temperature, reasoning effort levels, max tokens, stop sequences
|
||||
- **Conversation History**: Hybrid client/server approach with auto-generated IDs
|
||||
- **Context Management**: Token counting and window trimming
|
||||
- **Streaming**: Real-time SSE streaming with stop sequence and max token enforcement
|
||||
- **Error Handling**: Custom exception types (RateLimitError, ContextLengthError)
|
||||
- **Tool Calling**: PydanticAI tool integration with permanent tools
|
||||
- **Testing**: Comprehensive test suite with mocks and real Ollama integration
|
||||
|
||||
## FastAPI Best Practices
|
||||
|
||||
This project follows best practices from [github.com/zhanymkanov/fastapi-best-practices](https://github.com/zhanymkanov/fastapi-best-practices)
|
||||
|
||||
### Project Structure
|
||||
|
||||
**Domain-Based Organization**: Code is organized by domain/feature rather than by file type:
|
||||
|
||||
```
|
||||
**Correct Structure:**
|
||||
```text
|
||||
src/
|
||||
├── agents/ # Agent interface and implementations
|
||||
│ ├── base.py # Abstract AgentInterface
|
||||
│ ├── lorem_tester.py # Full-featured mock agent
|
||||
│ ├── tatlock.py # Placeholder for real agent
|
||||
│ └── registry.py # ModelRegistry for agent management
|
||||
├── responses/ # Responses API domain (PRIMARY)
|
||||
│ ├── router.py # POST /v1/responses endpoint
|
||||
│ ├── schemas.py # Request/response models with validators
|
||||
│ ├── service.py # Response generation logic
|
||||
│ ├── streaming.py # SSE streaming coordinator
|
||||
│ ├── history.py # Conversation history management
|
||||
│ └── context.py # Context window and token management
|
||||
├── chat/ # Chat Completions domain (WRAPPER)
|
||||
│ ├── router.py # POST /v1/chat/completions endpoint
|
||||
│ ├── schemas.py # Chat request/response models
|
||||
│ ├── service.py # Wraps Responses API, converts to <think> tags
|
||||
│ ├── constants.py # Chat constants (roles, finish reasons)
|
||||
│ └── __init__.py
|
||||
├── models/ # Models listing domain
|
||||
│ ├── router.py # GET /v1/models endpoint
|
||||
│ ├── schemas.py # Model schemas
|
||||
│ ├── service.py # Accesses ModelRegistry
|
||||
│ └── __init__.py
|
||||
├── core/ # Shared utilities
|
||||
│ ├── config.py # Global configuration (BaseSettings)
|
||||
│ ├── models.py # Custom base Pydantic models
|
||||
│ ├── exceptions.py # Custom exceptions (RateLimitError, etc.)
|
||||
│ ├── dependencies.py # Shared dependencies
|
||||
│ └── router.py # Core routes (health, root)
|
||||
├── ollama/ # Ollama client layer (not yet integrated)
|
||||
│ ├── client.py # Async Ollama HTTP client
|
||||
│ └── schemas.py # Ollama API models
|
||||
└── main.py # Application factory & configuration
|
||||
```
|
||||
|
||||
**Key Architectural Principles**:
|
||||
- **Single Source of Truth**: Responses API handles all generation logic
|
||||
- **Wrapper Pattern**: Chat Completions wraps Responses API without duplicating code
|
||||
- **Agent Abstraction**: AgentInterface defines contract for all models
|
||||
- **Domain Separation**: Each domain has its own router, schemas, service
|
||||
- **Service Layer**: Business logic in services, not routers
|
||||
- **Type Safety**: Pydantic models for ALL request/response validation
|
||||
- **Async First**: All I/O operations use async/await
|
||||
|
||||
### Async/Await Best Practices
|
||||
|
||||
**Critical Understanding**: FastAPI handles sync and async routes differently:
|
||||
|
||||
- **Async routes** (`async def`): Called directly in event loop
|
||||
- Use ONLY for non-blocking operations
|
||||
- Perfect for `await httpx.get()`, database queries, file I/O
|
||||
- **NEVER** use blocking calls like `time.sleep()` - this blocks entire server
|
||||
|
||||
- **Sync routes** (`def`): Run in thread pool
|
||||
- Use for CPU-intensive work or blocking SDKs
|
||||
- Blocking I/O won't freeze the event loop
|
||||
- Example: `time.sleep(10)` is safe here
|
||||
|
||||
**Example**:
|
||||
```python
|
||||
@router.get("/terrible")
|
||||
async def terrible():
|
||||
time.sleep(10) # ❌ BLOCKS ENTIRE SERVER
|
||||
|
||||
@router.get("/good")
|
||||
def good():
|
||||
time.sleep(10) # ✅ Runs in thread pool
|
||||
|
||||
@router.get("/perfect")
|
||||
async def perfect():
|
||||
await asyncio.sleep(10) # ✅ Non-blocking async
|
||||
```
|
||||
|
||||
**For CPU-intensive tasks**: Use separate worker processes (not threads) due to Python's GIL.
|
||||
|
||||
### Pydantic Configuration
|
||||
|
||||
**Custom Base Model**: All schemas inherit from `CustomBaseModel` for consistent behavior:
|
||||
|
||||
```python
|
||||
# src/core/models.py
|
||||
class CustomBaseModel(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
json_encoders={datetime: datetime_to_iso_str},
|
||||
populate_by_name=True,
|
||||
use_enum_values=True,
|
||||
validate_assignment=True,
|
||||
)
|
||||
|
||||
def serializable_dict(self, **kwargs):
|
||||
"""Return dict with only JSON-serializable fields."""
|
||||
return jsonable_encoder(self.model_dump(**kwargs))
|
||||
```
|
||||
|
||||
**Benefits**:
|
||||
- Consistent datetime serialization across all responses
|
||||
- Alias support for field name flexibility
|
||||
- Easy JSON encoding for logging/debugging
|
||||
|
||||
**Decoupled Settings**: Split configuration by domain instead of one monolithic file:
|
||||
|
||||
```python
|
||||
# src/core/config.py - Global settings
|
||||
class Config(BaseSettings):
|
||||
DATABASE_URL: PostgresDsn
|
||||
ENVIRONMENT: Environment
|
||||
|
||||
# src/chat/config.py - Chat-specific settings
|
||||
class ChatConfig(BaseSettings):
|
||||
MAX_TOKENS: int
|
||||
DEFAULT_TEMPERATURE: float
|
||||
```
|
||||
|
||||
### Dependency Injection Patterns
|
||||
|
||||
**Validation with Dependencies**: Use dependencies for complex validations:
|
||||
|
||||
```python
|
||||
async def valid_post_id(post_id: UUID4) -> dict:
|
||||
"""Validate post exists in database."""
|
||||
post = await service.get_by_id(post_id)
|
||||
if not post:
|
||||
raise PostNotFound()
|
||||
return post
|
||||
|
||||
@router.get("/posts/{post_id}")
|
||||
async def get_post(post: dict = Depends(valid_post_id)):
|
||||
return post # Already validated!
|
||||
```
|
||||
|
||||
**Chaining Dependencies**: Build reusable validation layers:
|
||||
|
||||
```python
|
||||
async def valid_owned_post(
|
||||
post: dict = Depends(valid_post_id),
|
||||
token_data: dict = Depends(parse_jwt_data),
|
||||
) -> dict:
|
||||
if post["creator_id"] != token_data["user_id"]:
|
||||
raise UserNotOwner()
|
||||
return post
|
||||
```
|
||||
|
||||
**Dependency Caching**: Dependencies are cached within request scope - FastAPI only executes each dependency once per request, even if used multiple times.
|
||||
|
||||
### Application Factory Pattern
|
||||
|
||||
Main.py uses factory pattern for testability and configuration:
|
||||
|
||||
```python
|
||||
def create_application() -> FastAPI:
|
||||
"""Create and configure FastAPI app."""
|
||||
app = FastAPI(title=config.APP_NAME)
|
||||
|
||||
# Add middleware
|
||||
app.add_middleware(CORSMiddleware, ...)
|
||||
|
||||
# Register exception handlers
|
||||
register_exception_handlers(app)
|
||||
|
||||
# Include routers
|
||||
app.include_router(chat_router, prefix="/v1")
|
||||
|
||||
return app
|
||||
|
||||
app = create_application()
|
||||
```
|
||||
|
||||
## Development Guidelines
|
||||
|
||||
### Git Workflow
|
||||
|
||||
**IMPORTANT**: Do NOT handle git commits or pushes automatically. Wait for explicit user instruction before:
|
||||
- Running `git add`
|
||||
- Running `git commit`
|
||||
- Running `git push`
|
||||
- Creating or pushing tags
|
||||
|
||||
The user will manage git operations themselves unless they specifically request assistance.
|
||||
|
||||
### Server Logs and Debugging
|
||||
|
||||
**Development Mode Logging**: When the server is started using `./wakeup.sh`, logs are written to `logs/server.log`. This file is:
|
||||
- Cleared on each server startup (fresh logs every time)
|
||||
- Written in real-time as the server runs
|
||||
- Already gitignored (won't be committed)
|
||||
|
||||
**Accessing Logs**: You can read the log file at any time while the server is running:
|
||||
```bash
|
||||
# View current logs
|
||||
cat logs/server.log
|
||||
|
||||
# Follow logs in real-time
|
||||
tail -f logs/server.log
|
||||
|
||||
# Search logs
|
||||
grep "ERROR" logs/server.log
|
||||
```
|
||||
|
||||
This is useful for debugging issues, monitoring API calls, and understanding server behavior during development.
|
||||
|
||||
### Code Structure Guidelines
|
||||
- Use async/await for ALL I/O operations (database, HTTP, file access)
|
||||
- Use sync (def) for blocking SDKs or CPU-intensive work
|
||||
- Implement proper error handling and logging
|
||||
- Follow dependency injection for validation and shared resources
|
||||
- Use Pydantic models for ALL request/response validation
|
||||
- Keep business logic in service modules, not routers
|
||||
- Domain-based project structure (not file-type based)
|
||||
|
||||
### Security Considerations
|
||||
- Validate all inputs using Pydantic models
|
||||
- Use environment variables for sensitive configuration
|
||||
- Keep dependencies updated and CVE-checked
|
||||
- Minor version locking for supply chain protection
|
||||
- Consider rate limiting for production deployment
|
||||
- Plan for authentication/API keys when needed
|
||||
|
||||
### Testing Approach
|
||||
- Write integration tests for API endpoints
|
||||
- Test streaming functionality with appropriate timeouts
|
||||
- Use pytest-asyncio for async test support
|
||||
- Validate OpenAI API compatibility in tests
|
||||
- Test both mock and real LLM integrations
|
||||
- Cover main application (CORS, exception handlers, lifespan)
|
||||
- Test wrapper layers (chat completions, etc.)
|
||||
- Include tool functionality tests
|
||||
|
||||
### Configuration Management
|
||||
- Use `.env` files for local development
|
||||
- Document all environment variables in README
|
||||
- Provide sensible defaults where possible
|
||||
- Use BaseSettings from pydantic-settings
|
||||
- Support both local and container-based configuration
|
||||
|
||||
## Common Patterns
|
||||
|
||||
### Streaming Response Pattern
|
||||
|
||||
Example from `src/chat/router.py`:
|
||||
|
||||
```python
|
||||
from sse_starlette.sse import EventSourceResponse
|
||||
from fastapi import FastAPI
|
||||
|
||||
async def event_generator():
|
||||
# Currently yields mock lorem ipsum chunks
|
||||
# Future: Stream from Ollama/PydanticAI
|
||||
yield {"data": chunk.model_dump_json()}
|
||||
yield {"data": "[DONE]"}
|
||||
|
||||
@app.post("/stream")
|
||||
async def stream():
|
||||
return EventSourceResponse(event_generator())
|
||||
```
|
||||
|
||||
### PydanticAI Agent Pattern
|
||||
|
||||
When implementing agents with PydanticAI and Ollama:
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
|
||||
agent = Agent(
|
||||
'ollama:mistral-nemo', # Target model
|
||||
# Configuration here
|
||||
)
|
||||
|
||||
# Use the agent
|
||||
result = await agent.run('Your prompt')
|
||||
```
|
||||
|
||||
### OpenAI-Compatible Response Format
|
||||
|
||||
Example schema from `src/chat/schemas.py`:
|
||||
|
||||
```python
|
||||
{
|
||||
"id": "chatcmpl-123",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1234567890,
|
||||
"model": "mistral-nemo:latest",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {"content": "response"},
|
||||
"finish_reason": None
|
||||
}]
|
||||
}
|
||||
```
|
||||
|
||||
### PydanticAI Tool Registration Pattern
|
||||
|
||||
Tools are registered with PydanticAI agents using decorators. See `src/agents/tatlock.py` for examples:
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent, RunContext
|
||||
|
||||
# After creating the agent
|
||||
@agent.tool
|
||||
def tool_name(ctx: RunContext[None], param: str) -> str:
|
||||
"""
|
||||
Tool description that the LLM sees.
|
||||
|
||||
Args:
|
||||
param: Parameter description
|
||||
|
||||
Returns:
|
||||
Result description
|
||||
"""
|
||||
return result
|
||||
```
|
||||
|
||||
**Tool Implementation Guidelines**:
|
||||
- Keep tools in `src/agents/tools.py` for reusability
|
||||
- Use clear, descriptive docstrings (LLM reads these)
|
||||
- Include parameter descriptions in docstrings
|
||||
- Handle errors gracefully and return error messages as strings
|
||||
- For async operations, declare the tool function as `async def`
|
||||
- Test tools independently before integration
|
||||
|
||||
**Example Tool Module** (`src/agents/tools.py`):
|
||||
```python
|
||||
def calculate(expression: str) -> str:
|
||||
"""Safe calculator implementation."""
|
||||
try:
|
||||
# Implementation
|
||||
return str(result)
|
||||
except Exception as e:
|
||||
return f"Error: {str(e)}"
|
||||
|
||||
async def search_web(query: str) -> str:
|
||||
"""Web search via SearXNG."""
|
||||
async with httpx.AsyncClient() as client:
|
||||
# Implementation
|
||||
return formatted_results
|
||||
```
|
||||
|
||||
## Update Policy
|
||||
|
||||
This document should be updated when:
|
||||
- New development patterns are established
|
||||
- Package versions are upgraded
|
||||
- Major architectural changes occur
|
||||
- New best practices are identified
|
||||
|
||||
Last updated: 2025-12-06 (Tools integration)
|
||||
├── auth/
|
||||
│ ├── router.py # Endpoints
|
||||
│ ├── schemas.py # Pydantic models
|
||||
│ ├── service.py # Business logic (CRUD, etc.)
|
||||
│ ├── dependencies.py# Module-specific dependencies
|
||||
│ └── config.py # Module-specific settings
|
||||
├── posts/
|
||||
│ ├── router.py
|
||||
│ └── ...
|
||||
└── main.py # App entry point
|
||||
+78
-1
@@ -7,6 +7,82 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
## [1.1.0] - 2025-12-11
|
||||
|
||||
### Added
|
||||
|
||||
#### Phase 3: Butler Orchestration (Multi-Agent Coordination)
|
||||
- **The Librarian Agent**: Expert agent for research and knowledge management
|
||||
- PydanticAI agent with specialized research assistant personality
|
||||
- Connects to library-desk API for HybridRAG capabilities
|
||||
- System prompt emphasizes fetching wiki pages before summarizing
|
||||
- Streaming support via `run_librarian_stream()`
|
||||
|
||||
- **Library-Desk API Client** (`src/agents/librarian/client.py`):
|
||||
- Async HTTP client with httpx for library-desk API integration
|
||||
- HybridRAG search (vector + graph + web search)
|
||||
- Wiki operations (search, get, list, create, update pages)
|
||||
- Smart page creation with HybridRAG research (`POST /wiki/pages/smart-create`)
|
||||
- Semantic vector search
|
||||
- Knowledge graph queries (Cypher execution)
|
||||
- Dossier (tag collection) browsing
|
||||
- Health check endpoint
|
||||
|
||||
- **Librarian Tools** (`src/agents/librarian/tools.py`):
|
||||
- Research tools:
|
||||
- `hybrid_search`: Combined vector, graph, and web search
|
||||
- `search_wiki`: Full-text wiki page search
|
||||
- `get_wiki_page`: Fetch full wiki page content by ID
|
||||
- `semantic_search`: Vector similarity search
|
||||
- `list_dossiers`: Browse knowledge collections
|
||||
- `get_dossier_pages`: Get pages in a dossier
|
||||
- `explore_knowledge_graph`: Entity and relationship discovery
|
||||
- `find_related_entities`: Find connected concepts
|
||||
- Write tools:
|
||||
- `smart_create_wiki_page`: Create page with automatic HybridRAG research (PREFERRED for topic-based creation)
|
||||
- `create_wiki_page`: Create page with user-provided content
|
||||
- `update_wiki_page`: Update existing page (partial updates supported)
|
||||
|
||||
- **Agent Communication Protocol** (`src/agents/protocol.py`):
|
||||
- `AgentRequest`: Standardized task request with context and constraints
|
||||
- `AgentResponse`: Response with result, reasoning, tool calls, confidence
|
||||
- `DelegationIntent`: Routing intent with target agent and reason
|
||||
- `CoordinationResult`: Aggregated multi-agent results
|
||||
- `DelegationReason` enum: domain expertise, tool access, resource efficiency, user preference
|
||||
- Error types: `AgentError`, `AgentTimeoutError`, `AgentUnavailableError`
|
||||
|
||||
- **Coordination Engine** (`src/agents/coordination.py`):
|
||||
- `CoordinationEngine`: Multi-agent task orchestration
|
||||
- Routing tasks to appropriate expert agents
|
||||
- Sequential and parallel execution support
|
||||
- Result aggregation from multiple agents
|
||||
- Graceful error handling and degradation
|
||||
- Streaming delegation support
|
||||
- Convenience functions: `delegate_to_librarian()`, `delegate_to_librarian_stream()`
|
||||
|
||||
- **Librarian Capability Registration**:
|
||||
- `LIBRARIAN_CAPABILITY` definition with research domains
|
||||
- Automatic registration on application startup
|
||||
- Integration with Household Registry
|
||||
|
||||
- **Configuration**:
|
||||
- `LIBRARY_DESK_HOST`: Library-desk API URL (default: `http://localhost:8089`)
|
||||
- `LIBRARY_DESK_API_KEY`: Optional API key for authentication
|
||||
- `LIBRARY_DESK_TIMEOUT`: Request timeout in seconds (default: 60)
|
||||
|
||||
- **Test Suite**:
|
||||
- 78 new tests for Phase 3 components
|
||||
- Protocol model tests (requests, responses, intents, errors)
|
||||
- Coordination engine tests (delegation, streaming, multi-agent)
|
||||
- Library-desk client tests (all endpoints with mocked HTTP)
|
||||
- Wiki write operation tests (update, smart-create)
|
||||
- Capability registration tests
|
||||
|
||||
### Changed
|
||||
- Application startup now registers The Librarian with Household Registry
|
||||
- Configuration expanded to support library-desk API integration
|
||||
- **Version loading**: APP_VERSION now dynamically loaded from pyproject.toml
|
||||
|
||||
## [1.0.0a] - 2025-12-11
|
||||
|
||||
### Added
|
||||
@@ -314,7 +390,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
- CORS middleware
|
||||
- Exception handlers (OpenAI-compatible error format)
|
||||
|
||||
[Unreleased]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v1.0.0a...main
|
||||
[Unreleased]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v1.1.0...main
|
||||
[1.1.0]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v1.0.0a...v1.1.0
|
||||
[1.0.0a]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v0.2.5...v1.0.0a
|
||||
[0.2.5]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v0.2.0...v0.2.5
|
||||
[0.2.0]: https://git.schweitz.net/jpmschweitzer/tatlock/compare/v0.1.1...v0.2.0
|
||||
|
||||
+1
-1
@@ -5,7 +5,7 @@ WORKDIR /app
|
||||
RUN apt-get update && apt-get install -y curl \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY requirements.txt .
|
||||
COPY requirements.txt pyproject.toml ./
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY src/ ./src/
|
||||
|
||||
@@ -0,0 +1,424 @@
|
||||
# Library-Desk API Requirements for Tatlock Integration
|
||||
|
||||
## Overview
|
||||
|
||||
The Librarian agent in Tatlock needs additional endpoints in library-desk to support wiki page editing and content management. Currently, the API provides read operations but The Librarian needs write capabilities for:
|
||||
|
||||
- Creating new wiki pages
|
||||
- Updating existing wiki pages (content, title, tags, description)
|
||||
|
||||
## Required Endpoints
|
||||
|
||||
### 1. Create Wiki Page (Already Exists)
|
||||
|
||||
**Endpoint:** `POST /wiki/pages`
|
||||
|
||||
This endpoint already exists and works correctly.
|
||||
|
||||
### 2. Update Wiki Page (Needs Enhancement)
|
||||
|
||||
**Endpoint:** `PUT /wiki/pages/{page_id}`
|
||||
|
||||
**Current Status:** May exist but needs verification that it supports partial updates.
|
||||
|
||||
**Required Behavior:**
|
||||
- Accept partial updates (only provided fields should be updated)
|
||||
- Support updating: `content`, `title`, `tags`, `description`
|
||||
- Auto-update vector embeddings after content changes
|
||||
- Auto-update knowledge graph after content changes
|
||||
|
||||
**Request Body:**
|
||||
```json
|
||||
{
|
||||
"content": "# New Content\n\nOptional - only if changing content",
|
||||
"title": "Optional - only if renaming",
|
||||
"tags": ["optional", "list", "of", "new", "tags"],
|
||||
"description": "Optional new description"
|
||||
}
|
||||
```
|
||||
|
||||
**Query Parameters:**
|
||||
- `user`: User identifier for multi-tenancy (required)
|
||||
|
||||
**Response:**
|
||||
```json
|
||||
{
|
||||
"id": 42,
|
||||
"path": "/projects/example",
|
||||
"title": "Updated Title",
|
||||
"description": "Updated description",
|
||||
"content": "# New Content...",
|
||||
"tags": ["updated", "tags"],
|
||||
"updated_at": "2024-01-15T10:30:00Z"
|
||||
}
|
||||
```
|
||||
|
||||
**Notes:**
|
||||
- Should trigger background tasks to re-index vectors and refresh graph entities
|
||||
- Should validate that user has access to the page (namespace check)
|
||||
- Should preserve fields that are not provided in the request
|
||||
|
||||
## Use Cases for The Librarian
|
||||
|
||||
### Adding New Knowledge
|
||||
When a user says "Add this to the wiki" or "Create a page about X":
|
||||
- Librarian uses `POST /wiki/pages` to create the page
|
||||
- Tags are assigned based on context (dossiers)
|
||||
|
||||
### Correcting Information
|
||||
When a user says "Update the page about X" or "Fix this fact":
|
||||
1. Librarian searches for the page with `GET /wiki/search`
|
||||
2. Fetches full content with `GET /wiki/pages/{id}`
|
||||
3. Updates with corrected content via `PUT /wiki/pages/{id}`
|
||||
|
||||
### Organizing Knowledge
|
||||
When a user says "Add this page to the projects dossier":
|
||||
- Librarian updates just the tags field via `PUT /wiki/pages/{id}`
|
||||
|
||||
## Integration Notes
|
||||
|
||||
- The Librarian will call these endpoints via HTTP from Tatlock
|
||||
- Authentication uses Bearer token (LIBRARY_DESK_API_KEY)
|
||||
- All operations are scoped to the user's namespace
|
||||
- Background processing (vectors, graph) should not block the response
|
||||
|
||||
## Testing Checklist
|
||||
|
||||
- [ ] `PUT /wiki/pages/{page_id}` accepts partial updates
|
||||
- [ ] Updating content triggers vector re-indexing
|
||||
- [ ] Updating content triggers graph entity extraction
|
||||
- [ ] Tags can be updated independently of content
|
||||
- [ ] Description can be updated independently
|
||||
- [ ] Title can be updated (with path remaining the same)
|
||||
- [ ] User namespace validation works correctly
|
||||
|
||||
|
||||
===== IMPLEMENTATION INSTRUCTIONS =========
|
||||
# Librarian Wiki Integration Guide
|
||||
|
||||
This document provides implementation instructions for integrating the library-desk wiki endpoints into the Librarian agent (Tatlock).
|
||||
|
||||
## Available Endpoints
|
||||
|
||||
### 1. Create Wiki Page
|
||||
|
||||
**Endpoint:** `POST /wiki/pages`
|
||||
|
||||
Use this for simple page creation when the Librarian already has the content.
|
||||
|
||||
```python
|
||||
async def create_wiki_page(
|
||||
title: str,
|
||||
path: str,
|
||||
content: str,
|
||||
tags: list[str],
|
||||
description: str = "",
|
||||
user: str = "default"
|
||||
) -> dict:
|
||||
"""Create a new wiki page."""
|
||||
response = await http_client.post(
|
||||
f"{LIBRARY_DESK_URL}/wiki/pages",
|
||||
headers={"Authorization": f"Bearer {LIBRARY_DESK_API_KEY}"},
|
||||
json={
|
||||
"title": title,
|
||||
"path": path,
|
||||
"content": content,
|
||||
"tags": tags,
|
||||
"description": description,
|
||||
"user": user
|
||||
}
|
||||
)
|
||||
return response.json()
|
||||
```
|
||||
|
||||
**When to use:**
|
||||
- User provides specific content to add
|
||||
- Librarian has already composed the content
|
||||
- Simple note-taking or quick additions
|
||||
|
||||
---
|
||||
|
||||
### 2. Smart Create Wiki Page (Recommended for Research)
|
||||
|
||||
**Endpoint:** `POST /wiki/pages/smart-create`
|
||||
|
||||
Use this when the Librarian should research a topic before creating the page. This endpoint:
|
||||
1. Searches existing wiki, knowledge graph, and web for context
|
||||
2. Uses LLM to synthesize findings into structured content
|
||||
3. Creates the page with proper attribution
|
||||
4. Automatically links entities bidirectionally
|
||||
|
||||
```python
|
||||
async def smart_create_wiki_page(
|
||||
topic: str,
|
||||
tags: list[str],
|
||||
user: str = "default",
|
||||
path: str | None = None,
|
||||
include_web_research: bool = True,
|
||||
include_wiki_search: bool = True
|
||||
) -> dict:
|
||||
"""Create a wiki page with HybridRAG research."""
|
||||
response = await http_client.post(
|
||||
f"{LIBRARY_DESK_URL}/wiki/pages/smart-create",
|
||||
headers={"Authorization": f"Bearer {LIBRARY_DESK_API_KEY}"},
|
||||
json={
|
||||
"topic": topic,
|
||||
"path": path, # Optional - auto-generated from topic if not provided
|
||||
"tags": tags,
|
||||
"user": user,
|
||||
"include_web_research": include_web_research,
|
||||
"include_wiki_search": include_wiki_search
|
||||
}
|
||||
)
|
||||
return response.json()
|
||||
```
|
||||
|
||||
**Response includes:**
|
||||
```json
|
||||
{
|
||||
"page": {
|
||||
"id": 123,
|
||||
"path": "/users/jpmschweitzer/technology/docker-orchestration",
|
||||
"title": "Docker orchestration",
|
||||
"content": "# Docker Orchestration\n\n...",
|
||||
"tags": ["technology", "devops"],
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z"
|
||||
},
|
||||
"research_summary": {
|
||||
"wiki_results": 3,
|
||||
"web_results": 8,
|
||||
"graph_entities": 5,
|
||||
"keywords_extracted": 12,
|
||||
"timing_ms": 4500
|
||||
},
|
||||
"sources_used": 11,
|
||||
"search_id": "uuid-for-reference",
|
||||
"entity_linking": {
|
||||
"forward_links": 5,
|
||||
"backward_links": 3,
|
||||
"pages_updated": 2
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**When to use:**
|
||||
- User says "Create a page about X"
|
||||
- User says "Add information about X to the wiki"
|
||||
- Librarian needs to research before writing
|
||||
- Topic benefits from context from existing knowledge
|
||||
|
||||
---
|
||||
|
||||
### 3. Update Wiki Page
|
||||
|
||||
**Endpoint:** `PUT /wiki/pages/{page_id}`
|
||||
|
||||
Use this for modifying existing pages. Supports partial updates.
|
||||
|
||||
```python
|
||||
async def update_wiki_page(
|
||||
page_id: int,
|
||||
user: str = "default",
|
||||
content: str | None = None,
|
||||
title: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
description: str | None = None
|
||||
) -> dict:
|
||||
"""Update an existing wiki page (partial updates supported)."""
|
||||
# Only include fields that are being updated
|
||||
update_data = {}
|
||||
if content is not None:
|
||||
update_data["content"] = content
|
||||
if title is not None:
|
||||
update_data["title"] = title
|
||||
if tags is not None:
|
||||
update_data["tags"] = tags
|
||||
if description is not None:
|
||||
update_data["description"] = description
|
||||
|
||||
response = await http_client.put(
|
||||
f"{LIBRARY_DESK_URL}/wiki/pages/{page_id}?user={user}",
|
||||
headers={"Authorization": f"Bearer {LIBRARY_DESK_API_KEY}"},
|
||||
json=update_data
|
||||
)
|
||||
return response.json()
|
||||
```
|
||||
|
||||
**When to use:**
|
||||
- User says "Update the page about X"
|
||||
- User says "Fix this information"
|
||||
- User says "Add this page to the projects dossier" (update tags only)
|
||||
- Correcting or enhancing existing content
|
||||
|
||||
---
|
||||
|
||||
### 4. Search Wiki Pages
|
||||
|
||||
**Endpoint:** `GET /wiki/search`
|
||||
|
||||
Use this to find existing pages before updating.
|
||||
|
||||
```python
|
||||
async def search_wiki(
|
||||
query: str,
|
||||
user: str = "default"
|
||||
) -> dict:
|
||||
"""Search wiki pages."""
|
||||
response = await http_client.get(
|
||||
f"{LIBRARY_DESK_URL}/wiki/search",
|
||||
headers={"Authorization": f"Bearer {LIBRARY_DESK_API_KEY}"},
|
||||
params={"q": query, "user": user}
|
||||
)
|
||||
return response.json()
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### 5. Get Wiki Page
|
||||
|
||||
**Endpoint:** `GET /wiki/pages/{page_id}`
|
||||
|
||||
Use this to fetch full page content before editing.
|
||||
|
||||
```python
|
||||
async def get_wiki_page(
|
||||
page_id: int,
|
||||
user: str = "default"
|
||||
) -> dict:
|
||||
"""Get a wiki page by ID."""
|
||||
response = await http_client.get(
|
||||
f"{LIBRARY_DESK_URL}/wiki/pages/{page_id}",
|
||||
headers={"Authorization": f"Bearer {LIBRARY_DESK_API_KEY}"},
|
||||
params={"user": user}
|
||||
)
|
||||
return response.json()
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Decision Flow for Librarian
|
||||
|
||||
```
|
||||
User Request
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────┐
|
||||
│ Does user want to CREATE or UPDATE a page? │
|
||||
└─────────────────────────────────────────────┘
|
||||
│ │
|
||||
▼ ▼
|
||||
CREATE UPDATE
|
||||
│ │
|
||||
▼ ▼
|
||||
┌─────────────────┐ ┌──────────────────────┐
|
||||
│ Does Librarian │ │ Search for the page │
|
||||
│ need to research│ │ GET /wiki/search │
|
||||
│ the topic? │ └──────────────────────┘
|
||||
└─────────────────┘ │
|
||||
│ │ ▼
|
||||
▼ ▼ ┌──────────────────────┐
|
||||
YES NO │ Get full page content│
|
||||
│ │ │ GET /wiki/pages/{id} │
|
||||
▼ ▼ └──────────────────────┘
|
||||
┌─────────┐ ┌─────────┐ │
|
||||
│ smart- │ │ POST │ ▼
|
||||
│ create │ │ /wiki/ │ ┌──────────────────────┐
|
||||
│ │ │ pages │ │ Update the page │
|
||||
└─────────┘ └─────────┘ │ PUT /wiki/pages/{id} │
|
||||
└──────────────────────┘
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Common Use Cases
|
||||
|
||||
### 1. "Create a page about Docker Compose"
|
||||
|
||||
```python
|
||||
# Use smart-create for research-backed content
|
||||
result = await smart_create_wiki_page(
|
||||
topic="Docker Compose",
|
||||
tags=["technology", "devops", "containers"],
|
||||
user="jpmschweitzer"
|
||||
)
|
||||
# Returns page with synthesized content from wiki + web research
|
||||
```
|
||||
|
||||
### 2. "Add this note to the wiki: Remember to renew SSL cert on Jan 15"
|
||||
|
||||
```python
|
||||
# Use simple create for user-provided content
|
||||
result = await create_wiki_page(
|
||||
title="SSL Certificate Renewal Reminder",
|
||||
path="/reminders/ssl-renewal",
|
||||
content="# SSL Certificate Renewal\n\nRemember to renew SSL cert on Jan 15",
|
||||
tags=["reminders", "infrastructure"],
|
||||
user="jpmschweitzer"
|
||||
)
|
||||
```
|
||||
|
||||
### 3. "Update the page about my home server to add the new IP"
|
||||
|
||||
```python
|
||||
# 1. Search for the page
|
||||
search_results = await search_wiki("home server", user="jpmschweitzer")
|
||||
page_id = search_results["results"][0]["id"]
|
||||
|
||||
# 2. Get current content
|
||||
page = await get_wiki_page(page_id, user="jpmschweitzer")
|
||||
|
||||
# 3. Modify content (Librarian edits the markdown)
|
||||
new_content = page["content"] + "\n\n## Updated IP\n\nNew IP: 192.168.1.100"
|
||||
|
||||
# 4. Update the page
|
||||
result = await update_wiki_page(
|
||||
page_id=page_id,
|
||||
content=new_content,
|
||||
user="jpmschweitzer"
|
||||
)
|
||||
```
|
||||
|
||||
### 4. "Add this page to the projects dossier"
|
||||
|
||||
```python
|
||||
# Update only tags (partial update)
|
||||
result = await update_wiki_page(
|
||||
page_id=page_id,
|
||||
tags=["projects", "existing-tag"], # Add "projects" tag
|
||||
user="jpmschweitzer"
|
||||
)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Background Processing
|
||||
|
||||
All write operations trigger background tasks that:
|
||||
|
||||
1. **Vector Indexing:** Chunks content and generates embeddings in Qdrant
|
||||
2. **Graph Extraction:** Extracts entities and creates Neo4j relationships
|
||||
3. **Entity Linking:** (smart-create only) Links entities bidirectionally
|
||||
|
||||
These run asynchronously and don't block the API response.
|
||||
|
||||
---
|
||||
|
||||
## Authentication
|
||||
|
||||
All endpoints require Bearer token authentication:
|
||||
|
||||
```
|
||||
Authorization: Bearer {LIBRARY_DESK_API_KEY}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Multi-Tenancy
|
||||
|
||||
All operations are scoped to the user's namespace:
|
||||
- Pages are stored under `/users/{user}/...`
|
||||
- Vector collections are per-user: `library_desk_{user}`
|
||||
- Graph nodes are labeled per-user: `User_{User}_Document`
|
||||
|
||||
Always pass the `user` parameter to ensure proper isolation.
|
||||
+1
-1
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "tatlock"
|
||||
version = "1.0.0a"
|
||||
version = "1.1.0"
|
||||
description = "OpenAI-compatible API with Ollama backend"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = []
|
||||
|
||||
@@ -0,0 +1,407 @@
|
||||
"""
|
||||
Multi-agent coordination engine.
|
||||
|
||||
Orchestrates delegation from Tatlock to expert agents (Librarian, etc.)
|
||||
based on Steward recommendations. Handles:
|
||||
- Routing tasks to appropriate agents
|
||||
- Parallel and sequential execution
|
||||
- Result aggregation
|
||||
- Error handling and graceful degradation
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Any, AsyncGenerator, Optional
|
||||
|
||||
from src.agents.librarian import run_librarian, run_librarian_stream
|
||||
from src.agents.protocol import (
|
||||
AgentError,
|
||||
AgentRequest,
|
||||
AgentResponse,
|
||||
AgentTimeoutError,
|
||||
AgentUnavailableError,
|
||||
CoordinationResult,
|
||||
DelegationIntent,
|
||||
DelegationReason,
|
||||
ToolCallRecord,
|
||||
)
|
||||
from src.core.household_registry import get_household_registry
|
||||
from src.core.logging_config import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# Agent execution functions registry
|
||||
AGENT_EXECUTORS: dict[str, Any] = {
|
||||
"librarian": run_librarian,
|
||||
}
|
||||
|
||||
AGENT_STREAM_EXECUTORS: dict[str, Any] = {
|
||||
"librarian": run_librarian_stream,
|
||||
}
|
||||
|
||||
|
||||
class CoordinationEngine:
|
||||
"""
|
||||
Coordinates multi-agent task execution.
|
||||
|
||||
Routes tasks from Tatlock to appropriate expert agents,
|
||||
handles execution, and aggregates results.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the coordination engine."""
|
||||
self.registry = get_household_registry()
|
||||
logger.info("coordination_engine_initialized")
|
||||
|
||||
def get_available_agents(self) -> list[str]:
|
||||
"""
|
||||
Get list of available expert agents.
|
||||
|
||||
Returns:
|
||||
List of agent names that can accept delegations
|
||||
"""
|
||||
available = []
|
||||
for name in self.registry.list_members():
|
||||
member = self.registry.get_member(name)
|
||||
if member and member.agent is not None:
|
||||
available.append(name)
|
||||
return available
|
||||
|
||||
def can_delegate_to(self, agent_name: str) -> bool:
|
||||
"""
|
||||
Check if delegation to an agent is possible.
|
||||
|
||||
Args:
|
||||
agent_name: Name of the target agent
|
||||
|
||||
Returns:
|
||||
True if agent is available and can accept tasks
|
||||
"""
|
||||
if agent_name not in AGENT_EXECUTORS:
|
||||
return False
|
||||
|
||||
member = self.registry.get_member(agent_name)
|
||||
return member is not None and member.agent is not None
|
||||
|
||||
async def execute_delegation(
|
||||
self,
|
||||
intent: DelegationIntent,
|
||||
context: str = "",
|
||||
message_history: Optional[list[Any]] = None,
|
||||
) -> AgentResponse:
|
||||
"""
|
||||
Execute a single delegation to an expert agent.
|
||||
|
||||
Args:
|
||||
intent: The delegation intent with task details
|
||||
context: Additional context for the agent
|
||||
message_history: Optional conversation history
|
||||
|
||||
Returns:
|
||||
AgentResponse with results
|
||||
|
||||
Raises:
|
||||
AgentUnavailableError: If agent is not available
|
||||
AgentTimeoutError: If execution times out
|
||||
AgentError: For other execution errors
|
||||
"""
|
||||
start_time = time.time()
|
||||
agent_name = intent.target_agent
|
||||
|
||||
logger.info(
|
||||
"delegation_started",
|
||||
agent=agent_name,
|
||||
task=intent.task[:100],
|
||||
reason=intent.reason.value,
|
||||
)
|
||||
|
||||
# Check if agent is available
|
||||
if not self.can_delegate_to(agent_name):
|
||||
raise AgentUnavailableError(
|
||||
f"Agent '{agent_name}' is not available for delegation",
|
||||
agent_name=agent_name,
|
||||
)
|
||||
|
||||
# Get the executor
|
||||
executor = AGENT_EXECUTORS.get(agent_name)
|
||||
if not executor:
|
||||
raise AgentUnavailableError(
|
||||
f"No executor found for agent '{agent_name}'",
|
||||
agent_name=agent_name,
|
||||
)
|
||||
|
||||
try:
|
||||
# Build the request
|
||||
request = AgentRequest(
|
||||
task=intent.task,
|
||||
context=context,
|
||||
delegation_reason=intent.reason,
|
||||
)
|
||||
|
||||
# Execute with timeout
|
||||
timeout = request.timeout_seconds or 60
|
||||
|
||||
result = await asyncio.wait_for(
|
||||
executor(
|
||||
task=request.task,
|
||||
context=request.context,
|
||||
message_history=message_history,
|
||||
),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
|
||||
logger.info(
|
||||
"delegation_completed",
|
||||
agent=agent_name,
|
||||
duration_ms=duration_ms,
|
||||
output_length=len(result),
|
||||
)
|
||||
|
||||
return AgentResponse(
|
||||
success=True,
|
||||
result=result,
|
||||
reasoning=f"Delegated to {agent_name}: {intent.expected_outcome}",
|
||||
duration_ms=duration_ms,
|
||||
)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
logger.error(
|
||||
"delegation_timeout",
|
||||
agent=agent_name,
|
||||
duration_ms=duration_ms,
|
||||
)
|
||||
raise AgentTimeoutError(
|
||||
f"Agent '{agent_name}' timed out after {duration_ms}ms",
|
||||
agent_name=agent_name,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
logger.error(
|
||||
"delegation_error",
|
||||
agent=agent_name,
|
||||
error=str(e),
|
||||
duration_ms=duration_ms,
|
||||
exc_info=True,
|
||||
)
|
||||
return AgentResponse(
|
||||
success=False,
|
||||
result="",
|
||||
error_message=str(e),
|
||||
duration_ms=duration_ms,
|
||||
)
|
||||
|
||||
async def execute_delegation_stream(
|
||||
self,
|
||||
intent: DelegationIntent,
|
||||
context: str = "",
|
||||
message_history: Optional[list[Any]] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Execute a delegation with streaming output.
|
||||
|
||||
Args:
|
||||
intent: The delegation intent with task details
|
||||
context: Additional context for the agent
|
||||
message_history: Optional conversation history
|
||||
|
||||
Yields:
|
||||
Text deltas from the agent
|
||||
|
||||
Raises:
|
||||
AgentUnavailableError: If agent is not available
|
||||
"""
|
||||
agent_name = intent.target_agent
|
||||
|
||||
logger.info(
|
||||
"delegation_stream_started",
|
||||
agent=agent_name,
|
||||
task=intent.task[:100],
|
||||
)
|
||||
|
||||
# Check if agent is available
|
||||
if agent_name not in AGENT_STREAM_EXECUTORS:
|
||||
raise AgentUnavailableError(
|
||||
f"Agent '{agent_name}' does not support streaming",
|
||||
agent_name=agent_name,
|
||||
)
|
||||
|
||||
executor = AGENT_STREAM_EXECUTORS[agent_name]
|
||||
|
||||
try:
|
||||
async for delta in executor(
|
||||
task=intent.task,
|
||||
context=context,
|
||||
message_history=message_history,
|
||||
):
|
||||
yield delta
|
||||
|
||||
logger.info("delegation_stream_completed", agent=agent_name)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"delegation_stream_error",
|
||||
agent=agent_name,
|
||||
error=str(e),
|
||||
exc_info=True,
|
||||
)
|
||||
yield f"\n\n[Error from {agent_name}: {str(e)}]"
|
||||
|
||||
async def coordinate(
|
||||
self,
|
||||
intents: list[DelegationIntent],
|
||||
context: str = "",
|
||||
message_history: Optional[list[Any]] = None,
|
||||
) -> CoordinationResult:
|
||||
"""
|
||||
Coordinate execution of multiple delegations.
|
||||
|
||||
Handles parallel execution for independent tasks and
|
||||
sequential execution for dependent tasks.
|
||||
|
||||
Args:
|
||||
intents: List of delegation intents to execute
|
||||
context: Shared context for all agents
|
||||
message_history: Optional conversation history
|
||||
|
||||
Returns:
|
||||
CoordinationResult with aggregated results
|
||||
"""
|
||||
start_time = time.time()
|
||||
agent_responses: dict[str, AgentResponse] = {}
|
||||
agents_consulted: list[str] = []
|
||||
|
||||
logger.info(
|
||||
"coordination_started",
|
||||
intent_count=len(intents),
|
||||
agents=[i.target_agent for i in intents],
|
||||
)
|
||||
|
||||
# Sort by priority
|
||||
sorted_intents = sorted(intents, key=lambda x: x.priority)
|
||||
|
||||
# Group by dependencies (simple version: sequential for now)
|
||||
# TODO: Implement parallel execution for independent tasks
|
||||
for intent in sorted_intents:
|
||||
try:
|
||||
response = await self.execute_delegation(
|
||||
intent=intent,
|
||||
context=context,
|
||||
message_history=message_history,
|
||||
)
|
||||
agent_responses[intent.target_agent] = response
|
||||
if response.success:
|
||||
agents_consulted.append(intent.target_agent)
|
||||
|
||||
except AgentError as e:
|
||||
agent_responses[intent.target_agent] = AgentResponse(
|
||||
success=False,
|
||||
result="",
|
||||
error_message=str(e),
|
||||
)
|
||||
|
||||
# Aggregate results
|
||||
successful_results = [
|
||||
r.result for r in agent_responses.values() if r.success and r.result
|
||||
]
|
||||
|
||||
final_response = "\n\n---\n\n".join(successful_results) if successful_results else ""
|
||||
|
||||
total_duration = int((time.time() - start_time) * 1000)
|
||||
|
||||
logger.info(
|
||||
"coordination_completed",
|
||||
total_duration_ms=total_duration,
|
||||
agents_consulted=agents_consulted,
|
||||
success_count=len(successful_results),
|
||||
)
|
||||
|
||||
return CoordinationResult(
|
||||
final_response=final_response,
|
||||
agent_responses=agent_responses,
|
||||
delegation_intents=intents,
|
||||
total_duration_ms=total_duration,
|
||||
agents_consulted=agents_consulted,
|
||||
)
|
||||
|
||||
|
||||
# Global coordination engine instance
|
||||
_coordination_engine: Optional[CoordinationEngine] = None
|
||||
|
||||
|
||||
def get_coordination_engine() -> CoordinationEngine:
|
||||
"""Get the global coordination engine instance."""
|
||||
global _coordination_engine
|
||||
if _coordination_engine is None:
|
||||
_coordination_engine = CoordinationEngine()
|
||||
return _coordination_engine
|
||||
|
||||
|
||||
async def delegate_to_librarian(
|
||||
task: str,
|
||||
context: str = "",
|
||||
reason: DelegationReason = DelegationReason.DOMAIN_EXPERTISE,
|
||||
message_history: Optional[list[Any]] = None,
|
||||
) -> AgentResponse:
|
||||
"""
|
||||
Convenience function to delegate a task to The Librarian.
|
||||
|
||||
Args:
|
||||
task: Research task description
|
||||
context: Additional context
|
||||
reason: Why delegating to Librarian
|
||||
message_history: Optional conversation history
|
||||
|
||||
Returns:
|
||||
AgentResponse with research results
|
||||
"""
|
||||
engine = get_coordination_engine()
|
||||
|
||||
intent = DelegationIntent(
|
||||
target_agent="librarian",
|
||||
task=task,
|
||||
reason=reason,
|
||||
expected_outcome="Research findings and relevant information",
|
||||
)
|
||||
|
||||
return await engine.execute_delegation(
|
||||
intent=intent,
|
||||
context=context,
|
||||
message_history=message_history,
|
||||
)
|
||||
|
||||
|
||||
async def delegate_to_librarian_stream(
|
||||
task: str,
|
||||
context: str = "",
|
||||
message_history: Optional[list[Any]] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Convenience function to delegate to Librarian with streaming.
|
||||
|
||||
Args:
|
||||
task: Research task description
|
||||
context: Additional context
|
||||
message_history: Optional conversation history
|
||||
|
||||
Yields:
|
||||
Text deltas from The Librarian
|
||||
"""
|
||||
engine = get_coordination_engine()
|
||||
|
||||
intent = DelegationIntent(
|
||||
target_agent="librarian",
|
||||
task=task,
|
||||
reason=DelegationReason.DOMAIN_EXPERTISE,
|
||||
expected_outcome="Research findings",
|
||||
)
|
||||
|
||||
async for delta in engine.execute_delegation_stream(
|
||||
intent=intent,
|
||||
context=context,
|
||||
message_history=message_history,
|
||||
):
|
||||
yield delta
|
||||
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
The Librarian - Expert agent for research and knowledge management.
|
||||
|
||||
Connects to the library-desk API to provide:
|
||||
- HybridRAG search (vector + graph + web)
|
||||
- Wiki.js operations
|
||||
- Knowledge graph queries
|
||||
- Semantic search
|
||||
"""
|
||||
from src.agents.librarian.agent import (
|
||||
get_librarian_agent,
|
||||
run_librarian,
|
||||
run_librarian_stream,
|
||||
)
|
||||
from src.agents.librarian.capability import (
|
||||
LIBRARIAN_CAPABILITY,
|
||||
get_librarian_capability,
|
||||
register_librarian,
|
||||
unregister_librarian,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"LIBRARIAN_CAPABILITY",
|
||||
"get_librarian_capability",
|
||||
"get_librarian_agent",
|
||||
"register_librarian",
|
||||
"unregister_librarian",
|
||||
"run_librarian",
|
||||
"run_librarian_stream",
|
||||
]
|
||||
@@ -0,0 +1,286 @@
|
||||
"""
|
||||
The Librarian - Expert agent for research and knowledge management.
|
||||
|
||||
A PydanticAI agent that provides research assistance through
|
||||
the library-desk API, offering:
|
||||
- HybridRAG search across all knowledge sources
|
||||
- Wiki and document management
|
||||
- Semantic search and knowledge graph exploration
|
||||
"""
|
||||
from typing import Any, Optional
|
||||
|
||||
from pydantic_ai import Agent
|
||||
|
||||
from src.agents.librarian.tools import (
|
||||
create_wiki_page,
|
||||
explore_knowledge_graph,
|
||||
find_related_entities,
|
||||
get_dossier_pages,
|
||||
get_wiki_page,
|
||||
hybrid_search,
|
||||
list_dossiers,
|
||||
search_wiki,
|
||||
semantic_search,
|
||||
smart_create_wiki_page,
|
||||
update_wiki_page,
|
||||
)
|
||||
from src.core.config import config
|
||||
from src.core.logging_config import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Librarian system prompt
|
||||
LIBRARIAN_SYSTEM_PROMPT = """You are The Librarian, an expert research assistant in the Tatlock household.
|
||||
|
||||
Your role is to help users find, understand, synthesize, and manage information from:
|
||||
- The personal wiki (Wiki.js) containing documentation and notes
|
||||
- The knowledge graph (Neo4j) with entities and relationships
|
||||
- Vector embeddings (Qdrant) for semantic search
|
||||
- Web search (SearXNG) for current information
|
||||
|
||||
## Your Personality
|
||||
- Scholarly and thorough in your research
|
||||
- Cite your sources and provide context
|
||||
- Organize information clearly
|
||||
- Suggest related topics when relevant
|
||||
- Acknowledge limitations when information is incomplete
|
||||
|
||||
## Your Tools
|
||||
|
||||
### Research Tools
|
||||
- **hybrid_search**: Your primary research tool - searches all sources at once
|
||||
- **search_wiki**: Find specific wiki pages by keyword
|
||||
- **semantic_search**: Find conceptually similar content
|
||||
- **explore_knowledge_graph** / **find_related_entities**: Discover connections
|
||||
- **list_dossiers** / **get_dossier_pages**: Browse knowledge collections
|
||||
|
||||
### Wiki Reading Tools
|
||||
- **get_wiki_page**: Read full content of a wiki page by ID
|
||||
- ALWAYS use this to fetch and read page content when summarizing
|
||||
- Use after search_wiki to get the full text of a specific page
|
||||
|
||||
### Wiki Writing Tools
|
||||
- **smart_create_wiki_page**: Create a page with automatic research (PREFERRED)
|
||||
- **This is the DEFAULT choice when user asks to create a wiki page about a topic**
|
||||
- When user says "Create a page about X" or "Add X to the wiki" without providing specific content, ALWAYS use this tool
|
||||
- Automatically researches the topic from wiki, graph, and web
|
||||
- Synthesizes content with proper source attribution
|
||||
- Creates bidirectional links in knowledge graph
|
||||
- **create_wiki_page**: Create a page with user-provided content
|
||||
- ONLY use when user provides specific text/content they want added verbatim
|
||||
- For simple notes, reminders, or quick additions with exact content
|
||||
- **update_wiki_page**: Update an existing page (partial updates)
|
||||
- Use when: "Update the page about X", "Fix this info", "Add to dossier"
|
||||
- First search_wiki to find the page, then get_wiki_page to read it
|
||||
- Only specify fields you want to change
|
||||
|
||||
## Research Approach
|
||||
1. Start with hybrid_search for broad queries
|
||||
2. Use search_wiki for specific document lookups
|
||||
3. **ALWAYS use get_wiki_page to fetch full content** before summarizing a page
|
||||
4. Use semantic_search when looking for conceptually similar content
|
||||
5. Explore the knowledge graph to find connections between concepts
|
||||
6. Synthesize and summarize findings clearly
|
||||
|
||||
## Writing Approach
|
||||
When asked to create or update wiki content:
|
||||
1. **"Create a page about X" (no specific content provided)**: Use smart_create_wiki_page
|
||||
- This is the PREFERRED tool for topic-based page creation
|
||||
- It researches first and creates comprehensive, well-sourced content
|
||||
2. **User provides exact text to add**: Use create_wiki_page with their content
|
||||
3. **Updating existing pages**:
|
||||
- Search for the page with search_wiki
|
||||
- Fetch full content with get_wiki_page
|
||||
- Make edits and use update_wiki_page
|
||||
4. **Organizing into dossiers**: Use update_wiki_page with just the tags field
|
||||
|
||||
## Response Format
|
||||
Your responses are returned to Tatlock (the butler) who will synthesize them into a final answer for the user. Keep this in mind:
|
||||
- Lead with the key findings or confirmation of action
|
||||
- Include relevant sources and citations
|
||||
- When summarizing wiki pages, fetch and read them first
|
||||
- Note any gaps in available information
|
||||
- Be concise but thorough - Tatlock will format the final response
|
||||
- Structure your findings clearly so they can be easily integrated with other responses
|
||||
"""
|
||||
|
||||
# Lazy initialization to avoid connection issues during imports
|
||||
_librarian_agent: Optional[Agent[None, str]] = None
|
||||
|
||||
|
||||
def _create_librarian_agent() -> Agent[None, str]:
|
||||
"""Create the Librarian PydanticAI agent."""
|
||||
# Import required classes for Ollama configuration
|
||||
from pydantic_ai.models.openai import OpenAIChatModel
|
||||
from pydantic_ai.providers.ollama import OllamaProvider
|
||||
|
||||
# PydanticAI expects Ollama base URL to end with /v1
|
||||
clean_host = str(config.OLLAMA_HOST).rstrip('/')
|
||||
base_url = f"{clean_host}/v1"
|
||||
|
||||
# Create Ollama model with provider
|
||||
model = OpenAIChatModel(
|
||||
model_name=config.OLLAMA_DEFAULT_MODEL,
|
||||
provider=OllamaProvider(base_url=base_url)
|
||||
)
|
||||
|
||||
agent: Agent[None, str] = Agent(
|
||||
model=model,
|
||||
system_prompt=LIBRARIAN_SYSTEM_PROMPT,
|
||||
retries=2,
|
||||
)
|
||||
|
||||
# Register research tools
|
||||
agent.tool_plain(hybrid_search)
|
||||
agent.tool_plain(search_wiki)
|
||||
agent.tool_plain(semantic_search)
|
||||
agent.tool_plain(list_dossiers)
|
||||
agent.tool_plain(get_dossier_pages)
|
||||
agent.tool_plain(explore_knowledge_graph)
|
||||
agent.tool_plain(find_related_entities)
|
||||
|
||||
# Register wiki read tools
|
||||
agent.tool_plain(get_wiki_page)
|
||||
|
||||
# Register wiki write tools
|
||||
agent.tool_plain(create_wiki_page)
|
||||
agent.tool_plain(update_wiki_page)
|
||||
agent.tool_plain(smart_create_wiki_page)
|
||||
|
||||
logger.info(
|
||||
"librarian_agent_created",
|
||||
model=config.OLLAMA_DEFAULT_MODEL,
|
||||
tool_count=11,
|
||||
)
|
||||
|
||||
return agent
|
||||
|
||||
|
||||
def get_librarian_agent() -> Agent[None, str]:
|
||||
"""
|
||||
Get the Librarian agent instance (lazy initialization).
|
||||
|
||||
Returns:
|
||||
PydanticAI Agent configured for research tasks
|
||||
"""
|
||||
global _librarian_agent
|
||||
if _librarian_agent is None:
|
||||
_librarian_agent = _create_librarian_agent()
|
||||
return _librarian_agent
|
||||
|
||||
|
||||
async def run_librarian(
|
||||
task: str,
|
||||
context: str = "",
|
||||
message_history: Optional[list[Any]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Execute a research task with The Librarian.
|
||||
|
||||
This is the main entry point for delegating research tasks
|
||||
to The Librarian from Tatlock or other agents.
|
||||
|
||||
Args:
|
||||
task: The research task or question
|
||||
context: Additional context from conversation
|
||||
message_history: Optional conversation history
|
||||
|
||||
Returns:
|
||||
Research results and findings
|
||||
|
||||
Example:
|
||||
result = await run_librarian(
|
||||
task="Find information about Docker networking",
|
||||
context="User is setting up a homelab",
|
||||
)
|
||||
"""
|
||||
agent = get_librarian_agent()
|
||||
|
||||
# Build prompt with context if provided
|
||||
prompt = task
|
||||
if context:
|
||||
prompt = f"Context: {context}\n\nTask: {task}"
|
||||
|
||||
logger.info(
|
||||
"librarian_task_started",
|
||||
task=task[:100],
|
||||
has_context=bool(context),
|
||||
has_history=bool(message_history),
|
||||
)
|
||||
|
||||
try:
|
||||
result = await agent.run(
|
||||
prompt,
|
||||
message_history=message_history,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"librarian_task_completed",
|
||||
task=task[:50],
|
||||
output_length=len(result.output),
|
||||
)
|
||||
|
||||
return result.output
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"librarian_task_error",
|
||||
task=task[:50],
|
||||
error=str(e),
|
||||
exc_info=True,
|
||||
)
|
||||
return f"The Librarian encountered an error: {str(e)}"
|
||||
|
||||
|
||||
async def run_librarian_stream(
|
||||
task: str,
|
||||
context: str = "",
|
||||
message_history: Optional[list[Any]] = None,
|
||||
):
|
||||
"""
|
||||
Execute a research task with streaming output.
|
||||
|
||||
Yields text deltas as The Librarian generates the response.
|
||||
|
||||
Args:
|
||||
task: The research task or question
|
||||
context: Additional context from conversation
|
||||
message_history: Optional conversation history
|
||||
|
||||
Yields:
|
||||
str: Text deltas from the response
|
||||
|
||||
Example:
|
||||
async for delta in run_librarian_stream("Find Docker docs"):
|
||||
print(delta, end="", flush=True)
|
||||
"""
|
||||
agent = get_librarian_agent()
|
||||
|
||||
# Build prompt with context if provided
|
||||
prompt = task
|
||||
if context:
|
||||
prompt = f"Context: {context}\n\nTask: {task}"
|
||||
|
||||
logger.info(
|
||||
"librarian_stream_started",
|
||||
task=task[:100],
|
||||
)
|
||||
|
||||
try:
|
||||
async with agent.run_stream(
|
||||
prompt,
|
||||
message_history=message_history,
|
||||
) as response:
|
||||
async for delta in response.stream_text(delta=True):
|
||||
yield delta
|
||||
|
||||
logger.info("librarian_stream_completed", task=task[:50])
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"librarian_stream_error",
|
||||
task=task[:50],
|
||||
error=str(e),
|
||||
exc_info=True,
|
||||
)
|
||||
yield f"\n\nThe Librarian encountered an error: {str(e)}"
|
||||
@@ -0,0 +1,86 @@
|
||||
"""
|
||||
Librarian capability registration for the Household Registry.
|
||||
|
||||
Defines The Librarian's capabilities and registers it as a
|
||||
household member for coordination by the Steward and Tatlock.
|
||||
"""
|
||||
from src.agents.librarian.agent import get_librarian_agent
|
||||
from src.agents.librarian.tools import LIBRARIAN_TOOLS
|
||||
from src.core.household_registry import (
|
||||
HouseholdCapability,
|
||||
get_household_registry,
|
||||
)
|
||||
from src.core.logging_config import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# The Librarian's capability summary for Steward coordination
|
||||
LIBRARIAN_CAPABILITY = HouseholdCapability(
|
||||
name="librarian",
|
||||
role="The Librarian",
|
||||
category="research",
|
||||
description=(
|
||||
"Research and wiki management: can CREATE wiki pages about topics "
|
||||
"(with automatic HybridRAG research), UPDATE existing pages, "
|
||||
"SEARCH wiki/knowledge graph/web, and synthesize information. "
|
||||
"Use for: 'create a page about X', 'update wiki', 'find info on X'"
|
||||
),
|
||||
domains=[
|
||||
"research",
|
||||
"knowledge",
|
||||
"information",
|
||||
"wiki",
|
||||
"documents",
|
||||
"search",
|
||||
"synthesis",
|
||||
"create",
|
||||
"write",
|
||||
"update",
|
||||
],
|
||||
cost="medium", # Multiple API calls to library-desk
|
||||
requires_network=True, # Needs library-desk API access
|
||||
)
|
||||
|
||||
|
||||
def get_librarian_capability() -> HouseholdCapability:
|
||||
"""Get The Librarian's capability definition."""
|
||||
return LIBRARIAN_CAPABILITY
|
||||
|
||||
|
||||
def register_librarian() -> None:
|
||||
"""
|
||||
Register The Librarian with the Household Registry.
|
||||
|
||||
This makes The Librarian available for:
|
||||
- Steward recommendations (via capability summary)
|
||||
- Tatlock delegation (via agent reference)
|
||||
- Tool scoping (via tool list)
|
||||
"""
|
||||
registry = get_household_registry()
|
||||
|
||||
# Check if already registered
|
||||
if "librarian" in registry:
|
||||
logger.debug("librarian_already_registered")
|
||||
return
|
||||
|
||||
registry.register(
|
||||
name="librarian",
|
||||
capability=LIBRARIAN_CAPABILITY,
|
||||
tools=LIBRARIAN_TOOLS,
|
||||
agent=get_librarian_agent(),
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"librarian_registered",
|
||||
role=LIBRARIAN_CAPABILITY.role,
|
||||
domains=LIBRARIAN_CAPABILITY.domains,
|
||||
tool_count=len(LIBRARIAN_TOOLS),
|
||||
)
|
||||
|
||||
|
||||
def unregister_librarian() -> None:
|
||||
"""Unregister The Librarian from the Household Registry."""
|
||||
registry = get_household_registry()
|
||||
registry.unregister("librarian")
|
||||
logger.info("librarian_unregistered")
|
||||
@@ -0,0 +1,685 @@
|
||||
"""
|
||||
HTTP client for the Library-Desk API.
|
||||
|
||||
Provides async methods for all relevant library-desk endpoints:
|
||||
- HybridRAG queries
|
||||
- Wiki operations
|
||||
- Vector search
|
||||
- Knowledge graph queries
|
||||
"""
|
||||
from typing import Any, Optional
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from src.core.config import config
|
||||
from src.core.logging_config import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Response Models
|
||||
# ============================================================================
|
||||
|
||||
class WikiPage(BaseModel):
|
||||
"""Wiki page from library-desk."""
|
||||
id: int
|
||||
path: str
|
||||
title: str
|
||||
description: Optional[str] = None
|
||||
content: Optional[str] = None
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
created_at: Optional[str] = None
|
||||
updated_at: Optional[str] = None
|
||||
|
||||
|
||||
class WikiSearchResult(BaseModel):
|
||||
"""Search result from wiki search."""
|
||||
id: int
|
||||
path: str
|
||||
title: str
|
||||
description: Optional[str] = None
|
||||
locale: Optional[str] = None
|
||||
|
||||
|
||||
class VectorSearchResult(BaseModel):
|
||||
"""Result from semantic vector search."""
|
||||
page_id: int
|
||||
page_path: str
|
||||
page_title: str
|
||||
chunk_text: str
|
||||
score: float
|
||||
chunk_index: int
|
||||
|
||||
|
||||
class HybridSearchResult(BaseModel):
|
||||
"""Result from HybridRAG search."""
|
||||
source: str # "vector", "graph", "web"
|
||||
title: str
|
||||
content: str
|
||||
url: Optional[str] = None
|
||||
score: float
|
||||
page_id: Optional[int] = None
|
||||
metadata: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class HybridRAGResponse(BaseModel):
|
||||
"""Full response from HybridRAG query."""
|
||||
results: list[HybridSearchResult] = Field(default_factory=list)
|
||||
keywords: list[str] = Field(default_factory=list)
|
||||
synonyms: list[str] = Field(default_factory=list)
|
||||
related_dossiers: list[str] = Field(default_factory=list)
|
||||
formatted_context: str = ""
|
||||
search_id: Optional[str] = None
|
||||
timing: dict[str, float] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class GraphNode(BaseModel):
|
||||
"""Node from knowledge graph."""
|
||||
id: str
|
||||
labels: list[str] = Field(default_factory=list)
|
||||
properties: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class Dossier(BaseModel):
|
||||
"""A dossier (tag-based collection)."""
|
||||
name: str
|
||||
page_count: int
|
||||
|
||||
|
||||
class ResearchSummary(BaseModel):
|
||||
"""Summary of research performed during smart-create."""
|
||||
wiki_results: int = 0
|
||||
web_results: int = 0
|
||||
graph_entities: int = 0
|
||||
keywords_extracted: int = 0
|
||||
timing_ms: int = 0
|
||||
|
||||
|
||||
class EntityLinking(BaseModel):
|
||||
"""Entity linking results from smart-create."""
|
||||
forward_links: int = 0
|
||||
backward_links: int = 0
|
||||
pages_updated: int = 0
|
||||
|
||||
|
||||
class SmartCreateResponse(BaseModel):
|
||||
"""Response from smart-create wiki page endpoint."""
|
||||
page: WikiPage
|
||||
research_summary: ResearchSummary = Field(default_factory=ResearchSummary)
|
||||
sources_used: int = 0
|
||||
search_id: Optional[str] = None
|
||||
entity_linking: EntityLinking = Field(default_factory=EntityLinking)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Client
|
||||
# ============================================================================
|
||||
|
||||
class LibraryDeskClient:
|
||||
"""
|
||||
Async HTTP client for Library-Desk API.
|
||||
|
||||
Usage:
|
||||
async with LibraryDeskClient() as client:
|
||||
results = await client.hybrid_search("docker kubernetes")
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
timeout: int = 60,
|
||||
):
|
||||
"""
|
||||
Initialize the client.
|
||||
|
||||
Args:
|
||||
base_url: Library-desk API URL (defaults to config)
|
||||
api_key: API key for authentication (defaults to config)
|
||||
timeout: Request timeout in seconds
|
||||
"""
|
||||
self.base_url = base_url or str(config.LIBRARY_DESK_HOST)
|
||||
self.api_key = api_key or config.LIBRARY_DESK_API_KEY
|
||||
self.timeout = timeout
|
||||
self._client: Optional[httpx.AsyncClient] = None
|
||||
|
||||
async def __aenter__(self) -> "LibraryDeskClient":
|
||||
"""Create HTTP client on context entry."""
|
||||
headers = {}
|
||||
if self.api_key:
|
||||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||||
|
||||
self._client = httpx.AsyncClient(
|
||||
base_url=self.base_url,
|
||||
headers=headers,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Close HTTP client on context exit."""
|
||||
if self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
def _ensure_client(self) -> httpx.AsyncClient:
|
||||
"""Ensure client is initialized."""
|
||||
if self._client is None:
|
||||
raise RuntimeError(
|
||||
"Client not initialized. Use 'async with LibraryDeskClient() as client:'"
|
||||
)
|
||||
return self._client
|
||||
|
||||
# ========================================================================
|
||||
# HybridRAG
|
||||
# ========================================================================
|
||||
|
||||
async def hybrid_search(
|
||||
self,
|
||||
query: str,
|
||||
user: str = "jpmschweitzer",
|
||||
vector_limit: int = 10,
|
||||
graph_limit: int = 10,
|
||||
web_limit: int = 5,
|
||||
enable_reranking: bool = True,
|
||||
final_result_count: int = 10,
|
||||
) -> HybridRAGResponse:
|
||||
"""
|
||||
Execute HybridRAG search combining vector, graph, and web results.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
user: User identifier for multi-tenancy
|
||||
vector_limit: Max results from vector search
|
||||
graph_limit: Max results from graph search
|
||||
web_limit: Max results from web search
|
||||
enable_reranking: Whether to rerank with LLM
|
||||
final_result_count: Number of final results after fusion
|
||||
|
||||
Returns:
|
||||
HybridRAGResponse with ranked results and context
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
payload = {
|
||||
"query": query,
|
||||
"config": {
|
||||
"vector_limit": vector_limit,
|
||||
"graph_limit": graph_limit,
|
||||
"web_limit": web_limit,
|
||||
"enable_reranking": enable_reranking,
|
||||
"final_result_count": final_result_count,
|
||||
},
|
||||
}
|
||||
|
||||
logger.info("library_desk_hybrid_search", query=query, user=user)
|
||||
|
||||
response = await client.post(
|
||||
"/query/hybrid",
|
||||
json=payload,
|
||||
params={"user": user},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
|
||||
# Parse results
|
||||
results = []
|
||||
for r in data.get("results", []):
|
||||
results.append(HybridSearchResult(
|
||||
source=r.get("source", "unknown"),
|
||||
title=r.get("title", ""),
|
||||
content=r.get("content", ""),
|
||||
url=r.get("url"),
|
||||
score=r.get("score", 0.0),
|
||||
page_id=r.get("page_id"),
|
||||
metadata=r.get("metadata", {}),
|
||||
))
|
||||
|
||||
return HybridRAGResponse(
|
||||
results=results,
|
||||
keywords=data.get("keywords", []),
|
||||
synonyms=data.get("synonyms", []),
|
||||
related_dossiers=data.get("related_dossiers", []),
|
||||
formatted_context=data.get("formatted_context", ""),
|
||||
search_id=data.get("search_id"),
|
||||
timing=data.get("timing", {}),
|
||||
)
|
||||
|
||||
# ========================================================================
|
||||
# Wiki Operations
|
||||
# ========================================================================
|
||||
|
||||
async def search_wiki(
|
||||
self,
|
||||
query: str,
|
||||
user: str = "jpmschweitzer",
|
||||
limit: int = 20,
|
||||
) -> list[WikiSearchResult]:
|
||||
"""
|
||||
Search wiki pages by text.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
user: User identifier
|
||||
limit: Maximum results
|
||||
|
||||
Returns:
|
||||
List of matching wiki pages
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
logger.debug("library_desk_wiki_search", query=query, user=user)
|
||||
|
||||
response = await client.get(
|
||||
"/wiki/search",
|
||||
params={"q": query, "user": user, "limit": limit},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
return [WikiSearchResult(**r) for r in data.get("results", [])]
|
||||
|
||||
async def get_wiki_page(
|
||||
self,
|
||||
page_id: int,
|
||||
user: str = "jpmschweitzer",
|
||||
) -> WikiPage:
|
||||
"""
|
||||
Get a wiki page by ID.
|
||||
|
||||
Args:
|
||||
page_id: Page ID
|
||||
user: User identifier
|
||||
|
||||
Returns:
|
||||
WikiPage with full content
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
response = await client.get(
|
||||
f"/wiki/pages/{page_id}",
|
||||
params={"user": user},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
return WikiPage(**response.json())
|
||||
|
||||
async def list_wiki_pages(
|
||||
self,
|
||||
user: str = "jpmschweitzer",
|
||||
tag: Optional[str] = None,
|
||||
limit: int = 50,
|
||||
) -> list[WikiPage]:
|
||||
"""
|
||||
List wiki pages, optionally filtered by tag.
|
||||
|
||||
Args:
|
||||
user: User identifier
|
||||
tag: Optional tag (dossier) to filter by
|
||||
limit: Maximum pages to return
|
||||
|
||||
Returns:
|
||||
List of wiki pages
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
params: dict[str, Any] = {"user": user, "limit": limit}
|
||||
if tag:
|
||||
params["tag"] = tag
|
||||
|
||||
response = await client.get("/wiki/pages", params=params)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
return [WikiPage(**p) for p in data.get("pages", [])]
|
||||
|
||||
async def create_wiki_page(
|
||||
self,
|
||||
title: str,
|
||||
path: str,
|
||||
content: str,
|
||||
user: str = "jpmschweitzer",
|
||||
description: str = "",
|
||||
tags: Optional[list[str]] = None,
|
||||
) -> WikiPage:
|
||||
"""
|
||||
Create a new wiki page.
|
||||
|
||||
Args:
|
||||
title: Page title
|
||||
path: Page path (e.g., "/projects/my-project")
|
||||
content: Markdown content
|
||||
user: User identifier
|
||||
description: Short description
|
||||
tags: List of tags (dossiers)
|
||||
|
||||
Returns:
|
||||
Created WikiPage
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
payload = {
|
||||
"title": title,
|
||||
"path": path,
|
||||
"content": content,
|
||||
"user": user,
|
||||
"description": description,
|
||||
"tags": tags or [],
|
||||
}
|
||||
|
||||
logger.info("library_desk_create_page", title=title, path=path)
|
||||
|
||||
response = await client.post("/wiki/pages", json=payload)
|
||||
response.raise_for_status()
|
||||
|
||||
return WikiPage(**response.json())
|
||||
|
||||
async def update_wiki_page(
|
||||
self,
|
||||
page_id: int,
|
||||
user: str = "jpmschweitzer",
|
||||
content: Optional[str] = None,
|
||||
title: Optional[str] = None,
|
||||
tags: Optional[list[str]] = None,
|
||||
description: Optional[str] = None,
|
||||
) -> WikiPage:
|
||||
"""
|
||||
Update an existing wiki page.
|
||||
|
||||
Supports partial updates - only provided fields are updated.
|
||||
Automatically triggers vector re-indexing and graph extraction.
|
||||
|
||||
Args:
|
||||
page_id: ID of the page to update
|
||||
user: User identifier
|
||||
content: New content (optional)
|
||||
title: New title (optional)
|
||||
tags: New tags list (optional)
|
||||
description: New description (optional)
|
||||
|
||||
Returns:
|
||||
Updated WikiPage
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
# Build update payload with only provided fields
|
||||
update_data: dict[str, Any] = {}
|
||||
if content is not None:
|
||||
update_data["content"] = content
|
||||
if title is not None:
|
||||
update_data["title"] = title
|
||||
if tags is not None:
|
||||
update_data["tags"] = tags
|
||||
if description is not None:
|
||||
update_data["description"] = description
|
||||
|
||||
logger.info(
|
||||
"library_desk_update_page",
|
||||
page_id=page_id,
|
||||
fields=list(update_data.keys()),
|
||||
)
|
||||
|
||||
response = await client.put(
|
||||
f"/wiki/pages/{page_id}",
|
||||
params={"user": user},
|
||||
json=update_data,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
return WikiPage(**response.json())
|
||||
|
||||
async def smart_create_wiki_page(
|
||||
self,
|
||||
topic: str,
|
||||
tags: list[str],
|
||||
user: str = "jpmschweitzer",
|
||||
path: Optional[str] = None,
|
||||
include_web_research: bool = True,
|
||||
include_wiki_search: bool = True,
|
||||
) -> SmartCreateResponse:
|
||||
"""
|
||||
Create a wiki page with HybridRAG research.
|
||||
|
||||
This endpoint:
|
||||
1. Searches existing wiki, knowledge graph, and web for context
|
||||
2. Uses LLM to synthesize findings into structured content
|
||||
3. Creates the page with proper attribution
|
||||
4. Automatically links entities bidirectionally
|
||||
|
||||
Args:
|
||||
topic: The topic to research and create a page about
|
||||
tags: List of tags (dossiers) for the page
|
||||
user: User identifier
|
||||
path: Optional custom path (auto-generated from topic if not provided)
|
||||
include_web_research: Whether to include web search results
|
||||
include_wiki_search: Whether to include existing wiki content
|
||||
|
||||
Returns:
|
||||
SmartCreateResponse with page and research metadata
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"topic": topic,
|
||||
"tags": tags,
|
||||
"user": user,
|
||||
"include_web_research": include_web_research,
|
||||
"include_wiki_search": include_wiki_search,
|
||||
}
|
||||
if path is not None:
|
||||
payload["path"] = path
|
||||
|
||||
logger.info(
|
||||
"library_desk_smart_create",
|
||||
topic=topic,
|
||||
tags=tags,
|
||||
include_web=include_web_research,
|
||||
)
|
||||
|
||||
response = await client.post("/wiki/pages/smart-create", json=payload)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
|
||||
# Parse nested response
|
||||
page = WikiPage(**data.get("page", {}))
|
||||
research_summary = ResearchSummary(**data.get("research_summary", {}))
|
||||
entity_linking = EntityLinking(**data.get("entity_linking", {}))
|
||||
|
||||
return SmartCreateResponse(
|
||||
page=page,
|
||||
research_summary=research_summary,
|
||||
sources_used=data.get("sources_used", 0),
|
||||
search_id=data.get("search_id"),
|
||||
entity_linking=entity_linking,
|
||||
)
|
||||
|
||||
async def list_dossiers(
|
||||
self,
|
||||
user: str = "jpmschweitzer",
|
||||
) -> list[Dossier]:
|
||||
"""
|
||||
List all dossiers (tag collections) for a user.
|
||||
|
||||
Args:
|
||||
user: User identifier
|
||||
|
||||
Returns:
|
||||
List of dossiers with page counts
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
response = await client.get(
|
||||
"/wiki/dossiers",
|
||||
params={"user": user},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
return [Dossier(**d) for d in data.get("dossiers", [])]
|
||||
|
||||
# ========================================================================
|
||||
# Vector Search
|
||||
# ========================================================================
|
||||
|
||||
async def semantic_search(
|
||||
self,
|
||||
query: str,
|
||||
user: str = "jpmschweitzer",
|
||||
limit: int = 10,
|
||||
score_threshold: float = 0.5,
|
||||
) -> list[VectorSearchResult]:
|
||||
"""
|
||||
Perform semantic (vector) search over documents.
|
||||
|
||||
Args:
|
||||
query: Natural language query
|
||||
user: User identifier
|
||||
limit: Maximum results
|
||||
score_threshold: Minimum similarity score
|
||||
|
||||
Returns:
|
||||
List of matching document chunks with scores
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
payload = {
|
||||
"query": query,
|
||||
"user": user,
|
||||
"limit": limit,
|
||||
"score_threshold": score_threshold,
|
||||
}
|
||||
|
||||
logger.debug("library_desk_semantic_search", query=query)
|
||||
|
||||
response = await client.post("/vector/search", json=payload)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
return [VectorSearchResult(**r) for r in data.get("results", [])]
|
||||
|
||||
# ========================================================================
|
||||
# Knowledge Graph
|
||||
# ========================================================================
|
||||
|
||||
async def query_graph(
|
||||
self,
|
||||
cypher_query: str,
|
||||
user: str = "jpmschweitzer",
|
||||
parameters: Optional[dict[str, Any]] = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Execute a Cypher query on the knowledge graph.
|
||||
|
||||
Note: Query is automatically scoped to user's data.
|
||||
|
||||
Args:
|
||||
cypher_query: Cypher query string
|
||||
user: User identifier
|
||||
parameters: Query parameters
|
||||
|
||||
Returns:
|
||||
List of result records
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
payload = {
|
||||
"query": cypher_query,
|
||||
"user": user,
|
||||
"parameters": parameters or {},
|
||||
}
|
||||
|
||||
logger.debug("library_desk_graph_query", query=cypher_query[:100])
|
||||
|
||||
response = await client.post("/graph/query", json=payload)
|
||||
response.raise_for_status()
|
||||
|
||||
return response.json().get("records", [])
|
||||
|
||||
async def list_graph_nodes(
|
||||
self,
|
||||
user: str = "jpmschweitzer",
|
||||
node_type: Optional[str] = None,
|
||||
limit: int = 100,
|
||||
) -> list[GraphNode]:
|
||||
"""
|
||||
List nodes in the knowledge graph.
|
||||
|
||||
Args:
|
||||
user: User identifier
|
||||
node_type: Optional filter by type (Document, Person, Concept, etc.)
|
||||
limit: Maximum nodes
|
||||
|
||||
Returns:
|
||||
List of graph nodes
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
params: dict[str, Any] = {"user": user, "limit": limit}
|
||||
if node_type:
|
||||
params["node_type"] = node_type
|
||||
|
||||
response = await client.get("/graph/nodes", params=params)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
return [GraphNode(**n) for n in data.get("nodes", [])]
|
||||
|
||||
async def get_graph_node(
|
||||
self,
|
||||
node_id: str,
|
||||
user: str = "jpmschweitzer",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get detailed information about a graph node.
|
||||
|
||||
Args:
|
||||
node_id: Node ID
|
||||
user: User identifier
|
||||
|
||||
Returns:
|
||||
Node with relationships and connected nodes
|
||||
"""
|
||||
client = self._ensure_client()
|
||||
|
||||
response = await client.get(
|
||||
f"/graph/nodes/{node_id}",
|
||||
params={"user": user},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
return response.json()
|
||||
|
||||
# ========================================================================
|
||||
# Health Check
|
||||
# ========================================================================
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""
|
||||
Check if library-desk is healthy.
|
||||
|
||||
Returns:
|
||||
True if healthy, False otherwise
|
||||
"""
|
||||
try:
|
||||
client = self._ensure_client()
|
||||
response = await client.get("/health")
|
||||
return response.status_code == 200
|
||||
except Exception as e:
|
||||
logger.warning("library_desk_health_check_failed", error=str(e))
|
||||
return False
|
||||
|
||||
|
||||
# Global client factory
|
||||
async def get_library_client() -> LibraryDeskClient:
|
||||
"""
|
||||
Get a library-desk client instance.
|
||||
|
||||
Usage:
|
||||
async with get_library_client() as client:
|
||||
results = await client.hybrid_search("query")
|
||||
"""
|
||||
return LibraryDeskClient()
|
||||
@@ -0,0 +1,701 @@
|
||||
"""
|
||||
Librarian tools for PydanticAI agent.
|
||||
|
||||
These tools wrap the library-desk API and are registered with
|
||||
The Librarian agent for research and knowledge management tasks.
|
||||
"""
|
||||
from src.agents.librarian.client import LibraryDeskClient
|
||||
from src.core.logging_config import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# HybridRAG Search
|
||||
# ============================================================================
|
||||
|
||||
async def hybrid_search(
|
||||
query: str,
|
||||
include_web: bool = True,
|
||||
) -> str:
|
||||
"""
|
||||
Search across all knowledge sources using HybridRAG.
|
||||
|
||||
This is the primary research tool, combining:
|
||||
- Vector search (semantic similarity over documents)
|
||||
- Knowledge graph (entities and relationships)
|
||||
- Web search (current information from SearXNG)
|
||||
|
||||
Results are fused and re-ranked by relevance.
|
||||
|
||||
Args:
|
||||
query: Natural language research query
|
||||
include_web: Whether to include web results (default: True)
|
||||
|
||||
Returns:
|
||||
Formatted search results with sources and context
|
||||
|
||||
Examples:
|
||||
hybrid_search("How does Docker orchestration work with Kubernetes?")
|
||||
hybrid_search("What projects use Neo4j?", include_web=False)
|
||||
"""
|
||||
try:
|
||||
async with LibraryDeskClient() as client:
|
||||
response = await client.hybrid_search(
|
||||
query=query,
|
||||
web_limit=5 if include_web else 0,
|
||||
)
|
||||
|
||||
if not response.results:
|
||||
return f"No results found for '{query}'"
|
||||
|
||||
# Format results
|
||||
output_parts = [f"## Search Results for: {query}\n"]
|
||||
|
||||
# Add keywords if extracted
|
||||
if response.keywords:
|
||||
output_parts.append(f"**Keywords:** {', '.join(response.keywords)}")
|
||||
|
||||
# Add related dossiers
|
||||
if response.related_dossiers:
|
||||
output_parts.append(
|
||||
f"**Related Dossiers:** {', '.join(response.related_dossiers)}"
|
||||
)
|
||||
|
||||
output_parts.append("")
|
||||
|
||||
# Add results
|
||||
for i, result in enumerate(response.results, 1):
|
||||
source_icon = {
|
||||
"vector": "📄",
|
||||
"graph": "🔗",
|
||||
"web": "🌐",
|
||||
}.get(result.source, "•")
|
||||
|
||||
output_parts.append(
|
||||
f"{i}. {source_icon} **{result.title}** (score: {result.score:.2f})"
|
||||
)
|
||||
if result.url:
|
||||
output_parts.append(f" URL: {result.url}")
|
||||
output_parts.append(f" {result.content[:300]}...")
|
||||
output_parts.append("")
|
||||
|
||||
logger.info(
|
||||
"librarian_hybrid_search",
|
||||
query=query,
|
||||
result_count=len(response.results),
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("librarian_hybrid_search_error", error=str(e), query=query)
|
||||
return f"Error searching: {str(e)}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Wiki Operations
|
||||
# ============================================================================
|
||||
|
||||
async def search_wiki(
|
||||
query: str,
|
||||
limit: int = 10,
|
||||
) -> str:
|
||||
"""
|
||||
Search the personal wiki for relevant pages.
|
||||
|
||||
Performs full-text search over wiki page titles, descriptions,
|
||||
and content. Use this for finding specific documents.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
limit: Maximum results (default: 10)
|
||||
|
||||
Returns:
|
||||
List of matching wiki pages with paths and descriptions
|
||||
|
||||
Examples:
|
||||
search_wiki("docker setup guide")
|
||||
search_wiki("architecture", limit=5)
|
||||
"""
|
||||
try:
|
||||
async with LibraryDeskClient() as client:
|
||||
results = await client.search_wiki(query=query, limit=limit)
|
||||
|
||||
if not results:
|
||||
return f"No wiki pages found for '{query}'"
|
||||
|
||||
output_parts = [f"## Wiki Search: {query}\n"]
|
||||
|
||||
for i, page in enumerate(results, 1):
|
||||
output_parts.append(f"{i}. **{page.title}**")
|
||||
output_parts.append(f" Path: {page.path}")
|
||||
if page.description:
|
||||
output_parts.append(f" {page.description}")
|
||||
output_parts.append("")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("librarian_wiki_search_error", error=str(e))
|
||||
return f"Error searching wiki: {str(e)}"
|
||||
|
||||
|
||||
async def get_wiki_page(
|
||||
page_id: int,
|
||||
) -> str:
|
||||
"""
|
||||
Get the full content of a wiki page.
|
||||
|
||||
Use this after searching to read the complete content
|
||||
of a specific page.
|
||||
|
||||
Args:
|
||||
page_id: The page ID from search results
|
||||
|
||||
Returns:
|
||||
Full page content including title, path, and markdown content
|
||||
|
||||
Examples:
|
||||
get_wiki_page(42)
|
||||
"""
|
||||
try:
|
||||
async with LibraryDeskClient() as client:
|
||||
page = await client.get_wiki_page(page_id=page_id)
|
||||
|
||||
output_parts = [
|
||||
f"# {page.title}",
|
||||
f"**Path:** {page.path}",
|
||||
]
|
||||
|
||||
if page.description:
|
||||
output_parts.append(f"**Description:** {page.description}")
|
||||
|
||||
if page.tags:
|
||||
output_parts.append(f"**Tags:** {', '.join(page.tags)}")
|
||||
|
||||
output_parts.append("")
|
||||
output_parts.append(page.content or "(No content)")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("librarian_get_page_error", error=str(e), page_id=page_id)
|
||||
return f"Error getting page {page_id}: {str(e)}"
|
||||
|
||||
|
||||
async def list_dossiers() -> str:
|
||||
"""
|
||||
List all research dossiers (tag collections).
|
||||
|
||||
Dossiers are collections of wiki pages grouped by tag.
|
||||
Use this to discover what knowledge collections exist.
|
||||
|
||||
Returns:
|
||||
List of dossiers with page counts
|
||||
|
||||
Examples:
|
||||
list_dossiers()
|
||||
"""
|
||||
try:
|
||||
async with LibraryDeskClient() as client:
|
||||
dossiers = await client.list_dossiers()
|
||||
|
||||
if not dossiers:
|
||||
return "No dossiers found"
|
||||
|
||||
output_parts = ["## Research Dossiers\n"]
|
||||
|
||||
for dossier in dossiers:
|
||||
output_parts.append(
|
||||
f"- **{dossier.name}** ({dossier.page_count} pages)"
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("librarian_list_dossiers_error", error=str(e))
|
||||
return f"Error listing dossiers: {str(e)}"
|
||||
|
||||
|
||||
async def get_dossier_pages(
|
||||
dossier_name: str,
|
||||
limit: int = 20,
|
||||
) -> str:
|
||||
"""
|
||||
Get all pages in a dossier.
|
||||
|
||||
Retrieves pages tagged with the specified dossier name.
|
||||
|
||||
Args:
|
||||
dossier_name: Name of the dossier/tag
|
||||
limit: Maximum pages to return
|
||||
|
||||
Returns:
|
||||
List of pages in the dossier
|
||||
|
||||
Examples:
|
||||
get_dossier_pages("projects")
|
||||
get_dossier_pages("architecture", limit=10)
|
||||
"""
|
||||
try:
|
||||
async with LibraryDeskClient() as client:
|
||||
pages = await client.list_wiki_pages(tag=dossier_name, limit=limit)
|
||||
|
||||
if not pages:
|
||||
return f"No pages found in dossier '{dossier_name}'"
|
||||
|
||||
output_parts = [f"## Dossier: {dossier_name}\n"]
|
||||
|
||||
for page in pages:
|
||||
output_parts.append(f"- **{page.title}** ({page.path})")
|
||||
if page.description:
|
||||
output_parts.append(f" {page.description}")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("librarian_get_dossier_error", error=str(e))
|
||||
return f"Error getting dossier: {str(e)}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Semantic Search
|
||||
# ============================================================================
|
||||
|
||||
async def semantic_search(
|
||||
query: str,
|
||||
limit: int = 10,
|
||||
) -> str:
|
||||
"""
|
||||
Perform semantic (vector) search over documents.
|
||||
|
||||
Finds documents similar in meaning to the query,
|
||||
even if they don't contain the exact words.
|
||||
|
||||
Args:
|
||||
query: Natural language query
|
||||
limit: Maximum results
|
||||
|
||||
Returns:
|
||||
Matching document chunks with similarity scores
|
||||
|
||||
Examples:
|
||||
semantic_search("containerization best practices")
|
||||
semantic_search("how to handle authentication")
|
||||
"""
|
||||
try:
|
||||
async with LibraryDeskClient() as client:
|
||||
results = await client.semantic_search(query=query, limit=limit)
|
||||
|
||||
if not results:
|
||||
return f"No semantically similar content found for '{query}'"
|
||||
|
||||
output_parts = [f"## Semantic Search: {query}\n"]
|
||||
|
||||
for i, result in enumerate(results, 1):
|
||||
output_parts.append(
|
||||
f"{i}. **{result.page_title}** (score: {result.score:.2f})"
|
||||
)
|
||||
output_parts.append(f" Path: {result.page_path}")
|
||||
output_parts.append(f" {result.chunk_text[:200]}...")
|
||||
output_parts.append("")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("librarian_semantic_search_error", error=str(e))
|
||||
return f"Error in semantic search: {str(e)}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Knowledge Graph
|
||||
# ============================================================================
|
||||
|
||||
async def explore_knowledge_graph(
|
||||
entity_type: str = "Document",
|
||||
limit: int = 20,
|
||||
) -> str:
|
||||
"""
|
||||
Explore entities in the knowledge graph.
|
||||
|
||||
Lists nodes of a specific type to understand what's
|
||||
in the knowledge base.
|
||||
|
||||
Args:
|
||||
entity_type: Type of entity (Document, Person, Project, Concept, Technology)
|
||||
limit: Maximum nodes to return
|
||||
|
||||
Returns:
|
||||
List of entities with their properties
|
||||
|
||||
Examples:
|
||||
explore_knowledge_graph("Person")
|
||||
explore_knowledge_graph("Technology", limit=50)
|
||||
"""
|
||||
try:
|
||||
async with LibraryDeskClient() as client:
|
||||
nodes = await client.list_graph_nodes(
|
||||
node_type=entity_type,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
if not nodes:
|
||||
return f"No {entity_type} nodes found in knowledge graph"
|
||||
|
||||
output_parts = [f"## Knowledge Graph: {entity_type} Entities\n"]
|
||||
|
||||
for node in nodes:
|
||||
name = node.properties.get("name", node.properties.get("title", node.id))
|
||||
output_parts.append(f"- **{name}**")
|
||||
|
||||
# Show a few key properties
|
||||
for key in ["description", "url", "path"]:
|
||||
if key in node.properties:
|
||||
output_parts.append(f" {key}: {node.properties[key]}")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("librarian_explore_graph_error", error=str(e))
|
||||
return f"Error exploring knowledge graph: {str(e)}"
|
||||
|
||||
|
||||
async def find_related_entities(
|
||||
entity_name: str,
|
||||
) -> str:
|
||||
"""
|
||||
Find entities related to a given concept or entity.
|
||||
|
||||
Queries the knowledge graph to find documents, people,
|
||||
and concepts connected to the specified entity.
|
||||
|
||||
Args:
|
||||
entity_name: Name of the entity to find relationships for
|
||||
|
||||
Returns:
|
||||
Related entities and their relationships
|
||||
|
||||
Examples:
|
||||
find_related_entities("Docker")
|
||||
find_related_entities("Kubernetes")
|
||||
"""
|
||||
try:
|
||||
async with LibraryDeskClient() as client:
|
||||
# Find entities mentioning or related to the search term
|
||||
cypher = """
|
||||
MATCH (n)
|
||||
WHERE toLower(n.name) CONTAINS toLower($name)
|
||||
OR toLower(n.title) CONTAINS toLower($name)
|
||||
OPTIONAL MATCH (n)-[r]-(related)
|
||||
RETURN n, collect(DISTINCT {type: type(r), node: related})[0..10] as relationships
|
||||
LIMIT 10
|
||||
"""
|
||||
|
||||
results = await client.query_graph(
|
||||
cypher,
|
||||
parameters={"name": entity_name},
|
||||
)
|
||||
|
||||
if not results:
|
||||
return f"No entities found related to '{entity_name}'"
|
||||
|
||||
output_parts = [f"## Entities Related to: {entity_name}\n"]
|
||||
|
||||
for record in results:
|
||||
node = record.get("n", {})
|
||||
relationships = record.get("relationships", [])
|
||||
|
||||
name = node.get("name", node.get("title", "Unknown"))
|
||||
labels = node.get("labels", [])
|
||||
|
||||
output_parts.append(f"### {name}")
|
||||
if labels:
|
||||
output_parts.append(f"Type: {', '.join(labels)}")
|
||||
|
||||
if relationships:
|
||||
output_parts.append("**Connections:**")
|
||||
for rel in relationships[:5]: # Limit to 5 relationships
|
||||
rel_type = rel.get("type", "RELATED_TO")
|
||||
related_node = rel.get("node", {})
|
||||
related_name = related_node.get(
|
||||
"name", related_node.get("title", "Unknown")
|
||||
)
|
||||
output_parts.append(f" - {rel_type} → {related_name}")
|
||||
|
||||
output_parts.append("")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("librarian_find_related_error", error=str(e))
|
||||
return f"Error finding related entities: {str(e)}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Wiki Write Operations
|
||||
# ============================================================================
|
||||
|
||||
async def update_wiki_page(
|
||||
page_id: int,
|
||||
content: str | None = None,
|
||||
title: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
description: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Update an existing wiki page.
|
||||
|
||||
Supports partial updates - only specify the fields you want to change.
|
||||
Changes trigger automatic vector re-indexing and knowledge graph updates.
|
||||
|
||||
Use this for:
|
||||
- Correcting information in a page
|
||||
- Adding content to an existing page
|
||||
- Updating tags to organize pages into dossiers
|
||||
- Fixing descriptions or titles
|
||||
|
||||
Args:
|
||||
page_id: ID of the page to update (get from search_wiki results)
|
||||
content: New markdown content (optional - only if changing content)
|
||||
title: New title (optional - only if renaming)
|
||||
tags: New tag list (optional - replaces existing tags)
|
||||
description: New description (optional)
|
||||
|
||||
Returns:
|
||||
Confirmation with updated page details
|
||||
|
||||
Examples:
|
||||
update_wiki_page(42, content="# Updated Content\\n\\nNew information here")
|
||||
update_wiki_page(42, tags=["projects", "devops"]) # Add to dossiers
|
||||
update_wiki_page(42, description="Updated description")
|
||||
"""
|
||||
try:
|
||||
async with LibraryDeskClient() as client:
|
||||
page = await client.update_wiki_page(
|
||||
page_id=page_id,
|
||||
content=content,
|
||||
title=title,
|
||||
tags=tags,
|
||||
description=description,
|
||||
)
|
||||
|
||||
# Build update summary
|
||||
updated_fields = []
|
||||
if content is not None:
|
||||
updated_fields.append("content")
|
||||
if title is not None:
|
||||
updated_fields.append("title")
|
||||
if tags is not None:
|
||||
updated_fields.append("tags")
|
||||
if description is not None:
|
||||
updated_fields.append("description")
|
||||
|
||||
output_parts = [
|
||||
f"## Page Updated: {page.title}",
|
||||
f"**Path:** {page.path}",
|
||||
f"**Updated fields:** {', '.join(updated_fields)}",
|
||||
]
|
||||
|
||||
if page.tags:
|
||||
output_parts.append(f"**Tags:** {', '.join(page.tags)}")
|
||||
|
||||
output_parts.append("\n*Vector embeddings and knowledge graph will be updated automatically.*")
|
||||
|
||||
logger.info(
|
||||
"librarian_update_page",
|
||||
page_id=page_id,
|
||||
updated_fields=updated_fields,
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("librarian_update_page_error", error=str(e), page_id=page_id)
|
||||
return f"Error updating page {page_id}: {str(e)}"
|
||||
|
||||
|
||||
async def create_wiki_page(
|
||||
title: str,
|
||||
path: str,
|
||||
content: str,
|
||||
tags: list[str],
|
||||
description: str = "",
|
||||
) -> str:
|
||||
"""
|
||||
Create a new wiki page with user-provided content.
|
||||
|
||||
Use this when:
|
||||
- User provides specific content to add
|
||||
- Creating simple notes or reminders
|
||||
- The content is already known/composed
|
||||
|
||||
For research-backed pages where you need to gather information first,
|
||||
use smart_create_wiki_page instead.
|
||||
|
||||
Args:
|
||||
title: Page title
|
||||
path: Page path (e.g., "/projects/my-project" or "/notes/meeting-2024")
|
||||
content: Markdown content for the page
|
||||
tags: List of tags/dossiers (e.g., ["projects", "devops"])
|
||||
description: Short description of the page
|
||||
|
||||
Returns:
|
||||
Confirmation with created page details
|
||||
|
||||
Examples:
|
||||
create_wiki_page(
|
||||
title="SSL Renewal Reminder",
|
||||
path="/reminders/ssl-renewal",
|
||||
content="# SSL Renewal\\n\\nRemember to renew SSL cert on Jan 15",
|
||||
tags=["reminders", "infrastructure"],
|
||||
description="Certificate renewal reminder"
|
||||
)
|
||||
"""
|
||||
try:
|
||||
async with LibraryDeskClient() as client:
|
||||
page = await client.create_wiki_page(
|
||||
title=title,
|
||||
path=path,
|
||||
content=content,
|
||||
tags=tags,
|
||||
description=description,
|
||||
)
|
||||
|
||||
output_parts = [
|
||||
f"## Page Created: {page.title}",
|
||||
f"**ID:** {page.id}",
|
||||
f"**Path:** {page.path}",
|
||||
]
|
||||
|
||||
if page.tags:
|
||||
output_parts.append(f"**Tags:** {', '.join(page.tags)}")
|
||||
|
||||
if page.description:
|
||||
output_parts.append(f"**Description:** {page.description}")
|
||||
|
||||
output_parts.append("\n*Vector embeddings and knowledge graph will be updated automatically.*")
|
||||
|
||||
logger.info(
|
||||
"librarian_create_page",
|
||||
page_id=page.id,
|
||||
title=title,
|
||||
path=path,
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("librarian_create_page_error", error=str(e), title=title)
|
||||
return f"Error creating page: {str(e)}"
|
||||
|
||||
|
||||
async def smart_create_wiki_page(
|
||||
topic: str,
|
||||
tags: list[str],
|
||||
path: str | None = None,
|
||||
include_web_research: bool = True,
|
||||
include_wiki_search: bool = True,
|
||||
) -> str:
|
||||
"""
|
||||
Create a wiki page with automatic research and content synthesis.
|
||||
|
||||
This is the RECOMMENDED way to create pages about topics. It will:
|
||||
1. Search existing wiki, knowledge graph, and web for relevant information
|
||||
2. Use an LLM to synthesize findings into well-structured content
|
||||
3. Create the page with proper source attribution
|
||||
4. Automatically link entities bidirectionally in the knowledge graph
|
||||
|
||||
Use this when:
|
||||
- User says "Create a page about X"
|
||||
- User says "Add information about X to the wiki"
|
||||
- You need to research a topic before writing
|
||||
- The topic would benefit from existing knowledge context
|
||||
|
||||
Args:
|
||||
topic: The topic to research and create a page about
|
||||
tags: List of tags/dossiers for categorization
|
||||
path: Optional custom path (auto-generated from topic if not provided)
|
||||
include_web_research: Whether to search the web (default: True)
|
||||
include_wiki_search: Whether to search existing wiki (default: True)
|
||||
|
||||
Returns:
|
||||
Summary of created page with research statistics
|
||||
|
||||
Examples:
|
||||
smart_create_wiki_page("Docker Compose", tags=["technology", "devops"])
|
||||
smart_create_wiki_page("Home network architecture", tags=["infrastructure"], include_web_research=False)
|
||||
"""
|
||||
try:
|
||||
async with LibraryDeskClient() as client:
|
||||
response = await client.smart_create_wiki_page(
|
||||
topic=topic,
|
||||
tags=tags,
|
||||
path=path,
|
||||
include_web_research=include_web_research,
|
||||
include_wiki_search=include_wiki_search,
|
||||
)
|
||||
|
||||
page = response.page
|
||||
research = response.research_summary
|
||||
linking = response.entity_linking
|
||||
|
||||
output_parts = [
|
||||
f"## Page Created: {page.title}",
|
||||
f"**ID:** {page.id}",
|
||||
f"**Path:** {page.path}",
|
||||
]
|
||||
|
||||
if page.tags:
|
||||
output_parts.append(f"**Tags:** {', '.join(page.tags)}")
|
||||
|
||||
# Research summary
|
||||
output_parts.append("\n### Research Summary")
|
||||
output_parts.append(f"- **Wiki results used:** {research.wiki_results}")
|
||||
output_parts.append(f"- **Web results used:** {research.web_results}")
|
||||
output_parts.append(f"- **Graph entities found:** {research.graph_entities}")
|
||||
output_parts.append(f"- **Keywords extracted:** {research.keywords_extracted}")
|
||||
output_parts.append(f"- **Total sources:** {response.sources_used}")
|
||||
output_parts.append(f"- **Research time:** {research.timing_ms}ms")
|
||||
|
||||
# Entity linking
|
||||
if linking.forward_links > 0 or linking.backward_links > 0:
|
||||
output_parts.append("\n### Knowledge Graph Updates")
|
||||
output_parts.append(f"- **Forward links created:** {linking.forward_links}")
|
||||
output_parts.append(f"- **Backward links created:** {linking.backward_links}")
|
||||
output_parts.append(f"- **Related pages updated:** {linking.pages_updated}")
|
||||
|
||||
logger.info(
|
||||
"librarian_smart_create",
|
||||
topic=topic,
|
||||
page_id=page.id,
|
||||
sources_used=response.sources_used,
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("librarian_smart_create_error", error=str(e), topic=topic)
|
||||
return f"Error creating page about '{topic}': {str(e)}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tool Collection for Registration
|
||||
# ============================================================================
|
||||
|
||||
# All tools available to The Librarian
|
||||
LIBRARIAN_TOOLS = [
|
||||
# Research tools
|
||||
hybrid_search,
|
||||
search_wiki,
|
||||
get_wiki_page,
|
||||
list_dossiers,
|
||||
get_dossier_pages,
|
||||
semantic_search,
|
||||
explore_knowledge_graph,
|
||||
find_related_entities,
|
||||
# Write tools
|
||||
create_wiki_page,
|
||||
update_wiki_page,
|
||||
smart_create_wiki_page,
|
||||
]
|
||||
@@ -0,0 +1,201 @@
|
||||
"""
|
||||
Agent communication protocol for multi-agent coordination.
|
||||
|
||||
Defines standardized request/response formats for communication between:
|
||||
- Steward (request analysis) → Tatlock (coordination)
|
||||
- Tatlock (coordination) → Expert agents (Librarian, Developer, etc.)
|
||||
"""
|
||||
from enum import Enum
|
||||
from typing import Any, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class DelegationReason(str, Enum):
|
||||
"""Why a task is being delegated to an expert agent."""
|
||||
DOMAIN_EXPERTISE = "domain_expertise" # Expert has specialized knowledge
|
||||
TOOL_ACCESS = "tool_access" # Expert has required tools
|
||||
RESOURCE_EFFICIENCY = "resource_efficiency" # Better handled by specialist
|
||||
USER_PREFERENCE = "user_preference" # User requested specific agent
|
||||
|
||||
|
||||
class TaskComplexity(str, Enum):
|
||||
"""Complexity estimate for task execution."""
|
||||
SIMPLE = "simple" # Single tool call, fast
|
||||
MODERATE = "moderate" # Multiple steps, moderate time
|
||||
COMPLEX = "complex" # Multi-agent, significant processing
|
||||
|
||||
|
||||
class AgentRequest(BaseModel):
|
||||
"""
|
||||
Request to an expert agent.
|
||||
|
||||
Contains everything the agent needs to execute a task,
|
||||
including context from the conversation and delegation intent.
|
||||
"""
|
||||
task: str = Field(
|
||||
...,
|
||||
description="Clear description of what the agent should do"
|
||||
)
|
||||
context: str = Field(
|
||||
default="",
|
||||
description="Relevant context from conversation history"
|
||||
)
|
||||
constraints: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="Any constraints or requirements for the task"
|
||||
)
|
||||
delegation_reason: DelegationReason = Field(
|
||||
default=DelegationReason.DOMAIN_EXPERTISE,
|
||||
description="Why this task was delegated to this agent"
|
||||
)
|
||||
user_id: str = Field(
|
||||
default="default",
|
||||
description="User identifier for multi-tenant operations"
|
||||
)
|
||||
max_tokens: Optional[int] = Field(
|
||||
default=None,
|
||||
description="Optional token limit for response"
|
||||
)
|
||||
timeout_seconds: Optional[int] = Field(
|
||||
default=60,
|
||||
description="Maximum time for task completion"
|
||||
)
|
||||
|
||||
|
||||
class ToolCallRecord(BaseModel):
|
||||
"""Record of a tool call made during execution."""
|
||||
tool_name: str
|
||||
arguments: dict[str, Any]
|
||||
result: str
|
||||
duration_ms: int
|
||||
|
||||
|
||||
class AgentResponse(BaseModel):
|
||||
"""
|
||||
Response from an expert agent.
|
||||
|
||||
Contains the result, reasoning, and metadata about execution.
|
||||
"""
|
||||
success: bool = Field(
|
||||
...,
|
||||
description="Whether the task completed successfully"
|
||||
)
|
||||
result: str = Field(
|
||||
...,
|
||||
description="The main output/answer from the agent"
|
||||
)
|
||||
reasoning: str = Field(
|
||||
default="",
|
||||
description="Agent's reasoning process (for transparency)"
|
||||
)
|
||||
tool_calls: list[ToolCallRecord] = Field(
|
||||
default_factory=list,
|
||||
description="Tools called during execution"
|
||||
)
|
||||
confidence: float = Field(
|
||||
default=1.0,
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
description="Agent's confidence in the result (0.0-1.0)"
|
||||
)
|
||||
sources: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="Sources or references used"
|
||||
)
|
||||
error_message: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Error details if success=False"
|
||||
)
|
||||
duration_ms: int = Field(
|
||||
default=0,
|
||||
description="Total execution time in milliseconds"
|
||||
)
|
||||
|
||||
|
||||
class DelegationIntent(BaseModel):
|
||||
"""
|
||||
Intent to delegate a task to an expert agent.
|
||||
|
||||
Created by Tatlock when deciding to delegate, based on
|
||||
Steward's recommendations.
|
||||
"""
|
||||
target_agent: str = Field(
|
||||
...,
|
||||
description="Name of the expert agent to delegate to"
|
||||
)
|
||||
task: str = Field(
|
||||
...,
|
||||
description="Task description for the agent"
|
||||
)
|
||||
reason: DelegationReason = Field(
|
||||
default=DelegationReason.DOMAIN_EXPERTISE,
|
||||
description="Why delegating to this agent"
|
||||
)
|
||||
expected_outcome: str = Field(
|
||||
default="",
|
||||
description="What we expect the agent to provide"
|
||||
)
|
||||
priority: int = Field(
|
||||
default=1,
|
||||
ge=1,
|
||||
le=10,
|
||||
description="Priority (1=highest, 10=lowest)"
|
||||
)
|
||||
depends_on: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="Other delegation IDs this depends on (for sequencing)"
|
||||
)
|
||||
|
||||
|
||||
class CoordinationResult(BaseModel):
|
||||
"""
|
||||
Result of multi-agent coordination.
|
||||
|
||||
Aggregates results from multiple expert agents into
|
||||
a single coherent response.
|
||||
"""
|
||||
final_response: str = Field(
|
||||
...,
|
||||
description="Synthesized response from all agents"
|
||||
)
|
||||
agent_responses: dict[str, AgentResponse] = Field(
|
||||
default_factory=dict,
|
||||
description="Individual responses keyed by agent name"
|
||||
)
|
||||
delegation_intents: list[DelegationIntent] = Field(
|
||||
default_factory=list,
|
||||
description="All delegations that were executed"
|
||||
)
|
||||
total_duration_ms: int = Field(
|
||||
default=0,
|
||||
description="Total coordination time"
|
||||
)
|
||||
agents_consulted: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="Names of agents that contributed"
|
||||
)
|
||||
|
||||
|
||||
class AgentError(Exception):
|
||||
"""Base exception for agent errors."""
|
||||
|
||||
def __init__(self, message: str, agent_name: str = "unknown"):
|
||||
self.message = message
|
||||
self.agent_name = agent_name
|
||||
super().__init__(f"[{agent_name}] {message}")
|
||||
|
||||
|
||||
class AgentTimeoutError(AgentError):
|
||||
"""Agent execution timed out."""
|
||||
pass
|
||||
|
||||
|
||||
class AgentUnavailableError(AgentError):
|
||||
"""Agent is not available or registered."""
|
||||
pass
|
||||
|
||||
|
||||
class DelegationError(AgentError):
|
||||
"""Error during task delegation."""
|
||||
pass
|
||||
@@ -48,7 +48,7 @@ AVAILABLE HOUSEHOLD CAPABILITIES:
|
||||
{capabilities_text}
|
||||
|
||||
YOUR TASK:
|
||||
Analyze the user's query and recommend which capabilities are needed.
|
||||
Analyze the user's query and recommend which capabilities are needed, with specific delegation instructions.
|
||||
{history_text}
|
||||
|
||||
USER QUERY: {query}
|
||||
@@ -57,18 +57,22 @@ GUIDELINES:
|
||||
- Be conservative - only recommend truly necessary capabilities
|
||||
- Simple greetings/chat → no capabilities needed (conversational response only)
|
||||
- Math/calculations → tatlock_core
|
||||
- Web searches → tatlock_core
|
||||
- Quick web searches → tatlock_core
|
||||
- Time/date queries → tatlock_core
|
||||
- Wiki creation ("create a page about X", "add X to wiki") → librarian with smart_create
|
||||
- Wiki updates ("update the page", "add to dossier") → librarian with update
|
||||
- Research queries ("find info", "what do we know about", "search for") → librarian with hybrid_search
|
||||
- In-depth research, knowledge synthesis, document lookup → librarian with hybrid_search
|
||||
- If conversation history is relevant, note which previous turns matter
|
||||
- Assess complexity: simple (1 tool), moderate (2-3 tools), complex (multiple steps)
|
||||
- If capabilities are missing, mention what would be needed
|
||||
|
||||
RESPOND WITH 2-3 SENTENCES:
|
||||
1. Which capabilities (if any) are needed and why
|
||||
2. Complexity assessment (simple/moderate/complex)
|
||||
3. Any conversation context or missing capabilities
|
||||
RESPOND WITH 2-4 SENTENCES:
|
||||
1. Capability needed: name the capability and the specific action (e.g., "librarian to create a wiki page about CI/CD using smart_create")
|
||||
2. Reason: brief explanation of why this capability handles the request
|
||||
3. Complexity assessment (simple/moderate/complex)
|
||||
4. Any relevant conversation context
|
||||
|
||||
Use capability names in your response (e.g., "tatlock_core for calculations").
|
||||
Be specific about what Tatlock should delegate - include the action verb (create, update, search, etc.).
|
||||
Plain text only - no JSON, no special formatting."""
|
||||
|
||||
|
||||
|
||||
+38
-1
@@ -4,11 +4,34 @@ Following best practice of splitting config across domains.
|
||||
"""
|
||||
from enum import Enum
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import Field, HttpUrl
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
def _get_version_from_pyproject() -> str:
|
||||
"""
|
||||
Load version from pyproject.toml.
|
||||
|
||||
Falls back to "unknown" if file cannot be read.
|
||||
"""
|
||||
try:
|
||||
# Find pyproject.toml relative to this file
|
||||
config_dir = Path(__file__).parent
|
||||
pyproject_path = config_dir.parent.parent / "pyproject.toml"
|
||||
|
||||
if pyproject_path.exists():
|
||||
content = pyproject_path.read_text()
|
||||
for line in content.splitlines():
|
||||
if line.strip().startswith("version"):
|
||||
# Parse: version = "1.0.0"
|
||||
return line.split("=", 1)[1].strip().strip('"').strip("'")
|
||||
except Exception:
|
||||
pass
|
||||
return "unknown"
|
||||
|
||||
|
||||
class Environment(str, Enum):
|
||||
"""Application environment."""
|
||||
DEVELOPMENT = "development"
|
||||
@@ -32,7 +55,7 @@ class Config(BaseSettings):
|
||||
|
||||
# Application
|
||||
APP_NAME: str = "OpenAI-Compatible API"
|
||||
APP_VERSION: str = "0.2.5"
|
||||
APP_VERSION: str = Field(default_factory=_get_version_from_pyproject)
|
||||
ENVIRONMENT: Environment = Environment.DEVELOPMENT
|
||||
DEBUG: bool = Field(default=False, description="Debug mode")
|
||||
|
||||
@@ -87,6 +110,20 @@ class Config(BaseSettings):
|
||||
description="Redis connection timeout in seconds"
|
||||
)
|
||||
|
||||
# Library-Desk Configuration (The Librarian backend)
|
||||
LIBRARY_DESK_HOST: HttpUrl = Field(
|
||||
default="http://localhost:8089",
|
||||
description="Library-Desk API URL"
|
||||
)
|
||||
LIBRARY_DESK_API_KEY: str = Field(
|
||||
default="",
|
||||
description="API key for Library-Desk authentication"
|
||||
)
|
||||
LIBRARY_DESK_TIMEOUT: int = Field(
|
||||
default=60,
|
||||
description="Library-Desk request timeout in seconds"
|
||||
)
|
||||
|
||||
# Logging
|
||||
LOG_LEVEL: str = Field(default="INFO", description="Logging level")
|
||||
ENABLE_BENCHMARKS: bool = Field(default=True, description="Enable performance benchmarking")
|
||||
|
||||
@@ -4,6 +4,7 @@ Request preprocessing pipeline.
|
||||
Analyzes requests via the Steward and creates scoped toolsets for Tatlock.
|
||||
"""
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any, Optional
|
||||
|
||||
from src.agents.steward import analyze_request, format_steward_note
|
||||
@@ -14,6 +15,23 @@ from src.core.logging_config import get_logger
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _inject_temporal_context(request: str) -> str:
|
||||
"""
|
||||
Append current time context to user request.
|
||||
|
||||
Provides Tatlock with temporal awareness for time-sensitive queries.
|
||||
|
||||
Args:
|
||||
request: Original user request
|
||||
|
||||
Returns:
|
||||
Request with appended time context
|
||||
"""
|
||||
now = datetime.now()
|
||||
time_str = now.strftime("%Y-%m-%d %H:%M")
|
||||
return f"{request}\n\n[Current time: {time_str}]"
|
||||
|
||||
|
||||
@dataclass
|
||||
class EnrichedRequest:
|
||||
"""
|
||||
@@ -65,6 +83,9 @@ async def preprocess_request(
|
||||
>>> print(len(enriched.scoped_tools))
|
||||
5 # All tatlock_core tools
|
||||
"""
|
||||
# Inject temporal context for time-aware processing
|
||||
enriched_request = _inject_temporal_context(user_request)
|
||||
|
||||
logger.info(
|
||||
"preprocessing_request",
|
||||
request_preview=user_request[:100],
|
||||
@@ -74,7 +95,7 @@ async def preprocess_request(
|
||||
|
||||
# Call Steward with full conversation history
|
||||
recommendation = await analyze_request(
|
||||
user_request,
|
||||
enriched_request,
|
||||
conversation_history=conversation_history,
|
||||
conversation_id=conversation_id,
|
||||
)
|
||||
@@ -97,7 +118,7 @@ async def preprocess_request(
|
||||
)
|
||||
|
||||
return EnrichedRequest(
|
||||
original_request=user_request,
|
||||
original_request=enriched_request,
|
||||
steward_note=steward_note,
|
||||
scoped_tools=scoped_tools,
|
||||
recommendation=recommendation,
|
||||
|
||||
+12
-5
@@ -5,6 +5,7 @@ Handles initialization of household registry and other startup tasks.
|
||||
This module should be called during application startup to register
|
||||
all household members.
|
||||
"""
|
||||
from src.agents.librarian import register_librarian
|
||||
from src.agents.tatlock_core import TATLOCK_CORE_CAPABILITY, tatlock_core_tools
|
||||
from src.core.household_registry import get_household_registry
|
||||
from src.core.logging_config import get_logger
|
||||
@@ -21,11 +22,7 @@ def register_household_members():
|
||||
|
||||
Currently registers:
|
||||
- tatlock_core: Butler's core tools (calculator, datetime, web search)
|
||||
|
||||
Future phases will add:
|
||||
- librarian: Research and knowledge management
|
||||
- developer: Software development assistance
|
||||
- etc.
|
||||
- librarian: Research and knowledge management (Phase 3)
|
||||
"""
|
||||
registry = get_household_registry()
|
||||
|
||||
@@ -45,6 +42,16 @@ def register_household_members():
|
||||
tool_count=len(tatlock_core_tools),
|
||||
)
|
||||
|
||||
# Register The Librarian (Phase 3)
|
||||
try:
|
||||
register_librarian()
|
||||
except Exception as e:
|
||||
# Don't fail startup if Librarian registration fails
|
||||
logger.warning(
|
||||
"librarian_registration_failed",
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"household_registration_complete",
|
||||
total_members=len(registry),
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for The Librarian agent."""
|
||||
@@ -0,0 +1,129 @@
|
||||
"""
|
||||
Tests for Librarian capability registration.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from src.agents.librarian.capability import (
|
||||
LIBRARIAN_CAPABILITY,
|
||||
get_librarian_capability,
|
||||
register_librarian,
|
||||
unregister_librarian,
|
||||
)
|
||||
from src.core.household_registry import HouseholdCapability
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestLibrarianCapability:
|
||||
"""Tests for the Librarian capability definition."""
|
||||
|
||||
def test_capability_is_household_capability(self):
|
||||
"""Test capability is correct type."""
|
||||
assert isinstance(LIBRARIAN_CAPABILITY, HouseholdCapability)
|
||||
|
||||
def test_capability_name(self):
|
||||
"""Test capability has correct name."""
|
||||
assert LIBRARIAN_CAPABILITY.name == "librarian"
|
||||
|
||||
def test_capability_role(self):
|
||||
"""Test capability has correct role."""
|
||||
assert LIBRARIAN_CAPABILITY.role == "The Librarian"
|
||||
|
||||
def test_capability_category(self):
|
||||
"""Test capability is in research category."""
|
||||
assert LIBRARIAN_CAPABILITY.category == "research"
|
||||
|
||||
def test_capability_domains(self):
|
||||
"""Test capability covers expected domains."""
|
||||
domains = LIBRARIAN_CAPABILITY.domains
|
||||
|
||||
assert "research" in domains
|
||||
assert "knowledge" in domains
|
||||
assert "wiki" in domains
|
||||
assert "search" in domains
|
||||
|
||||
def test_capability_requires_network(self):
|
||||
"""Test capability requires network access."""
|
||||
assert LIBRARIAN_CAPABILITY.requires_network is True
|
||||
|
||||
def test_get_librarian_capability(self):
|
||||
"""Test getter returns same capability."""
|
||||
cap = get_librarian_capability()
|
||||
|
||||
assert cap is LIBRARIAN_CAPABILITY
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestLibrarianRegistration:
|
||||
"""Tests for Librarian registration functions."""
|
||||
|
||||
def test_register_librarian(self):
|
||||
"""Test registering librarian with registry."""
|
||||
mock_registry = MagicMock()
|
||||
mock_registry.__contains__ = MagicMock(return_value=False)
|
||||
|
||||
with patch(
|
||||
"src.agents.librarian.capability.get_household_registry",
|
||||
return_value=mock_registry,
|
||||
):
|
||||
with patch(
|
||||
"src.agents.librarian.capability.get_librarian_agent"
|
||||
) as mock_get_agent:
|
||||
mock_agent = MagicMock()
|
||||
mock_get_agent.return_value = mock_agent
|
||||
|
||||
register_librarian()
|
||||
|
||||
mock_registry.register.assert_called_once()
|
||||
call_kwargs = mock_registry.register.call_args[1]
|
||||
|
||||
assert call_kwargs["name"] == "librarian"
|
||||
assert call_kwargs["capability"] is LIBRARIAN_CAPABILITY
|
||||
assert call_kwargs["agent"] is mock_agent
|
||||
|
||||
def test_register_librarian_already_registered(self):
|
||||
"""Test registering when already registered does nothing."""
|
||||
mock_registry = MagicMock()
|
||||
mock_registry.__contains__ = MagicMock(return_value=True)
|
||||
|
||||
with patch(
|
||||
"src.agents.librarian.capability.get_household_registry",
|
||||
return_value=mock_registry,
|
||||
):
|
||||
register_librarian()
|
||||
|
||||
# Should not call register since already registered
|
||||
mock_registry.register.assert_not_called()
|
||||
|
||||
def test_unregister_librarian(self):
|
||||
"""Test unregistering librarian from registry."""
|
||||
mock_registry = MagicMock()
|
||||
|
||||
with patch(
|
||||
"src.agents.librarian.capability.get_household_registry",
|
||||
return_value=mock_registry,
|
||||
):
|
||||
unregister_librarian()
|
||||
|
||||
mock_registry.unregister.assert_called_once_with("librarian")
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCapabilityDescription:
|
||||
"""Tests for capability description."""
|
||||
|
||||
def test_description_mentions_wiki_capabilities(self):
|
||||
"""Test description mentions wiki read/write capabilities."""
|
||||
desc = LIBRARIAN_CAPABILITY.description.lower()
|
||||
assert "create" in desc
|
||||
assert "update" in desc
|
||||
assert "search" in desc
|
||||
|
||||
def test_description_mentions_search(self):
|
||||
"""Test description mentions search capability."""
|
||||
assert "search" in LIBRARIAN_CAPABILITY.description.lower()
|
||||
|
||||
def test_description_mentions_wiki(self):
|
||||
"""Test description mentions wiki access."""
|
||||
assert "wiki" in LIBRARIAN_CAPABILITY.description.lower()
|
||||
@@ -0,0 +1,598 @@
|
||||
"""
|
||||
Tests for the Library-Desk HTTP client.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
import httpx
|
||||
|
||||
from src.agents.librarian.client import (
|
||||
LibraryDeskClient,
|
||||
HybridRAGResponse,
|
||||
HybridSearchResult,
|
||||
WikiPage,
|
||||
WikiSearchResult,
|
||||
VectorSearchResult,
|
||||
GraphNode,
|
||||
Dossier,
|
||||
SmartCreateResponse,
|
||||
ResearchSummary,
|
||||
EntityLinking,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_httpx_client():
|
||||
"""Create a mock httpx client."""
|
||||
return AsyncMock(spec=httpx.AsyncClient)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client_with_mock(mock_httpx_client):
|
||||
"""Create a LibraryDeskClient with mocked httpx client."""
|
||||
client = LibraryDeskClient(
|
||||
base_url="http://test:8089",
|
||||
api_key="test-key",
|
||||
)
|
||||
client._client = mock_httpx_client
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestLibraryDeskClientInit:
|
||||
"""Tests for client initialization."""
|
||||
|
||||
def test_default_initialization(self):
|
||||
"""Test client initializes with defaults from config."""
|
||||
client = LibraryDeskClient()
|
||||
|
||||
assert client.base_url is not None
|
||||
assert client.timeout == 60
|
||||
assert client._client is None
|
||||
|
||||
def test_custom_initialization(self):
|
||||
"""Test client with custom parameters."""
|
||||
client = LibraryDeskClient(
|
||||
base_url="http://custom:9000",
|
||||
api_key="my-api-key",
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
assert client.base_url == "http://custom:9000"
|
||||
assert client.api_key == "my-api-key"
|
||||
assert client.timeout == 120
|
||||
|
||||
def test_ensure_client_not_initialized(self):
|
||||
"""Test _ensure_client raises when not in context."""
|
||||
client = LibraryDeskClient()
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
client._ensure_client()
|
||||
|
||||
assert "not initialized" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestContextManager:
|
||||
"""Tests for async context manager."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_manager_creates_client(self):
|
||||
"""Test context manager creates httpx client."""
|
||||
async with LibraryDeskClient(
|
||||
base_url="http://test:8089",
|
||||
api_key="test-key",
|
||||
) as client:
|
||||
assert client._client is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_manager_closes_client(self):
|
||||
"""Test context manager closes client on exit."""
|
||||
client = LibraryDeskClient(base_url="http://test:8089")
|
||||
|
||||
async with client:
|
||||
assert client._client is not None
|
||||
|
||||
# After exit, client should be None
|
||||
assert client._client is None
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestHybridSearch:
|
||||
"""Tests for hybrid search."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hybrid_search_success(self, client_with_mock, mock_httpx_client):
|
||||
"""Test successful hybrid search."""
|
||||
# Mock response
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"results": [
|
||||
{
|
||||
"source": "vector",
|
||||
"title": "Docker Guide",
|
||||
"content": "Docker networking basics...",
|
||||
"score": 0.95,
|
||||
"page_id": 123,
|
||||
}
|
||||
],
|
||||
"keywords": ["docker", "networking"],
|
||||
"synonyms": ["container"],
|
||||
"formatted_context": "Context here",
|
||||
"timing": {"total": 1.5},
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.post.return_value = mock_response
|
||||
|
||||
result = await client_with_mock.hybrid_search(
|
||||
query="Docker networking",
|
||||
user="testuser",
|
||||
)
|
||||
|
||||
assert isinstance(result, HybridRAGResponse)
|
||||
assert len(result.results) == 1
|
||||
assert result.results[0].title == "Docker Guide"
|
||||
assert result.results[0].source == "vector"
|
||||
assert "docker" in result.keywords
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hybrid_search_empty_results(
|
||||
self, client_with_mock, mock_httpx_client
|
||||
):
|
||||
"""Test hybrid search with no results."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"results": [],
|
||||
"keywords": [],
|
||||
"formatted_context": "",
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.post.return_value = mock_response
|
||||
|
||||
result = await client_with_mock.hybrid_search("nonexistent query")
|
||||
|
||||
assert len(result.results) == 0
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestWikiOperations:
|
||||
"""Tests for wiki operations."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_wiki(self, client_with_mock, mock_httpx_client):
|
||||
"""Test wiki search."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"results": [
|
||||
{
|
||||
"id": 1,
|
||||
"path": "/docs/docker",
|
||||
"title": "Docker Documentation",
|
||||
"description": "Docker docs",
|
||||
}
|
||||
]
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.get.return_value = mock_response
|
||||
|
||||
results = await client_with_mock.search_wiki("docker")
|
||||
|
||||
assert len(results) == 1
|
||||
assert isinstance(results[0], WikiSearchResult)
|
||||
assert results[0].title == "Docker Documentation"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_wiki_page(self, client_with_mock, mock_httpx_client):
|
||||
"""Test getting a wiki page."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"id": 123,
|
||||
"path": "/docs/docker",
|
||||
"title": "Docker Guide",
|
||||
"content": "# Docker\n\nFull content here...",
|
||||
"tags": ["docker", "devops"],
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.get.return_value = mock_response
|
||||
|
||||
page = await client_with_mock.get_wiki_page(123)
|
||||
|
||||
assert isinstance(page, WikiPage)
|
||||
assert page.id == 123
|
||||
assert page.title == "Docker Guide"
|
||||
assert "docker" in page.tags
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_wiki_pages(self, client_with_mock, mock_httpx_client):
|
||||
"""Test listing wiki pages."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"pages": [
|
||||
{"id": 1, "path": "/page1", "title": "Page 1"},
|
||||
{"id": 2, "path": "/page2", "title": "Page 2"},
|
||||
]
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.get.return_value = mock_response
|
||||
|
||||
pages = await client_with_mock.list_wiki_pages()
|
||||
|
||||
assert len(pages) == 2
|
||||
assert pages[0].title == "Page 1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_dossiers(self, client_with_mock, mock_httpx_client):
|
||||
"""Test listing dossiers."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"dossiers": [
|
||||
{"name": "docker", "page_count": 10},
|
||||
{"name": "kubernetes", "page_count": 5},
|
||||
]
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.get.return_value = mock_response
|
||||
|
||||
dossiers = await client_with_mock.list_dossiers()
|
||||
|
||||
assert len(dossiers) == 2
|
||||
assert isinstance(dossiers[0], Dossier)
|
||||
assert dossiers[0].name == "docker"
|
||||
assert dossiers[0].page_count == 10
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSemanticSearch:
|
||||
"""Tests for semantic/vector search."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_search(self, client_with_mock, mock_httpx_client):
|
||||
"""Test semantic search."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"results": [
|
||||
{
|
||||
"page_id": 1,
|
||||
"page_path": "/docs/networking",
|
||||
"page_title": "Networking Guide",
|
||||
"chunk_text": "Container networking...",
|
||||
"score": 0.92,
|
||||
"chunk_index": 0,
|
||||
}
|
||||
]
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.post.return_value = mock_response
|
||||
|
||||
results = await client_with_mock.semantic_search("container networking")
|
||||
|
||||
assert len(results) == 1
|
||||
assert isinstance(results[0], VectorSearchResult)
|
||||
assert results[0].score == 0.92
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestGraphOperations:
|
||||
"""Tests for knowledge graph operations."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_graph(self, client_with_mock, mock_httpx_client):
|
||||
"""Test executing a Cypher query."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"records": [
|
||||
{"name": "Docker", "type": "Technology"},
|
||||
{"name": "Kubernetes", "type": "Technology"},
|
||||
]
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.post.return_value = mock_response
|
||||
|
||||
records = await client_with_mock.query_graph(
|
||||
"MATCH (n:Technology) RETURN n.name as name, n.type as type"
|
||||
)
|
||||
|
||||
assert len(records) == 2
|
||||
assert records[0]["name"] == "Docker"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_graph_nodes(self, client_with_mock, mock_httpx_client):
|
||||
"""Test listing graph nodes."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"nodes": [
|
||||
{
|
||||
"id": "node1",
|
||||
"labels": ["Technology"],
|
||||
"properties": {"name": "Docker"},
|
||||
}
|
||||
]
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.get.return_value = mock_response
|
||||
|
||||
nodes = await client_with_mock.list_graph_nodes()
|
||||
|
||||
assert len(nodes) == 1
|
||||
assert isinstance(nodes[0], GraphNode)
|
||||
assert nodes[0].id == "node1"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestHealthCheck:
|
||||
"""Tests for health check."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_healthy(self, client_with_mock, mock_httpx_client):
|
||||
"""Test health check returns true when healthy."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_httpx_client.get.return_value = mock_response
|
||||
|
||||
result = await client_with_mock.health_check()
|
||||
|
||||
assert result is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_unhealthy(self, client_with_mock, mock_httpx_client):
|
||||
"""Test health check returns false on error."""
|
||||
mock_httpx_client.get.side_effect = httpx.ConnectError("Connection refused")
|
||||
|
||||
result = await client_with_mock.health_check()
|
||||
|
||||
assert result is False
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestResponseModels:
|
||||
"""Tests for response model validation."""
|
||||
|
||||
def test_wiki_page_model(self):
|
||||
"""Test WikiPage model."""
|
||||
page = WikiPage(
|
||||
id=1,
|
||||
path="/test",
|
||||
title="Test Page",
|
||||
content="Content here",
|
||||
tags=["tag1"],
|
||||
)
|
||||
|
||||
assert page.id == 1
|
||||
assert page.title == "Test Page"
|
||||
|
||||
def test_wiki_page_optional_fields(self):
|
||||
"""Test WikiPage with minimal fields."""
|
||||
page = WikiPage(id=1, path="/test", title="Test")
|
||||
|
||||
assert page.content is None
|
||||
assert page.tags == []
|
||||
|
||||
def test_hybrid_search_result_model(self):
|
||||
"""Test HybridSearchResult model."""
|
||||
result = HybridSearchResult(
|
||||
source="vector",
|
||||
title="Title",
|
||||
content="Content",
|
||||
score=0.9,
|
||||
)
|
||||
|
||||
assert result.source == "vector"
|
||||
assert result.url is None
|
||||
assert result.metadata == {}
|
||||
|
||||
def test_vector_search_result_model(self):
|
||||
"""Test VectorSearchResult model."""
|
||||
result = VectorSearchResult(
|
||||
page_id=1,
|
||||
page_path="/doc",
|
||||
page_title="Doc",
|
||||
chunk_text="Text chunk",
|
||||
score=0.85,
|
||||
chunk_index=0,
|
||||
)
|
||||
|
||||
assert result.score == 0.85
|
||||
assert result.chunk_index == 0
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestUpdateWikiPage:
|
||||
"""Tests for update_wiki_page method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_wiki_page_content(self, client_with_mock, mock_httpx_client):
|
||||
"""Test updating wiki page content."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"id": 42,
|
||||
"path": "/docs/test",
|
||||
"title": "Test Page",
|
||||
"content": "# Updated\n\nNew content",
|
||||
"tags": ["test"],
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.put.return_value = mock_response
|
||||
|
||||
page = await client_with_mock.update_wiki_page(
|
||||
page_id=42,
|
||||
content="# Updated\n\nNew content",
|
||||
)
|
||||
|
||||
assert isinstance(page, WikiPage)
|
||||
assert page.id == 42
|
||||
assert "Updated" in page.content
|
||||
mock_httpx_client.put.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_wiki_page_tags_only(self, client_with_mock, mock_httpx_client):
|
||||
"""Test updating only tags (partial update)."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"id": 42,
|
||||
"path": "/docs/test",
|
||||
"title": "Test Page",
|
||||
"tags": ["projects", "devops"],
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.put.return_value = mock_response
|
||||
|
||||
page = await client_with_mock.update_wiki_page(
|
||||
page_id=42,
|
||||
tags=["projects", "devops"],
|
||||
)
|
||||
|
||||
assert page.tags == ["projects", "devops"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_wiki_page_multiple_fields(
|
||||
self, client_with_mock, mock_httpx_client
|
||||
):
|
||||
"""Test updating multiple fields at once."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"id": 42,
|
||||
"path": "/docs/test",
|
||||
"title": "New Title",
|
||||
"description": "New description",
|
||||
"tags": ["updated"],
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.put.return_value = mock_response
|
||||
|
||||
page = await client_with_mock.update_wiki_page(
|
||||
page_id=42,
|
||||
title="New Title",
|
||||
description="New description",
|
||||
tags=["updated"],
|
||||
)
|
||||
|
||||
assert page.title == "New Title"
|
||||
assert page.description == "New description"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSmartCreateWikiPage:
|
||||
"""Tests for smart_create_wiki_page method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_smart_create_basic(self, client_with_mock, mock_httpx_client):
|
||||
"""Test basic smart create."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"page": {
|
||||
"id": 123,
|
||||
"path": "/users/test/technology/docker-compose",
|
||||
"title": "Docker Compose",
|
||||
"content": "# Docker Compose\n\nContent...",
|
||||
"tags": ["technology", "devops"],
|
||||
},
|
||||
"research_summary": {
|
||||
"wiki_results": 3,
|
||||
"web_results": 8,
|
||||
"graph_entities": 5,
|
||||
"keywords_extracted": 12,
|
||||
"timing_ms": 4500,
|
||||
},
|
||||
"sources_used": 11,
|
||||
"search_id": "uuid-123",
|
||||
"entity_linking": {
|
||||
"forward_links": 5,
|
||||
"backward_links": 3,
|
||||
"pages_updated": 2,
|
||||
},
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.post.return_value = mock_response
|
||||
|
||||
result = await client_with_mock.smart_create_wiki_page(
|
||||
topic="Docker Compose",
|
||||
tags=["technology", "devops"],
|
||||
)
|
||||
|
||||
assert isinstance(result, SmartCreateResponse)
|
||||
assert result.page.id == 123
|
||||
assert result.page.title == "Docker Compose"
|
||||
assert result.sources_used == 11
|
||||
assert result.research_summary.wiki_results == 3
|
||||
assert result.research_summary.web_results == 8
|
||||
assert result.entity_linking.forward_links == 5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_smart_create_with_options(self, client_with_mock, mock_httpx_client):
|
||||
"""Test smart create with custom options."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"page": {
|
||||
"id": 456,
|
||||
"path": "/custom/path",
|
||||
"title": "Custom Topic",
|
||||
"tags": ["custom"],
|
||||
},
|
||||
"research_summary": {
|
||||
"wiki_results": 5,
|
||||
"web_results": 0, # Web disabled
|
||||
"timing_ms": 2000,
|
||||
},
|
||||
"sources_used": 5,
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_httpx_client.post.return_value = mock_response
|
||||
|
||||
result = await client_with_mock.smart_create_wiki_page(
|
||||
topic="Custom Topic",
|
||||
tags=["custom"],
|
||||
path="/custom/path",
|
||||
include_web_research=False,
|
||||
)
|
||||
|
||||
assert result.page.path == "/custom/path"
|
||||
assert result.research_summary.web_results == 0
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestNewResponseModels:
|
||||
"""Tests for new response models."""
|
||||
|
||||
def test_research_summary_model(self):
|
||||
"""Test ResearchSummary model."""
|
||||
summary = ResearchSummary(
|
||||
wiki_results=3,
|
||||
web_results=5,
|
||||
graph_entities=2,
|
||||
keywords_extracted=10,
|
||||
timing_ms=3000,
|
||||
)
|
||||
|
||||
assert summary.wiki_results == 3
|
||||
assert summary.timing_ms == 3000
|
||||
|
||||
def test_research_summary_defaults(self):
|
||||
"""Test ResearchSummary default values."""
|
||||
summary = ResearchSummary()
|
||||
|
||||
assert summary.wiki_results == 0
|
||||
assert summary.timing_ms == 0
|
||||
|
||||
def test_entity_linking_model(self):
|
||||
"""Test EntityLinking model."""
|
||||
linking = EntityLinking(
|
||||
forward_links=5,
|
||||
backward_links=3,
|
||||
pages_updated=2,
|
||||
)
|
||||
|
||||
assert linking.forward_links == 5
|
||||
assert linking.pages_updated == 2
|
||||
|
||||
def test_smart_create_response_model(self):
|
||||
"""Test SmartCreateResponse model."""
|
||||
page = WikiPage(id=1, path="/test", title="Test")
|
||||
response = SmartCreateResponse(
|
||||
page=page,
|
||||
sources_used=10,
|
||||
search_id="uuid-456",
|
||||
)
|
||||
|
||||
assert response.page.id == 1
|
||||
assert response.sources_used == 10
|
||||
assert response.search_id == "uuid-456"
|
||||
@@ -0,0 +1,339 @@
|
||||
"""
|
||||
Tests for multi-agent coordination engine.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from src.agents.coordination import (
|
||||
CoordinationEngine,
|
||||
get_coordination_engine,
|
||||
delegate_to_librarian,
|
||||
)
|
||||
from src.agents.protocol import (
|
||||
AgentResponse,
|
||||
AgentUnavailableError,
|
||||
DelegationIntent,
|
||||
DelegationReason,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def coordination_engine():
|
||||
"""Create a fresh coordination engine for testing."""
|
||||
return CoordinationEngine()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_registry():
|
||||
"""Mock the household registry."""
|
||||
with patch("src.agents.coordination.get_household_registry") as mock:
|
||||
registry = MagicMock()
|
||||
mock.return_value = registry
|
||||
yield registry
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def librarian_intent():
|
||||
"""Create a standard librarian delegation intent."""
|
||||
return DelegationIntent(
|
||||
target_agent="librarian",
|
||||
task="Find information about Docker networking",
|
||||
reason=DelegationReason.DOMAIN_EXPERTISE,
|
||||
expected_outcome="Documentation and examples",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCoordinationEngine:
|
||||
"""Tests for CoordinationEngine class."""
|
||||
|
||||
def test_initialization(self, coordination_engine):
|
||||
"""Test engine initializes correctly."""
|
||||
assert coordination_engine is not None
|
||||
assert coordination_engine.registry is not None
|
||||
|
||||
def test_get_available_agents_empty(self, mock_registry):
|
||||
"""Test getting available agents when none have agents."""
|
||||
mock_registry.list_members.return_value = ["tatlock_core"]
|
||||
mock_member = MagicMock()
|
||||
mock_member.agent = None # No agent
|
||||
mock_registry.get_member.return_value = mock_member
|
||||
|
||||
engine = CoordinationEngine()
|
||||
available = engine.get_available_agents()
|
||||
|
||||
assert available == []
|
||||
|
||||
def test_get_available_agents_with_librarian(self, mock_registry):
|
||||
"""Test getting available agents with librarian registered."""
|
||||
mock_registry.list_members.return_value = ["tatlock_core", "librarian"]
|
||||
|
||||
# tatlock_core has no agent
|
||||
core_member = MagicMock()
|
||||
core_member.agent = None
|
||||
|
||||
# librarian has an agent
|
||||
librarian_member = MagicMock()
|
||||
librarian_member.agent = MagicMock()
|
||||
|
||||
def get_member_side_effect(name):
|
||||
if name == "tatlock_core":
|
||||
return core_member
|
||||
elif name == "librarian":
|
||||
return librarian_member
|
||||
return None
|
||||
|
||||
mock_registry.get_member.side_effect = get_member_side_effect
|
||||
|
||||
engine = CoordinationEngine()
|
||||
available = engine.get_available_agents()
|
||||
|
||||
assert "librarian" in available
|
||||
assert "tatlock_core" not in available
|
||||
|
||||
def test_can_delegate_to_unknown_agent(self, mock_registry):
|
||||
"""Test checking delegation to unknown agent."""
|
||||
mock_registry.get_member.return_value = None
|
||||
|
||||
engine = CoordinationEngine()
|
||||
|
||||
assert engine.can_delegate_to("unknown_agent") is False
|
||||
|
||||
def test_can_delegate_to_librarian(self, mock_registry):
|
||||
"""Test checking delegation to librarian."""
|
||||
mock_member = MagicMock()
|
||||
mock_member.agent = MagicMock() # Has an agent
|
||||
mock_registry.get_member.return_value = mock_member
|
||||
|
||||
engine = CoordinationEngine()
|
||||
|
||||
assert engine.can_delegate_to("librarian") is True
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestDelegationExecution:
|
||||
"""Tests for delegation execution."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_delegation_unavailable_agent(
|
||||
self, mock_registry, librarian_intent
|
||||
):
|
||||
"""Test delegation fails for unavailable agent."""
|
||||
mock_registry.get_member.return_value = None
|
||||
|
||||
engine = CoordinationEngine()
|
||||
|
||||
with pytest.raises(AgentUnavailableError) as exc_info:
|
||||
await engine.execute_delegation(librarian_intent)
|
||||
|
||||
assert "librarian" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_delegation_success(
|
||||
self, mock_registry, librarian_intent
|
||||
):
|
||||
"""Test successful delegation execution."""
|
||||
# Setup mock member with agent
|
||||
mock_member = MagicMock()
|
||||
mock_member.agent = MagicMock()
|
||||
mock_registry.get_member.return_value = mock_member
|
||||
|
||||
# Mock the executor
|
||||
with patch(
|
||||
"src.agents.coordination.AGENT_EXECUTORS",
|
||||
{"librarian": AsyncMock(return_value="Research results here")},
|
||||
):
|
||||
engine = CoordinationEngine()
|
||||
response = await engine.execute_delegation(librarian_intent)
|
||||
|
||||
assert response.success is True
|
||||
assert response.result == "Research results here"
|
||||
# Duration might be 0 for very fast mock execution
|
||||
assert response.duration_ms >= 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_delegation_error(
|
||||
self, mock_registry, librarian_intent
|
||||
):
|
||||
"""Test delegation handles executor errors."""
|
||||
mock_member = MagicMock()
|
||||
mock_member.agent = MagicMock()
|
||||
mock_registry.get_member.return_value = mock_member
|
||||
|
||||
# Mock executor that raises
|
||||
async def failing_executor(**kwargs):
|
||||
raise ValueError("API connection failed")
|
||||
|
||||
with patch(
|
||||
"src.agents.coordination.AGENT_EXECUTORS",
|
||||
{"librarian": failing_executor},
|
||||
):
|
||||
engine = CoordinationEngine()
|
||||
response = await engine.execute_delegation(librarian_intent)
|
||||
|
||||
assert response.success is False
|
||||
assert "API connection failed" in response.error_message
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCoordinate:
|
||||
"""Tests for multi-agent coordination."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_coordinate_single_intent(self, mock_registry, librarian_intent):
|
||||
"""Test coordinating a single delegation."""
|
||||
mock_member = MagicMock()
|
||||
mock_member.agent = MagicMock()
|
||||
mock_registry.get_member.return_value = mock_member
|
||||
|
||||
with patch(
|
||||
"src.agents.coordination.AGENT_EXECUTORS",
|
||||
{"librarian": AsyncMock(return_value="Found docs")},
|
||||
):
|
||||
engine = CoordinationEngine()
|
||||
result = await engine.coordinate([librarian_intent])
|
||||
|
||||
assert result.final_response == "Found docs"
|
||||
assert "librarian" in result.agents_consulted
|
||||
# Duration might be 0 for very fast mock execution
|
||||
assert result.total_duration_ms >= 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_coordinate_empty_intents(self, mock_registry):
|
||||
"""Test coordinating with no intents."""
|
||||
engine = CoordinationEngine()
|
||||
result = await engine.coordinate([])
|
||||
|
||||
assert result.final_response == ""
|
||||
assert result.agents_consulted == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_coordinate_multiple_intents(self, mock_registry):
|
||||
"""Test coordinating multiple delegations."""
|
||||
mock_member = MagicMock()
|
||||
mock_member.agent = MagicMock()
|
||||
mock_registry.get_member.return_value = mock_member
|
||||
|
||||
intents = [
|
||||
DelegationIntent(
|
||||
target_agent="librarian",
|
||||
task="Task 1",
|
||||
reason=DelegationReason.DOMAIN_EXPERTISE,
|
||||
expected_outcome="Result 1",
|
||||
priority=1,
|
||||
),
|
||||
DelegationIntent(
|
||||
target_agent="librarian",
|
||||
task="Task 2",
|
||||
reason=DelegationReason.DOMAIN_EXPERTISE,
|
||||
expected_outcome="Result 2",
|
||||
priority=2,
|
||||
),
|
||||
]
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def mock_executor(**kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
return f"Result {call_count}"
|
||||
|
||||
with patch(
|
||||
"src.agents.coordination.AGENT_EXECUTORS",
|
||||
{"librarian": mock_executor},
|
||||
):
|
||||
engine = CoordinationEngine()
|
||||
result = await engine.coordinate(intents)
|
||||
|
||||
# Both intents were executed (check agents_consulted count)
|
||||
assert len(result.agents_consulted) == 2
|
||||
# Current implementation replaces same-agent responses in dict
|
||||
# So final_response has the last result (or combined if different agents)
|
||||
assert len(result.final_response) > 0
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestDelegateToLibrarian:
|
||||
"""Tests for convenience delegation function."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delegate_to_librarian(self, mock_registry):
|
||||
"""Test the delegate_to_librarian helper."""
|
||||
mock_member = MagicMock()
|
||||
mock_member.agent = MagicMock()
|
||||
mock_registry.get_member.return_value = mock_member
|
||||
|
||||
with patch(
|
||||
"src.agents.coordination.AGENT_EXECUTORS",
|
||||
{"librarian": AsyncMock(return_value="Wiki search results")},
|
||||
):
|
||||
# Reset global engine
|
||||
with patch(
|
||||
"src.agents.coordination._coordination_engine",
|
||||
None,
|
||||
):
|
||||
response = await delegate_to_librarian(
|
||||
task="Search for Docker docs",
|
||||
context="Setting up homelab",
|
||||
)
|
||||
|
||||
assert response.success is True
|
||||
assert response.result == "Wiki search results"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestGetCoordinationEngine:
|
||||
"""Tests for engine singleton."""
|
||||
|
||||
def test_get_coordination_engine_singleton(self):
|
||||
"""Test engine is singleton."""
|
||||
with patch("src.agents.coordination._coordination_engine", None):
|
||||
engine1 = get_coordination_engine()
|
||||
engine2 = get_coordination_engine()
|
||||
|
||||
# Should be same instance
|
||||
assert engine1 is engine2
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestDelegationStreaming:
|
||||
"""Tests for streaming delegation."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_delegation_stream_unavailable(
|
||||
self, mock_registry, librarian_intent
|
||||
):
|
||||
"""Test streaming fails for unavailable agent."""
|
||||
engine = CoordinationEngine()
|
||||
|
||||
# Change target to an agent that doesn't have a stream executor
|
||||
librarian_intent.target_agent = "nonexistent_agent"
|
||||
|
||||
with pytest.raises(AgentUnavailableError):
|
||||
async for _ in engine.execute_delegation_stream(librarian_intent):
|
||||
pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_delegation_stream_success(
|
||||
self, mock_registry, librarian_intent
|
||||
):
|
||||
"""Test successful streaming delegation."""
|
||||
mock_member = MagicMock()
|
||||
mock_member.agent = MagicMock()
|
||||
mock_registry.get_member.return_value = mock_member
|
||||
|
||||
async def mock_stream(**kwargs):
|
||||
yield "Hello "
|
||||
yield "world"
|
||||
|
||||
with patch(
|
||||
"src.agents.coordination.AGENT_STREAM_EXECUTORS",
|
||||
{"librarian": mock_stream},
|
||||
):
|
||||
engine = CoordinationEngine()
|
||||
chunks = []
|
||||
async for chunk in engine.execute_delegation_stream(librarian_intent):
|
||||
chunks.append(chunk)
|
||||
|
||||
assert chunks == ["Hello ", "world"]
|
||||
@@ -0,0 +1,256 @@
|
||||
"""
|
||||
Tests for agent communication protocol.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from src.agents.protocol import (
|
||||
AgentError,
|
||||
AgentRequest,
|
||||
AgentResponse,
|
||||
AgentTimeoutError,
|
||||
AgentUnavailableError,
|
||||
CoordinationResult,
|
||||
DelegationIntent,
|
||||
DelegationReason,
|
||||
ToolCallRecord,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAgentRequest:
|
||||
"""Tests for AgentRequest model."""
|
||||
|
||||
def test_basic_request(self):
|
||||
"""Test creating a basic agent request."""
|
||||
request = AgentRequest(task="Find information about Docker")
|
||||
|
||||
assert request.task == "Find information about Docker"
|
||||
assert request.context == ""
|
||||
assert request.timeout_seconds == 60
|
||||
|
||||
def test_request_with_context(self):
|
||||
"""Test request with additional context."""
|
||||
request = AgentRequest(
|
||||
task="Find Docker networking docs",
|
||||
context="User is setting up a homelab",
|
||||
delegation_reason=DelegationReason.DOMAIN_EXPERTISE,
|
||||
)
|
||||
|
||||
assert request.task == "Find Docker networking docs"
|
||||
assert request.context == "User is setting up a homelab"
|
||||
assert request.delegation_reason == DelegationReason.DOMAIN_EXPERTISE
|
||||
|
||||
def test_request_serialization(self):
|
||||
"""Test request can be serialized to dict."""
|
||||
request = AgentRequest(
|
||||
task="Research task",
|
||||
context="Some context",
|
||||
)
|
||||
|
||||
data = request.model_dump()
|
||||
|
||||
assert data["task"] == "Research task"
|
||||
assert data["context"] == "Some context"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAgentResponse:
|
||||
"""Tests for AgentResponse model."""
|
||||
|
||||
def test_successful_response(self):
|
||||
"""Test creating a successful response."""
|
||||
response = AgentResponse(
|
||||
success=True,
|
||||
result="Here are the findings...",
|
||||
reasoning="Searched wiki and found relevant docs",
|
||||
duration_ms=1500,
|
||||
)
|
||||
|
||||
assert response.success is True
|
||||
assert response.result == "Here are the findings..."
|
||||
assert response.reasoning == "Searched wiki and found relevant docs"
|
||||
assert response.duration_ms == 1500
|
||||
assert response.error_message is None
|
||||
|
||||
def test_failed_response(self):
|
||||
"""Test creating a failed response."""
|
||||
response = AgentResponse(
|
||||
success=False,
|
||||
result="",
|
||||
error_message="Connection timeout",
|
||||
duration_ms=30000,
|
||||
)
|
||||
|
||||
assert response.success is False
|
||||
assert response.result == ""
|
||||
assert response.error_message == "Connection timeout"
|
||||
|
||||
def test_response_with_tool_calls(self):
|
||||
"""Test response tracking tool calls."""
|
||||
tool_call = ToolCallRecord(
|
||||
tool_name="hybrid_search",
|
||||
arguments={"query": "Docker networking"},
|
||||
result="Found 5 results",
|
||||
duration_ms=500,
|
||||
)
|
||||
|
||||
response = AgentResponse(
|
||||
success=True,
|
||||
result="Based on search...",
|
||||
tool_calls=[tool_call],
|
||||
)
|
||||
|
||||
assert len(response.tool_calls) == 1
|
||||
assert response.tool_calls[0].tool_name == "hybrid_search"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestDelegationIntent:
|
||||
"""Tests for DelegationIntent model."""
|
||||
|
||||
def test_basic_intent(self):
|
||||
"""Test creating a basic delegation intent."""
|
||||
intent = DelegationIntent(
|
||||
target_agent="librarian",
|
||||
task="Research Docker networking",
|
||||
reason=DelegationReason.DOMAIN_EXPERTISE,
|
||||
expected_outcome="Documentation and examples",
|
||||
)
|
||||
|
||||
assert intent.target_agent == "librarian"
|
||||
assert intent.task == "Research Docker networking"
|
||||
assert intent.reason == DelegationReason.DOMAIN_EXPERTISE
|
||||
assert intent.priority == 1 # Default
|
||||
|
||||
def test_intent_with_priority(self):
|
||||
"""Test intent with custom priority."""
|
||||
intent = DelegationIntent(
|
||||
target_agent="librarian",
|
||||
task="Urgent research",
|
||||
reason=DelegationReason.RESOURCE_EFFICIENCY,
|
||||
expected_outcome="Quick answer",
|
||||
priority=1,
|
||||
)
|
||||
|
||||
assert intent.priority == 1
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestDelegationReason:
|
||||
"""Tests for DelegationReason enum."""
|
||||
|
||||
def test_all_reasons_have_values(self):
|
||||
"""Test all delegation reasons are defined."""
|
||||
reasons = list(DelegationReason)
|
||||
|
||||
assert DelegationReason.DOMAIN_EXPERTISE in reasons
|
||||
assert DelegationReason.TOOL_ACCESS in reasons
|
||||
assert DelegationReason.RESOURCE_EFFICIENCY in reasons
|
||||
assert DelegationReason.USER_PREFERENCE in reasons
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCoordinationResult:
|
||||
"""Tests for CoordinationResult model."""
|
||||
|
||||
def test_single_agent_result(self):
|
||||
"""Test coordination with single agent."""
|
||||
agent_response = AgentResponse(
|
||||
success=True,
|
||||
result="Research findings",
|
||||
duration_ms=1000,
|
||||
)
|
||||
|
||||
intent = DelegationIntent(
|
||||
target_agent="librarian",
|
||||
task="Research task",
|
||||
reason=DelegationReason.DOMAIN_EXPERTISE,
|
||||
expected_outcome="Findings",
|
||||
)
|
||||
|
||||
result = CoordinationResult(
|
||||
final_response="Research findings",
|
||||
agent_responses={"librarian": agent_response},
|
||||
delegation_intents=[intent],
|
||||
total_duration_ms=1200,
|
||||
agents_consulted=["librarian"],
|
||||
)
|
||||
|
||||
assert result.final_response == "Research findings"
|
||||
assert len(result.agent_responses) == 1
|
||||
assert result.agents_consulted == ["librarian"]
|
||||
|
||||
def test_empty_result(self):
|
||||
"""Test coordination with no delegations."""
|
||||
result = CoordinationResult(
|
||||
final_response="",
|
||||
agent_responses={},
|
||||
delegation_intents=[],
|
||||
total_duration_ms=0,
|
||||
agents_consulted=[],
|
||||
)
|
||||
|
||||
assert result.final_response == ""
|
||||
assert len(result.agents_consulted) == 0
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAgentErrors:
|
||||
"""Tests for agent error types."""
|
||||
|
||||
def test_agent_error(self):
|
||||
"""Test base AgentError."""
|
||||
error = AgentError("Something went wrong")
|
||||
|
||||
assert "Something went wrong" in str(error)
|
||||
assert error.agent_name == "unknown"
|
||||
|
||||
def test_agent_timeout_error(self):
|
||||
"""Test AgentTimeoutError."""
|
||||
error = AgentTimeoutError(
|
||||
"Timed out after 60s",
|
||||
agent_name="librarian",
|
||||
)
|
||||
|
||||
assert "Timed out" in str(error)
|
||||
assert error.agent_name == "librarian"
|
||||
|
||||
def test_agent_unavailable_error(self):
|
||||
"""Test AgentUnavailableError."""
|
||||
error = AgentUnavailableError(
|
||||
"Agent not registered",
|
||||
agent_name="unknown_agent",
|
||||
)
|
||||
|
||||
assert "not registered" in str(error)
|
||||
assert error.agent_name == "unknown_agent"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestToolCallRecord:
|
||||
"""Tests for ToolCallRecord model."""
|
||||
|
||||
def test_tool_call_record(self):
|
||||
"""Test creating a tool call record."""
|
||||
record = ToolCallRecord(
|
||||
tool_name="semantic_search",
|
||||
arguments={"query": "networking concepts", "limit": 10},
|
||||
result="Found 10 relevant documents",
|
||||
duration_ms=250,
|
||||
)
|
||||
|
||||
assert record.tool_name == "semantic_search"
|
||||
assert record.arguments["query"] == "networking concepts"
|
||||
assert record.duration_ms == 250
|
||||
|
||||
def test_tool_call_with_empty_result(self):
|
||||
"""Test tool call with empty result."""
|
||||
record = ToolCallRecord(
|
||||
tool_name="query_graph",
|
||||
arguments={"cypher": "MATCH (n) RETURN n"},
|
||||
result="",
|
||||
duration_ms=100,
|
||||
)
|
||||
|
||||
assert record.result == ""
|
||||
@@ -43,8 +43,8 @@ LOG_FILE="$LOGS_DIR/server.log"
|
||||
echo -e "${YELLOW}Logs will be written to: ${LOG_FILE}${NC}"
|
||||
|
||||
# Start the server
|
||||
echo -e "${GREEN}Starting uvicorn server on http://localhost:8000${NC}"
|
||||
echo -e "${GREEN}Starting uvicorn server on http://localhost:8123${NC}"
|
||||
echo -e "${YELLOW}Press Ctrl+C to stop the server${NC}"
|
||||
echo ""
|
||||
|
||||
uvicorn src.main:app --reload --host 0.0.0.0 --port 8000 2>&1 | tee "$LOG_FILE"
|
||||
uvicorn src.main:app --reload --host 0.0.0.0 --port 8123 2>&1 | tee "$LOG_FILE"
|
||||
|
||||
Reference in New Issue
Block a user