patientzero 0.1.0__py3-none-any.whl

backend/README.md

@@ -0,0 +1,71 @@
+# PatientZero Backend
+
+FastAPI server with SQLite database and abstracted LLM provider layer.
+
+## Setup
+
+```bash
+uv sync
+cp ../.env.example ../.env  # if not done already
+```
+
+## Running
+
+```bash
+uv run uvicorn api.main:app --reload
+
+```
+
+Server runs at http://localhost:8000
+
+## Project Structure
+
+```
+backend/
+├── api/
+│   ├── main.py              # FastAPI app, CORS, lifespan
+│   ├── dependencies.py      # Shared db + provider instances
+│   └── routes/
+│       └── chat.py          # Chat + session endpoints
+├── config/
+│   └── settings.py          # Environment config
+├── db/
+│   ├── database.py          # Database class (SQLite, raw queries)
+│   ├── schema.sql           # Table definitions
+│   └── queries/
+│       └── sessions.py      # Session + turn CRUD
+├── llm/
+│   ├── base.py              # Abstract LLMProvider
+│   ├── mock.py              # Mock provider (testing)
+│   └── factory.py           # Provider factory
+└── pyproject.toml
+```
+
+## API Endpoints
+
+| Method | Path | Description |
+|--------|------|-------------|
+| `POST` | `/api/sessions` | Create a new chat session |
+| `GET` | `/api/sessions` | List all sessions |
+| `GET` | `/api/sessions/{id}` | Get session with turns |
+| `POST` | `/api/chat` | Send message, receive SSE stream |
+
+## LLM Providers
+
+Set `LLM_PROVIDER` in `.env`:
+
+| Provider | Value | Status |
+|----------|-------|--------|
+| Mock | `mock` | Available |
+| OpenAI | `openai` | Planned |
+| Claude | `claude` | Planned |
+| Local | `local` | Planned |
+
+## Database
+
+SQLite with WAL mode. Tables:
+
+- `sessions` — chat sessions (id, title, created_at)
+- `turns` — individual messages (session_id, role, content, turn_number)
+
+Database file is created automatically on first run.

backend/api/dependencies.py

@@ -0,0 +1,8 @@
+from core.config.settings import DB_PATH
+from core.db.database import Database
+from core.repositories import RepoSet
+from core.logger import SimulationLogger
+
+db = Database(DB_PATH)
+repos = RepoSet.for_db(db)
+logger = SimulationLogger()

backend/api/main.py

@@ -0,0 +1,81 @@
+import logging
+import traceback
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI, Request
+from fastapi.exceptions import HTTPException, RequestValidationError
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
+logger = logging.getLogger("patientzero.api")
+
+from core import Experiment
+from core.config.settings import FRONTEND_URL
+from core.examples.medical.config import MEDICAL_EXAMPLE_CONFIG
+from backend.api.dependencies import db, repos
+from backend.api.routes.agents import router as agents_router
+from backend.api.routes.analysis import router as analysis_router
+from backend.api.routes.chat import router as chat_router
+from backend.api.routes.distributions import router as distributions_router
+from backend.api.routes.experiments import router as experiments_router
+from backend.api.routes.settings import router as settings_router
+from backend.api.routes.simulate import router as simulate_router
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    db.init()
+    if not repos.experiments.list_all():
+        Experiment(MEDICAL_EXAMPLE_CONFIG, repos)
+    yield
+    db.close()
+
+
+app = FastAPI(title="PatientZero", lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=[FRONTEND_URL],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.exception_handler(HTTPException)
+async def http_exception_handler(request: Request, exc: HTTPException):
+    if exc.status_code >= 400:
+        print(
+            f"\033[33m[{exc.status_code}]\033[0m {request.method} {request.url.path} "
+            f"→ {exc.detail}",
+            flush=True,
+        )
+    return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})
+
+
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request: Request, exc: RequestValidationError):
+    print(
+        f"\033[33m[422]\033[0m {request.method} {request.url.path} → validation error:\n"
+        f"  {exc.errors()}",
+        flush=True,
+    )
+    return JSONResponse(status_code=422, content={"detail": exc.errors()})
+
+
+@app.exception_handler(Exception)
+async def unhandled_exception_handler(request: Request, exc: Exception):
+    print(
+        f"\033[31m[500]\033[0m {request.method} {request.url.path} → {type(exc).__name__}: {exc}",
+        flush=True,
+    )
+    traceback.print_exc()
+    return JSONResponse(status_code=500, content={"detail": f"{type(exc).__name__}: {exc}"})
+
+
+app.include_router(chat_router, prefix="/api")
+app.include_router(simulate_router, prefix="/api")
+app.include_router(analysis_router, prefix="/api")
+app.include_router(settings_router, prefix="/api")
+app.include_router(experiments_router, prefix="/api")
+app.include_router(distributions_router, prefix="/api")
+app.include_router(agents_router, prefix="/api")

