Legal Research OpenEnv ⚖️

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 # Ensure root is in sys.path _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 # Global State _env: Optional[LegalResearchEnv] = None _is_ready: bool = False _init_error: Optional[str] = None # UI Content LANDING_HTML = """\ Legal Research OpenEnv ⚖️

📘 API Docs

Interactive Swagger UI.

POST /reset

POST /step

GET /state

Open Swagger UI →

💚 Health Status

Live environment health check.

Status: {{STATUS_TEXT}}

Version: 2.0.0

Check health →

""" # Background Initialization 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=["*"], ) # --- Endpoints --- @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" # docs_loaded = len(_env.search_engine.case_ids) if _env and _env.search_engine else 0 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": ""} ``` 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()

Teach AI Agentsto Research the Law

Teach AI Agents
to Research the Law