main.py

"""FastAPI application exposing /ask and a minimal browser UI."""
from __future__ import annotations

from pathlib import Path
from typing import Any

from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from src.agent import AgentResult, TextToSqlAgent
from src.config import Settings
from src.db import ReadOnlyDB


# ---------- App + DI ----------

settings = Settings.load()
db = ReadOnlyDB(settings.db_path)
agent = TextToSqlAgent(settings=settings, db=db)

app = FastAPI(title="Agentic Text-to-SQL", version="0.1.0")

STATIC_DIR = Path(__file__).resolve().parent.parent / "static"


# ---------- Models ----------

class AskRequest(BaseModel):
    question: str = Field(min_length=1, max_length=2000)
    previous_response_id: str | None = Field(default=None, max_length=200)


class AskResponse(BaseModel):
    answer: str
    sql: str
    reasoning: str
    columns: list[str]
    rows: list[list[Any]]
    row_count_total: int
    truncated: bool
    turns_used: int
    trace: list[dict[str, Any]]
    chart: dict[str, Any] | None = None
    response_id: str | None = None
    web_sources: list[dict[str, str]] = []


# ---------- Routes ----------

@app.get("/health")
def health() -> dict[str, Any]:
    return {
        "status": "ok",
        "model": settings.openai_model,
        "tables": db.list_tables(),
    }


@app.get("/schema")
def schema() -> list[dict[str, Any]]:
    return db.full_schema()


@app.post("/ask", response_model=AskResponse)
def ask(req: AskRequest) -> AskResponse:
    try:
        result: AgentResult = agent.ask(
            req.question,
            previous_response_id=req.previous_response_id,
        )
    except RuntimeError as e:
        raise HTTPException(status_code=422, detail=str(e)) from e
    return AskResponse(
        answer=result.answer,
        sql=result.sql,
        reasoning=result.reasoning,
        columns=result.columns,
        rows=result.rows,
        row_count_total=result.row_count_total,
        truncated=result.truncated,
        turns_used=result.turns_used,
        trace=[step.model_dump() for step in result.trace],
        chart=result.chart,
        response_id=result.response_id,
        web_sources=result.web_sources,
    )


# Static UI (must be mounted last so /ask, /health etc. win route resolution)
if STATIC_DIR.exists():
    app.mount("/", StaticFiles(directory=str(STATIC_DIR), html=True), name="static")

    @app.get("/")
    def index() -> FileResponse:  # explicit redirect-friendly handler
        return FileResponse(STATIC_DIR / "index.html")