Differentiated graders, LLM-driven inference, and API improvements

README.md

@@ -9,10 +9,102 @@ pinned: false
 
 # Contract Negotiation Environment
 
-An AI-powered OpenEnv environment for contract negotiation using hybrid rule-based and LLM-driven decision making.
+An AI-powered OpenEnv environment for evaluating contract-negotiation agents
+through hybrid rule-based and LLM-driven decision making.
 
 ## Overview
-This project simulates real-world contract negotiation with structured tasks, grading, and agent evaluation.
+
+This project simulates real-world contract negotiation scenarios where an AI
+agent must:
+
+1. **Analyse** contract clauses to identify legal risks (unlimited liability,
+   hidden traps, one-sided IP terms, etc.).
+2. **Decide** on the best negotiation action: flag the risk, edit the clause,
+   propose a counter-offer, reject, or accept.
+3. **Generate** safer clause rewrites that protect the customer while keeping
+   commercially reasonable terms.
+
+Agents are scored on three dimensions:
+- **Correctness** — how well the agent identifies risky language.
+- **Improvement** — how much the proposed edits reduce risk.
+- **Risk alignment** — whether the chosen action matches the actual risk level.
+
+## Tasks
+
+| ID | Difficulty | Clause Type | Industry |
+|----|-----------|-------------|----------|
+| `easy_unlimited_liability` | EASY | Liability | SaaS B2B |
+| `medium_auto_renewal` | MEDIUM | Term/Renewal | SaaS B2B |
+| `hard_conflicting_obligations` | HARD | Performance/Changes | Professional Services |
+| `easy_compliance_agreement` | EASY+ | Compliance | SaaS B2B |
+| `hard_intellectual_property` | HARD+ | IP Ownership | Professional Services |
+
+Each task has a **dedicated grader** with difficulty-specific scoring adjustments
+(e.g., harder tasks penalise unresolved hidden traps).
 
 ## Structure
-- `contract_env/` → core environment, server, inference
+
+```
+contract_env/
+├── env/
+│   ├── environment.py   # ContractEnv — the main OpenEnv environment
+│   ├── graders.py       # Task-specific grading functions
+│   ├── models.py        # Pydantic models (Action, Reward, Observation)
+│   └── tasks.py         # Task definitions and metadata
+├── server/
+│   └── app.py           # FastAPI server exposing /reset, /step, /state, /tasks
+├── tests/               # Unit tests for API, graders, and environment
+└── scripts/             # Helper scripts for local/Docker runs
+inference.py             # LLM-driven inference agent
+openenv.yaml             # OpenEnv manifest
+Dockerfile               # Production container definition
+```
+
+## Quick Start
+
+### Local development
+
+```bash
+pip install -e ".[dev]"
+python -m pytest contract_env/tests/ -v
+```
+
+### Run the server
+
+```bash
+uvicorn contract_env.server.app:app --host 0.0.0.0 --port 7860
+```
+
+### Run inference
+
+```bash
+export HF_TOKEN="your-huggingface-token"
+python inference.py --episodes 5
+python inference.py --benchmark   # one episode per task
+```
+
+### Docker
+
+```bash
+docker build -t contract-negotiation-env .
+docker run -p 7860:7860 contract-negotiation-env
+```
+
+## API Endpoints
+
+| Method | Path | Description |
+|--------|------|-------------|
+| `GET` | `/health` | Health check |
+| `GET` | `/tasks` | List all tasks with metadata |
+| `GET` | `/state` | Current environment state |
+| `POST` | `/reset` | Reset and get first observation |
+| `POST` | `/step` | Submit an action, receive reward |
+
+## Environment Variables
+
+| Variable | Required | Default | Description |
+|----------|----------|---------|-------------|
+| `API_BASE_URL` | No | `https://router.huggingface.co/v1` | LLM API endpoint |
+| `MODEL_NAME` | No | `Qwen/Qwen2.5-72B-Instruct` | Model identifier |
+| `HF_TOKEN` | Yes | — | HuggingFace API token |
+| `PORT` | No | `7860` | Server port |

contract_env/env/environment.py

