Compare commits
4 Commits
71fd1a0a7c
...
864dfdc4e6
| Author | SHA1 | Date | |
|---|---|---|---|
| 864dfdc4e6 | |||
| 0d16729036 | |||
| 82669d3704 | |||
| 4d0917f5df |
28
.env.example
Normal file
28
.env.example
Normal file
@@ -0,0 +1,28 @@
|
||||
# ── Application ──────────────────────────────────────────────────────────────
|
||||
ENV=dev
|
||||
|
||||
# ── Database ──────────────────────────────────────────────────────────────────
|
||||
DATABASE_URL=postgresql+asyncpg://postgres:postgres@localhost:5432/adiuva
|
||||
|
||||
# ── Auth ──────────────────────────────────────────────────────────────────────
|
||||
JWT_SECRET=replace-with-a-long-random-secret
|
||||
JWT_ALGORITHM=HS256
|
||||
JWT_ACCESS_TOKEN_EXPIRE_MINUTES=30
|
||||
JWT_REFRESH_TOKEN_EXPIRE_DAYS=30
|
||||
|
||||
# ── OpenAI ────────────────────────────────────────────────────────────────────
|
||||
OPENAI_API_KEY=sk-...
|
||||
|
||||
# ── Stripe ────────────────────────────────────────────────────────────────────
|
||||
STRIPE_SECRET_KEY=sk_test_...
|
||||
STRIPE_WEBHOOK_SECRET=whsec_...
|
||||
|
||||
# ── AWS / S3 ──────────────────────────────────────────────────────────────────
|
||||
S3_BUCKET=adiuva-backups
|
||||
S3_REGION=us-east-1
|
||||
AWS_ACCESS_KEY_ID=AKIA...
|
||||
AWS_SECRET_ACCESS_KEY=...
|
||||
|
||||
# ── CORS ──────────────────────────────────────────────────────────────────────
|
||||
# Comma-separated list parsed by Settings (override default if needed)
|
||||
# CORS_ORIGINS=["app://.","http://localhost:3000"]
|
||||
33
.gitignore
vendored
Normal file
33
.gitignore
vendored
Normal file
@@ -0,0 +1,33 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*.egg-info/
|
||||
dist/
|
||||
build/
|
||||
|
||||
# Virtual environment
|
||||
.venv/
|
||||
venv/
|
||||
env/
|
||||
|
||||
# Environment variables
|
||||
.env
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
|
||||
# Testing / coverage
|
||||
.pytest_cache/
|
||||
htmlcov/
|
||||
.coverage
|
||||
|
||||
# Docker
|
||||
*.log
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Claude Code
|
||||
.claude/
|
||||
@@ -68,9 +68,9 @@ adiuva-api/
|
||||
|
||||
## Step-by-Step Implementation
|
||||
|
||||
### Step 1 — Project scaffolding
|
||||
- [ ] Initialize repo with the directory structure above
|
||||
- [ ] Write `requirements.txt`:
|
||||
### Step 1 — Project scaffolding ✅
|
||||
- [x] Initialize repo with the directory structure above
|
||||
- [x] Write `requirements.txt`:
|
||||
```
|
||||
fastapi>=0.115.0
|
||||
uvicorn[standard]>=0.34.0
|
||||
@@ -91,15 +91,15 @@ adiuva-api/
|
||||
pytest>=8.0.0
|
||||
pytest-asyncio>=0.24.0
|
||||
```
|
||||
- [ ] Write `app/main.py`: FastAPI app with CORS (allow `app://`, `http://localhost:*`), lifespan (init DB pool, init agent registry), include all routers under `/api/v1`
|
||||
- [ ] Write `app/config/settings.py`: `Settings(BaseSettings)` with fields: `DATABASE_URL`, `JWT_SECRET`, `JWT_ALGORITHM` (default HS256), `STRIPE_SECRET_KEY`, `STRIPE_WEBHOOK_SECRET`, `S3_BUCKET`, `S3_REGION`, `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `OPENAI_API_KEY`, `CORS_ORIGINS`, `ENV` (dev/prod)
|
||||
- [ ] Write `Dockerfile`: Python 3.12 slim, multi-stage (builder + runtime), non-root user
|
||||
- [ ] Write `docker-compose.yml`: app, postgres:16, optional redis
|
||||
- [ ] Write `.env.example`
|
||||
- [x] Write `app/main.py`: FastAPI app with CORS (allow `app://`, `http://localhost:*`), lifespan (init DB pool, init agent registry), include all routers under `/api/v1`
|
||||
- [x] Write `app/config/settings.py`: `Settings(BaseSettings)` with fields: `DATABASE_URL`, `JWT_SECRET`, `JWT_ALGORITHM` (default HS256), `STRIPE_SECRET_KEY`, `STRIPE_WEBHOOK_SECRET`, `S3_BUCKET`, `S3_REGION`, `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `OPENAI_API_KEY`, `CORS_ORIGINS`, `ENV` (dev/prod)
|
||||
- [x] Write `Dockerfile`: Python 3.12 slim, multi-stage (builder + runtime), non-root user
|
||||
- [x] Write `docker-compose.yml`: app, postgres:16, optional redis
|
||||
- [x] Write `.env.example`
|
||||
- **Outcome:** Runnable FastAPI skeleton (returns 404 on all routes).
|
||||
|
||||
### Step 2 — Pydantic schemas (API contracts)
|
||||
- [ ] Create `app/schemas.py` (mirrors `src/shared/api-types.ts` from Electron repo):
|
||||
### Step 2 — Pydantic schemas (API contracts) ✅
|
||||
- [x] Create `app/schemas.py` (mirrors `src/shared/api-types.ts` from Electron repo):
|
||||
- `ChatRequest`: `message: str`, `context: ChatContext`, `execution_mode: Literal['direct', 'plan']`
|
||||
- `ChatContext`: `user_profile: dict`, `relevant_documents: list[str]`, `recent_tasks: list[dict]`, `conversation_history: list[dict]`
|
||||
- `ChatResponse`: `response: str`, `actions: list[PlanAction]`
|
||||
@@ -112,8 +112,8 @@ adiuva-api/
|
||||
- `UserProfile`: `id: str`, `email: str`, `tier: BillingTier`
|
||||
- **Outcome:** All request/response models defined and validated.
|
||||
|
||||
### Step 3 — Agent Registry + base classes
|
||||
- [ ] `app/core/agent_registry.py`:
|
||||
### Step 3 — Agent Registry + base classes ✅
|
||||
- [x] `app/core/agent_registry.py`:
|
||||
- `BaseAgent(ABC)`:
|
||||
- `user_id: str`, `shared_memory: dict`, `vector_store_context: list[str]`, `skills: list[str]`
|
||||
- Abstract `get_name() -> str`, `get_description() -> str`
|
||||
@@ -127,7 +127,7 @@ adiuva-api/
|
||||
- `get(name) -> ChatAgent`
|
||||
- `list_agents() -> list[dict]` — returns `[{name, description}]` for orchestrator prompt
|
||||
- `async call_agent(name, query, context) -> str` — for inter-agent calls
|
||||
- [ ] Unit tests: register, get, list, call_agent with mock
|
||||
- [x] Unit tests: register, get, list, call_agent with mock
|
||||
- **Outcome:** Pluggable agent framework.
|
||||
|
||||
### Step 4 — Orchestrator
|
||||
@@ -356,3 +356,4 @@ adiuva-api/
|
||||
4. **Type hints everywhere.** All functions have full type annotations.
|
||||
5. **Test every agent.** Each chat agent has unit tests with mocked LLM responses.
|
||||
6. **Structured logging.** JSON logs with request ID correlation.
|
||||
7. **One step at a time.** Implement one numbered step per session. When the step is fully done, mark all its checkboxes as `[x]` in this file and commit with message `step N complete: <outcome line>`.
|
||||
|
||||
31
Dockerfile
Normal file
31
Dockerfile
Normal file
@@ -0,0 +1,31 @@
|
||||
# ── builder ──────────────────────────────────────────────────────────────────
|
||||
FROM python:3.12-slim AS builder
|
||||
|
||||
WORKDIR /build
|
||||
|
||||
COPY requirements.txt .
|
||||
RUN pip install --upgrade pip && \
|
||||
pip install --no-cache-dir --prefix=/install -r requirements.txt
|
||||
|
||||
# ── runtime ──────────────────────────────────────────────────────────────────
|
||||
FROM python:3.12-slim AS runtime
|
||||
|
||||
# Non-root user
|
||||
RUN addgroup --system appgroup && adduser --system --ingroup appgroup appuser
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Copy installed packages from builder
|
||||
COPY --from=builder /install /usr/local
|
||||
|
||||
# Copy application source
|
||||
COPY app/ app/
|
||||
|
||||
# Ensure appuser owns the working directory
|
||||
RUN chown -R appuser:appgroup /app
|
||||
|
||||
USER appuser
|
||||
|
||||
EXPOSE 8000
|
||||
|
||||
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "2"]
|
||||
0
app/__init__.py
Normal file
0
app/__init__.py
Normal file
0
app/agents/__init__.py
Normal file
0
app/agents/__init__.py
Normal file
0
app/api/__init__.py
Normal file
0
app/api/__init__.py
Normal file
0
app/api/middleware/__init__.py
Normal file
0
app/api/middleware/__init__.py
Normal file
0
app/api/routes/__init__.py
Normal file
0
app/api/routes/__init__.py
Normal file
0
app/billing/__init__.py
Normal file
0
app/billing/__init__.py
Normal file
0
app/config/__init__.py
Normal file
0
app/config/__init__.py
Normal file
31
app/config/settings.py
Normal file
31
app/config/settings.py
Normal file
@@ -0,0 +1,31 @@
|
||||
from typing import Literal
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
DATABASE_URL: str = "postgresql+asyncpg://postgres:postgres@localhost:5432/adiuva"
|
||||
JWT_SECRET: str = "change-me-in-production"
|
||||
JWT_ALGORITHM: str = "HS256"
|
||||
JWT_ACCESS_TOKEN_EXPIRE_MINUTES: int = 30
|
||||
JWT_REFRESH_TOKEN_EXPIRE_DAYS: int = 30
|
||||
|
||||
STRIPE_SECRET_KEY: str = ""
|
||||
STRIPE_WEBHOOK_SECRET: str = ""
|
||||
|
||||
S3_BUCKET: str = ""
|
||||
S3_REGION: str = "us-east-1"
|
||||
AWS_ACCESS_KEY_ID: str = ""
|
||||
AWS_SECRET_ACCESS_KEY: str = ""
|
||||
|
||||
OPENAI_API_KEY: str = ""
|
||||
|
||||
CORS_ORIGINS: list[str] = ["app://.", "http://localhost:3000", "http://localhost:5173"]
|
||||
|
||||
ENV: Literal["dev", "prod"] = "dev"
|
||||
|
||||
class Config:
|
||||
env_file = ".env"
|
||||
env_file_encoding = "utf-8"
|
||||
|
||||
|
||||
settings = Settings()
|
||||
0
app/core/__init__.py
Normal file
0
app/core/__init__.py
Normal file
137
app/core/agent_registry.py
Normal file
137
app/core/agent_registry.py
Normal file
@@ -0,0 +1,137 @@
|
||||
"""Agent Registry — base classes and singleton registry for chat agents."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
|
||||
class BaseAgent(ABC):
|
||||
"""Common base for all agents."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
user_id: str = "",
|
||||
shared_memory: dict[str, Any] | None = None,
|
||||
vector_store_context: list[str] | None = None,
|
||||
) -> None:
|
||||
self.user_id = user_id
|
||||
self.shared_memory: dict[str, Any] = shared_memory or {}
|
||||
self.vector_store_context: list[str] = vector_store_context or []
|
||||
|
||||
@abstractmethod
|
||||
def get_name(self) -> str: ...
|
||||
|
||||
@abstractmethod
|
||||
def get_description(self) -> str: ...
|
||||
|
||||
@property
|
||||
def skills(self) -> list[str]:
|
||||
"""Override in subclasses to advertise capabilities."""
|
||||
return []
|
||||
|
||||
|
||||
class ChatAgent(BaseAgent):
|
||||
"""Base class for LLM-powered chat agents."""
|
||||
|
||||
@abstractmethod
|
||||
async def handle(self, query: str, context: dict[str, Any]) -> str:
|
||||
"""Process a user query and return a text response."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def get_tools(self) -> list[Any]:
|
||||
"""Return LangChain tool definitions available to this agent."""
|
||||
...
|
||||
|
||||
async def _tool_loop(
|
||||
self,
|
||||
llm: Any,
|
||||
messages: list[Any],
|
||||
tools: list[Any],
|
||||
max_iter: int = 5,
|
||||
) -> str:
|
||||
"""Shared tool-calling loop.
|
||||
|
||||
Binds *tools* to *llm*, invokes iteratively until the model stops
|
||||
requesting tool calls or *max_iter* is reached, and returns the
|
||||
final text response.
|
||||
"""
|
||||
from langchain_core.messages import AIMessage, ToolMessage
|
||||
|
||||
llm_with_tools = llm.bind_tools(tools) if tools else llm
|
||||
|
||||
for _ in range(max_iter):
|
||||
response: AIMessage = await llm_with_tools.ainvoke(messages)
|
||||
messages.append(response)
|
||||
|
||||
if not response.tool_calls:
|
||||
return str(response.content)
|
||||
|
||||
# Execute each requested tool call
|
||||
tool_map = {t.name: t for t in tools}
|
||||
for call in response.tool_calls:
|
||||
tool_fn = tool_map.get(call["name"])
|
||||
if tool_fn is None:
|
||||
result = f"Unknown tool: {call['name']}"
|
||||
else:
|
||||
result = await tool_fn.ainvoke(call["args"])
|
||||
messages.append(
|
||||
ToolMessage(content=str(result), tool_call_id=call["id"])
|
||||
)
|
||||
|
||||
# Exhausted iterations — ask model for a final answer without tools
|
||||
response = await llm.ainvoke(messages)
|
||||
return str(response.content)
|
||||
|
||||
|
||||
class AgentRegistry:
|
||||
"""Singleton registry for ChatAgent subclasses."""
|
||||
|
||||
_instance: AgentRegistry | None = None
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._agents: dict[str, type[ChatAgent]] = {}
|
||||
|
||||
def __new__(cls) -> AgentRegistry:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._agents = {}
|
||||
return cls._instance
|
||||
|
||||
# ── public API ───────────────────────────────────────────────────
|
||||
|
||||
def register(self, agent_class: type[ChatAgent]) -> type[ChatAgent]:
|
||||
"""Class decorator — registers an agent by its name."""
|
||||
instance = agent_class()
|
||||
name = instance.get_name()
|
||||
self._agents[name] = agent_class
|
||||
return agent_class
|
||||
|
||||
def get(self, name: str) -> ChatAgent:
|
||||
"""Return a fresh instance of the named agent."""
|
||||
cls = self._agents.get(name)
|
||||
if cls is None:
|
||||
raise KeyError(f"Agent not found: {name}")
|
||||
return cls()
|
||||
|
||||
def list_agents(self) -> list[dict[str, str]]:
|
||||
"""Return ``[{name, description}]`` for the orchestrator prompt."""
|
||||
result: list[dict[str, str]] = []
|
||||
for cls in self._agents.values():
|
||||
inst = cls()
|
||||
result.append(
|
||||
{"name": inst.get_name(), "description": inst.get_description()}
|
||||
)
|
||||
return result
|
||||
|
||||
async def call_agent(
|
||||
self, name: str, query: str, context: dict[str, Any]
|
||||
) -> str:
|
||||
"""Instantiate the named agent and call its ``handle`` method."""
|
||||
agent = self.get(name)
|
||||
return await agent.handle(query, context)
|
||||
|
||||
|
||||
# Module-level singleton
|
||||
registry = AgentRegistry()
|
||||
52
app/main.py
Normal file
52
app/main.py
Normal file
@@ -0,0 +1,52 @@
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
from app.config.settings import settings
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
# Startup: initialise DB connection pool and agent registry
|
||||
from app.core.agent_registry import registry # noqa: F401 — triggers module load
|
||||
import app.agents # noqa: F401 — triggers @registry.register decorators
|
||||
|
||||
yield
|
||||
|
||||
# Shutdown: nothing to clean up for now
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
app = FastAPI(
|
||||
title="Adiuva Cloud API",
|
||||
version="0.1.0",
|
||||
docs_url="/docs" if settings.ENV == "dev" else None,
|
||||
redoc_url=None,
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=settings.CORS_ORIGINS,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Routers (registered when implemented)
|
||||
# from app.api.routes import auth, chat, plans, backup, billing
|
||||
# app.include_router(auth.router, prefix="/api/v1")
|
||||
# app.include_router(chat.router, prefix="/api/v1")
|
||||
# app.include_router(plans.router, prefix="/api/v1")
|
||||
# app.include_router(backup.router, prefix="/api/v1")
|
||||
# app.include_router(billing.router, prefix="/api/v1")
|
||||
|
||||
@app.get("/api/v1/health", tags=["health"])
|
||||
async def health() -> dict:
|
||||
return {"status": "ok", "version": app.version}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
app = create_app()
|
||||
84
app/schemas.py
Normal file
84
app/schemas.py
Normal file
@@ -0,0 +1,84 @@
|
||||
"""Pydantic schemas — API request/response contracts.
|
||||
|
||||
Mirrors the TypeScript types from the Electron app (src/shared/api-types.ts).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
# ── Billing ──────────────────────────────────────────────────────────
|
||||
|
||||
BillingTier = Literal["free", "pro", "power", "team"]
|
||||
|
||||
|
||||
# ── Auth ─────────────────────────────────────────────────────────────
|
||||
|
||||
class AuthTokens(BaseModel):
|
||||
access_token: str
|
||||
refresh_token: str
|
||||
expires_at: int
|
||||
|
||||
|
||||
class UserProfile(BaseModel):
|
||||
id: str
|
||||
email: str
|
||||
tier: BillingTier
|
||||
|
||||
|
||||
# ── Chat ─────────────────────────────────────────────────────────────
|
||||
|
||||
class ChatContext(BaseModel):
|
||||
user_profile: dict[str, Any] = Field(default_factory=dict)
|
||||
relevant_documents: list[str] = Field(default_factory=list)
|
||||
recent_tasks: list[dict[str, Any]] = Field(default_factory=list)
|
||||
conversation_history: list[dict[str, Any]] = Field(default_factory=list)
|
||||
|
||||
|
||||
class PlanAction(BaseModel):
|
||||
type: Literal[
|
||||
"create_record",
|
||||
"update_record",
|
||||
"delete_record",
|
||||
"index_document",
|
||||
"send_notification",
|
||||
]
|
||||
table: str | None = None
|
||||
data: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class ChatRequest(BaseModel):
|
||||
message: str
|
||||
context: ChatContext = Field(default_factory=ChatContext)
|
||||
execution_mode: Literal["direct", "plan"] = "direct"
|
||||
|
||||
|
||||
class ChatResponse(BaseModel):
|
||||
response: str
|
||||
actions: list[PlanAction] = Field(default_factory=list)
|
||||
|
||||
|
||||
# ── Execution Plans ──────────────────────────────────────────────────
|
||||
|
||||
class PlanStep(BaseModel):
|
||||
action: str
|
||||
prompt_template: str | None = None
|
||||
variables: dict[str, Any] | None = None
|
||||
data_from_step: int | None = None
|
||||
|
||||
|
||||
class ExecutionPlan(BaseModel):
|
||||
agent: str
|
||||
steps: list[PlanStep] = Field(default_factory=list)
|
||||
|
||||
|
||||
# ── Backup ───────────────────────────────────────────────────────────
|
||||
|
||||
class BackupMetadata(BaseModel):
|
||||
version: int
|
||||
timestamp: int
|
||||
checksum: str
|
||||
chunk_count: int
|
||||
38
docker-compose.yml
Normal file
38
docker-compose.yml
Normal file
@@ -0,0 +1,38 @@
|
||||
version: "3.9"
|
||||
|
||||
services:
|
||||
app:
|
||||
build: .
|
||||
ports:
|
||||
- "8000:8000"
|
||||
env_file:
|
||||
- .env
|
||||
environment:
|
||||
DATABASE_URL: postgresql+asyncpg://postgres:postgres@db:5432/adiuva
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
restart: unless-stopped
|
||||
|
||||
db:
|
||||
image: postgres:16-alpine
|
||||
environment:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: adiuva
|
||||
volumes:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U postgres"]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
restart: unless-stopped
|
||||
|
||||
# Optional Redis for future rate-limit or caching needs
|
||||
# redis:
|
||||
# image: redis:7-alpine
|
||||
# restart: unless-stopped
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
19
requirements.txt
Normal file
19
requirements.txt
Normal file
@@ -0,0 +1,19 @@
|
||||
fastapi>=0.115.0
|
||||
uvicorn[standard]>=0.34.0
|
||||
langchain>=0.3.0
|
||||
langchain-openai>=0.3.0
|
||||
pydantic>=2.10.0
|
||||
pydantic-settings>=2.7.0
|
||||
python-jose[cryptography]>=3.3.0
|
||||
stripe>=11.0.0
|
||||
boto3>=1.35.0
|
||||
slowapi>=0.1.9
|
||||
sqlalchemy>=2.0.0
|
||||
asyncpg>=0.30.0
|
||||
alembic>=1.14.0
|
||||
bcrypt>=4.2.0
|
||||
python-dotenv>=1.0.0
|
||||
httpx>=0.28.0
|
||||
websockets>=14.0
|
||||
pytest>=8.0.0
|
||||
pytest-asyncio>=0.24.0
|
||||
0
tests/__init__.py
Normal file
0
tests/__init__.py
Normal file
214
tests/test_agent_registry.py
Normal file
214
tests/test_agent_registry.py
Normal file
@@ -0,0 +1,214 @@
|
||||
"""Unit tests for the agent registry, base classes, and tool loop."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from app.core.agent_registry import AgentRegistry, ChatAgent
|
||||
|
||||
|
||||
# ── Helpers ──────────────────────────────────────────────────────────
|
||||
|
||||
class _StubAgent(ChatAgent):
|
||||
"""Minimal concrete agent for testing."""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "stub"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "A stub agent for tests"
|
||||
|
||||
def get_tools(self) -> list[Any]:
|
||||
return []
|
||||
|
||||
async def handle(self, query: str, context: dict[str, Any]) -> str:
|
||||
return f"echo: {query}"
|
||||
|
||||
|
||||
class _AnotherAgent(ChatAgent):
|
||||
def get_name(self) -> str:
|
||||
return "another"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "Another stub"
|
||||
|
||||
def get_tools(self) -> list[Any]:
|
||||
return []
|
||||
|
||||
async def handle(self, query: str, context: dict[str, Any]) -> str:
|
||||
return "another"
|
||||
|
||||
|
||||
# ── Fixtures ─────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _fresh_registry():
|
||||
"""Reset the singleton between tests."""
|
||||
AgentRegistry._instance = None
|
||||
yield
|
||||
AgentRegistry._instance = None
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def reg() -> AgentRegistry:
|
||||
return AgentRegistry()
|
||||
|
||||
|
||||
# ── Tests ────────────────────────────────────────────────────────────
|
||||
|
||||
class TestRegisterAndGet:
|
||||
def test_register_decorator(self, reg: AgentRegistry) -> None:
|
||||
reg.register(_StubAgent)
|
||||
agent = reg.get("stub")
|
||||
assert isinstance(agent, _StubAgent)
|
||||
|
||||
def test_get_unknown_raises(self, reg: AgentRegistry) -> None:
|
||||
with pytest.raises(KeyError, match="not found"):
|
||||
reg.get("nonexistent")
|
||||
|
||||
def test_register_multiple(self, reg: AgentRegistry) -> None:
|
||||
reg.register(_StubAgent)
|
||||
reg.register(_AnotherAgent)
|
||||
assert reg.get("stub").get_name() == "stub"
|
||||
assert reg.get("another").get_name() == "another"
|
||||
|
||||
|
||||
class TestListAgents:
|
||||
def test_empty(self, reg: AgentRegistry) -> None:
|
||||
assert reg.list_agents() == []
|
||||
|
||||
def test_list_after_register(self, reg: AgentRegistry) -> None:
|
||||
reg.register(_StubAgent)
|
||||
agents = reg.list_agents()
|
||||
assert len(agents) == 1
|
||||
assert agents[0] == {"name": "stub", "description": "A stub agent for tests"}
|
||||
|
||||
def test_list_multiple(self, reg: AgentRegistry) -> None:
|
||||
reg.register(_StubAgent)
|
||||
reg.register(_AnotherAgent)
|
||||
names = {a["name"] for a in reg.list_agents()}
|
||||
assert names == {"stub", "another"}
|
||||
|
||||
|
||||
class TestCallAgent:
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_agent(self, reg: AgentRegistry) -> None:
|
||||
reg.register(_StubAgent)
|
||||
result = await reg.call_agent("stub", "hello", {})
|
||||
assert result == "echo: hello"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_unknown_raises(self, reg: AgentRegistry) -> None:
|
||||
with pytest.raises(KeyError):
|
||||
await reg.call_agent("nope", "hi", {})
|
||||
|
||||
|
||||
class TestSingleton:
|
||||
def test_singleton_identity(self) -> None:
|
||||
a = AgentRegistry()
|
||||
b = AgentRegistry()
|
||||
assert a is b
|
||||
|
||||
|
||||
class TestToolLoop:
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_tool_calls(self) -> None:
|
||||
"""When the LLM responds without tool calls, return content directly."""
|
||||
agent = _StubAgent()
|
||||
|
||||
ai_msg = MagicMock()
|
||||
ai_msg.content = "final answer"
|
||||
ai_msg.tool_calls = []
|
||||
|
||||
llm = AsyncMock()
|
||||
llm.bind_tools = MagicMock(return_value=llm)
|
||||
llm.ainvoke = AsyncMock(return_value=ai_msg)
|
||||
|
||||
result = await agent._tool_loop(llm, [], [])
|
||||
assert result == "final answer"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_then_answer(self) -> None:
|
||||
"""LLM requests one tool call, gets result, then answers."""
|
||||
agent = _StubAgent()
|
||||
|
||||
# First response: tool call
|
||||
tool_call_msg = MagicMock()
|
||||
tool_call_msg.content = ""
|
||||
tool_call_msg.tool_calls = [
|
||||
{"id": "call_1", "name": "my_tool", "args": {"x": 1}}
|
||||
]
|
||||
|
||||
# Second response: final answer
|
||||
final_msg = MagicMock()
|
||||
final_msg.content = "done"
|
||||
final_msg.tool_calls = []
|
||||
|
||||
llm = AsyncMock()
|
||||
llm.bind_tools = MagicMock(return_value=llm)
|
||||
llm.ainvoke = AsyncMock(side_effect=[tool_call_msg, final_msg])
|
||||
|
||||
# Mock tool
|
||||
tool = AsyncMock()
|
||||
tool.name = "my_tool"
|
||||
tool.ainvoke = AsyncMock(return_value="tool_result")
|
||||
|
||||
result = await agent._tool_loop(llm, [], [tool])
|
||||
assert result == "done"
|
||||
tool.ainvoke.assert_called_once_with({"x": 1})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_tool_handled(self) -> None:
|
||||
"""Unknown tool names produce an error message instead of crashing."""
|
||||
agent = _StubAgent()
|
||||
|
||||
tool_call_msg = MagicMock()
|
||||
tool_call_msg.content = ""
|
||||
tool_call_msg.tool_calls = [
|
||||
{"id": "call_1", "name": "missing", "args": {}}
|
||||
]
|
||||
|
||||
final_msg = MagicMock()
|
||||
final_msg.content = "recovered"
|
||||
final_msg.tool_calls = []
|
||||
|
||||
llm = AsyncMock()
|
||||
llm.bind_tools = MagicMock(return_value=llm)
|
||||
llm.ainvoke = AsyncMock(side_effect=[tool_call_msg, final_msg])
|
||||
|
||||
result = await agent._tool_loop(llm, [], [])
|
||||
assert result == "recovered"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_iter_reached(self) -> None:
|
||||
"""When max iterations are exhausted, a final no-tools call is made."""
|
||||
agent = _StubAgent()
|
||||
|
||||
# Every response requests a tool call
|
||||
loop_msg = MagicMock()
|
||||
loop_msg.content = ""
|
||||
loop_msg.tool_calls = [
|
||||
{"id": "call_x", "name": "t", "args": {}}
|
||||
]
|
||||
|
||||
final_msg = MagicMock()
|
||||
final_msg.content = "gave up"
|
||||
final_msg.tool_calls = []
|
||||
|
||||
tool = AsyncMock()
|
||||
tool.name = "t"
|
||||
tool.ainvoke = AsyncMock(return_value="ok")
|
||||
|
||||
llm_with_tools = AsyncMock()
|
||||
llm_with_tools.ainvoke = AsyncMock(return_value=loop_msg)
|
||||
|
||||
llm = AsyncMock()
|
||||
llm.bind_tools = MagicMock(return_value=llm_with_tools)
|
||||
llm.ainvoke = AsyncMock(return_value=final_msg)
|
||||
|
||||
result = await agent._tool_loop(llm, [], [tool], max_iter=2)
|
||||
assert result == "gave up"
|
||||
assert llm_with_tools.ainvoke.call_count == 2
|
||||
Reference in New Issue
Block a user