M21: API Design & Deployment

# src/models.py
# WHAT: Typed request and response contracts for the agent API
# WHY: Pydantic catches malformed inputs before they reach the agent,
#      and FastAPI generates OpenAPI docs from these classes automatically
# GOTCHA: Use Optional[str] not str | None for Python < 3.10 compatibility

from pydantic import BaseModel, Field, field_validator
from typing import Optional, List, Any
import uuid

class AgentRequest(BaseModel):
    """Incoming agent invocation request."""
    query: str = Field(
        ...,                              # required
        min_length=1,
        max_length=4096,
        description="The question or task to send to the agent",
    )
    session_id: Optional[str] = Field(
        default=None,
        description="Resume a prior conversation. None = fresh context.",
    )
    max_iterations: int = Field(
        default=8,
        ge=1,
        le=20,
        description="Max ReAct loop iterations before forced stop",
    )
    stream: bool = Field(
        default=False,
        description="If true, return SSE stream instead of JSON blob",
    )

    @field_validator("query")
    @classmethod
    def strip_query(cls, v: str) -> str:
        v = v.strip()
        if not v:
            raise ValueError("query must not be empty after stripping whitespace")
        return v


class ToolCallRecord(BaseModel):
    """One tool invocation captured from the agent loop."""
    tool_name: str
    input_summary: str       # truncated for log safety
    output_summary: str
    duration_ms: int


class AgentResponse(BaseModel):
    """Successful synchronous agent response."""
    result: str
    session_id: str = Field(
        default_factory=lambda: str(uuid.uuid4()),
        description="Use this in the next request to continue the conversation",
    )
    iterations: int          = Field(description="ReAct loop iterations used")
    tool_calls: List[ToolCallRecord] = Field(default_factory=list)
    latency_ms: int          = Field(description="Total wall-clock time")
    model: str               = Field(description="Ollama model used")


class ErrorResponse(BaseModel):
    """Standard error envelope returned on 4xx/5xx."""
    error: str
    detail: Optional[str] = None
    request_id: Optional[str] = None

# src/main.py
# WHAT: FastAPI app with CORS, auth, error middleware, and health check
# WHY: Each concern is a separate middleware layer so they can be
#      tested and replaced independently
# GOTCHA: Place the error handler FIRST (outermost) so it catches
#         exceptions from ALL inner middleware, not just route handlers

import os, time, uuid, logging
from fastapi import FastAPI, Request, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from .models import AgentRequest, AgentResponse, ErrorResponse
from .agent import run_agent   # your ReAct loop from Capstone C3

logger = logging.getLogger(__name__)

API_KEY = os.environ.get("API_KEY")
if not API_KEY:
    raise RuntimeError("API_KEY environment variable must be set")

OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://localhost:11434")

app = FastAPI(
    title="Entity Resolution Agent API",
    version="1.0.0",
    docs_url="/docs",          # Swagger UI
    redoc_url="/redoc",
)

# ── CORS ──────────────────────────────────────────────────────────────
# WHAT: Allow browser frontends to call this API directly
# GOTCHA: In production, replace allow_origins=["*"] with your
#         actual frontend domain to prevent cross-origin abuse
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],       # tighten in production
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)

# ── API KEY AUTH ───────────────────────────────────────────────────────
async def verify_api_key(request: Request) -> str:
    """Dependency: extract and validate the Bearer token."""
    auth = request.headers.get("Authorization", "")
    if not auth.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing Bearer token")
    token = auth.removeprefix("Bearer ").strip()
    if token != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")
    return token

# ── GLOBAL ERROR HANDLER ───────────────────────────────────────────────
@app.exception_handler(Exception)
async def global_error_handler(request: Request, exc: Exception):
    request_id = request.headers.get("X-Request-ID", str(uuid.uuid4()))
    logger.exception("Unhandled error on %s (req=%s)", request.url, request_id)
    return JSONResponse(
        status_code=500,
        content=ErrorResponse(
            error="internal_server_error",
            detail=str(exc) if os.environ.get("DEBUG") else None,
            request_id=request_id,
        ).model_dump(),
    )

