simulator.py

#!/usr/bin/env python3
"""
Contract Negotiation Simulator
Author: Marco De Roni | Legal AI
GitHub: https://github.com/marcoderoni/contract-negotiation-simulator
"""

import os
import sys
import json
import time
import textwrap
from datetime import datetime
from anthropic import Anthropic

# ── Colours ────────────────────────────────────────────────────────────────
GREEN   = "\033[92m"
YELLOW  = "\033[93m"
RED     = "\033[91m"
CYAN    = "\033[96m"
BOLD    = "\033[1m"
DIM     = "\033[2m"
RESET   = "\033[0m"

# ── Scenarios ───────────────────────────────────────────────────────────────
SCENARIOS = {
    "1": {
        "name": "SaaS MSA — Limitation of Liability",
        "type": "MSA",
        "your_role": "Vendor Counsel",
        "counterparty_role": "Customer Counsel",
        "description": "Customer wants uncapped liability; you must defend a 12-month ACV cap.",
        "opening_clause": (
            "Except for breaches of confidentiality, IP infringement, or wilful misconduct, "
            "each party's total liability shall not exceed the fees paid in the 12 months "
            "preceding the claim."
        ),
        "playbook_guidance": (
            "Acceptable: up to 24-month cap for regulated industries. "
            "Never accept: unlimited liability or revenue-based caps. "
            "Fallback: carve out data breaches at 2× ACV."
        ),
        "difficulty": "Medium",
        "turns": 6,
    },
    "2": {
        "name": "DPA — Sub-processor Approval",
        "type": "DPA",
        "your_role": "Vendor Counsel",
        "counterparty_role": "Customer DPO/Counsel",
        "description": (
            "Customer demands individual prior written consent for every sub-processor change. "
            "You must protect operational flexibility while maintaining GDPR compliance."
        ),
        "opening_clause": (
            "Vendor may engage sub-processors listed in Annex A. Vendor shall notify Customer "
            "30 days in advance of any new sub-processor. Customer may object within 14 days; "
            "if no objection is raised, the sub-processor is deemed approved."
        ),
        "playbook_guidance": (
            "Acceptable: 14-day objection window, published sub-processor list. "
            "Never accept: prior individual written consent per sub-processor. "
            "Fallback: termination right if objection is reasonable and vendor cannot remedy."
        ),
        "difficulty": "Hard",
        "turns": 6,
    },
    "3": {
        "name": "CSA — IP Ownership & Work for Hire",
        "type": "CSA",
        "your_role": "Vendor Counsel",
        "counterparty_role": "Customer Counsel",
        "description": (
            "Customer demands ownership of all deliverables including pre-existing IP. "
            "You must protect the vendor's core technology and IP."
        ),
        "opening_clause": (
            "All deliverables created under this Agreement shall be considered works made for hire. "
            "To the extent any deliverable does not qualify, Vendor hereby assigns all rights, "
            "title and interest to Customer, excluding Vendor's Pre-Existing IP as defined herein."
        ),
        "playbook_guidance": (
            "Acceptable: assignment of bespoke deliverables excluding pre-existing IP and general know-how. "
            "Never accept: assignment of platform IP or any broadly defined 'improvements'. "
            "Fallback: licence-back of deliverables incorporating Vendor IP."
        ),
        "difficulty": "Medium",
        "turns": 5,
    },
    "4": {
        "name": "MSA — Termination for Convenience",
        "type": "MSA",
        "your_role": "Vendor Counsel",
        "counterparty_role": "Customer Counsel",
        "description": (
            "Customer insists on termination for convenience at any time with 30-day notice. "
            "You need to protect multi-year revenue commitments."
        ),
        "opening_clause": (
            "This Agreement may be terminated by either party for any reason upon 90 days' "
            "written notice, provided that any outstanding fees for the remainder of the "
            "current subscription term shall become immediately due and payable upon termination."
        ),
        "playbook_guidance": (
            "Acceptable: T4C after first anniversary with 90-day notice and pro-rata refund. "
            "Never accept: T4C with full refund within first 12 months. "
            "Fallback: T4C with 6-month notice after year 1 and no refund of past fees."
        ),
        "difficulty": "Easy",
        "turns": 5,
    },
}

# ── Counterparty system prompt ───────────────────────────────────────────────
def build_counterparty_prompt(scenario: dict) -> str:
    return f"""You are an experienced {scenario['counterparty_role']} negotiating a {scenario['type']} contract clause.

SCENARIO: {scenario['description']}

YOUR MANDATE:
- Protect your client's interests aggressively but professionally
- Push back on the vendor's opening position
- Propose concrete redlines and alternative language — always use proper legal drafting
- Vary your tactics: sometimes concede minor points to gain major ones
- You may escalate pressure or show flexibility depending on how well the user argues
- Stay in character at all times — respond as a real counterparty counsel would in a live negotiation

CURRENT CLAUSE UNDER NEGOTIATION:
{scenario['opening_clause']}

TONE: Professional, direct, occasionally firm. Use concise legal language. 
LENGTH: Keep responses to 3–6 sentences unless proposing redline language, in which case draft the full clause.
FORMAT: Speak as yourself (the counterparty counsel). Do not narrate or explain what you are doing."""

