Long Document Processing.

import hashlib
from typing import TypedDict

class Chunk(TypedDict):
    chunk_id: str
    page: int
    text: str

def chunk_document(pages: list[str], target_tokens: int = 1200, overlap: float = 0.15) -> list[Chunk]:
    """Semantic chunking with overlap. Split at section, paragraph, sentence."""
    chunks = []
    buffer = ""
    page_buffer_started = 1

    for page_num, page_text in enumerate(pages, start=1):
        for paragraph in split_into_paragraphs(page_text):
            # If adding this paragraph would exceed target, emit current buffer
            if approx_tokens(buffer + paragraph) > target_tokens and buffer:
                chunks.append({
                    "chunk_id": hashlib.md5(buffer.encode()).hexdigest()[:12],
                    "page": page_buffer_started,
                    "text": buffer.strip(),
                })
                # Carry forward the last 15% as overlap
                tail = buffer[-int(len(buffer) * overlap):]
                buffer = tail + paragraph + "\n\n"
                page_buffer_started = page_num
            else:
                buffer += paragraph + "\n\n"

    if buffer.strip():
        chunks.append({
            "chunk_id": hashlib.md5(buffer.encode()).hexdigest()[:12],
            "page": page_buffer_started,
            "text": buffer.strip(),
        })
    return chunks

def split_into_paragraphs(page_text: str) -> list[str]:
    return [p for p in page_text.split("\n\n") if p.strip()]

def approx_tokens(text: str) -> int:
    return len(text) // 4  # rule of thumb

# embed_and_index.py
import voyageai

vo = voyageai.Client()  # picks up VOYAGE_API_KEY

def embed_and_index(chunks: list[Chunk], collection: str):
    """Embed every chunk; store in vector DB keyed by chunk_id."""
    texts = [c["text"] for c in chunks]
    # Voyage supports batched embedding. Much cheaper than per-chunk
    embeddings = vo.embed(texts, model="voyage-3", input_type="document").embeddings

    # Pseudo-code for vector store; real impl uses pgvector / Pinecone / FAISS
    for chunk, embedding in zip(chunks, embeddings):
        vector_store.upsert(
            collection=collection,
            id=chunk["chunk_id"],
            vector=embedding,
            metadata={"page": chunk["page"], "text_preview": chunk["text"][:200]},
        )
    return len(embeddings)

# Re-embedding cache: skip if chunk content hash hasn't changed
def smart_reindex(doc_id: str, new_chunks: list[Chunk]):
    existing = vector_store.list_chunks(collection=doc_id)
    existing_ids = {c["id"] for c in existing}
    new_ids = {c["chunk_id"] for c in new_chunks}
    to_delete = existing_ids - new_ids
    to_add = [c for c in new_chunks if c["chunk_id"] not in existing_ids]
    vector_store.delete_many(doc_id, list(to_delete))
    embed_and_index(to_add, doc_id)

def retrieve_top_k(question: str, doc_id: str, k: int = 5) -> list[dict]:
    """Top-K retrieval. Full doc never enters context."""
    q_embed = vo.embed([question], model="voyage-3", input_type="query").embeddings[0]
    candidates = vector_store.search(collection=doc_id, vector=q_embed, top_k=20)
    # Optional re-rank with Haiku for precision
    if len(candidates) > k:
        candidates = rerank_with_haiku(question, candidates)[:k]
    return [
        {
            "chunk_id": c["id"],
            "page": c["metadata"]["page"],
            "text": c["metadata"]["text_full"],  # rehydrate from chunk store
            "score": c["score"],
        }
        for c in candidates[:k]
    ]

def rerank_with_haiku(question: str, candidates: list[dict]) -> list[dict]:
    """Use Haiku to re-rank embedding candidates (cheap, focused)."""
    prompt = (
        f"Question: {question}\n\n"
        + "\n\n".join(f"[{i}] {c['metadata']['text_preview']}"
                       for i, c in enumerate(candidates))
        + "\n\nReturn the indices of the top 5 most relevant chunks, JSON only:"
    )
    resp = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=64,
        messages=[{"role": "user", "content": prompt}],
    )
    indices = json.loads(resp.content[0].text)
    return [candidates[i] for i in indices]

