api/server.py

import os
import time
import asyncio
from typing import Optional
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, Response
from pydantic import BaseModel, ValidationError

from env.environment import environment
from env.models import (
    Action, Observation, EpisodeState,
    DifficultyLevel, ActionType,
    StepResponse, ResetResponse, TaskListResponse,
    BaselineResponse, BaselineResult,
    GraderRequest, GraderResponse,
    HealthResponse, TaskInfo, ProgressResponse
)
from env.tasks import task_manager, ACTION_SCHEMA
from env.graders import grade, grade_db_action, _is_scenario_task, _get_scenario


# ─────────────────────────────────────────────
#  STARTUP / SHUTDOWN
# ─────────────────────────────────────────────

_startup_time = time.time()

@asynccontextmanager
async def lifespan(app: FastAPI):
    environment.reset(difficulty="easy")
    yield


# ─────────────────────────────────────────────
#  APP DEFINITION
# ─────────────────────────────────────────────

app = FastAPI(
    title       = "SQL Database Engineer Agent — OpenEnv Environment",
    description = (
        "An OpenEnv-compliant reinforcement learning environment where AI agents "
        "learn to act like senior database engineers. "
        "The agent manages a simulated production database over 50+ steps: "
        "inspecting slow queries, creating indexes, rewriting queries, partitioning tables. "
        "Built for the META x PyTorch x SST OpenEnv Hackathon Finals — April 25-26, Bangalore."
    ),
    version     = "2.0.0",
    lifespan    = lifespan,
    docs_url    = "/docs",
    redoc_url   = "/redoc",
)

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


# ─────────────────────────────────────────────
#  GLOBAL EXCEPTION HANDLER
# ─────────────────────────────────────────────

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    return JSONResponse(
        status_code = 500,
        content     = {"error": str(exc), "type": type(exc).__name__}
    )


# ─────────────────────────────────────────────
#  FAVICON
# ─────────────────────────────────────────────

@app.get("/favicon.ico", include_in_schema=False)
async def favicon():
    return Response(status_code=204)


# ─────────────────────────────────────────────
#  1. /health — GET
# ─────────────────────────────────────────────

@app.get("/health", response_model=HealthResponse, tags=["System"])
async def health():
    """Liveness check. Always returns 200."""
    return HealthResponse(
        status  = "ok",
        version = "2.0.0",
        uptime  = round(time.time() - _startup_time, 2)
    )


# ─────────────────────────────────────────────
#  2. /reset — POST
# ─────────────────────────────────────────────

class ResetBody(BaseModel):
    difficulty: Optional[str] = None
    task_id:    Optional[str] = None

@app.post("/reset", response_model=Observation, tags=["Environment"])
async def reset(body: ResetBody = ResetBody()):
    """
    Starts a fresh episode. Initializes DatabaseSimulator.
    Returns the initial Observation with DB state and slow queries.
    """
    try:
        obs = environment.reset(
            difficulty = body.difficulty,
            task_id    = body.task_id
        )
        return obs
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Reset failed: {str(e)}")


# ─────────────────────────────────────────────
#  3. /step — POST
# ─────────────────────────────────────────────