# ── Scoring prompt ───────────────────────────────────────────────────────────
def build_scoring_prompt(scenario: dict, history: list) -> str:
    convo = "\n".join([
        f"{'YOU' if m['role'] == 'user' else 'COUNTERPARTY'}: {m['content']}"
        for m in history
    ])
    return f"""You are a senior legal negotiation coach evaluating a contract negotiation exercise.

SCENARIO: {scenario['name']}
USER ROLE: {scenario['your_role']}
PLAYBOOK GUIDANCE: {scenario['playbook_guidance']}

NEGOTIATION TRANSCRIPT:
{convo}

Evaluate the user's negotiation performance and respond ONLY with a JSON object in exactly this format:
{{
  "overall_score": <integer 0-100>,
  "grade": "<A / B / C / D / F>",
  "dimensions": {{
    "playbook_adherence": {{
      "score": <integer 0-100>,
      "comment": "<one sentence>"
    }},
    "legal_precision": {{
      "score": <integer 0-100>,
      "comment": "<one sentence>"
    }},
    "negotiation_strategy": {{
      "score": <integer 0-100>,
      "comment": "<one sentence>"
    }},
    "commercial_awareness": {{
      "score": <integer 0-100>,
      "comment": "<one sentence>"
    }}
  }},
  "strengths": ["<strength 1>", "<strength 2>"],
  "improvements": ["<improvement 1>", "<improvement 2>"],
  "best_move": "<The single best thing the user did or said>",
  "missed_opportunity": "<The key point the user failed to press or exploit>",
  "suggested_closing_position": "<How a senior counsel would have closed this clause>"
}}"""

# ── Helpers ──────────────────────────────────────────────────────────────────
def wrap(text: str, width: int = 82) -> str:
    return "\n".join(
        textwrap.fill(line, width) if line.strip() else line
        for line in text.split("\n")
    )

def print_header():
    print(f"\n{BOLD}{CYAN}")
    print("╔══════════════════════════════════════════════════════════════════════╗")
    print("║        CONTRACT NEGOTIATION SIMULATOR  ·  Legal AI by Marco De Roni  ║")
    print("╚══════════════════════════════════════════════════════════════════════╝")
    print(RESET)

def print_scenario_menu():
    print(f"{BOLD}Select a scenario:{RESET}\n")
    for key, s in SCENARIOS.items():
        diff_colour = GREEN if s["difficulty"] == "Easy" else (YELLOW if s["difficulty"] == "Medium" else RED)
        print(f"  {CYAN}[{key}]{RESET} {s['name']}")
        print(f"       {DIM}{s['description'][:72]}...{RESET}")
        print(f"       Type: {s['type']}  ·  Turns: {s['turns']}  ·  Difficulty: {diff_colour}{s['difficulty']}{RESET}\n")

def print_clause_box(label: str, text: str):
    print(f"\n{BOLD}{YELLOW}{'─'*72}")
    print(f"  {label}")
    print(f"{'─'*72}{RESET}")
    print(wrap(f"  {text}"))
    print(f"{YELLOW}{'─'*72}{RESET}\n")

def print_counterparty(text: str):
    print(f"\n{RED}{BOLD}⚖  COUNTERPARTY:{RESET}")
    print(wrap(f"   {text}"))