def build_system_prompt(case_facts: dict, retrieved_chunks: list[dict]) -> str:
    """Pin CASE_FACTS at the top; retrieved chunks below; conversation last."""
    chunks_text = "\n\n".join(
        f"### CHUNK {c['chunk_id']} (page {c['page']})\n{c['text']}"
        for c in retrieved_chunks
    )
    return f"""You are a long-document extraction agent.

CASE_FACTS (immutable; re-read every turn; values are EXACT, never paraphrased):
- doc_id: {case_facts['doc_id']}
- doc_type: {case_facts.get('doc_type', 'unknown')}
- extracted_count: {case_facts.get('extracted_count', 0)}
- last_clause_id: {case_facts.get('last_clause_id', 'none')}
- policy_cap: ${case_facts.get('policy_cap', 0):,.2f}

Constraints:
- Cite every extraction with chunk_id + page from the chunks below.
- Never paraphrase exact values from CASE_FACTS or chunk text.
- Branch on stop_reason. On max_tokens, save state. The harness will resume.

RETRIEVED CHUNKS (top-K; only these are in context):
{chunks_text}"""

def update_case_facts(case_facts: dict, new_extraction: dict) -> dict:
    """Hook-style update; preserves all values verbatim."""
    return {
        **case_facts,
        "extracted_count": case_facts.get("extracted_count", 0) + 1,
        "last_clause_id": new_extraction["clause_id"],
    }

import json
from datetime import datetime

CHECKPOINT_DIR = ".checkpoints"

def save_checkpoint(case_facts: dict, last_record: dict, position: dict):
    """Persist state on max_tokens. Idempotent on doc_id."""
    path = f"{CHECKPOINT_DIR}/{case_facts['doc_id']}.json"
    with open(path, "w") as f:
        json.dump({
            "case_facts": case_facts,
            "last_record": last_record,
            "position": position,  # {chunk_id, paragraph_offset}
            "saved_at": datetime.utcnow().isoformat() + "Z",
        }, f)
    return path

def load_checkpoint(doc_id: str) -> dict | None:
    path = f"{CHECKPOINT_DIR}/{doc_id}.json"
    if not os.path.exists(path):
        return None
    with open(path) as f:
        return json.load(f)

def extract_with_resume(doc_id: str, question: str, max_iter: int = 50):
    """Top-level loop with automatic checkpoint-and-resume."""
    checkpoint = load_checkpoint(doc_id) or {}
    case_facts = checkpoint.get("case_facts") or {"doc_id": doc_id}

    for iteration in range(max_iter):
        chunks = retrieve_top_k(question, doc_id, k=5)
        resp = client.messages.create(
            model="claude-sonnet-4.5",
            max_tokens=4096,
            system=build_system_prompt(case_facts, chunks),
            tools=[EXTRACT_CLAUSE_TOOL],
            messages=[{"role": "user", "content": question}],
        )
        if resp.stop_reason == "end_turn":
            return {"status": "complete", "case_facts": case_facts}
        if resp.stop_reason == "max_tokens":
            # Save and continue in a fresh session
            save_checkpoint(case_facts, last_record={}, position={})
            continue  # next iteration starts a fresh session
        if resp.stop_reason == "tool_use":
            tool_use = next(b for b in resp.content if b.type == "tool_use")
            case_facts = update_case_facts(case_facts, tool_use.input)
    return {"status": "iteration_cap", "case_facts": case_facts}

