app/main.py

"""
OpenEnv — Multi-Domain AI Testing Environment
FastAPI server implementing the full OpenEnv spec across 4 domains.
"""

from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
from typing import Dict, Any, Optional
import os, sys, statistics

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.models import Action, Observation, Reward, AgentAction
from tasks.task1_classification import ClassificationTask
from tasks.task2_drafting import ResponseDraftingTask
from tasks.task3_queue import QueueManagementTask
from tasks.task4_multiturn import MultiTurnTask
from tasks.legal_tasks import LegalIdentifyClauseTask, LegalRiskFlagTask, LegalRedlineTask
from tasks.clinical_tasks import ClinicalTriageTask, ClinicalESITask, ClinicalNoteTask
from tasks.pr_tasks import PRTypeTask, PRBugIdentifyTask, PRReviewTask

# ─── Task Registry ────────────────────────────────────────────────────────────

TASKS = {
    "ticket_classification":        ClassificationTask,
    "response_drafting":            ResponseDraftingTask,
    "queue_management":             QueueManagementTask,
    "multi_turn_conversation":      MultiTurnTask,
    "legal_clause_identification":  LegalIdentifyClauseTask,
    "legal_risk_flagging":          LegalRiskFlagTask,
    "legal_clause_redlining":       LegalRedlineTask,
    "clinical_triage_classification": ClinicalTriageTask,
    "clinical_esi_assignment":      ClinicalESITask,
    "clinical_triage_note":         ClinicalNoteTask,
    "pr_type_classification":       PRTypeTask,
    "pr_bug_identification":        PRBugIdentifyTask,
    "pr_review_comment":            PRReviewTask,
}

_active: Dict[str, Any] = {}
_current_task_id: str = "ticket_classification"


def get_task(task_id: str = None):
    tid = task_id or _current_task_id
    if tid not in _active:
        raise HTTPException(
            status_code=400,
            detail=f"No active episode for task '{tid}'. Call /reset first."
        )
    return _active[tid]


# ─── App ──────────────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    for tid, TaskClass in TASKS.items():
        instance = TaskClass()
        instance.reset()
        _active[tid] = instance
    yield


app = FastAPI(
    title="OpenEnv — Multi-Domain AI Testing Environment",
    description=(
        "Real-world OpenEnv environment testing AI agents across "
        "Customer Support, Legal Review, Clinical Triage, and Software Engineering."
    ),
    version="2.0.0",
    lifespan=lifespan,
)

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


# ─── Core OpenEnv Endpoints ───────────────────────────────────────────────────

@app.post("/reset")
async def reset(task_id: Optional[str] = None) -> Dict[str, Any]:
    global _current_task_id
    tid = task_id or _current_task_id
    if tid not in TASKS:
        raise HTTPException(
            status_code=400,
            detail=f"Unknown task '{tid}'. Valid: {list(TASKS.keys())}"
        )
    _current_task_id = tid
    instance = TASKS[tid]()
    _active[tid] = instance
    obs = instance.reset()
    return {"task_id": tid, "observation": obs.model_dump()}


@app.post("/step")
async def step(action: Action, task_id: Optional[str] = None) -> Dict[str, Any]:
    task = get_task(task_id)
    try:
        obs, reward, done, info = task.step(action)
    except Exception as e:
        raise HTTPException(status_code=422, detail=f"Step failed: {str(e)}")
    return {
        "observation": obs.model_dump(),
        "reward": reward.model_dump(),
        "done": done,
        "info": info,
    }


@app.get("/state")
async def state(task_id: Optional[str] = None) -> Dict[str, Any]:
    task = get_task(task_id)
    return task.state()


# ─── Extended Endpoints ───────────────────────────────────────────────────────

