Add FastAPI backend wrapping ML inference pipeline

app/api/routes.py

@@ -0,0 +1,288 @@
+"""
+FastAPI routes for Agent Trace Anomaly Detection.
+
+All ML logic lives in scripts/inference.py (partner's code).
+This file only handles HTTP ↔ inference translation.
+
+Endpoints
+---------
+GET  /health                 – service health & loaded model info
+POST /models/load            – load a model (xgboost or distilbert)
+POST /predict                – predict anomaly for a single trace
+POST /predict/batch          – predict anomalies for multiple traces
+POST /predict/compare        – run both models on same trace, compare
+POST /pipeline/train         – trigger the full training pipeline
+GET  /pipeline/status        – check if models dir has trained models
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import logging
+from typing import Any
+
+from fastapi import APIRouter, HTTPException, Query
+
+from app.schemas import (
+    PredictRequest,
+    PredictBatchRequest,
+    PredictResponse,
+    PredictBatchResponse,
+    CompareResponse,
+    ModelLoadRequest,
+    HealthResponse,
+)
+
+logger = logging.getLogger(__name__)
+router = APIRouter()
+
+
+# ── In-memory state ──────────────────────────────────────────────────────────
+
+_state: dict[str, Any] = {
+    "detector": None,          # TraceAnomalyDetector instance
+    "model_type": None,        # "xgboost" or "distilbert"
+    "model_dir": "models",     # path to saved models
+}
+
+
+def _get_detector():
+    """Return the loaded detector or raise 409."""
+    if _state["detector"] is None:
+        raise HTTPException(
+            status_code=409,
+            detail="No model loaded. POST /models/load first.",
+        )
+    return _state["detector"]
+
+
+def _check_available_models(model_dir: str) -> list[str]:
+    """Check which trained models exist on disk."""
+    available = []
+    if os.path.exists(os.path.join(model_dir, "xgboost_model.joblib")):
+        available.append("xgboost")
+    if os.path.exists(os.path.join(model_dir, "distilbert_trace", "trace_config.json")):
+        available.append("distilbert")
+    if os.path.exists(os.path.join(model_dir, "naive_baseline.joblib")):
+        available.append("naive_baseline")
+    return available
+
+
+# ── Health ───────────────────────────────────────────────────────────────────
+
+@router.get("/health", response_model=HealthResponse, tags=["System"])
+def health_check():
+    model_dir = _state["model_dir"]
+    return HealthResponse(
+        status="ok",
+        loaded_model=_state["model_type"],
+        available_models=_check_available_models(model_dir),
+        model_dir=model_dir,
+    )
+
+
+# ── Model Loading ───────────────────────────────────────────────────────────
+
+@router.post("/models/load", tags=["Models"])
+def load_model(req: ModelLoadRequest):
+    """
+    Load a trained model into memory for inference.
+    Models must be trained first via the pipeline (setup.py).
+    """
+    from scripts.inference import TraceAnomalyDetector
+
+    available = _check_available_models(req.model_dir)
+    if req.model_type not in available:
+        raise HTTPException(
+            status_code=404,
+            detail=(
+                f"Model '{req.model_type}' not found in '{req.model_dir}/'. "
+                f"Available models: {available}. "
+                f"Run the training pipeline first: python setup.py"
+            ),
+        )
+
+    try:
+        detector = TraceAnomalyDetector(
+            model_dir=req.model_dir,
+            model_type=req.model_type,
+        )
+        _state["detector"] = detector
+        _state["model_type"] = req.model_type
+        _state["model_dir"] = req.model_dir
+
+        return {
+            "message": f"Model '{req.model_type}' loaded successfully",
+            "model_type": req.model_type,
+            "model_dir": req.model_dir,
+        }
+    except Exception as exc:
+        logger.exception("Failed to load model")
+        raise HTTPException(status_code=500, detail=f"Failed to load model: {exc}")
+
+
+# ── Prediction ───────────────────────────────────────────────────────────────
+
+@router.post("/predict", response_model=PredictResponse, tags=["Prediction"])
+def predict(req: PredictRequest):
+    """
+    Predict whether a single agent trace is anomalous.
+
+    Pass conversations in ShareGPT/ToolBench format:
+    [{"from": "user", "value": "..."}, {"from": "assistant", "value": "..."}, ...]
+    """
+    detector = _get_detector()
+
+    try:
+        result = detector.predict(req.conversations)
+    except Exception as exc:
+        logger.exception("Prediction failed")
+        raise HTTPException(status_code=500, detail=f"Prediction error: {exc}")
+
+    return PredictResponse(
+        is_anomalous=result["is_anomalous"],
+        confidence=result["confidence"],
+        label=result["label"],
+        anomaly_signals=result.get("anomaly_signals", []),
+        model_used=_state["model_type"],
+        features=result.get("features"),
+    )
+
+
+@router.post("/predict/batch", response_model=PredictBatchResponse, tags=["Prediction"])
+def predict_batch(req: PredictBatchRequest):
+    """Predict anomalies for multiple traces at once."""
+    detector = _get_detector()
+
+    try:
+        results = detector.predict_batch(req.traces)
+    except Exception as exc:
+        logger.exception("Batch prediction failed")
+        raise HTTPException(status_code=500, detail=f"Batch prediction error: {exc}")
+
+    predictions = [
+        PredictResponse(
+            is_anomalous=r["is_anomalous"],
+            confidence=r["confidence"],
+            label=r["label"],
+            anomaly_signals=r.get("anomaly_signals", []),
+            model_used=_state["model_type"],
+            features=r.get("features"),
+        )
+        for r in results
+    ]
+
+    return PredictBatchResponse(
+        predictions=predictions,
+        anomaly_count=sum(1 for p in predictions if p.is_anomalous),
+        total=len(predictions),
+    )
+
+
+@router.post("/predict/compare", response_model=CompareResponse, tags=["Prediction"])
+def predict_compare(req: PredictRequest):
+    """
+    Run both XGBoost and DistilBERT on the same trace and compare results.
+    Both models must be trained and available in the models directory.
+    """
+    from scripts.inference import TraceAnomalyDetector
+
+    model_dir = _state["model_dir"]
+    available = _check_available_models(model_dir)
+    results = {}
+
+    for model_type in ["xgboost", "distilbert"]:
+        if model_type in available:
+            try:
+                det = TraceAnomalyDetector(model_dir=model_dir, model_type=model_type)
+                r = det.predict(req.conversations)
+                results[model_type] = PredictResponse(
+                    is_anomalous=r["is_anomalous"],
+                    confidence=r["confidence"],
+                    label=r["label"],
+                    anomaly_signals=r.get("anomaly_signals", []),
+                    model_used=model_type,
+                    features=r.get("features"),
+                )
+            except Exception as exc:
+                logger.warning("Compare: %s failed: %s", model_type, exc)
+
+    if not results:
+        raise HTTPException(
+            status_code=404,
+            detail=f"No trained models found in '{model_dir}/'. Run the training pipeline first.",
+        )
+
+    xgb = results.get("xgboost")
+    bert = results.get("distilbert")
+    agreement = True
+    if xgb and bert:
+        agreement = xgb.label == bert.label
+
+    return CompareResponse(
+        xgboost=xgb,
+        distilbert=bert,
+        agreement=agreement,
+    )
+
+
+# ── Training Pipeline ────────────────────────────────────────────────────────
+
+@router.post("/pipeline/train", tags=["Pipeline"])
+def trigger_training(
+    max_samples: int | None = Query(None, ge=100, description="Cap on dataset rows"),
+    model: str = Query("all", description="Which model to train: all, naive, classical, deep"),
+):
+    """
+    Trigger the full training pipeline (setup.py).
+
+    This runs data download → feature extraction → model training.
+    May take several minutes depending on dataset size and model choice.
+    """
+    cmd = ["python", "setup.py"]
+    if max_samples:
+        cmd.extend(["--max_samples", str(max_samples)])
+    if model != "all":
+        cmd.extend(["--step", "train", "--model", model])
+
+    try:
+        result = subprocess.run(
+            cmd, capture_output=True, text=True, timeout=1800,  # 30 min max
+        )
+        return {
+            "message": "Training pipeline completed" if result.returncode == 0 else "Pipeline failed",
+            "returncode": result.returncode,
+            "stdout": result.stdout[-3000:] if result.stdout else "",  # last 3000 chars
+            "stderr": result.stderr[-1000:] if result.stderr else "",
+        }
+    except subprocess.TimeoutExpired:
+        raise HTTPException(status_code=504, detail="Training pipeline timed out (30 min limit)")
+    except FileNotFoundError:
+        raise HTTPException(
+            status_code=500,
+            detail="setup.py not found. Make sure you're running from the OffRails project root.",
+        )
+
+
+@router.get("/pipeline/status", tags=["Pipeline"])
+def pipeline_status():
+    """Check what trained models and data files are available."""
+    model_dir = _state["model_dir"]
+    data_dir = "data/processed"
+
+    data_files = {}
+    if os.path.isdir(data_dir):
+        for f in os.listdir(data_dir):
+            path = os.path.join(data_dir, f)
+            data_files[f] = {
+                "size_mb": round(os.path.getsize(path) / 1_048_576, 2),
+            }
+
+    return {
+        "available_models": _check_available_models(model_dir),
+        "loaded_model": _state["model_type"],
+        "model_dir": model_dir,
+        "data_files": data_files,
+    }

app/main.py

@@ -0,0 +1,84 @@
+"""
+Agent Trace Anomaly Detection — FastAPI Backend
+
+This is the API layer that wraps the ML pipeline built in scripts/.
+All model training, feature extraction, and inference logic lives
+in the partner's code (scripts/inference.py). This file just serves it.
+
+Run from the OffRails project root:
+    uvicorn app.main:app --reload --host 0.0.0.0 --port 8000
+
+Interactive docs:
+    http://localhost:8000/docs
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import logging
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+# ── Make partner's scripts/ importable ───────────────────────────────────────
+# inference.py does `from model import ...` and `from build_features import ...`
+# so we need scripts/ on sys.path.
+SCRIPTS_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "scripts")
+if SCRIPTS_DIR not in sys.path:
+    sys.path.insert(0, SCRIPTS_DIR)
+
+from app.api.routes import router
+
+# ── Logging ──────────────────────────────────────────────────────────────────
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s  %(levelname)-8s  %(name)s — %(message)s",
+    datefmt="%H:%M:%S",
+)
+
+# ── App ──────────────────────────────────────────────────────────────────────
+
+app = FastAPI(
+    title="Agent Trace Anomaly Detection API",
+    description=(
+        "Detects anomalous agent execution traces — unnecessary tool calls, "
+        "circular reasoning, and goal drift.\n\n"
+        "**ML models** (XGBoost, DistilBERT) are trained via the pipeline in `scripts/`.\n"
+        "**This API** serves predictions from those trained models.\n\n"
+        "## Workflow\n"
+        "1. Train models: `python setup.py` (or `POST /pipeline/train`)\n"
+        "2. Load a model: `POST /models/load`\n"
+        "3. Predict: `POST /predict`\n"
+        "4. Compare models: `POST /predict/compare`\n"
+    ),
+    version="1.0.0",
+)
+
+# Allow Gradio / any frontend to call the API
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+app.include_router(router)
+
+
+# ── Root ─────────────────────────────────────────────────────────────────────
+
+@app.get("/", include_in_schema=False)
+def root():
+    return {
+        "service": "Agent Trace Anomaly Detection API",
+        "docs": "/docs",
+        "workflow": [
+            "1. Train models: python setup.py",
+            "2. POST /models/load  (load xgboost or distilbert)",
+            "3. POST /predict      (classify a trace)",
+            "4. POST /predict/compare (run both models)",
+        ],
+    }

