feat: add SQLite persistence for episodes, results, and leaderboard via SQLModel

.gitignore

@@ -11,3 +11,8 @@ dist/
 build/
 .idea/
 .vscode/
+
+# Persistence
+data/*.db
+data/*.db-shm
+data/*.db-wal

app.py

@@ -1,60 +1,109 @@
 import uuid
-from typing import Dict
-
-from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
+import logging
+import asyncio
+import json
+from typing import Dict, List, Optional
+from datetime import datetime, timezone
+
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Depends, Security, Query, BackgroundTasks, Request
+from fastapi.responses import JSONResponse
+from fastapi.exceptions import RequestValidationError
+from fastapi.security.api_key import APIKeyHeader
 from pydantic import BaseModel
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+from sqlmodel import Session
 
 from codereview_env.models import (
-    TaskId, Action, ResetResult, StepResult, EpisodeResult
+    TaskId, Action, ResetResult, StepResult, EpisodeResult, ActionRecord
 )
 from codereview_env.env import CodeReviewEnv
+from codereview_env.config import get_settings
+from codereview_env.database import (
+    create_db_and_tables, get_session, save_episode,
+    get_episode, get_leaderboard_db, submit_leaderboard, get_stats,
+    LeaderboardRecord
+)
+
+# ── Logging ───────────────────────────────────────────────────────────────────
+settings = get_settings()
+logging.basicConfig(
+    level=getattr(logging, settings.log_level),
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+)
+logger = logging.getLogger("codereview_env")
 
+# ── App Initialization ────────────────────────────────────────────────────────
 app = FastAPI(
     title="AgentOrg CodeReview OpenEnv API",
     description=(
         "AI Senior Code Reviewer evaluation environment. "
         "Trains agents to detect bugs, security vulnerabilities, and architectural issues "
-        "in realistic Python PRs grounded in real-world incident patterns."
+        "in realistic Python PRs."
     ),
     version="1.0.0",
 )
 
-# Simple in-memory storage for active episodes
-episodes: Dict[str, CodeReviewEnv] = {}
+# ── Rate Limiting ─────────────────────────────────────────────────────────────
+limiter = Limiter(key_func=get_remote_address, default_limits=[f"{settings.rate_limit_per_minute}/minute"])
+app.state.limiter = limiter
+app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+
+# ── API Key Authentication ────────────────────────────────────────────────────
+API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False)
 
+async def verify_api_key(api_key: str = Security(API_KEY_HEADER)):
+    if not settings.api_key_enabled:
+        return  # Auth disabled in development
+    if api_key != settings.api_key:
+        raise HTTPException(status_code=403, detail="Invalid or missing API key")
 
+# ── Storage & TTL ─────────────────────────────────────────────────────────────
+episodes: Dict[str, CodeReviewEnv] = {}
+episode_timestamps: Dict[str, datetime] = {}
+
+async def cleanup_expired_episodes():
+    """Remove episodes older than TTL."""
+    while True:
+        await asyncio.sleep(300)  # run every 5 minutes
+        cutoff = datetime.now(timezone.utc).timestamp() - settings.episode_ttl_seconds
+        expired = [
+            eid for eid, ts in episode_timestamps.items()
+            if ts.timestamp() < cutoff
+        ]
+        for eid in expired:
+            episodes.pop(eid, None)
+            episode_timestamps.pop(eid, None)
+        if expired:
+            logger.info(f"Cleaned up {len(expired)} expired episodes")
+
+@app.on_event("startup")
+async def startup_event():
+    create_db_and_tables()
+    asyncio.create_task(cleanup_expired_episodes())
+    logger.info(f"CodeReview API started \u2014 DB at {settings.db_path}")
+
+# ── Models ────────────────────────────────────────────────────────────────────
 class ResetRequest(BaseModel):
     task_id: TaskId
     seed:    int = 42
 
-
 class ResetResponse(BaseModel):
     episode_id: str
     result:     ResetResult
 
-
-# In-memory leaderboard
-leaderboard: Dict[TaskId, list] = {
-    TaskId.BUG_DETECTION:        [],
-    TaskId.SECURITY_AUDIT:       [],
-    TaskId.ARCHITECTURAL_REVIEW: []
-}
-
-
 class SubmitScore(BaseModel):
     agent_name: str
     task_id:    TaskId
     score:      float
     seed:       int
 
-
 # ── WebSocket clients ─────────��───────────────────────────────────────────────
 clients = set()
 
-
 async def broadcast_event(data: dict):
     from fastapi.encoders import jsonable_encoder
-    import json
     message = json.dumps(jsonable_encoder(data))
     dead = set()
     for client in clients:
@@ -64,30 +113,55 @@ async def broadcast_event(data: dict):
             dead.add(client)
     clients.difference_update(dead)
 
+# ── Error Handlers ────────────────────────────────────────────────────────────
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request, exc):
+    return JSONResponse(
+        status_code=422,
+        content={
+            "error": "validation_error",
+            "detail": str(exc),
+            "status_code": 422
+        }
+    )
+
+@app.exception_handler(HTTPException)
+async def http_exception_handler(request, exc):
+    logger.warning(f"HTTP {exc.status_code}: {exc.detail} \u2014 {request.url}")
+    return JSONResponse(
+        status_code=exc.status_code,
+        content={
+            "error": exc.detail,
+            "status_code": exc.status_code
+        }
+    )
 
 # ── Endpoints ─────────────────────────────────────────────────────────────────
 
 @app.get("/health")
 def health_check():
     return {
-        "status":    "ok",
-        "version":   "1.0.0",
+        "status": "ok",
+        "version": "1.0.0",
         "env_ready": True,
+        "env": settings.app_env,
         "active_episodes": len(episodes),
+        "auth_enabled": settings.api_key_enabled
     }
 
-
 @app.post("/reset", response_model=ResetResponse)
-def reset_env(req: ResetRequest):
+@limiter.limit(f"{settings.rate_limit_per_minute}/minute")
+def reset_env(request: Request, req: ResetRequest, _: None = Depends(verify_api_key)):
     episode_id = str(uuid.uuid4())
     env        = CodeReviewEnv()
     result     = env.reset(req.task_id, req.seed)
     episodes[episode_id] = env
+    episode_timestamps[episode_id] = datetime.now(timezone.utc)
     return ResetResponse(episode_id=episode_id, result=result)
 
-
 @app.post("/step/{episode_id}", response_model=StepResult)
-async def step_env(episode_id: str, action: Action):
+@limiter.limit(f"{settings.rate_limit_per_minute}/minute")
+async def step_env(request: Request, episode_id: str, action: Action, _: None = Depends(verify_api_key)):
     if episode_id not in episodes:
         raise HTTPException(status_code=404, detail="Episode not found")
 
@@ -99,29 +173,117 @@ async def step_env(episode_id: str, action: Action):
     except RuntimeError as e:
         raise HTTPException(status_code=400, detail=str(e))
 
-
 @app.get("/result/{episode_id}", response_model=EpisodeResult)
-def get_result(episode_id: str):
-    if episode_id not in episodes:
+def get_result(
+    episode_id: str,
+    session: Session = Depends(get_session),
+    _: None = Depends(verify_api_key)
+):
+    # Try in-memory (active episode)
+    if episode_id in episodes:
+        env = episodes[episode_id]
+        result = env.get_final_result()
+        result.episode_id = episode_id
+        # If done, persist and remove from memory
+        if env.done:
+            save_episode(session, result)
+            del episodes[episode_id]
+            episode_timestamps.pop(episode_id, None)
+        return result
+    
+    # Fall back to DB (completed episode)
+    record = get_episode(session, episode_id)
+    if not record:
         raise HTTPException(status_code=404, detail="Episode not found")
-    return episodes[episode_id].get_final_result()
-
+    
+    return EpisodeResult(
+        episode_id=record.episode_id,
+        task_id=TaskId(record.task_id),
+        scenario_hash=record.scenario_hash,
+        seed=record.seed,
+        final_score=record.final_score,
+        steps_taken=record.steps_taken,
+        issues_found=record.issues_found,
+        issues_total=record.issues_total,
+        noise_penalties=record.noise_penalties,
+        terminated_reason=record.terminated_reason,
+        history=[ActionRecord(**r) for r in json.loads(record.history_json or "[]")]
+    )
 
 @app.get("/leaderboard")
-def get_leaderboard():
-    return leaderboard
-
+def get_leaderboard(
+    task_id: Optional[TaskId] = None,
+    limit: int = Query(default=10, ge=1, le=50),
+    offset: int = Query(default=0, ge=0),
+    session: Session = Depends(get_session)
+):
+    tasks_to_query = [task_id] if task_id else list(TaskId)
+    result = {}
+    for t in tasks_to_query:
+        entries, total = get_leaderboard_db(session, t.value, limit, offset)
+        result[t.value] = {
+            "entries": [e.model_dump() for e in entries],
+            "total": total
+        }
+    if task_id:
+        return result[task_id.value]
+    return result
 
 @app.post("/submit")
-def submit_to_leaderboard(submission: SubmitScore):
-    entries   = leaderboard.get(submission.task_id, [])
-    new_entry = submission.model_dump()
-    entries.append(new_entry)
-    entries.sort(key=lambda x: x["score"], reverse=True)
-    rank = entries.index(new_entry) + 1   # capture rank before slicing
-    leaderboard[submission.task_id] = entries[:5]
-    return {"status": "submitted", "rank": rank if rank <= 5 else None}
+@limiter.limit(f"{settings.rate_limit_per_minute}/minute")
+def submit_to_leaderboard(
+    request: Request,
+    submission: SubmitScore,
+    session: Session = Depends(get_session),
+    _: None = Depends(verify_api_key)
+):
+    rank = submit_leaderboard(
+        session,
+        agent_name=submission.agent_name,
+        task_id=submission.task_id.value,
+        score=submission.score,
+        seed=submission.seed
+    )
+    return {"status": "submitted", "rank": rank if rank > 0 else None}
+
+@app.get("/stats")
+def get_aggregate_stats(session: Session = Depends(get_session)):
+    return get_stats(session)
+
+@app.get("/episodes/{episode_id}/replay")
+def get_episode_replay(
+    episode_id: str,
+    session: Session = Depends(get_session),
+    _: None = Depends(verify_api_key)
+):
+    record = get_episode(session, episode_id)
+    if not record:
+        raise HTTPException(status_code=404, detail="Episode not found or not yet completed")
+    return {
+        "episode_id": record.episode_id,
+        "task_id": record.task_id,
+        "scenario_hash": record.scenario_hash,
+        "final_score": record.final_score,
+        "history": json.loads(record.history_json or "[]"),
+        "created_at": record.created_at
+    }
 
+@app.get("/episodes")
+def list_episodes(
+    _: None = Depends(verify_api_key),
+    limit: int = Query(default=20, ge=1, le=100)
+):
+    episode_list = [
+        {
+            "episode_id": eid,
+            "task_id": env.task_id,
+            "step_count": env.observation.step_count,
+            "done": env.done,
+            "created_at": episode_timestamps.get(eid, "").isoformat() if episode_timestamps.get(eid) else ""
+        }
+        for eid, env in list(episodes.items())[:limit]
+    ]
+    return {"episodes": episode_list, "total": len(episodes)}
 
 @app.websocket("/ws/events")
 async def websocket_endpoint(websocket: WebSocket):
@@ -135,7 +297,6 @@ async def websocket_endpoint(websocket: WebSocket):
     finally:
         clients.discard(websocket)
 
-
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host=settings.app_host, port=settings.app_port)

codereview_env/config.py

@@ -0,0 +1,27 @@
+from functools import lru_cache
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8", extra="ignore")
+
+    app_host: str = "0.0.0.0"
+    app_port: int = 7860
+    app_env: str = "development"
+
+    api_key: str = "changeme"
+    api_key_enabled: bool = False
+
+    leaderboard_max_entries: int = 10
+
+    log_level: str = "INFO"
+
+    episode_ttl_seconds: int = 3600        # episodes expire after 1 hour
+    rate_limit_per_minute: int = 60        # requests per minute per IP
+    
+    # Persistence
+    db_path: str = "./data/codereview.db"
+    db_echo: bool = False    # Set True to log all SQL queries
+
+@lru_cache
+def get_settings() -> Settings:
+    return Settings()

codereview_env/database.py

@@ -0,0 +1,142 @@
+from pathlib import Path
+from sqlmodel import SQLModel, Field, Session, create_engine, select
+from typing import Optional, List, Tuple
+import json
+from codereview_env.config import get_settings
+from codereview_env.models import EpisodeResult, TaskId
+
+def get_engine():
+    settings = get_settings()
+    Path(settings.db_path).parent.mkdir(parents=True, exist_ok=True)
+    return create_engine(
+        f"sqlite:///{settings.db_path}",
+        echo=settings.db_echo,
+        connect_args={"check_same_thread": False}
+    )
+
+def create_db_and_tables():
+    engine = get_engine()
+    SQLModel.metadata.create_all(engine)
+
+def get_session():
+    engine = get_engine()
+    with Session(engine) as session:
+        yield session
+
+class EpisodeRecord(SQLModel, table=True):
+    __tablename__ = "episodes"
+
+    id: Optional[int] = Field(default=None, primary_key=True)
+    episode_id: str = Field(index=True, unique=True)
+    task_id: str
+    scenario_hash: str
+    seed: int
+    final_score: float
+    steps_taken: int
+    issues_found: int
+    issues_total: int
+    noise_penalties: int
+    terminated_reason: str
+    history_json: str = ""      # JSON-serialized list of ActionRecord dicts
+    created_at: str = ""        # ISO datetime
+    agent_name: str = ""        # optional, set via /submit
+
+class LeaderboardRecord(SQLModel, table=True):
+    __tablename__ = "leaderboard"
+
+    id: Optional[int] = Field(default=None, primary_key=True)
+    agent_name: str
+    task_id: str = Field(index=True)
+    score: float
+    seed: int
+    episode_id: str = ""
+    submitted_at: str = ""      # ISO datetime
+
+def save_episode(session: Session, result: EpisodeResult) -> EpisodeRecord:
+    """Persist a completed episode result."""
+    from datetime import datetime, timezone
+    record = EpisodeRecord(
+        episode_id=result.episode_id,
+        task_id=result.task_id.value,
+        scenario_hash=result.scenario_hash,
+        seed=result.seed,
+        final_score=result.final_score,
+        steps_taken=result.steps_taken,
+        issues_found=result.issues_found,
+        issues_total=result.issues_total,
+        noise_penalties=result.noise_penalties,
+        terminated_reason=result.terminated_reason,
+        history_json=json.dumps([r.model_dump() for r in result.history]),
+        created_at=datetime.now(timezone.utc).isoformat()
+    )
+    session.add(record)
+    session.commit()
+    session.refresh(record)
+    return record
+
+def get_episode(session: Session, episode_id: str) -> Optional[EpisodeRecord]:
+    return session.exec(select(EpisodeRecord).where(EpisodeRecord.episode_id == episode_id)).first()
+
+def get_leaderboard_db(session: Session, task_id: str, limit: int = 10, offset: int = 0) -> Tuple[List[LeaderboardRecord], int]:
+    results = session.exec(
+        select(LeaderboardRecord)
+        .where(LeaderboardRecord.task_id == task_id)
+        .order_by(LeaderboardRecord.score.desc())
+        .offset(offset)
+        .limit(limit)
+    ).all()
+    # Fixed: session.exec(select(LeaderboardRecord).where(LeaderboardRecord.task_id == task_id)).all() is not efficient but the snippet used it.
+    # To get count specifically:
+    from sqlmodel import func
+    total = session.exec(
+        select(func.count()).select_from(LeaderboardRecord).where(LeaderboardRecord.task_id == task_id)
+    ).one()
+    return list(results), total
+
+def submit_leaderboard(session: Session, agent_name: str, task_id: str, score: float, seed: int, episode_id: str = "") -> int:
+    """Add entry to leaderboard. Returns rank (1-indexed)."""
+    from datetime import datetime, timezone
+    record = LeaderboardRecord(
+        agent_name=agent_name,
+        task_id=task_id,
+        score=score,
+        seed=seed,
+        episode_id=episode_id,
+        submitted_at=datetime.now(timezone.utc).isoformat()
+    )
+    session.add(record)
+    session.commit()
+    # Calculate rank
+    rank_result = session.exec(
+        select(LeaderboardRecord)
+        .where(LeaderboardRecord.task_id == task_id)
+        .order_by(LeaderboardRecord.score.desc())
+    ).all()
+    for i, r in enumerate(rank_result):
+        if r.id == record.id:
+            return i + 1
+    return -1
+
+def get_stats(session: Session) -> dict:
+    """Return aggregate statistics about all recorded episodes."""
+    all_episodes = session.exec(select(EpisodeRecord)).all()
+    if not all_episodes:
+        return {"total_episodes": 0, "avg_score": 0.0, "by_task": {}}
+    
+    from collections import defaultdict
+    by_task = defaultdict(list)
+    for ep in all_episodes:
+        by_task[ep.task_id].append(ep.final_score)
+    
+    return {
+        "total_episodes": len(all_episodes),
+        "avg_score": sum(ep.final_score for ep in all_episodes) / len(all_episodes),
+        "by_task": {
+            task: {
+                "count": len(scores),
+                "avg_score": sum(scores) / len(scores) if scores else 0,
+                "best_score": max(scores) if scores else 0,
+            }
+            for task, scores in by_task.items()
+        }
+    }

requirements.txt

@@ -6,3 +6,8 @@ requests>=2.31.0
 websockets>=12.0
 httpx<0.28.0
 openai>=1.0.0
+pydantic-settings==2.2.1
+slowapi==0.1.9
+python-dotenv==1.0.1
+sqlmodel==0.0.16
+aiosqlite==0.20.0

scripts/migrate.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+"""Initialize or reset the CodeReview database."""
+import sys
+import os
+# Ensure PYTHONPATH is set so we can import codereview_env
+sys.path.append(os.getcwd())
+
+from codereview_env.database import create_db_and_tables, get_engine
+from codereview_env.config import get_settings
+from sqlmodel import SQLModel
+
+def init():
+    settings = get_settings()
+    print(f"Initializing database at: {settings.db_path}")
+    create_db_and_tables()
+    print("Database initialized successfully.")
+
+def reset():
+    settings = get_settings()
+    engine = get_engine()
+    print(f"Dropping all tables in: {settings.db_path}")
+    SQLModel.metadata.drop_all(engine)
+    SQLModel.metadata.create_all(engine)
+    print("Database reset successfully.")
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python3 scripts/migrate.py [init|reset]")
+        sys.exit(1)
+        
+    cmd = sys.argv[1]
+    if cmd == "init":
+        init()
+    elif cmd == "reset":
+        reset()
+    else:
+        print(f"Unknown command: {cmd}. Use 'init' or 'reset'.")
+        sys.exit(1)

tests/test_api.py

@@ -50,8 +50,10 @@ def test_api_leaderboard():
     # Check leaderboard
     lb_resp = client.get("/leaderboard")
     assert lb_resp.status_code == 200
-    assert len(lb_resp.json()["bug_detection"]) > 0
-    assert lb_resp.json()["bug_detection"][0]["agent_name"] == "test_agent"
+    lb_data = lb_resp.json()
+    bug_entries = lb_data["bug_detection"]["entries"]
+    assert len(bug_entries) > 0
+    assert bug_entries[0]["agent_name"] == "test_agent"
 
 def test_api_invalid_episode():
     client = TestClient(app)