backend/api/routes/agents.py

@@ -0,0 +1,24 @@
+from fastapi import APIRouter, HTTPException
+
+from backend.api.dependencies import repos
+
+router = APIRouter()
+
+
+@router.get("/experiments/{exp_id}/agents")
+def get_experiment_agents(exp_id: str):
+    experiment = repos.experiments.get(exp_id)
+    if experiment is None:
+        raise HTTPException(status_code=404, detail="Experiment not found")
+    config = experiment.config
+    return {
+        "agents": [
+            {"name": a.name, "prompt": a.prompt, "model": a.model}
+            for a in config.agents
+        ],
+        "judge": {
+            "rubric": dict(config.judge.rubric),
+            "instructions": config.judge.instructions,
+            "model": config.judge.model,
+        },
+    }

backend/api/routes/analysis.py

@@ -0,0 +1,138 @@
+"""
+Experiment analysis — aggregates judge scores across each agent's sampled traits.
+
+The shape is derived from the experiment's actual distributions: group by
+(agent_name, trait, value) and report per-metric mean/std/n per bucket.
+No hardcoded dimension names — whatever traits the distributions declare
+are what you get.
+"""
+
+import math
+
+from fastapi import APIRouter, HTTPException
+
+from backend.api.dependencies import repos
+
+router = APIRouter()
+
+
+# ── Stats helpers ─────────────────────────────────────────────────────────────
+
+
+def _mean(vals: list[float]) -> float | None:
+    return sum(vals) / len(vals) if vals else None
+
+
+def _std(vals: list[float]) -> float | None:
+    if len(vals) < 2:
+        return None
+    m = _mean(vals)
+    assert m is not None
+    return math.sqrt(sum((v - m) ** 2 for v in vals) / (len(vals) - 1))
+
+
+def _score_stats(per_dim_values: dict[str, list[float]]) -> dict:
+    out: dict[str, dict] = {}
+    for dim, vals in per_dim_values.items():
+        m = _mean(vals)
+        s = _std(vals)
+        out[dim] = {
+            "mean": round(m, 2) if m is not None else None,
+            "std": round(s, 2) if s is not None else None,
+            "n": len(vals),
+        }
+    return out
+
+
+def _row_dims(row: dict, dims: set[str]) -> dict[str, list[float]]:
+    return {dim: [row[dim]] if row.get(dim) is not None else [] for dim in dims}
+
+
+# ── Row building ──────────────────────────────────────────────────────────────
+
+
+def _build_rows(experiment_id: str) -> tuple[list[dict], set[str]]:
+    """Join completed simulations with their evaluations and flatten per-dim scores.
+
+    Each returned row has: ``profiles`` (dict[agent → trait dict]) plus every
+    judge rubric dimension at the top level as a single float (mean across
+    judge_results for that dimension on that evaluation).
+    Returns ``(rows, judge_dimensions)``.
+    """
+    pairs = repos.evaluations.list_completed_with_evaluations_for_experiment(experiment_id)
+    rows: list[dict] = []
+    dims: set[str] = set()
+    for sim, ev in pairs:
+        row: dict = {"profiles": sim.config.profiles}
+        for judge in ev.judge_results:
+            for name, value in judge.scores.items():
+                dims.add(name)
+        for dim in dims:
+            vals = [
+                j.scores.get(dim)
+                for j in ev.judge_results
+                if j.scores.get(dim) is not None
+            ]
+            row[dim] = (sum(vals) / len(vals)) if vals else None
+        rows.append(row)
+    return rows, dims
+
+
+# ── Grouping ──────────────────────────────────────────────────────────────────
+
+
+def _aggregate_per_dim(rows: list[dict], dims: set[str]) -> dict[str, list[float]]:
+    out: dict[str, list[float]] = {dim: [] for dim in dims}
+    for row in rows:
+        for dim in dims:
+            v = row.get(dim)
+            if v is not None:
+                out[dim].append(v)
+    return out
+
+
+def _group_by_trait(
+    rows: list[dict], dims: set[str], agent_name: str, trait: str
+) -> dict[str, dict]:
+    buckets: dict[str, list[dict]] = {}
+    for row in rows:
+        value = row.get("profiles", {}).get(agent_name, {}).get(trait)
+        if value is None:
+            continue
+        buckets.setdefault(value, []).append(row)
+    return {
+        value: _score_stats(_aggregate_per_dim(bucket_rows, dims))
+        for value, bucket_rows in buckets.items()
+    }
+
+
+# ── Endpoints ─────────────────────────────────────────────────────────────────
+
+
+@router.get("/experiments/{exp_id}/analysis")
+def get_experiment_analysis(exp_id: str):
+    experiment = repos.experiments.get(exp_id)
+    if experiment is None:
+        raise HTTPException(status_code=404, detail="Experiment not found")
+
+    rows, dims = _build_rows(exp_id)
+    if not rows:
+        return {
+            "total_evaluations": 0,
+            "overall": {},
+            "by_trait": {},
+        }
+
+    overall = _score_stats(_aggregate_per_dim(rows, dims))
+
+    by_trait: dict[str, dict[str, dict]] = {}
+    for agent in experiment.config.agents:
+        for trait in agent.distribution.topo_order:
+            key = f"{agent.name}.{trait}"
+            by_trait[key] = _group_by_trait(rows, dims, agent.name, trait)
+
+    return {
+        "total_evaluations": len(rows),
+        "overall": overall,
+        "by_trait": by_trait,
+    }