app/schemas.py

@@ -0,0 +1,123 @@
+"""
+Pydantic schemas for the Agent Trace Anomaly Detection API.
+
+Matches the interface defined in scripts/inference.py:
+    TraceAnomalyDetector.predict() returns:
+        is_anomalous: bool
+        confidence:   float
+        label:        int (0=normal, 1=anomalous)
+        anomaly_signals: list[str]
+        features:     dict (xgboost only)
+"""
+
+from pydantic import BaseModel, Field
+from typing import Optional
+
+
+# ── Request Schemas ──────────────────────────────────────────────────────────
+
+class TraceMessage(BaseModel):
+    """
+    A single message in an agent execution trace.
+    Accepts both ToolBench/ShareGPT format ('from'/'value')
+    and OpenAI format ('role'/'content').
+    """
+    role: Optional[str] = Field(None, alias="from", description="Message role (OpenAI format)")
+    value: Optional[str] = Field(None, description="Message content (ShareGPT format)")
+    content: Optional[str] = Field(None, description="Message content (OpenAI format)")
+
+    model_config = {"populate_by_name": True}
+
+    def to_dict(self) -> dict:
+        """Normalize to the format inference.py expects."""
+        d = {}
+        if self.role is not None:
+            d["from"] = self.role
+        if self.value is not None:
+            d["value"] = self.value
+        if self.content is not None:
+            d["value"] = self.content  # map content → value for ToolBench compat
+            if "from" not in d and self.role:
+                d["from"] = self.role
+        return d
+
+
+class PredictRequest(BaseModel):
+    """Request body for single-trace anomaly prediction."""
+    conversations: list[dict] = Field(
+        ...,
+        description=(
+            "List of message dicts in ShareGPT/ToolBench format. "
+            "Each dict should have 'from' (role) and 'value' (content) keys."
+        ),
+    )
+
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "conversations": [
+                        {"from": "user", "value": "Find me flights from NYC to London"},
+                        {"from": "assistant", "value": "I'll search for flights using the travel API."},
+                        {"from": "function", "value": '{"flights": [{"price": 450}]}'},
+                        {"from": "assistant", "value": "I found flights starting at $450."},
+                    ]
+                }
+            ]
+        }
+    }
+
+
+class PredictBatchRequest(BaseModel):
+    """Request body for batch prediction on multiple traces."""
+    traces: list[list[dict]] = Field(
+        ..., description="List of traces, each trace is a list of message dicts"
+    )
+
+
+class ModelLoadRequest(BaseModel):
+    """Request to load a specific model type."""
+    model_type: str = Field(
+        "xgboost",
+        description="Model to load: 'xgboost' or 'distilbert'",
+        pattern="^(xgboost|distilbert)$",
+    )
+    model_dir: str = Field("models", description="Path to saved models directory")
+
+
+# ── Response Schemas ─────────────────────────────────────────────────────────
+
+class PredictResponse(BaseModel):
+    """Response from the anomaly detector — mirrors TraceAnomalyDetector.predict() output."""
+    is_anomalous: bool = Field(..., description="True if the trace is predicted anomalous")
+    confidence: float = Field(..., ge=0.0, le=1.0, description="Probability of anomaly")
+    label: int = Field(..., description="0 = normal, 1 = anomalous")
+    anomaly_signals: list[str] = Field(
+        default_factory=list,
+        description="Human-readable explanations of detected anomaly patterns",
+    )
+    model_used: str = Field(..., description="Which model produced this prediction")
+    features: Optional[dict] = Field(
+        None, description="Extracted feature values (xgboost only)"
+    )
+
+
+class PredictBatchResponse(BaseModel):
+    """Response for batch predictions."""
+    predictions: list[PredictResponse]
+    anomaly_count: int
+    total: int
+
+
+class CompareResponse(BaseModel):
+    """Side-by-side prediction from both models on the same trace."""
+    xgboost: Optional[PredictResponse] = None
+    distilbert: Optional[PredictResponse] = None
+    agreement: bool = Field(..., description="Whether both models agree on the label")
+
+
+class HealthResponse(BaseModel):
+    status: str
+    loaded_model: Optional[str] = None
+    available_models: list[str]
+    model_dir: str