# ── HEALTH CHECK ───────────────────────────────────────────────────────
@app.get("/health", tags=["ops"])
async def health():
    """Load balancer probe — returns 200 if the service is alive."""
    import httpx
    try:
        r = httpx.get(f"{OLLAMA_HOST}/api/tags", timeout=3)
        ollama_ok = r.status_code == 200
    except Exception:
        ollama_ok = False
    return {
        "status": "ok" if ollama_ok else "degraded",
        "ollama": ollama_ok,
        "version": app.version,
    }

# ── AGENT ROUTE ────────────────────────────────────────────────────────
@app.post(
    "/agent/run",
    response_model=AgentResponse,
    responses={422: {"model": ErrorResponse}, 429: {"model": ErrorResponse}},
    dependencies=[Depends(verify_api_key)],
    tags=["agent"],
)
async def run_agent_endpoint(body: AgentRequest, request: Request):
    """Synchronous agent invocation. For long queries use /agent/jobs."""
    start = time.perf_counter()
    result_data = await run_agent(
        query=body.query,
        session_id=body.session_id,
        max_iterations=body.max_iterations,
        ollama_host=OLLAMA_HOST,
    )
    latency_ms = int((time.perf_counter() - start) * 1000)
    return AgentResponse(latency_ms=latency_ms, **result_data)

# src/streaming.py
# WHAT: SSE endpoint that emits thought / tool_call / result events
# WHY: ReAct agents can take 30-90s; streaming keeps the UI alive
# GOTCHA: async generators must use "yield" not "return"; returning
#         from an async generator closes the stream cleanly

import json, time
from fastapi import APIRouter, Request, Depends
from fastapi.responses import StreamingResponse
from .models import AgentRequest
from .main import verify_api_key, OLLAMA_HOST

router = APIRouter()

async def agent_event_generator(body: AgentRequest):
    """
    WHAT: Converts the ReAct loop into an async SSE generator.
    WHY:  FastAPI's StreamingResponse consumes this generator and
          writes each yielded string directly to the HTTP response
          without buffering the whole output.
    GOTCHA: Every chunk MUST end with \n\n (two newlines). One newline
            splits fields within an event; two newlines end the event.
    """
    from .agent import stream_agent   # yields step dicts

    try:
        async for step in stream_agent(
            query=body.query,
            session_id=body.session_id,
            max_iterations=body.max_iterations,
            ollama_host=OLLAMA_HOST,
        ):
            step_type = step.get("type", "unknown")
            payload = json.dumps(step, ensure_ascii=False)

            # SSE wire format: "event: {type}\ndata: {json}\n\n"
            yield f"event: {step_type}\ndata: {payload}\n\n"

    except Exception as e:
        # Send a final error event before closing the stream
        error_payload = json.dumps({"error": str(e)})
        yield f"event: error\ndata: {error_payload}\n\n"


@router.post(
    "/agent/stream",
    dependencies=[Depends(verify_api_key)],
    tags=["agent"],
    summary="Stream agent steps as SSE",
)
async def stream_agent_endpoint(body: AgentRequest, request: Request):
    """
    WHAT: Wraps the async generator in FastAPI's StreamingResponse
    GOTCHA: headers must include Cache-Control: no-cache or reverse
            proxies (Nginx, Caddy) will buffer the stream
    """
    return StreamingResponse(
        agent_event_generator(body),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",   # disables Nginx proxy buffering
            "Connection": "keep-alive",
        },
    )

// client/stream.js  (browser)
// WHAT: POST to /agent/stream then read SSE events progressively
// WHY: Native EventSource only does GET; we use fetch + ReadableStream
//      for POST bodies so we can send { query, session_id }
// GOTCHA: TextDecoderStream is available in all modern browsers;
//         for Safari < 14.1, polyfill with text-encoding-utf-8