def print_score(result: dict):
    grade_colours = {"A": GREEN, "B": GREEN, "C": YELLOW, "D": YELLOW, "F": RED}
    g = result.get("grade", "C")
    c = grade_colours.get(g, YELLOW)
    score = result.get("overall_score", 0)

    print(f"\n{BOLD}{CYAN}{'═'*72}")
    print(f"  FINAL SCORE: {c}{score}/100  ·  Grade: {g}{RESET}{BOLD}{CYAN}")
    print(f"{'═'*72}{RESET}\n")

    dims = result.get("dimensions", {})
    labels = {
        "playbook_adherence":    "Playbook Adherence  ",
        "legal_precision":       "Legal Precision     ",
        "negotiation_strategy":  "Negotiation Strategy",
        "commercial_awareness":  "Commercial Awareness",
    }
    for key, label in labels.items():
        d = dims.get(key, {})
        s = d.get("score", 0)
        bar = "█" * (s // 10) + "░" * (10 - s // 10)
        colour = GREEN if s >= 75 else (YELLOW if s >= 50 else RED)
        print(f"  {label}  {colour}{bar} {s:>3}{RESET}  {DIM}{d.get('comment','')}{RESET}")

    print(f"\n{BOLD}✅ Strengths:{RESET}")
    for s in result.get("strengths", []):
        print(f"   • {s}")

    print(f"\n{BOLD}🔧 Areas to improve:{RESET}")
    for i in result.get("improvements", []):
        print(f"   • {i}")

    print(f"\n{BOLD}🏆 Best move:{RESET}")
    print(wrap(f"   {result.get('best_move','—')}"))

    print(f"\n{BOLD}⚠️  Missed opportunity:{RESET}")
    print(wrap(f"   {result.get('missed_opportunity','—')}"))

    print(f"\n{BOLD}{GREEN}💡 How a senior counsel would have closed it:{RESET}")
    print(wrap(f"   {result.get('suggested_closing_position','—')}"))
    print()

def save_session(scenario: dict, history: list, score: dict):
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    fname = f"session_{scenario['type']}_{ts}.json"
    data = {
        "scenario": scenario["name"],
        "timestamp": ts,
        "score": score,
        "transcript": history,
    }
    with open(fname, "w") as f:
        json.dump(data, f, indent=2)
    print(f"{DIM}Session saved → {fname}{RESET}\n")

# ── Main loop ────────────────────────────────────────────────────────────────
def run():
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        print(f"{RED}Error: ANTHROPIC_API_KEY not set.{RESET}")
        sys.exit(1)

    client = Anthropic(api_key=api_key)

    print_header()
    print_scenario_menu()

    # Scenario selection
    while True:
        choice = input(f"{BOLD}Enter scenario number (1-{len(SCENARIOS)}) or Q to quit: {RESET}").strip()
        if choice.upper() == "Q":
            print("Goodbye.")
            sys.exit(0)
        if choice in SCENARIOS:
            scenario = SCENARIOS[choice]
            break
        print(f"{RED}Invalid choice. Try again.{RESET}")

    print(f"\n{BOLD}{GREEN}Starting: {scenario['name']}{RESET}")
    print(f"{DIM}You are: {scenario['your_role']}  ·  {scenario['turns']} turns  ·  Difficulty: {scenario['difficulty']}{RESET}")
    print(f"\n{DIM}Tip: Propose redline language, cite playbook limits, use commercial arguments.{RESET}")
    print(f"{DIM}Type 'score' at any point to end the session and receive your evaluation.{RESET}\n")
    input(f"Press {BOLD}ENTER{RESET} to begin...\n")

    # Show opening clause
    print_clause_box("📄 OPENING CLAUSE (your position)", scenario["opening_clause"])

    # Build message history for counterparty AI
    system_prompt = build_counterparty_prompt(scenario)
    history = []      # stores {"role": "user"/"assistant", "content": "..."}
    ui_history = []   # same structure, used for scoring

    # Counterparty opens
    opening_msg = (
        f"I've reviewed your proposed {scenario['type']} clause. "
        f"Let me explain why this doesn't work for us and what we need instead."
    )
    init_response = client.messages.create(
        model="claude-opus-4-5",
        max_tokens=512,
        system=system_prompt,
        messages=[{"role": "user", "content": opening_msg}],
    )
    cp_opening = init_response.content[0].text
    history.append({"role": "user", "content": opening_msg})
    history.append({"role": "assistant", "content": cp_opening})
    ui_history.append({"role": "assistant", "content": cp_opening})
    print_counterparty(cp_opening)

    # ── Negotiation turns ────────────────────────────────────────────────────
    turn = 0
    while turn < scenario["turns"]:
        turn += 1
        prompt = f"{BOLD}[Turn {turn}/{scenario['turns']}] Your response:{RESET}\n> "
        try:
            user_input = input(prompt).strip()
        except (KeyboardInterrupt, EOFError):
            print("\nSession interrupted.")
            break

        if not user_input:
            continue

        if user_input.lower() in ("score", "exit", "quit", "q"):
            break

        ui_history.append({"role": "user", "content": user_input})

        # Append to history for counterparty context
        history.append({"role": "user", "content": user_input})

        # Get counterparty response
        print(f"{DIM}Counterparty is responding...{RESET}")
        response = client.messages.create(
            model="claude-opus-4-5",
            max_tokens=512,
            system=system_prompt,
            messages=history,
        )
        cp_reply = response.content[0].text
        history.append({"role": "assistant", "content": cp_reply})
        ui_history.append({"role": "assistant", "content": cp_reply})
        print_counterparty(cp_reply)

        if turn == scenario["turns"]:
            print(f"\n{YELLOW}{BOLD}⏱  Final turn reached. Generating your score...{RESET}\n")
            time.sleep(1)

    # ── Scoring ──────────────────────────────────────────────────────────────
    if not ui_history:
        print(f"{YELLOW}No turns completed — no score generated.{RESET}")
        return

    print(f"\n{DIM}Analysing your negotiation...{RESET}")
    scoring_prompt = build_scoring_prompt(scenario, ui_history)
    score_response = client.messages.create(
        model="claude-opus-4-5",
        max_tokens=1024,
        messages=[{"role": "user", "content": scoring_prompt}],
    )
    raw = score_response.content[0].text.strip()

    # Strip markdown fences if present
    if raw.startswith("```"):
        raw = raw.split("```")[1]
        if raw.startswith("json"):
            raw = raw[4:]

    try:
        score_data = json.loads(raw)
    except json.JSONDecodeError:
        print(f"{RED}Could not parse score. Raw output:{RESET}\n{raw}")
        return

    print_score(score_data)
    save_session(scenario, ui_history, score_data)

    # Replay option
    again = input(f"{BOLD}Run another scenario? (Y/N): {RESET}").strip().upper()
    if again == "Y":
        run()

if __name__ == "__main__":
    run()