pipeline_api.py

"""HTTP wrapper around the LangGraph pipeline (pipeline.py).

Drives the full SCAN -> TRIAGE -> ENSEMBLE -> SYNTHESISE -> HUMAN CHECKPOINT loop
from the browser:

    POST /pipeline/start   {question, question_type} -> runs to the checkpoint
    POST /pipeline/resume  {thread_id, resume}       -> "" accepts, "override 0.XX"

Robustness:
* the ensemble is computed IN-PROCESS (no HTTP self-call), avoiding
  localhost/IPv6 and nested-request issues.
* any failure is returned as a readable {"status":"error","error":...} with
  HTTP 200, so the UI shows the message instead of a raw 500.
"""

from __future__ import annotations

import traceback
import uuid
from typing import Literal, Optional

from fastapi import APIRouter
from pydantic import BaseModel, Field

import pipeline

router = APIRouter(prefix="/pipeline", tags=["pipeline"])

_SNAPSHOT_KEYS = [
    "question", "question_type", "horizon", "market_context", "context_source",
    "eligible_market", "eligibility_status", "hits", "scored", "top", "arbitrage",
    "ensemble", "use_forecast", "ensemble_trend", "brier_estimate",
    "market_consensus", "final_probability", "rationale", "history_source",
    "history_status", "decision_note", "forecast_id",
]

_GRAPH = None
_PATCHED = False


def _graph():
    global _GRAPH
    if _GRAPH is None:
        _GRAPH = pipeline.build_graph()
    return _GRAPH


def _use_inprocess_forecast():
    """Replace the pipeline's HTTP forecast call with a direct in-process call."""
    global _PATCHED
    if _PATCHED:
        return

    def _inprocess(prices, question_type, horizon):
        import main  # function-level import avoids a circular import at load
        req = main.ForecastRequest(
            prices=list(prices or []),
            question_type=question_type,
            horizon=horizon,
        )
        return main.forecast(req)

    pipeline.call_forecast = _inprocess
    _PATCHED = True


def _snapshot(graph, config) -> dict:
    try:
        values = graph.get_state(config).values or {}
    except Exception:
        values = {}
    return {k: values.get(k) for k in _SNAPSHOT_KEYS if k in values}


def _interrupt_payload(result):
    if isinstance(result, dict) and result.get("__interrupt__"):
        try:
            return result["__interrupt__"][0].value
        except Exception:
            return None
    return None


def _error(thread_id, graph, config, exc) -> dict:
    return {
        "thread_id": thread_id,
        "status": "error",
        "state": _snapshot(graph, config) if graph is not None else {},
        "checkpoint": None,
        "error": f"{type(exc).__name__}: {exc}",
        "trace": traceback.format_exc()[-1500:],
    }


class StartRequest(BaseModel):
    question: str = Field(min_length=3, max_length=500)
    question_type: Literal["numeric", "event"] = "event"
    horizon: int = Field(default=5, ge=1, le=64)
    market_context: Optional[dict] = None


class ResumeRequest(BaseModel):
    thread_id: str
    resume: str = ""


@router.post("/start")
def start(req: StartRequest) -> dict:
    _use_inprocess_forecast()
    thread_id = f"web-{uuid.uuid4().hex[:10]}"
    try:
        graph = _graph()
    except Exception as exc:
        return _error(thread_id, None, None, exc)

    config = {"configurable": {"thread_id": thread_id}}
    initial = {
        "question": req.question,
        "question_type": req.question_type,
        "horizon": req.horizon,
    }
    if req.market_context is not None:
        initial["market_context"] = req.market_context
    try:
        result = graph.invoke(initial, config=config)
    except Exception as exc:
        return _error(thread_id, graph, config, exc)

    payload = _interrupt_payload(result)
    return {
        "thread_id": thread_id,
        "status": "checkpoint" if payload is not None else "done",
        "state": _snapshot(graph, config),
        "checkpoint": payload,
    }


@router.post("/resume")
def resume(req: ResumeRequest) -> dict:
    from langgraph.types import Command

    try:
        graph = _graph()
    except Exception as exc:
        return _error(req.thread_id, None, None, exc)

    config = {"configurable": {"thread_id": req.thread_id}}
    try:
        existing = graph.get_state(config)
        if existing is None or not existing.values:
            return {
                "thread_id": req.thread_id, "status": "error", "state": {},
                "checkpoint": None, "error": "unknown or expired thread_id",
            }
        result = graph.invoke(Command(resume=req.resume), config=config)
    except Exception as exc:
        return _error(req.thread_id, graph, config, exc)

    payload = _interrupt_payload(result)
    return {
        "thread_id": req.thread_id,
        "status": "checkpoint" if payload is not None else "done",
        "state": _snapshot(graph, config),
        "checkpoint": payload,
    }