backend/api/routes/chat.py

@@ -0,0 +1,109 @@
+import json
+
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel, Field
+from sse_starlette.sse import EventSourceResponse
+
+from backend.api.dependencies import db
+from core.config.settings import AVAILABLE_MODELS
+from core.db.queries.sessions import (
+    create_session,
+    create_turn,
+    delete_session,
+    get_session,
+    get_turn_count,
+    get_turns,
+    list_sessions,
+    update_session_model,
+    update_session_title,
+)
+from core.llm.factory import parse_provider_model
+
+router = APIRouter()
+
+
+class CreateSessionRequest(BaseModel):
+    model: str = "mock:default"
+
+
+class ChatRequest(BaseModel):
+    session_id: str
+    message: str = Field(min_length=1, max_length=10000)
+
+
+class UpdateSessionRequest(BaseModel):
+    model: str
+
+
+@router.get("/models")
+def get_available_models():
+    return AVAILABLE_MODELS
+
+
+@router.post("/sessions")
+def create_new_session(request: CreateSessionRequest):
+    return create_session(db, request.model).to_dict()
+
+
+@router.get("/sessions")
+def get_all_sessions():
+    return [s.to_dict() for s in list_sessions(db)]
+
+
+@router.get("/sessions/{session_id}")
+def get_session_detail(session_id: str):
+    session = get_session(db, session_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+    turns = get_turns(db, session_id)
+    return {**session.to_dict(), "turns": [t.to_dict() for t in turns]}
+
+
+@router.patch("/sessions/{session_id}")
+def update_session(session_id: str, request: UpdateSessionRequest):
+    session = get_session(db, session_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+    update_session_model(db, session_id, request.model)
+    return get_session(db, session_id).to_dict()
+
+
+@router.delete("/sessions/{session_id}")
+def delete_session_endpoint(session_id: str):
+    session = get_session(db, session_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+    delete_session(db, session_id)
+    return {"ok": True}
+
+
+@router.post("/chat")
+async def chat(request: ChatRequest):
+    session = get_session(db, request.session_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+
+    turn_number = get_turn_count(db, request.session_id)
+    create_turn(db, request.session_id, "user", request.message, turn_number)
+
+    if turn_number == 0:
+        title = request.message[:50] + ("..." if len(request.message) > 50 else "")
+        update_session_title(db, request.session_id, title)
+
+    turns = get_turns(db, request.session_id)
+    messages = [{"role": t.role, "content": t.content} for t in turns]
+
+    provider, model = parse_provider_model(session.model)
+
+    async def generate():
+        full_response = ""
+        try:
+            async for chunk in provider.stream(messages, model):
+                full_response += chunk
+                yield {"data": json.dumps({"token": chunk})}
+            create_turn(db, request.session_id, "assistant", full_response, turn_number + 1)
+            yield {"event": "done", "data": ""}
+        except Exception as e:
+            yield {"event": "error", "data": json.dumps({"error": str(e)})}
+
+    return EventSourceResponse(generate())

backend/api/routes/distributions.py

@@ -0,0 +1,18 @@
+from fastapi import APIRouter, HTTPException
+
+from backend.api.dependencies import repos
+from core.distribution import distribution_to_dict
+
+router = APIRouter()
+
+
+@router.get("/experiments/{exp_id}/distributions/{agent_name}")
+def get_agent_distribution(exp_id: str, agent_name: str):
+    experiment = repos.experiments.get(exp_id)
+    if experiment is None:
+        raise HTTPException(status_code=404, detail="Experiment not found")
+    try:
+        agent = experiment.config.agent(agent_name)
+    except KeyError:
+        raise HTTPException(status_code=404, detail=f"Agent {agent_name!r} not found in experiment")
+    return {"distribution": distribution_to_dict(agent.distribution)}

backend/api/routes/experiments.py

@@ -0,0 +1,150 @@
+import threading
+
+from fastapi import APIRouter, HTTPException, Query, Response
+from pydantic import BaseModel, Field
+
+from backend.api.dependencies import repos
+from core import Experiment
+from core.analysis.coverage import compute_coverage
+from core.config.settings import APP_SETTINGS
+from core.examples.medical.config import MEDICAL_EXAMPLE_CONFIG
+from core.services.feedback import FeedbackService
+
+router = APIRouter()
+
+_optimize_semaphore = threading.Semaphore(APP_SETTINGS.max_concurrent_optimizations)
+
+
+class CreateExperimentRequest(BaseModel):
+    name: str
+
+
+class SetCurrentOptimizationTargetBody(BaseModel):
+    optimization_target_id: str
+
+
+def _experiment_or_404(exp_id: str):
+    exp = repos.experiments.get(exp_id)
+    if not exp:
+        raise HTTPException(status_code=404, detail="Experiment not found")
+    return exp
+
+
+@router.post("/experiments")
+def post_experiment(request: CreateExperimentRequest):
+    config = MEDICAL_EXAMPLE_CONFIG
+    # Shallow-override the name for this instance.
+    from dataclasses import replace
+    named_config = replace(config, name=request.name)
+    try:
+        exp = Experiment(named_config, repos)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    record = exp.record
+    return record.to_dict(counts=repos.experiments.counts_for(record.id))
+
+
+@router.get("/experiments")
+def get_experiments():
+    exps = repos.experiments.list_all()
+    counts_by_id = repos.experiments.counts_all()
+    return [e.to_dict(counts=counts_by_id.get(e.id)) for e in exps]
+
+
+@router.get("/experiments/{exp_id}")
+def get_experiment_by_id(exp_id: str):
+    exp = _experiment_or_404(exp_id)
+    return exp.to_dict(counts=repos.experiments.counts_for(exp_id))
+
+
+@router.delete("/experiments/{exp_id}")
+def delete_experiment_by_id(exp_id: str):
+    _experiment_or_404(exp_id)
+    repos.experiments.delete(exp_id)
+    return Response(status_code=204)
+
+
+@router.patch("/experiments/{exp_id}")
+def patch_experiment(exp_id: str):
+    """Currently only supports resetting the sample draw index."""
+    _experiment_or_404(exp_id)
+    repos.experiments.reset_sample_draw_index(exp_id)
+    updated = repos.experiments.get(exp_id)
+    assert updated is not None
+    return updated.to_dict(counts=repos.experiments.counts_for(exp_id))
+
+
+# ── Experiment-scoped lists ─────────────────────────────────────────────────
+
+
+@router.get("/experiments/{exp_id}/simulations")
+def list_experiment_simulations(exp_id: str):
+    _experiment_or_404(exp_id)
+    return [s.to_dict() for s in repos.simulations.list_for_experiment(exp_id)]
+
+
+@router.get("/experiments/{exp_id}/evaluations")
+def list_experiment_evaluations(exp_id: str):
+    _experiment_or_404(exp_id)
+    return [e.to_dict() for e in repos.evaluations.list_for_experiment(exp_id)]
+
+
+# ── Optimization targets ────────────────────────────────────────────────────
+
+
+@router.get("/experiments/{exp_id}/optimization-targets")
+def list_experiment_optimization_targets(exp_id: str):
+    _experiment_or_404(exp_id)
+    return [t.to_dict() for t in repos.optimization_targets.list_for_experiment(exp_id)]
+
+
+@router.post("/experiments/{exp_id}/optimization-target/current")
+def set_experiment_current_optimization_target(
+    exp_id: str, body: SetCurrentOptimizationTargetBody
+):
+    _experiment_or_404(exp_id)
+    target = repos.optimization_targets.get(body.optimization_target_id)
+    if target is None or target.experiment_id != exp_id:
+        raise HTTPException(status_code=404, detail="Optimization target not found for this experiment")
+    repos.experiments.set_current_optimization_target(exp_id, body.optimization_target_id)
+    updated = repos.experiments.get(exp_id)
+    assert updated is not None
+    return updated.to_dict(counts=repos.experiments.counts_for(exp_id))
+
+
+# ── Coverage ────────────────────────────────────────────────────────────────
+
+
+@router.get("/experiments/{exp_id}/coverage")
+def get_experiment_coverage(
+    exp_id: str,
+    mc_samples: int = Query(100_000, ge=5_000, le=500_000),
+):
+    exp = _experiment_or_404(exp_id)
+    sims = repos.simulations.list_for_experiment(exp_id)
+    return compute_coverage(sims, exp.config.agents, samples=mc_samples).to_dict()
+
+
+# ── Optimize ────────────────────────────────────────────────────────────────
+
+
+@router.post("/experiments/{exp_id}/optimize")
+async def optimize_experiment(exp_id: str):
+    _experiment_or_404(exp_id)
+
+    if not _optimize_semaphore.acquire(blocking=False):
+        raise HTTPException(
+            status_code=409,
+            detail=(
+                "Another optimization run is in progress "
+                f"(max_concurrent_optimizations={APP_SETTINGS.max_concurrent_optimizations})"
+            ),
+        )
+    try:
+        try:
+            result = await FeedbackService(repos).optimize(exp_id)
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+        return result.to_dict()
+    finally:
+        _optimize_semaphore.release()

backend/api/routes/settings.py

@@ -0,0 +1,12 @@
+from dataclasses import asdict
+
+from fastapi import APIRouter
+
+from core.config.settings import APP_SETTINGS
+
+router = APIRouter()
+
+
+@router.get("/settings")
+def get_settings():
+    return asdict(APP_SETTINGS)