async function streamAgentQuery(query, sessionId = null) {
  const outputDiv = document.getElementById('agent-output');
  outputDiv.innerHTML = '';   // clear previous run

  const response = await fetch('/agent/stream', {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'Authorization': `Bearer ${window.API_KEY}`,
    },
    body: JSON.stringify({ query, session_id: sessionId, max_iterations: 8 }),
  });

  if (!response.ok) {
    outputDiv.textContent = `Error: ${response.status}`;
    return;
  }

  const reader = response.body
    .pipeThrough(new TextDecoderStream())
    .getReader();

  let buffer = '';

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;

    buffer += value;
    // SSE events are terminated by double newline
    const events = buffer.split('\n\n');
    buffer = events.pop();   // keep incomplete event in buffer

    for (const raw of events) {
      if (!raw.trim()) continue;

      // Parse "event: TYPE\ndata: JSON"
      const lines  = raw.split('\n');
      const evType = lines.find(l => l.startsWith('event:'))?.slice(6).trim() ?? 'unknown';
      const dataLine = lines.find(l => l.startsWith('data:'))?.slice(5).trim() ?? '{}';

      let step;
      try { step = JSON.parse(dataLine); } catch { step = { raw: dataLine }; }

      renderStep(evType, step, outputDiv);
    }
  }
}

function renderStep(type, step, container) {
  const div = document.createElement('div');
  div.className = `sse-event sse-${type}`;

  switch (type) {
    case 'thought':
      div.innerHTML = `Thought: ${escapeHtml(step.content ?? '')}`;
      break;
    case 'tool_call':
      div.innerHTML = `Tool: ${escapeHtml(step.tool ?? '')} → ${escapeHtml(step.summary ?? '')}`;
      break;
    case 'result':
      div.innerHTML = `Result: ${escapeHtml(step.content ?? '')}`;
      div.style.borderColor = '#10B981';
      break;
    case 'error':
      div.innerHTML = `Error: ${escapeHtml(step.error ?? '')}`;
      div.style.borderColor = '#F43F5E';
      break;
    default:
      div.textContent = JSON.stringify(step);
  }

  container.appendChild(div);
  div.scrollIntoView({ behavior: 'smooth', block: 'nearest' });
}

