| import asyncio |
| import os |
| import sys |
| import time |
| import subprocess |
| from pathlib import Path |
| from contextlib import asynccontextmanager |
| from typing import Optional |
|
|
| from fastapi import FastAPI, HTTPException, Request, Body |
| from fastapi.middleware.cors import CORSMiddleware |
| from fastapi.responses import HTMLResponse |
|
|
| |
| _root = str(Path(__file__).resolve().parent.parent) |
| if _root not in sys.path: |
| sys.path.insert(0, _root) |
|
|
| from env import LegalResearchEnv |
| from models import Observation, State, Action, StepResponse, ResetRequest, ResetResponse |
|
|
| |
| _env: Optional[LegalResearchEnv] = None |
| _is_ready: bool = False |
| _init_error: Optional[str] = None |
|
|
| |
|
|
| LANDING_HTML = """\ |
| <!DOCTYPE html> |
| <html lang="en"> |
| <head> |
| <meta charset="UTF-8"> |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> |
| <title>Legal Research OpenEnv ⚖️</title> |
| <link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&display=swap" rel="stylesheet"> |
| <style> |
| :root { |
| --bg: #0d1117; |
| --card-bg: #161b22; |
| --accent: #58a6ff; |
| --text-main: #c9d1d9; |
| --text-dim: #8b949e; |
| --success: #3fb950; |
| } |
| * { box-sizing: border-box; margin: 0; padding: 0; } |
| body { |
| font-family: 'Inter', sans-serif; |
| background-color: var(--bg); |
| color: var(--text-main); |
| line-height: 1.6; |
| display: flex; |
| flex-direction: column; |
| align-items: center; |
| min-height: 100vh; |
| padding: 2rem; |
| } |
| .header { text-align: center; margin-bottom: 3rem; max-width: 800px; } |
| .badge { |
| display: inline-block; |
| background: rgba(88, 166, 255, 0.1); |
| color: var(--accent); |
| padding: 0.4rem 1rem; |
| border-radius: 2rem; |
| font-size: 0.85rem; |
| font-weight: 600; |
| margin-bottom: 1rem; |
| border: 1px solid rgba(88, 166, 255, 0.2); |
| } |
| h1 { font-size: 3rem; font-weight: 800; margin-bottom: 1rem; color: #fff; } |
| .description { font-size: 1.1rem; color: var(--text-dim); margin-bottom: 2rem; } |
| |
| .features { |
| display: grid; |
| grid-template-columns: repeat(2, 1fr); |
| gap: 1.5rem; |
| margin-bottom: 3rem; |
| } |
| .feature-item { |
| display: flex; |
| align-items: center; |
| gap: 0.75rem; |
| font-weight: 500; |
| color: var(--accent); |
| } |
| .feature-item::before { content: '✓'; color: var(--success); font-weight: 900; } |
| |
| .cards { |
| display: grid; |
| grid-template-columns: repeat(auto-fit, minmax(350px, 1fr)); |
| gap: 2rem; |
| width: 100%; |
| max-width: 1000px; |
| } |
| .card { |
| background: var(--card-bg); |
| border: 1px solid #30363d; |
| border-radius: 12px; |
| padding: 2rem; |
| transition: transform 0.2s, border-color 0.2s; |
| } |
| .card:hover { transform: translateY(-4px); border-color: var(--accent); } |
| .card-title { |
| display: flex; |
| align-items: center; |
| gap: 0.75rem; |
| font-size: 1.5rem; |
| margin-bottom: 1rem; |
| color: #fff; |
| } |
| .endpoint { |
| font-family: 'Courier New', monospace; |
| background: #0d1117; |
| padding: 0.5rem; |
| border-radius: 6px; |
| margin-bottom: 0.5rem; |
| font-size: 0.9rem; |
| color: #79c0ff; |
| } |
| .stat-line { |
| display: flex; |
| justify-content: space-between; |
| margin-bottom: 0.5rem; |
| font-size: 0.95rem; |
| } |
| .stat-val { color: var(--success); font-weight: 700; } |
| |
| .btn { |
| display: inline-block; |
| margin-top: 1.5rem; |
| background: var(--accent); |
| color: #fff; |
| padding: 0.75rem 1.5rem; |
| border-radius: 8px; |
| text-decoration: none; |
| font-weight: 600; |
| } |
| .footer { margin-top: 4rem; color: var(--text-dim); font-size: 0.9rem; } |
| .status-ready { color: var(--success); } |
| .status-loading { color: #f2cc60; animation: pulse 2s infinite; } |
| @keyframes pulse { 0% { opacity: 1; } 50% { opacity: 0.5; } 100% { opacity: 1; } } |
| </style> |
| </head> |
| <body> |
| <div class="header"> |
| <div class="badge">⚖️ OpenEnv · Legal Research RL Environment</div> |
| <h1>Teach AI Agents<br>to Research the Law</h1> |
| <p class="description"> |
| An OpenEnv-compatible reinforcement learning environment where AI agents learn to search, read, extract and reason about real Supreme Court judgments over 26,000+ cases. |
| </p> |
| <div class="features"> |
| <div class="feature-item">26,000+ Judgments indexed</div> |
| <div class="feature-item">3 Graders</div> |
| <div class="feature-item">BM25 + TF-IDF</div> |
| <div class="feature-item">Continuous Reward signal</div> |
| </div> |
| </div> |
| |
| <div class="cards"> |
| <div class="card"> |
| <div class="card-title">📘 API Docs</div> |
| <p style="margin-bottom: 1rem; color: var(--text-dim);"> |
| Interactive Swagger UI. |
| </p> |
| <div class="endpoint">POST /reset</div> |
| <div class="endpoint">POST /step</div> |
| <div class="endpoint">GET /state</div> |
| <a href="/docs" class="btn">Open Swagger UI →</a> |
| </div> |
| |
| <div class="card"> |
| <div class="card-title">💚 Health Status</div> |
| <p style="margin-bottom: 1rem; color: var(--text-dim);"> |
| Live environment health check. |
| </p> |
| <div class="stat-line"><span>Status:</span> <span class="stat-val {{STATUS_CLASS}}">{{STATUS_TEXT}}</span></div> |
| <div class="stat-line"><span>Version:</span> <span class="stat-val">2.0.0</span></div> |
| <a href="/health" class="btn">Check health →</a> |
| </div> |
| </div> |
| |
| <div class="footer"> |
| GitHub · Swagger UI · OpenEnv Compliant |
| </div> |
| </body> |
| </html> |
| """ |
|
|
| |
| async def bg_init_corpus(): |
| global _env, _is_ready, _init_error |
| try: |
| start_t = time.time() |
| print("[INIT] Running corpus preparation check...") |
| subprocess.run(["python", "prepare_corpus.py"], check=False) |
|
|
| corpus_dir = os.getenv("CORPUS_DIR", "data/judgments") |
| print(f"[INIT] Loading LegalResearchEnv from: {corpus_dir}") |
| |
| loop = asyncio.get_event_loop() |
| _env = await loop.run_in_executor(None, lambda: LegalResearchEnv(corpus_dir=corpus_dir)) |
| |
| _is_ready = True |
| duration = time.time() - start_t |
| docs = len(_env.search_engine.case_ids) if _env.search_engine else 0 |
| print(f"[INIT] Environment ready in {duration:.2f}s. {docs} documents loaded.") |
| except Exception as e: |
| _init_error = str(e) |
| print(f"[ERROR] Initialization failed: {e}") |
|
|
| @asynccontextmanager |
| async def lifespan(app: FastAPI): |
| print("[SERVER] Startup complete. Spawning background loader...") |
| asyncio.create_task(bg_init_corpus()) |
| yield |
| print("[SERVER] Shutting down.") |
|
|
| app = FastAPI( |
| title="Legal Research OpenEnv ⚖️", |
| version="2.0.0", |
| description="OpenEnv environment for legal research on Indian Supreme Court judgments. Actions: search, read, extract, assess, submit.", |
| lifespan=lifespan |
| ) |
|
|
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
| |
| @app.get("/", response_class=HTMLResponse, include_in_schema=False) |
| async def root(): |
| status_text = "ok" if _is_ready else "loading..." |
| status_class = "status-ready" if _is_ready else "status-loading" |
| |
| |
| html = LANDING_HTML.replace("{{STATUS_TEXT}}", status_text) \ |
| .replace("{{STATUS_CLASS}}", status_class) \ |
| |
| return HTMLResponse(content=html) |
|
|
| @app.post("/reset", response_model=ResetResponse) |
| async def reset( |
| request: Optional[ResetRequest] = Body( |
| None, |
| openapi_examples={ |
| "easy": { |
| "summary": "Easy task (constitutional law)", |
| "value": {"difficulty": "easy"} |
| }, |
| "medium": { |
| "summary": "Medium task (criminal law)", |
| "value": {"difficulty": "medium"} |
| }, |
| "hard": { |
| "summary": "Hard task (corporate law, multi-step)", |
| "value": {"difficulty": "hard"} |
| }, |
| "empty": { |
| "summary": "No body (defaults to easy)", |
| "value": {} |
| } |
| } |
| ) |
| ): |
| """ Reset environment and start a new episode. |
| **Body (optional JSON):** |
| - `difficulty`: `"easy"` | `"medium"` | `"hard"` (default: `"easy"`) |
| - Sending an empty body `{}` or no body is also valid. |
| **Returns:** Observation with task description, available actions, and empty search results.""" |
| if not _is_ready: |
| if _init_error: raise HTTPException(status_code=500, detail=_init_error) |
| raise HTTPException(status_code=503, detail="Corpus still loading...") |
| |
| difficulty = "easy" |
| if request and request.difficulty: |
| difficulty = request.difficulty |
| |
| obs = _env.reset(difficulty=difficulty) |
| return ResetResponse(observation=obs) |
|
|
| @app.post("/step", response_model=StepResponse) |
| async def step(action: Action): |
| """ |
| **Always call `/reset` first to start an episode.** |
| **Action examples** — paste one of these into the request body: |
| 1. **Search** (use first — gets you case IDs): |
| ```json |
| {"action_type": "search", "query": "constitutional fundamental rights article"} |
| ``` |
| 2. **Read** (copy a `case_id` from search results above): |
| ```json |
| {"action_type": "read", "case_id": "<paste case_id from search results>"} |
| ``` |
| 3. **Extract** (after reading a document): |
| ```json |
| {"action_type": "extract", "extraction_type": "legal_principle"} |
| ``` |
| 4. **Assess** (hard task only — use string values, not objects): |
| ```json |
| {"action_type": "assess", "assessment": {"precedent_application": "The cases establish director liability under fiduciary duty.", "conclusion": "Director is personally liable for breach of duty of care."}} |
| ``` |
| 5. **Submit** (ends the episode and returns final score): |
| ```json |
| {"action_type": "submit", "submission": {"summary": "Research complete", "cases_found": 3}} |
| ``` |
| """ |
| if not _is_ready: |
| raise HTTPException(status_code=503, detail="Corpus still loading...") |
| return _env.step(action) |
|
|
| @app.get("/state") |
| async def get_state(): |
| """Get current environment state.""" |
| if not _is_ready: return {"status": "loading"} |
| return _env.state() |
|
|
| @app.get("/health") |
| async def health(): |
| """Health check endpoint.""" |
| if not _is_ready: return {"status": "loading", "init_error": _init_error} |
| docs = len(_env.search_engine.case_ids) if _env and _env.search_engine else 0 |
| return { |
| "status": "ok", |
| "environment": "legal_research_env", |
| "tasks": ["easy", "medium", "hard"] |
| } |
|
|
| @app.get("/validate") |
| async def validate(): |
| """Endpoint for OpenEnv automated task discovery.""" |
| return { |
| "tasks": [ |
| {"id": "easy", "name": "Constitutional Retrieval", "difficulty": "easy", "reward_range": [0.01, 0.99]}, |
| {"id": "medium", "name": "Criminal Law Analysis", "difficulty": "medium", "reward_range": [0.01, 0.99]}, |
| {"id": "hard", "name": "Corporate Liability Research", "difficulty": "hard", "reward_range": [0.01, 0.99]} |
| ], |
| "grader_type": "heuristic_legal_grader", |
| "spec_version": "2.0.0" |
| } |
|
|
| def main(): |
| import uvicorn |
| port = int(os.getenv("PORT", "7860")) |
| uvicorn.run("server.app:app", host="0.0.0.0", port=port, workers=1, reload=False) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
