backend/main.py

"""
SQL Agent OpenEnv — FastAPI entry point.

Start with:
  uvicorn main:app --reload --port 8000

Environment variables:
  API_BASE_URL   — OpenAI-compatible base URL
  MODEL_NAME     — model name
  HF_TOKEN       — API key / bearer token
  DATA_DIR       — override data directory (default: ./data)
"""

from __future__ import annotations

import logging
import os
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles

from api.demo import router as demo_router
from api.openenv import router as openenv_router, ResetRequest, StepRequest, env_reset, env_step, env_state
from env.database import ensure_seeded

app = FastAPI(
    title="Self-Improving SQL Agent",
    description=(
        "A SQL generation environment that learns from its own mistakes. "
        "Powered by a LinUCB contextual bandit for repair strategy selection "
        "and GEPA prompt evolution for continuous self-improvement."
    ),
    version="1.0.0",
)

# ─── CORS ────────────────────────────────────────────────────────

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

# ─── Routers ─────────────────────────────────────────────────────

app.include_router(demo_router, prefix="/api", tags=["demo"])
app.include_router(openenv_router, prefix="/env", tags=["openenv"])

# ─── Top-level OpenEnv aliases (required by openenv validate + pre-validation) ─
# The validator pings POST <url>/reset — these mirror /env/* without the prefix.

@app.post("/reset", tags=["openenv"])
async def root_reset(req: ResetRequest = None):
    return await env_reset(req or ResetRequest())


@app.post("/step", tags=["openenv"])
async def root_step(req: StepRequest = None):
    return await env_step(req or StepRequest())


@app.get("/state", tags=["openenv"])
async def root_state():
    return await env_state()


# ─── Health check ────────────────────────────────────────────────

@app.get("/health", tags=["system"])
async def health():
    return {"status": "healthy", "service": "sql-agent-openenv"}


# ─── Startup ─────────────────────────────────────────────────────

@app.on_event("startup")
async def startup_event():
    """Seed the database on first startup."""
    try:
        ensure_seeded()
    except Exception as e:
        print(f"Warning: database seed failed: {e}")

    # Log LLM config so it's visible in container logs
    token = os.environ.get("HF_TOKEN")  # no default
    api_base = os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1")
    model = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
    token_status = f"set ({len(token)} chars)" if token else "NOT SET"
    print(f"[startup] LLM config: base_url={api_base} model={model} HF_TOKEN={token_status}", flush=True)


# ─── Static files (frontend) — mount last ─────────────────────────

_frontend_dist = Path(__file__).parent.parent / "frontend" / "dist"
if _frontend_dist.exists():
    app.mount(
        "/",
        StaticFiles(directory=str(_frontend_dist), html=True),
        name="frontend",
    )
else:
    @app.get("/", tags=["system"])
    async def root():
        return {
            "message": "SQL Agent OpenEnv API",
            "docs": "/docs",
            "health": "/health",
            "env_info": "/env/info",
        }