@app.post("/step", response_model=StepResponse, tags=["Environment"])
async def step(action: Action):
    """
    Submits an action to the environment.
    Round 2 actions: inspect_query, create_index, rewrite_query,
    partition_table, analyze_statistics, analyze_indexes, submit_report.
    Returns (observation, reward, done, info) with DB performance delta.
    """
    try:
        response = environment.step(action)
        return response
    except ValidationError as e:
        from env.models import Reward
        return StepResponse(
            observation = environment._build_observation(),
            reward      = Reward(
                score     = 0.001,
                breakdown = {"validation_error": 0.001},
                feedback  = f"Malformed action: {str(e)}"
            ),
            done = False,
            info = {"error": "validation_error", "detail": str(e)}
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Step failed: {str(e)}")


# ─────────────────────────────────────────────
#  4. /state — GET
# ─────────────────────────────────────────────

@app.get("/state", response_model=EpisodeState, tags=["Environment"])
async def state():
    """Returns full current environment state including performance history."""
    return environment.state()


# ─────────────────────────────────────────────
#  5. /tasks — GET
# ─────────────────────────────────────────────

@app.get("/tasks", response_model=TaskListResponse, tags=["Tasks"])
async def tasks():
    """
    Lists all 30 tasks (15 Round 2 scenarios + 15 Round 1 cases).
    Includes complete action schema for all 15 action types.
    """
    all_tasks = task_manager.list_all_tasks()
    return TaskListResponse(
        tasks        = all_tasks,
        total        = len(all_tasks),
        action_types = [a.value for a in ActionType]
    )


# ─────────────────────────────────────────────
#  6. /grader — POST  (FIXED)
# ─────────────────────────────────────────────

@app.post("/grader", response_model=GraderResponse, tags=["Grading"])
async def grader(request: GraderRequest):
    """
    Grades an action for a given task_id. STATELESS — does not change episode state.

    Routing:
      Round 2 scenario IDs (easy_s001, medium_s002, hard_s003):
        - submit_report   → computes score from current DB performance delta
        - all other types → grade_db_action() scores action quality vs scenario

      Round 1 task IDs (easy_001, medium_001, hard_001):
        → grade() → grade_easy/medium/hard() (original Round 1 graders)

    Score is ALWAYS strictly between 0.001 and 0.999.
    NEVER crashes — all exceptions caught and returned as 0.001.

    FIXES applied vs original:
      - Round 2 non-terminal actions now route to grade_db_action() instead of
        grade_easy() which was looking for "fixed_query" in Round 2 payloads
        and returning 0.001 for every create_index / analyze_indexes / inspect_query
      - submit_report score now uses db_simulator state from environment directly
        instead of brittle action_counts dict lookup which could be empty or stale
    """
    try:
        if request.action is None:
            return GraderResponse(
                score     = 0.001,
                feedback  = "No action provided for grading.",
                breakdown = {"error": "null_action"}
            )

        task_id     = request.task_id or ""
        action_type = (
            request.action.action_type.value
            if hasattr(request.action.action_type, "value")
            else str(request.action.action_type)
        )

        # ── ROUND 2: DB ENGINEERING SCENARIO ─────────────────────
        if _is_scenario_task(task_id):

            # submit_report: use live DB state from environment simulator
            if action_type == "submit_report":
                return _grade_submit_report(request, task_id)

            # All other Round 2 actions: stateless scenario-aware grading
            score, breakdown, feedback = grade_db_action(request.action, task_id)
            score = max(0.001, min(0.999, score))
            return GraderResponse(score=score, feedback=feedback, breakdown=breakdown)

        # ── ROUND 1: SQL DEBUGGING TASK ───────────────────────────
        score, breakdown, feedback = grade(request.action, task_id)
        score = max(0.001, min(0.999, score))
        return GraderResponse(score=score, feedback=feedback, breakdown=breakdown)

    except Exception as e:
        return GraderResponse(
            score     = 0.001,
            feedback  = f"Grader error: {str(e)}",
            breakdown = {"error": str(e)}
        )


def _grade_submit_report(request: GraderRequest, task_id: str) -> GraderResponse:
    """
    Grade a submit_report action for a Round 2 scenario.

    Score components:
      60% — performance improvement (baseline → current)
      20% — step efficiency (fewer steps = higher bonus)
      10% — base credit for submitting
      10% — report summary quality

    Falls back gracefully if DB simulator state is unavailable.
    """
    try:
        ep_state = environment.state()

        # Get performance data from environment state
        # Use action_counts as the store (set by environment.py during steps)
        ac          = ep_state.action_counts or {}
        perf_history = ac.get("_perf_history", [])
        baseline     = float(ac.get("_baseline_score", 0.0))
        current      = float(perf_history[-1]) if perf_history else baseline
        steps_used   = ep_state.step_count
        max_steps    = 50  # Round 2 default

        # If no perf history (called before reset, or env in wrong state):
        # fall back to scenario-based quality score
        if not perf_history or baseline == 0.0:
            scenario = _get_scenario(task_id)
            if scenario:
                baseline = float(scenario.get("performance_score_baseline", 0.0))
                target   = float(scenario.get("target_score", 85.0))
                # Score based on report quality only
                summary    = str((request.action.payload or {}).get("summary", ""))
                base_score = 0.15 + min(len(summary) / 400, 0.25)
                return GraderResponse(
                    score    = round(max(0.001, min(0.999, base_score)), 4),
                    feedback = (
                        f"Report graded on quality only (episode state unavailable). "
                        f"Run a full episode via /reset then /step to get performance-based score."
                    ),
                    breakdown = {"report_quality": round(base_score, 4), "note": "no_episode_state"}
                )

        max_possible     = max(1.0, 100.0 - baseline)
        perf_improvement = max(0.0, (current - baseline) / max_possible)
        step_efficiency  = max(0.0, 1.0 - (steps_used / max(1, max_steps)))
        summary          = str((request.action.payload or {}).get("summary", ""))
        report_quality   = min(len(summary) / 300, 0.10) if summary else 0.0

        raw_score = (
            perf_improvement * 0.60
            + step_efficiency * 0.20
            + 0.10                  # base credit
            + report_quality        # up to 0.10
        )
        score = round(max(0.001, min(0.999, raw_score)), 4)

        return GraderResponse(
            score    = score,
            feedback = (
                f"DB performance: {baseline:.1f} → {current:.1f} "
                f"(improvement: {perf_improvement*100:.1f}%). "
                f"Steps used: {steps_used}/{max_steps}. "
                f"Efficiency: {step_efficiency*100:.1f}%."
            ),
            breakdown = {
                "perf_improvement": round(perf_improvement, 4),
                "step_efficiency":  round(step_efficiency,  4),
                "base_credit":      0.10,
                "report_quality":   round(report_quality,   4),
            }
        )

    except Exception as e:
        # Last resort — don't return an error, return a low but non-zero score
        return GraderResponse(
            score     = 0.10,
            feedback  = f"Submit report scored with fallback (error: {str(e)}).",
            breakdown = {"fallback": 0.10, "error": str(e)}
        )


# ─────────────────────────────────────────────
#  7. /baseline — POST
# ─────────────────────────────────────────────

@app.post("/baseline", response_model=BaselineResponse, tags=["Baseline"])
async def baseline():
    """
    Runs the baseline agent against all difficulty levels.
    Must complete within 60 seconds.
    """
    try:
        import baseline as baseline_module
        results = await asyncio.wait_for(
            asyncio.to_thread(baseline_module.run_baseline),
            timeout=55.0
        )
        return results
    except asyncio.TimeoutError:
        return BaselineResponse(
            results=[BaselineResult(
                task_id="timeout", difficulty=DifficultyLevel.EASY,
                score=0.0, steps=0, feedback="Baseline timed out."
            )],
            average_score=0.0
        )
    except Exception as e:
        return BaselineResponse(
            results=[BaselineResult(
                task_id="error", difficulty=DifficultyLevel.EASY,
                score=0.0, steps=0, feedback=f"Baseline error: {str(e)}"
            )],
            average_score=0.0
        )


# ─────────────────────────────────────────────
#  8. /progress — GET  (Round 2)
# ─────────────────────────────────────────────

@app.get("/progress", response_model=ProgressResponse, tags=["Training"])
async def progress():
    """
    Returns DB performance history for training visualization.
    Used by evaluate_agent.py to generate reward curves.
    Shows improvement from baseline to current score.
    """
    ep_state     = environment.state()
    ac           = ep_state.action_counts or {}
    perf_history = ac.get("_perf_history", [])
    milestones   = ac.get("_milestones", [])
    baseline     = ac.get("_baseline_score", 0.0)
    target       = ac.get("_target_score", 85.0)
    best         = ac.get("_best_score", 0.0)
    current      = perf_history[-1] if perf_history else 0.0

    return ProgressResponse(
        scenario_id         = ep_state.task_id,
        performance_score   = current,
        baseline_score      = baseline,
        target_score        = target,
        improvement_history = perf_history,
        milestones_earned   = milestones,
        best_score          = best,
        steps_used          = ep_state.step_count,
        budget_remaining    = max(0, 50 - ep_state.step_count),
        total_reward        = ep_state.total_reward,
    )


# ─────────────────────────────────────────────
#  ROOT
# ─────────────────────────────────────────────

@app.get("/", tags=["System"])
async def root():
    return {
        "name":        "SQL Database Engineer Agent — OpenEnv Environment",
        "version":     "2.0.0",
        "tagline":     "Training LLMs to act like senior database engineers",
        "docs":        "/docs",
        "health":      "/health",
        "endpoints":   ["/reset", "/step", "/state", "/tasks", "/grader", "/baseline", "/progress", "/health"],
        "hackathon":   "META x PyTorch x SST OpenEnv Hackathon — Finals April 25-26 Bangalore",
        "domain":      "Long-Horizon Database Engineering",
        "tasks_count": 30,
        "max_steps":   50,
        "themes":      ["Long-Horizon Planning", "World Modeling", "Self-Improvement", "Wildcard"],
    }