function escapeHtml(s) {
  return s.replace(/&/g,'&').replace(//g,'>');
}

# src/jobs.py
# WHAT: In-process asyncio job queue — simple version for dev/staging
# WHY: Zero dependencies; good enough for low-volume or demo deployments
# GOTCHA: State is in-memory. On container restart, pending jobs are lost.
#         For production, replace job_store with Redis and use rq or Celery.

import asyncio, uuid, time
from enum import Enum
from typing import Any, Optional, Dict
from fastapi import APIRouter, BackgroundTasks, Depends
from .models import AgentRequest
from .main import verify_api_key, OLLAMA_HOST
from .agent import run_agent

router = APIRouter()

class JobStatus(str, Enum):
    PENDING   = "pending"
    RUNNING   = "running"
    COMPLETE  = "complete"
    FAILED    = "failed"

# WHAT: In-memory job store — replace with Redis hashes in production
job_store: Dict[str, Dict[str, Any]] = {}


async def _execute_job(job_id: str, body: AgentRequest) -> None:
    """Background coroutine that runs the agent and updates job_store."""
    job_store[job_id]["status"] = JobStatus.RUNNING
    job_store[job_id]["started_at"] = time.time()
    try:
        result = await run_agent(
            query=body.query,
            session_id=body.session_id,
            max_iterations=body.max_iterations,
            ollama_host=OLLAMA_HOST,
        )
        job_store[job_id].update({
            "status": JobStatus.COMPLETE,
            "result": result,
            "finished_at": time.time(),
        })
    except Exception as e:
        job_store[job_id].update({
            "status": JobStatus.FAILED,
            "error": str(e),
            "finished_at": time.time(),
        })


@router.post(
    "/agent/jobs",
    status_code=202,
    dependencies=[Depends(verify_api_key)],
    tags=["agent"],
    summary="Submit a long-running agent job",
)
async def submit_job(body: AgentRequest, background_tasks: BackgroundTasks):
    """
    WHAT: Enqueue job, return job_id immediately with 202 Accepted.
    GOTCHA: BackgroundTasks runs after the response is sent but WITHIN
            the same process. For true isolation, use rq or Celery.
    """
    job_id = str(uuid.uuid4())
    job_store[job_id] = {
        "job_id": job_id,
        "status": JobStatus.PENDING,
        "query": body.query,
        "created_at": time.time(),
    }
    background_tasks.add_task(_execute_job, job_id, body)
    return {"job_id": job_id, "status": JobStatus.PENDING, "poll_url": f"/jobs/{job_id}"}


@router.get(
    "/jobs/{job_id}",
    dependencies=[Depends(verify_api_key)],
    tags=["agent"],
    summary="Poll job status and retrieve result",
)
async def get_job(job_id: str):
    """
    WHAT: Returns current job state. Client polls every 2s until complete.
    GOTCHA: Return 404 for unknown IDs so clients detect bad job_ids
            rather than polling forever on a silent empty response.
    """
    job = job_store.get(job_id)
    if not job:
        from fastapi import HTTPException
        raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
    return job

# src/rate_limit.py
# WHAT: In-process token bucket rate limiter, 10 req/min per API key
# WHY: Protect Ollama from inference overload caused by a noisy caller
# GOTCHA: Use a threading.Lock (or asyncio.Lock) around bucket state;
#         async routes run concurrently so unsynchronized reads are a race

import time, math
from threading import Lock
from typing import Dict
from fastapi import HTTPException, Request

class TokenBucket:
    """Thread-safe token bucket. Capacity=10, refill=10/min."""

    def __init__(self, capacity: int = 10, refill_rate: float = 10 / 60):
        self.capacity    = capacity
        self.refill_rate = refill_rate          # tokens per second
        self._buckets: Dict[str, Dict] = {}
        self._lock = Lock()

    def _get_or_create(self, key: str) -> Dict:
        if key not in self._buckets:
            self._buckets[key] = {"tokens": self.capacity, "last": time.monotonic()}
        return self._buckets[key]

    def consume(self, key: str) -> float:
        """
        WHAT: Try to consume one token. Returns seconds until retry if refused.
        WHY:  Returns a float so the caller can set Retry-After header precisely.
        GOTCHA: math.floor for tokens avoids sub-token accumulation drift.
        """
        with self._lock:
            bucket = self._get_or_create(key)
            now    = time.monotonic()
            elapsed = now - bucket["last"]

            # Refill tokens earned since last call
            new_tokens = elapsed * self.refill_rate
            bucket["tokens"] = min(self.capacity, bucket["tokens"] + new_tokens)
            bucket["last"]   = now

            if bucket["tokens"] >= 1:
                bucket["tokens"] -= 1
                return 0.0          # allowed
            else:
                # seconds until one token is available
                return (1 - bucket["tokens"]) / self.refill_rate


limiter = TokenBucket(capacity=10, refill_rate=10 / 60)


async def rate_limit_dependency(request: Request) -> None:
    """
    WHAT: FastAPI dependency that rejects over-rate requests with 429.
    GOTCHA: Reads the Authorization header to identify the key; use
            a different identifier (IP, user ID) if keys are shared.
    """
    auth = request.headers.get("Authorization", "")
    key  = auth.removeprefix("Bearer ").strip() or request.client.host

    retry_after = limiter.consume(key)
    if retry_after > 0:
        raise HTTPException(
            status_code=429,
            detail={
                "error": "rate_limit_exceeded",
                "message": f"10 req/min limit reached. Retry in {retry_after:.1f}s.",
                "retry_after_seconds": math.ceil(retry_after),
            },
            headers={"Retry-After": str(math.ceil(retry_after))},
        )

# Dockerfile  (Python FastAPI)
# Stage 1: builder — installs deps with build tools available
FROM python:3.12-slim AS builder
WORKDIR /build

# WHAT: Install only what's needed to compile native extensions (e.g. httptools)
# GOTCHA: --no-install-recommends saves ~80MB; removes man pages, etc.
RUN apt-get update && apt-get install -y --no-install-recommends \
    gcc g++ && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .
RUN pip install --user --no-cache-dir -r requirements.txt

# Stage 2: runtime — lean image with only the compiled packages
FROM python:3.12-slim AS runtime
WORKDIR /app

# WHAT: Copy compiled site-packages from builder without gcc
COPY --from=builder /root/.local /root/.local
COPY src/ ./src/
COPY .env.example ./.env.example

# WHAT: Create non-root user to follow least-privilege principle
# GOTCHA: Files copied before this RUN are owned by root; chown them
RUN useradd --no-create-home appuser && chown -R appuser /app
USER appuser

# WHAT: HEALTHCHECK tells Docker and docker-compose when the app is ready
# WHY: depends_on: condition: service_healthy waits for this to pass
HEALTHCHECK --interval=15s --timeout=5s --start-period=20s --retries=3 \
  CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" \
  || exit 1

ENV PATH=/root/.local/bin:$PATH \
    PYTHONUNBUFFERED=1

EXPOSE 8000
CMD ["uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "2"]

# docker-compose.yml
# WHAT: Orchestrate api + ollama sidecar on a private bridge network
# WHY: The private network means ollama is NOT exposed to the internet
#      unless you explicitly bind its port — only the api container talks to it
# GOTCHA: The model must be pulled once: docker exec -it ollama ollama pull mistral

services:
  api:
    build:
      context: .
      target: runtime          # use the lean runtime stage
    ports:
      - "8000:8000"
    environment:
      API_KEY:      ${API_KEY}
      OLLAMA_HOST:  http://ollama:11434   # DNS name = service name inside compose
      DEBUG:        "false"
    depends_on:
      ollama:
        condition: service_healthy        # waits for ollama HEALTHCHECK to pass
    networks:
      - agent-network
    restart: unless-stopped
    stop_grace_period: 30s               # WHAT: graceful shutdown window
    # GOTCHA: send SIGTERM, not SIGKILL; uvicorn needs time to finish
    #         in-flight requests. 30s covers even slow Ollama inference.

  ollama:
    image: ollama/ollama:latest
    ports:
      - "11434:11434"                    # expose only for local CLI pulls
    environment:
      OLLAMA_HOST:         "0.0.0.0"    # listen on all interfaces, not just 127.0.0.1
      OLLAMA_NUM_THREADS:  "2"
      OLLAMA_MAX_LOADED_MODELS: "1"
    volumes:
      - ollama-models:/root/.ollama     # persist pulled models across restarts
    healthcheck:
      test: ["CMD", "curl", "-sf", "http://localhost:11434/api/tags"]
      interval: 15s
      timeout: 5s
      retries: 5
      start_period: 30s                  # Ollama needs ~15s to initialize
    networks:
      - agent-network
    restart: unless-stopped

volumes:
  ollama-models:                         # named volume; survives docker-compose down

networks:
  agent-network:
    driver: bridge

#!/bin/bash
# scripts/provision.sh — run once on a fresh Ubuntu 22.04 server
# WHAT: Install Docker, Caddy, pull repo, inject secrets, start stack
# GOTCHA: Run as root or prefix commands with sudo

set -euo pipefail

# 1. Install Docker
curl -fsSL https://get.docker.com | sh
usermod -aG docker $USER

# 2. Install Caddy (reverse proxy with auto-TLS via Let's Encrypt)
apt-get install -y debian-keyring debian-archive-keyring apt-transport-https
curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/gpg.key' | gpg --dearmor -o /usr/share/keyrings/caddy-stable-archive-keyring.gpg
curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/debian.deb.txt' | tee /etc/apt/sources.list.d/caddy-stable.list
apt-get update && apt-get install -y caddy

# 3. Clone repo and create .env
git clone https://github.com/your-org/entity-agent.git /opt/agent
cd /opt/agent

# 4. Inject secrets (never commit these to git!)
cat > .env <