{"repo": "aayushbaniya2006/contract-data-extraction", "files": {"llm.py": {"77df07b56fee2007d90b743627a46965c31ee34d": {"t": 1785674432.5156386, "n": 7316, "h": "362447fab0f242cf", "c": "\"\"\"\nllm.py — the \"extract\" stage: document -> structured fields, via Claude.\n\nTwo tiers (this is the cost/accuracy cascade from the memo):\n  - extract_fields()        Tier 1: one CHEAP call over the OCR'd text. Handles the\n                            easy ~80% (parties, contract type, clean dates).\n  - extract_from_pdf_pages()Tier 2: only for the few fields the text pass couldn't\n                            read, re-read the cited pages as actual IMAGES with the\n                            strong model — so garbled OCR no longer hides the value.\n\nThis is the only file that talks to the model. Everything model-specific lives here,\nso the rest of the pipeline never changes if we swap models or add a tier.\n\"\"\"\n\nimport base64\nimport json\nimport os\nimport time\n\nimport anthropic\n\nfrom schema import (build_json_schema, build_system_prompt, build_vision_system_prompt,\n                    build_term_schema, TERM_DERIVATION_PROMPT)\n\n# The pipeline is model-AGNOSTIC. Two tiers, each configurable via env vars:\n#   BASE_MODEL   — cheap first pass over the whole document text.\n#   VISION_MODEL — re-reads hard pages as page IMAGES (multimodal model required).\n# Defaults are Haiku, which every key (including the limited OAuth token) can reach.\n# To \"go up\", set BASE_MODEL / VISION_MODEL to a stronger model (e.g. claude-opus-4-8)\n# when you have a standard sk-ant-api key that can reach it — no code change needed.\nDEFAULT_MODEL = \"claude-haiku-4-5\"\nBASE_MODEL = os.environ.get(\"BASE_MODEL\", DEFAULT_MODEL)\nVISION_MODEL = os.environ.get(\"VISION_MODEL\", DEFAULT_MODEL)\n# If a chosen model is not reachable on the current key, we fall back to this one so\n# the pipeline still runs end-to-end instead of hard-failing.\nFALLBACK_MODEL = os.environ.get(\"FALLBACK_MODEL\", DEFAULT_MODEL)\n\n_fallback_announced = set()  # so we log each \"X unavailable, using Y\" note only once\n\n# Built lazily (on first use) so a missing key gives a clear message, not a stack trace.\n_client = None\n\n\ndef _get_client():\n    global _client\n    if _client is None:\n        key = os.environ.get(\"ANTHROPIC_API_KEY\")\n        if not key:\n            raise SystemExit(\n                \"ANTHROPIC_API_KEY is not set.\\n\"\n                \"Put it in a .env file, or:  export ANTHROPIC_API_KEY=sk-ant-...\"\n            )\n        # OAuth tokens (sk-ant-oat...) authenticate with a Bearer token; classic API\n        # keys (sk-ant-api...) use the x-api-key header. Pick the matching one.\n        if key.startswith(\"sk-ant-oat\"):\n            _client = anthropic.Anthropic(auth_token=key)\n        else:\n            _client = anthropic.Anthropic(api_key=key)\n    return _client\n\n\ndef _create(model, _attempt=0, **kwargs):\n    \"\"\"Call the Messages API with `model`, but if that model isn't reachable on the\n    current key (403 no-access / 404 unknown), transparently fall back to FALLBACK_MODEL.\n    This is what lets ONE codebase run on a Haiku-only OAuth token AND scale up to\n    Opus/Sonnet on a full API key, with no edits.\n\n    A 429 is treated as a *transient rate limit first*: we back off and retry the SAME\n    model a few times before ever downgrading. Silently swapping a strong model for Haiku\n    on a momentary rate spike would degrade quality on legal data — the opposite of what\n    a rate limit should trigger. Only a persistent 429 (e.g. a token not entitled to the\n    model) falls through to the fallback, so the run still completes end-to-end.\"\"\"\n    client = _get_client()\n    try:\n        return client.messages.create(model=model, **kwargs)\n    except anthropic.APIStatusError as e:\n        # Some models (e.g. Opus 4.8) reject `temperature` -> drop it and retry, so\n        # the same call works whether the model supports sampling params or not.\n        if e.status_code == 400 and \"temperature\" in kwargs and \"temperature\" in str(e).lower():\n            kwargs.pop(\"temperature\")\n            return _create(model, _attempt=_attempt, **kwargs)\n        # 429 rate-limited: back off (1s, 2s, 4s) and retry the SAME model before downgrading.\n        if e.status_code == 429 and _attempt < 3:\n            time.sleep(2 ** _attempt)\n            return _create(model, _attempt=_attempt + 1, **kwargs)\n        # Model not reachable on this key (no access / unknown / persistently rate-limited)\n        # -> fall back so the run still completes end-to-end instead of hard-failing.\n        if model == FALLBACK_MODEL or e.status_code not in (403, 404, 429):\n            raise  # already on fallback, or a real error — surface it\n        if model not in _fallback_announced:\n            print(f\"    note: '{model}' unavailable on this key \"\n                  f\"(HTTP {e.status_code}); falling back to '{FALLBACK_MODEL}'\")\n            _fallback_announced.add(model)\n        return _create(FALLBACK_MODEL, **kwargs)\n\n\ndef extract_fields(document_text, fields):\n    \"\"\"Tier 1: one cheap Claude call over the document TEXT -> validated dict of fields.\"\"\"\n    response = _create(\n        BASE_MODEL,\n        max_tokens=3000,                        # plenty for a small JSON answer\n        temperature=0,                          # reproducible runs (auto-dropped if a model rejects it)\n        system=build_system_prompt(fields),\n        messages=[{\"role\": \"user\", \"content\": document_text}],\n        # structured output: the response is guaranteed to be JSON matching the schema\n        output_config={\"format\": {\"type\": \"json_schema\", \"schema\": build_json_schema(fields)}},\n    )\n    text = next(block.text for block in response.content if block.type == \"text\")\n    return json.loads(text)\n\n\ndef extract_term_components(document_text):\n    \"\"\"Focused pass: pull the term length + any commencement anchor/cap FROM THE LEASE\n    (over the whole text) as structured JSON. Python does the date math, not the model.\"\"\"\n    response = _create(\n        BASE_MODEL,\n        max_tokens=800,\n        temperature=0,\n        system=TERM_DERIVATION_PROMPT,\n        messages=[{\"role\": \"user\", \"content\": document_text}],\n        output_config={\"format\": {\"type\": \"json_schema\", \"schema\": build_term_schema()}},\n    )\n    text = next(block.text for block in response.content if block.type == \"text\")\n    return json.loads(text)\n\n\ndef extract_from_pdf_pages(pdf_bytes, fields):\n    \"\"\"Tier 2: re-read the given pages as IMAGES with the strong model.\n\n    `pdf_bytes` is a small PDF holding only the pages the text pass couldn't read.\n    Sent as a `document` block, Claude sees the real page image, so OCR garbling no\n    longer hides the value. Returns a dict for just `fields`.\n    \"\"\"\n    b64 = base64.standard_b64encode(pdf_bytes).decode()\n    response = _create(\n        VISION_MODEL,\n        max_tokens=3000,\n        temperature=0,                          # reproducible runs (auto-dropped if a model rejects it)\n        system=build_vision_system_prompt(fields),\n        messages=[{\"role\": \"user\", \"content\": [\n            {\"type\": \"document\",\n             \"source\": {\"type\": \"base64\", \"media_type\": \"application/pdf\", \"data\": b64}},\n            {\"type\": \"text\", \"text\": \"Re-read these page images and extract the fields.\"},\n        ]}],\n        output_config={\"format\": {\"type\": \"json_schema\", \"schema\": build_json_schema(fields)}},\n    )\n    text = next(block.text for block in response.content if block.type == \"text\")\n    return json.loads(text)\n"}}}}