@app.get("/tasks")
async def list_tasks() -> Dict[str, Any]:
    return {
        "tasks": [
            # ── Customer Support ──
            {
                "id": "ticket_classification",
                "name": "Ticket Classification & Routing",
                "difficulty": "easy",
                "description": "Classify incoming support tickets by category and priority.",
                "action_schema": {
                    "action_type": "classify",
                    "ticket_id": "string",
                    "category": "billing|technical|account|feature_request|abuse|unknown",
                    "priority": "P1|P2|P3|P4",
                },
            },
            {
                "id": "response_drafting",
                "name": "Response Drafting & Quality",
                "difficulty": "medium",
                "description": "Draft professional customer-facing responses using KB context.",
                "action_schema": {
                    "action_type": "draft_response",
                    "ticket_id": "string",
                    "response_text": "string (100-300 words)",
                },
            },
            {
                "id": "queue_management",
                "name": "SLA Queue Management",
                "difficulty": "hard",
                "description": "Manage 20 tickets across 3 agents. Maximize SLA compliance and FCR.",
                "action_schema": {
                    "action_type": "assign_ticket|resolve|escalate|close|no_op",
                    "ticket_id": "string",
                    "target_agent_id": "agent_billing|agent_tech|agent_general",
                },
            },
            {
                "id": "multi_turn_conversation",
                "name": "Multi-Turn De-escalation",
                "difficulty": "very_hard",
                "description": "Handle a multi-turn frustrated customer conversation requiring context retention.",
                "action_schema": {
                    "action_type": "draft_response|escalate|resolve|close",
                    "response_text": "string",
                },
            },
            # ── Legal ──
            {
                "id": "legal_clause_identification",
                "name": "Legal Clause Identification",
                "difficulty": "easy",
                "description": "Identify the type of each contractual clause.",
                "action_schema": {
                    "action_type": "identify_clause",
                    "clause_type": "indemnity|liability|ip|termination|unknown",
                },
            },
            {
                "id": "legal_risk_flagging",
                "name": "Legal Risk Flagging",
                "difficulty": "medium",
                "description": "Assess and flag risk levels for contractual clauses.",
                "action_schema": {
                    "action_type": "flag_risk",
                    "risk_level": "low|medium|high|critical",
                    "reasoning": "string",
                },
            },
            {
                "id": "legal_clause_redlining",
                "name": "Legal Clause Redlining",
                "difficulty": "hard",
                "description": "Rewrite risky contract clauses with safer alternative language.",
                "action_schema": {
                    "action_type": "redline",
                    "redline_text": "string (improved clause text)",
                },
            },
            # ── Clinical ──
            {
                "id": "clinical_triage_classification",
                "name": "Clinical Triage Body System",
                "difficulty": "easy",
                "description": "Classify patient chief complaints into affected body systems.",
                "action_schema": {
                    "action_type": "classify_triage",
                    "body_system": "cardiac|respiratory|neurologic|gi|musculoskeletal|other",
                },
            },
            {
                "id": "clinical_esi_assignment",
                "name": "Clinical ESI Assignment",
                "difficulty": "medium",
                "description": "Assign Emergency Severity Index (ESI) 1-5 triage levels.",
                "action_schema": {
                    "action_type": "assign_esi",
                    "esi_level": "integer 1-5",
                    "reasoning": "string",
                },
            },
            {
                "id": "clinical_triage_note",
                "name": "Clinical Triage Note Generation",
                "difficulty": "hard",
                "description": "Write a complete clinical triage note for ED patients.",
                "action_schema": {
                    "action_type": "write_triage_note",
                    "triage_note": "string (clinical note)",
                },
            },
            # ── Engineering ──
            {
                "id": "pr_type_classification",
                "name": "PR Type Classification",
                "difficulty": "easy",
                "description": "Classify pull request diffs by type.",
                "action_schema": {
                    "action_type": "classify_pr",
                    "pr_type": "bug_fix|feature|refactor|security",
                },
            },
            {
                "id": "pr_bug_identification",
                "name": "PR Bug Identification",
                "difficulty": "medium",
                "description": "Identify the specific bug or security issue in a code diff.",
                "action_schema": {
                    "action_type": "identify_bug",
                    "bug_description": "string",
                },
            },
            {
                "id": "pr_review_comment",
                "name": "PR Code Review",
                "difficulty": "hard",
                "description": "Write expert code review comments approving or blocking PRs.",
                "action_schema": {
                    "action_type": "review_pr",
                    "review_comment": "string",
                },
            },
        ]
    }


@app.post("/grader")
async def grader(task_id: Optional[str] = None) -> Dict[str, Any]:
    task = get_task(task_id)
    result = task.grader_score()

    # Global reward shaping bonuses
    reward_history = getattr(task, "reward_history", [])
    if reward_history and len(reward_history) >= 3:
        first_half = sum(reward_history[:len(reward_history) // 2])
        second_half = sum(reward_history[len(reward_history) // 2:])
        traj_bonus = 0.05 if second_half > first_half else 0.0
        try:
            std = statistics.stdev(reward_history)
            cons_penalty = -0.05 if std > 0.4 else 0.0
        except Exception:
            cons_penalty = 0.0
        result["final_score"] = max(0.001, min(0.999, result["final_score"] + traj_bonus + cons_penalty))
        result.setdefault("metrics", {})
        result["metrics"]["trajectory_bonus"] = traj_bonus
        result["metrics"]["consistency_penalty"] = cons_penalty

    # FORCE clamp — validator requires strictly (0, 1) exclusive
    result["final_score"] = max(0.001, min(0.999, float(result["final_score"])))
    
    return result


@app.post("/baseline")
async def baseline() -> Dict[str, Any]:
    """Run heuristic baseline agent across all tasks. No API key needed."""
    from graders.baseline_agent import run_baseline_all_tasks
    results = await run_baseline_all_tasks()
    # Clamp all task scores
    for tid in results.get("tasks", {}):
        s = results["tasks"][tid].get("final_score", 0.5)
        results["tasks"][tid]["final_score"] = max(0.001, min(0.999, float(s)))
    overall = results.get("overall_score", 0.5)
    results["overall_score"] = max(0.001, min(0.999, float(overall)))
    return results


@app.get("/health")
async def health():
    return {"status": "ok", "tasks": list(TASKS.keys())}


@app.get("/", response_class=FileResponse)
async def root():
    root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    static_path = os.path.join(root_dir, "static", "index.html")
    return FileResponse(static_path)