server/main.py

"""FastAPI app — JSON API for ProcureMind (Hugging Face Space / local dev)."""

from __future__ import annotations

import os
from contextlib import asynccontextmanager
from typing import Any

from fastapi import FastAPI
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from server.agent import coerce_payload, run_agent
from server.catalog import connect, get_commodity, summarize_row
from server.form_schema import ensure_form_schema_table, get_or_create_schema
from server.pr_lines import build_pr_rows


@asynccontextmanager
async def lifespan(_: FastAPI):
    from pathlib import Path

    p = os.environ.get("UNSPSC_DB_PATH", "")
    if p and not Path(p).exists():
        print(f"WARNING: catalogue database path does not exist: {p}")
    try:
        conn = connect()
        ensure_form_schema_table(conn)
        conn.close()
    except Exception as ex:
        print(f"WARNING: could not init auxiliary tables: {ex}")
    yield


app = FastAPI(title="ProcureMind API", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class ChatRequest(BaseModel):
    message: str = ""
    selected_commodity_code: int | None = Field(
        default=None,
        description="Skip LLM and lock this catalogue commodity (disambiguation choice)",
    )


class BuildPrRequest(BaseModel):
    commodity_code: int
    dynamic_values: dict[str, Any] = Field(default_factory=dict)
    deliveries: int = 4
    interval: str = "Quarterly"
    other_spec: str = ""
    year: int = 2026


@app.get("/api/health")
def health():
    return {"status": "ok"}


def _enrich_found(conn, payload: dict) -> dict:
    if payload.get("status") != "found":
        return payload
    code = payload.get("commodity_code")
    if code is None or payload.get("selected_details"):
        return payload
    row = get_commodity(conn, int(code))
    if row:
        payload["selected_details"] = summarize_row(row)
    return payload


@app.get("/api/form-schema/{commodity_code}")
def form_schema(commodity_code: int):
    conn = connect()
    try:
        return get_or_create_schema(conn, commodity_code)
    finally:
        conn.close()


@app.post("/api/build-pr")
def build_pr(req: BuildPrRequest):
    conn = connect()
    try:
        schema = get_or_create_schema(conn, req.commodity_code)
        fields = schema.get("fields") or []
        rows = build_pr_rows(
            mat_grp=req.commodity_code,
            dynamic_fields_ordered=fields,
            dynamic_values=req.dynamic_values,
            deliveries=req.deliveries,
            interval=req.interval,
            other_spec=req.other_spec,
            year=req.year,
        )
        return {"rows": rows}
    finally:
        conn.close()


@app.post("/api/chat")
def chat(req: ChatRequest):
    msg = (req.message or "").strip()
    if not msg and req.selected_commodity_code is None:
        return JSONResponse(
            {"detail": "Provide message or selected_commodity_code"},
            status_code=400,
        )
    conn = connect()
    try:
        payload = run_agent(
            msg,
            conn=conn,
            selected_code=req.selected_commodity_code,
        )
        payload = _enrich_found(conn, payload)
        return coerce_payload(payload)
    finally:
        conn.close()