api.py

import logging
import uuid

from fastapi import FastAPI, HTTPException, Request, Security
from fastapi.security import APIKeyHeader
from pydantic import BaseModel, Field

from src.config import ConfigError, get_settings
from src.graph.workflow import run_task
from src.job_store import count_jobs, create_job, get_job, list_jobs
from src.ratelimit import get_rate_limiter

logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s %(levelname)s %(name)s %(message)s")
log = logging.getLogger("coding-agent.api")

app = FastAPI(title="Coding Assistant Agent API", version="2.0.0")

_api_key_header = APIKeyHeader(name="X-API-Key", auto_error=True)
_limiter = None


@app.on_event("startup")
def _startup() -> None:
    settings = get_settings()
    global _limiter
    _limiter = get_rate_limiter(settings.rate_limit, settings.redis_url)
    if not settings.is_production and settings.api_key == "dev-key":
        log.warning("Using default API_KEY='dev-key' — fine for local, NOT production.")
    log.info("Startup ok env=%s model=%s rate=%s redis=%s",
             settings.environment, settings.groq_model,
             settings.rate_limit, bool(settings.redis_url))


def _verify_key(api_key: str = Security(_api_key_header)) -> str:
    if api_key != get_settings().api_key:
        raise HTTPException(status_code=403, detail="Invalid API key")
    return api_key


def _check_rate(request: Request) -> None:
    if _limiter is None:
        return
    client = request.client.host if request.client else "unknown"
    if not _limiter.allow(client):
        raise HTTPException(status_code=429, detail="Rate limit exceeded")


class SolveRequest(BaseModel):
    task: str = Field(..., min_length=3, max_length=2000)
    thread_id: str | None = Field(None, description="Session ID for conversation memory")


# ── Health ────────────────────────────────────────────────────────────────

@app.get("/health")
def health() -> dict:
    try:
        get_settings()
    except ConfigError as exc:
        return {"status": "error", "detail": str(exc)}
    return {"status": "ok"}


# ── Synchronous solve ─────────────────────────────────────────────────────

@app.post("/solve")
def solve(body: SolveRequest, request: Request,
          _key: str = Security(_verify_key)) -> dict:
    """Solve a coding task synchronously. Blocks until the agent finishes."""
    _check_rate(request)
    log.info("solve task=%r thread=%s", body.task[:80], body.thread_id)
    try:
        result = run_task(body.task.strip(), thread_id=body.thread_id)
    except Exception as exc:
        log.exception("solve failed")
        raise HTTPException(status_code=500, detail=f"Agent error: {exc}")
    return {
        "task": result["task"],
        "final_answer": result.get("final_answer", ""),
        "plan": result.get("plan", []),
        "verdict": result.get("verdict", ""),
        "critique": result.get("critique", ""),
        "code_runs": result.get("code_runs", []),
        "fix_attempts": result.get("fix_attempts", 0),
    }


# ── Async solve (Celery queue) ────────────────────────────────────────────

@app.post("/solve/async", status_code=202)
def solve_async(body: SolveRequest, request: Request,
                _key: str = Security(_verify_key)) -> dict:
    """Queue a coding task. Returns job_id immediately; poll /jobs/{id}."""
    _check_rate(request)
    try:
        from src.tasks import solve_task
    except Exception as exc:
        raise HTTPException(status_code=503,
                            detail=f"Celery/Redis not available: {exc}")
    job_id = str(uuid.uuid4())
    create_job(job_id, body.task.strip())
    solve_task.delay(job_id, body.task.strip())
    log.info("queued job=%s task=%r", job_id, body.task[:80])
    return {"job_id": job_id, "status": "queued"}


@app.get("/jobs/{job_id}")
def get_job_status(job_id: str, _key: str = Security(_verify_key)) -> dict:
    job = get_job(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail="Job not found")
    return {"job_id": job_id, **job}


@app.get("/jobs")
def get_all_jobs(
    limit: int = 20,
    offset: int = 0,
    _key: str = Security(_verify_key),
) -> dict:
    """List jobs with pagination. Default: 20 per page."""
    if limit < 1 or limit > 100:
        raise HTTPException(status_code=400, detail="limit must be 1–100")
    jobs = list_jobs(limit=limit, offset=offset)
    return {
        "jobs": jobs,
        "total": count_jobs(),
        "limit": limit,
        "offset": offset,
    }


# ── Human-in-the-loop review ──────────────────────────────────────────────

class ReviewRequest(BaseModel):
    action: str = Field(..., pattern="^(approve|revise)$")
    feedback: str = Field("", max_length=2000)


@app.post("/jobs/{job_id}/review")
def review_job(job_id: str, body: ReviewRequest,
               _key: str = Security(_verify_key)) -> dict:
    """Inject a human verdict into a HITL-interrupted job.

    Requires the job to have been started via /solve/async and the graph
    to be paused at the critic interrupt. Returns the completed result.
    """
    from src.graph.workflow import resume_task

    job = get_job(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail="Job not found")
    if job.get("status") != "awaiting_review":
        raise HTTPException(
            status_code=409,
            detail=f"Job status is '{job['status']}', expected 'awaiting_review'",
        )
    log.info("review job=%s action=%s", job_id, body.action)
    try:
        result = resume_task(job_id, body.action, feedback=body.feedback)
    except Exception as exc:
        log.exception("review failed job=%s", job_id)
        raise HTTPException(status_code=500, detail=f"Resume error: {exc}")

    from src.job_store import update_job
    update_job(job_id, status="complete", result=result.get("final_answer", ""))
    return {
        "job_id": job_id,
        "action": body.action,
        "final_answer": result.get("final_answer", ""),
        "verdict": result.get("verdict", ""),
    }