EXTRACT_CLAUSE_TOOL = {
    "name": "extract_clause",
    "description": (
        "Extract a contractual clause from the retrieved chunks.\n"
        "Use this when the user asks for a specific clause type.\n"
        "Edge cases: if the clause type is not present in any retrieved chunk, "
        "emit clause_text='not_found' with empty citations.\n"
        "ALWAYS cite chunk_id and page for every extraction."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "clause_id": {"type": "string"},
            "clause_type": {
                "type": "string",
                "enum": ["indemnification", "termination", "payment", "ip", "other"],
            },
            "clause_text": {"type": "string"},
            "citations": {
                "type": "array",
                "minItems": 1,  # require at least one citation
                "items": {
                    "type": "object",
                    "properties": {
                        "chunk_id": {"type": "string"},
                        "page": {"type": "integer", "minimum": 1},
                        "span": {
                            "type": "string",
                            "description": "character offsets within the chunk, e.g. '128-340'",
                        },
                    },
                    "required": ["chunk_id", "page"],
                },
            },
        },
        "required": ["clause_id", "clause_type", "clause_text", "citations"],
    },
}

def submit_bulk_extraction(docs: list[dict], clause_type: str) -> str:
    """Submit a batch of clause-extraction requests for overnight processing."""
    requests = []
    for doc in docs:
        chunks = retrieve_top_k(f"Find {clause_type} clauses", doc["id"], k=5)
        case_facts = load_checkpoint(doc["id"]) or {"doc_id": doc["id"]}
        requests.append({
            "custom_id": f"clause-{doc['id']}-{clause_type}",
            "params": {
                "model": "claude-sonnet-4.5",
                "max_tokens": 2048,
                "system": [{
                    "type": "text",
                    "text": build_system_prompt(case_facts, chunks),
                    "cache_control": {"type": "ephemeral"},  # cache the system prompt
                }],
                "tools": [
                    {**EXTRACT_CLAUSE_TOOL, "cache_control": {"type": "ephemeral"}},
                ],
                "tool_choice": {"type": "tool", "name": "extract_clause"},
                "messages": [{"role": "user", "content": f"Extract all {clause_type} clauses."}],
            },
        })
    batch = client.messages.batches.create(requests=requests)
    return batch.id

# Next morning. Fetch + harvest
def harvest_bulk(batch_id: str):
    results = client.messages.batches.results(batch_id)
    accepted, retry_queue = [], []
    for r in results:
        if r.result.type == "succeeded":
            tu = next(b for b in r.result.message.content if b.type == "tool_use")
            accepted.append(tu.input)
        else:
            retry_queue.append(r.custom_id)
    return {"accepted": accepted, "retry": retry_queue}

from collections import defaultdict

def stratified_accuracy(extractions: list[dict]) -> dict:
    """Pass rate by doc_type × clause_type × page-section."""
    buckets = defaultdict(lambda: {"pass": 0, "fail": 0})
    for e in extractions:
        section = ("front" if e["citations"][0]["page"] <= 30
                   else "back" if e["citations"][0]["page"] > 150
                   else "middle")
        key = (e["doc_type"], e["clause_type"], section)
        bucket = "pass" if validate_extraction(e) else "fail"
        buckets[key][bucket] += 1

    report = {}
    for (doc_type, clause_type, section), counts in buckets.items():
        total = counts["pass"] + counts["fail"]
        report[f"{doc_type}/{clause_type}/{section}"] = {
            "total": total,
            "pass_rate": counts["pass"] / total if total else 0,
        }
    return dict(sorted(report.items(), key=lambda kv: kv[1]["pass_rate"]))

def adversarial_silent_source_test() -> float:
    """50 docs where the requested clause is GENUINELY absent."""
    correct = 0
    for doc in load_silent_source_docs():  # known to NOT contain the clause
        result = extract_with_resume(doc["id"], "Find the indemnification clause")
        # Right: clause_text='not_found' with empty citations
        # Wrong: ANY invented clause text
        if result.get("case_facts", {}).get("last_extraction", {}).get("clause_text") == "not_found":
            correct += 1
    return correct / 50  # target: ≥ 95%