Merge pull request #28 from ArshVermaGit/main

feat: implement SQLite persistence and database management

Author: DsThakurRawat

.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()