app/api.py

"""FastAPI service - the production interface.

Exposes the CodeAssistant over REST so any client (web app, CI bot, IDE plugin)
can call it. Run locally:
    uvicorn app.api:app --host 0.0.0.0 --port 8000
Then POST to /generate. Interactive docs at /docs.

The model loads once at startup (lifespan) and is reused across requests.
"""
from __future__ import annotations

import os
import sys
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI
from pydantic import BaseModel, Field

sys.path.append(str(Path(__file__).resolve().parents[1]))

# Allow tests / CI to skip the heavy model load.
_STATE: dict = {"assistant": None}


def get_assistant():
    return _STATE["assistant"]


@asynccontextmanager
async def lifespan(app: FastAPI):
    if os.environ.get("CGA_SKIP_MODEL") != "1":
        from src.rag.generator import CodeAssistant
        _STATE["assistant"] = CodeAssistant.from_config(with_index=True)
    yield
    _STATE["assistant"] = None


app = FastAPI(title="Code Generation Assistant", version="1.0", lifespan=lifespan)


class GenerateRequest(BaseModel):
    intent: str = Field(..., description="Natural-language description of the code")
    mode: str = Field("rag", description="'rag' or 'baseline'")
    repair: bool = Field(False, description="Run the agentic repair loop")
    return_sources: bool = Field(True)


class GenerateResponse(BaseModel):
    code: str
    mode: str
    sources: list = []


@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": _STATE["assistant"] is not None}


@app.post("/generate", response_model=GenerateResponse)
def generate(req: GenerateRequest):
    assistant = get_assistant()
    if assistant is None:
        return GenerateResponse(code="# model not loaded", mode=req.mode)

    if req.repair:
        from src.agent.repair_loop import make_repair_generator, repair_loop
        gen_fn = make_repair_generator(assistant)
        trace = repair_loop(req.intent, gen_fn,
                            check_program_fn=lambda c: c, max_iters=3)
        return GenerateResponse(code=trace.final_code, mode=f"{req.mode}+repair")

    result = assistant.generate(req.intent, mode=req.mode,
                                return_sources=req.return_sources)
    if isinstance(result, tuple):
        code, sources = result
        return GenerateResponse(code=code, mode=req.mode, sources=sources)
    return GenerateResponse(code=result, mode=req.mode)