@@ -102,11 +102,17 @@ def step(self, action: Action) -> Tuple[Observation, float, bool, dict[str, Any]
         contract_before = self.state_data["contract_text"]
         proposed = build_proposed_contract_for_step(contract_before, action)
 
-        reward_obj, grade_info = evaluate_action(
-            self.current_task, contract_before, action, proposed
-        )
+        # Use the task-specific grader when available
+        task = self.current_task
+        if task.has_grader():
+            reward_obj = task.grade(contract_before, action, proposed)
+            # Collect grade info from evaluate_action for transparency
+            _, grade_info = evaluate_action(task, contract_before, action, proposed)
+        else:
+            reward_obj, grade_info = evaluate_action(
+                task, contract_before, action, proposed
+            )
 
-        # ✅ FIX: convert Reward → float
         reward = float(reward_obj.score)
 
         info.update(grade_info)

contract_env/env/graders.py

@@ -182,25 +182,68 @@ def grade_action(
     return reward
 
 
-# Specific graders for each task
+# ============ TASK-SPECIFIC GRADERS ============
+# Each grader applies difficulty-specific adjustments on top of the base evaluation.
+
+# -- Grading multipliers (named constants for clarity) --
+_EASY_SAFE_EDIT_BONUS = 1.08          # +8 % for well-matched safe edits
+_MEDIUM_PREMATURE_ACCEPT_PENALTY = 0.65  # −35 % for accepting risky terms
+_HARD_UNRESOLVED_TRAP_PENALTY = 0.5   # −50 % when hidden traps remain
+_EASY_PLUS_NOTIFICATION_BONUS = 1.06  # +6 % for breach-notification language
+_HARD_PLUS_TRAP_PENALTY = 0.55        # −45 % for unresolved IP traps
+_HARD_PLUS_OWNERSHIP_BONUS = 1.07     # +7 % for explicit customer-ownership
+
+
 def grade_easy(task: NegotiationTask, contract_before: str, action: Action, proposed_contract_text: str) -> Reward:
-    return grade_action(task, contract_before, action, proposed_contract_text)
+    """Grade easy tasks with a bias toward accepting safe-looking clauses quickly."""
+    reward, _ = evaluate_action(task, contract_before, action, proposed_contract_text)
+    if action.action_type in ("EDIT_CLAUSE", "PROPOSE_COUNTER"):
+        safe = _safe_overlap(
+            (action.content or "").strip(), task.safe_keywords, task.expected_safe_edit
+        )
+        if safe > 0.5:
+            reward.score = max(0.001, min(0.999, reward.score * _EASY_SAFE_EDIT_BONUS))
+    return reward
 
 
 def grade_medium(task: NegotiationTask, contract_before: str, action: Action, proposed_contract_text: str) -> Reward:
-    return grade_action(task, contract_before, action, proposed_contract_text)
+    """Grade medium tasks, penalising premature acceptance of risky auto-renewal terms."""
+    reward, _ = evaluate_action(task, contract_before, action, proposed_contract_text)
+    if action.action_type == "ACCEPT":
+        risk = _weighted_risk_hits(proposed_contract_text, task.risk_keywords)
+        if risk >= 0.3:
+            reward.score = max(0.001, min(0.999, reward.score * _MEDIUM_PREMATURE_ACCEPT_PENALTY))
+    return reward
 
 
 def grade_hard(task: NegotiationTask, contract_before: str, action: Action, proposed_contract_text: str) -> Reward:
-    return grade_action(task, contract_before, action, proposed_contract_text)
+    """Grade hard tasks with trap-resolution checking and heavier penalty for missed traps."""
+    reward, _ = evaluate_action(task, contract_before, action, proposed_contract_text)
+    if action.action_type in ("ACCEPT", "EDIT_CLAUSE", "PROPOSE_COUNTER"):
+        if trap_unresolved(task, proposed_contract_text):
+            reward.score = max(0.001, min(0.999, reward.score * _HARD_UNRESOLVED_TRAP_PENALTY))
+    return reward
 
 
 def grade_easy_plus(task: NegotiationTask, contract_before: str, action: Action, proposed_contract_text: str) -> Reward:
-    return grade_action(task, contract_before, action, proposed_contract_text)
+    """Grade easy-plus compliance tasks, rewarding mention of notification obligations."""
+    reward, _ = evaluate_action(task, contract_before, action, proposed_contract_text)
+    content = (action.content or "").strip().lower()
+    if action.action_type in ("EDIT_CLAUSE", "PROPOSE_COUNTER"):
+        if any(kw in content for kw in ("notify", "notification", "promptly inform")):
+            reward.score = max(0.001, min(0.999, reward.score * _EASY_PLUS_NOTIFICATION_BONUS))
+    return reward
 
 
 def grade_hard_plus(task: NegotiationTask, contract_before: str, action: Action, proposed_contract_text: str) -> Reward:
-    return grade_action(task, contract_before, action, proposed_contract_text)
+    """Grade hard-plus IP tasks with trap-resolution + ownership-clarity checks."""
+    reward, _ = evaluate_action(task, contract_before, action, proposed_contract_text)
+    content = (action.content or "").strip().lower()
+    if trap_unresolved(task, proposed_contract_text):
+        reward.score = max(0.001, min(0.999, reward.score * _HARD_PLUS_TRAP_PENALTY))
+    if any(kw in content for kw in ("customer owns", "customer-owned", "owned by customer")):
+        reward.score = max(0.001, min(0.999, reward.score * _HARD_PLUS_OWNERSHIP_BONUS))
+    return reward
 
 
 # ============ GRADER REGISTRY ============

contract_env/server/app.py

@@ -11,13 +11,19 @@
 from pydantic import ValidationError
 
 from contract_env.env.environment import ContractEnv
+from contract_env.env.graders import TASK_GRADERS, NUM_GRADED_TASKS
 from contract_env.env.models import Action, StepRequest
+from contract_env.env.tasks import TASKS
 
 _env = ContractEnv()
 
 app = FastAPI(
     title="Contract Negotiation OpenEnv",
-    version="1.1.0",
+    description=(
+        "AI-driven environment for evaluating contract-negotiation agents. "
+        "Agents analyse clauses, identify risks, and propose safer alternatives."
+    ),
+    version="1.2.0",
 )
 
 app.add_middleware(
@@ -29,13 +35,13 @@
 )
 
 
-# ---------------- ROOT ----------------
+# ── ROOT ────────────────────────────────────────────────────────────────
 @app.get("/")
 def root():
     return {"status": "ok", "service": "contract-negotiation-env"}
 
 
-# ---------------- ERROR HANDLERS ----------------
+# ── ERROR HANDLERS ──────────────────────────────────────────────────────
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request: Request, exc: HTTPException):
     return JSONResponse(
@@ -60,43 +66,59 @@ async def global_exception_handler(request: Request, exc: Exception):
     )
 
 
-# ---------------- HEALTH ----------------
+# ── HEALTH ──────────────────────────────────────────────────────────────
 @app.get("/health")
 def health():
     return {"status": "ok"}
 
 
-# ---------------- STATE ----------------
+# ── TASKS LISTING ───────────────────────────────────────────────────────
+@app.get("/tasks")
+def list_tasks():
+    """Return metadata for every registered task."""
+    return {
+        "total": len(TASKS),
+        "graded": NUM_GRADED_TASKS,
+        "tasks": [
+            {
+                "id": t.id,
+                "name": t.name,
+                "clause_type": t.clause_type,
+                "risk_level": t.risk_level,
+                "industry_context": t.industry_context,
+                "has_grader": t.id in TASK_GRADERS,
+            }
+            for t in TASKS
+        ],
+    }
+
+
+# ── STATE ───────────────────────────────────────────────────────────────
 @app.get("/state")
 def get_state():
     return _env.state()
 
 
-# ---------------- RESET (FIXED) ----------------
+# ── RESET ───────────────────────────────────────────────────────────────
 @app.post("/reset")
 def reset():
     try:
         obs = _env.reset()
-
-        return {
-            "observation": obs.model_dump()
-        }
-
+        return {"observation": obs.model_dump()}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 
-# ---------------- STEP (FIXED) ----------------
+# ── STEP ────────────────────────────────────────────────────────────────
 @app.post("/step")
 def step(req: StepRequest):
     try:
         action = Action(action_type=req.action_type, content=req.content)
-
         obs, reward, done, info = _env.step(action)
 
         return {
             "observation": obs.model_dump(),
-            "reward": {"score": reward},  # CRITICAL FIX
+            "reward": {"score": reward},
             "done": done,
             "info": info,
         }
@@ -105,7 +127,7 @@ def step(req: StepRequest):
         raise HTTPException(status_code=422, detail=e.errors())
 
 
-# ---------------- MAIN ----------------
+# ── MAIN ────────────────────────────────────────────────────────────────
 def main():
     import uvicorn
 

contract_env/tests/test_api.py

@@ -43,6 +43,18 @@ def test_state_endpoint(self) -> None:
         self.assertEqual(r.status_code, 200)
         self.assertIn("current_step", r.json())
 
+    def test_tasks_endpoint(self) -> None:
+        r = self.client.get("/tasks")
+        self.assertEqual(r.status_code, 200)
+        data = r.json()
+        self.assertGreaterEqual(data["total"], 5)
+        self.assertGreaterEqual(data["graded"], 3)
+        self.assertEqual(len(data["tasks"]), data["total"])
+        for t in data["tasks"]:
+            self.assertIn("id", t)
+            self.assertIn("clause_type", t)
+            self.assertIn("has_grader", t)
+
 
 if __name__ == "__main__":
     unittest.main()