Evaluate and Improve Structured Extraction

%pip install -q "langfuse>=4.14" "openai>=3" pydantic pandas

import os

# Langfuse: project Settings → API Keys → https://cloud.langfuse.com
os.environ.setdefault("LANGFUSE_PUBLIC_KEY", "pk-lf-...")
os.environ.setdefault("LANGFUSE_SECRET_KEY", "sk-lf-...")
os.environ.setdefault("LANGFUSE_BASE_URL", "https://cloud.langfuse.com")  # 🇪🇺 EU region
# 🇺🇸 US region: https://us.cloud.langfuse.com

# OpenAI: https://platform.openai.com/api-keys
os.environ.setdefault("OPENAI_API_KEY", "sk-proj-...");

from langfuse import get_client

langfuse = get_client()
assert langfuse.auth_check(), "Langfuse auth failed. Check your keys and host."
print("Langfuse auth OK")

from __future__ import annotations
from datetime import date
from typing import List, Literal, Optional
from pydantic import BaseModel, ConfigDict, Field

# ---- shared enum aliases (defined once -> both variants stay identical) -------
Channel = Literal["Email", "Phone", "Portal", "In-Person"]
CoverageType = Literal["Property", "Auto", "Liability", "Health", "Travel", "Other"]
ObjectType = Literal["Vehicle", "Building", "Person", "Other"]
IncidentType = Literal[
    "rear_end_collision", "side_impact_collision", "head_on_collision", "parking_lot_collision",
    "house_fire", "kitchen_fire", "electrical_fire", "burst_pipe_flood", "storm_damage", "roof_leak",
    "slip_and_fall", "property_injury", "product_liability", "theft_burglary", "vandalism",
]
LocationType = Literal[
    "intersection", "highway", "parking_lot", "driveway", "residential_street",
    "residence_interior", "residence_exterior", "commercial_property", "public_property",
]
_STRICT = ConfigDict(extra="forbid")

# ---- BARE variant: structure + constraints only ------------------------------
class ClaimHeaderBare(BaseModel):
    model_config = _STRICT
    claim_id: str
    report_date: date
    incident_date: date
    reported_by: str = Field(..., min_length=1)
    channel: Channel

class PolicyDetailsBare(BaseModel):
    model_config = _STRICT
    policy_number: str
    policyholder_name: str = Field(..., min_length=1)
    coverage_type: CoverageType
    effective_date: date
    expiration_date: date

class InsuredObjectBare(BaseModel):
    model_config = _STRICT
    object_id: str
    object_type: ObjectType
    make_model: Optional[str] = None
    year: Optional[int] = None
    location_address: Optional[str] = None
    estimated_value: Optional[int] = None

class IncidentDescriptionBare(BaseModel):
    model_config = _STRICT
    incident_type: IncidentType
    location_type: LocationType
    estimated_damage_amount: Optional[int] = None
    police_report_number: Optional[str] = None

class InsuranceClaimBare(BaseModel):
    model_config = _STRICT
    header: ClaimHeaderBare
    policy_details: Optional[PolicyDetailsBare] = None
    insured_objects: Optional[List[InsuredObjectBare]] = None
    incident_description: IncidentDescriptionBare

# ---- DESCRIPTIVE variant: same structure + rich descriptions (the "spec") ----
class ClaimHeader(BaseModel):
    model_config = _STRICT
    claim_id: str = Field(..., description="Claim ID in format CLM-XXXXXX, where X is a digit")
    report_date: date = Field(..., description="Date claim was reported (ISO 8601)")
    incident_date: date = Field(..., description="Date incident occurred (ISO 8601)")
    reported_by: str = Field(..., min_length=1, description="Full name of person reporting the claim")
    channel: Channel = Field(..., description="Channel used to report the claim")

class PolicyDetails(BaseModel):
    model_config = _STRICT
    policy_number: str = Field(..., description="Policy number in format POL-XXXXXXXXX, where X is a digit")
    policyholder_name: str = Field(..., min_length=1, description="Full legal name on the policy")
    coverage_type: CoverageType = Field(..., description="Type of insurance coverage")
    effective_date: date = Field(..., description="Policy effective start date (ISO 8601)")
    expiration_date: date = Field(..., description="Policy expiration end date (ISO 8601)")

class InsuredObject(BaseModel):
    model_config = _STRICT
    object_id: str = Field(..., description=(
        "Unique identifier for the insured object. For vehicles use VIN format "
        "(e.g., VIN12345678901234567). For buildings use PROP-XXXXXX. For liability use "
        "LIAB-XXXXXX. Otherwise use OBJ-XXXXXX, where X is a digit"))
    object_type: ObjectType = Field(..., description="Type of insured object")
    make_model: Optional[str] = Field(None, description=(
        "Make and model for vehicles (use standardized manufacturer names and models), "
        "or building type for property"))
    year: Optional[int] = Field(None, description="Year for vehicles or year built for buildings")
    location_address: Optional[str] = Field(None, description="Full street address where object is located")
    estimated_value: Optional[int] = Field(None, description="Estimated value in USD, no currency symbol")

class IncidentDescription(BaseModel):
    model_config = _STRICT
    incident_type: IncidentType = Field(..., description="Specific standardized incident type")
    location_type: LocationType = Field(..., description="Standardized location type where it occurred")
    estimated_damage_amount: Optional[int] = Field(None, description="Estimated damage in USD, no symbol")
    police_report_number: Optional[str] = Field(None, description="Police report number if applicable")

class InsuranceClaim(BaseModel):
    model_config = _STRICT
    header: ClaimHeader = Field(..., description="Basic claim information")
    policy_details: Optional[PolicyDetails] = Field(None, description="Policy information if available")
    insured_objects: Optional[List[InsuredObject]] = Field(None, description="Insured objects, if applicable")
    incident_description: IncidentDescription = Field(..., description="Structured incident details")

SCHEMAS = {"bare": InsuranceClaimBare, "descriptive": InsuranceClaim}

# Peek at the strict JSON schema OpenAI will enforce.
from openai.lib._pydantic import to_strict_json_schema
strict = to_strict_json_schema(InsuranceClaim)
print("sections:", list(strict["properties"].keys()))
print("incident_type enum:", len(IncidentType.__args__), "values | location_type:", len(LocationType.__args__))

import ast
import pandas as pd

DATA_URL = "https://huggingface.co/datasets/Cleanlab/insurance-claims-extraction/resolve/main/insurance_claims_extraction.csv"
df = pd.read_csv(DATA_URL)
print("rows:", len(df), "| columns:", list(df.columns))
example_gt = ast.literal_eval(df["ground_truth"].iloc[0])   # ground_truth is a stringified dict
print("ground-truth sections:", list(example_gt.keys()))

DATASET_NAME = "insurance-claims"

langfuse.create_dataset(
    name=DATASET_NAME,
    description="Nested insurance-claim extraction (Cleanlab). 30 synthetic claims; verify before production.",
    metadata={"source": "Cleanlab/insurance-claims-extraction", "n_items": len(items)})
for it in items:
    langfuse.create_dataset_item(
        dataset_name=DATASET_NAME, id=f"{DATASET_NAME}:{it['id']}",
        input=it["input"], expected_output=it["expected_output"], metadata=it["metadata"])
langfuse.flush()
dataset = langfuse.get_dataset(DATASET_NAME)
print("uploaded & fetched:", len(dataset.items), "items")

from langfuse import Evaluation

def extraction_evaluator(*, output, expected_output, **kwargs):
    if not isinstance(output, dict):
        return [Evaluation(name="valid_json", value=False, data_type="BOOLEAN")]

    breakdown = compare_claim(output.get("prediction"), expected_output)
    evals = [Evaluation(name="valid_json", value=bool(breakdown.get("valid_json")), data_type="BOOLEAN")]

    for path, ok in breakdown.get("leaf_results") or []:
        evals.append(Evaluation(name=f"field.{path}", value=bool(ok), data_type="BOOLEAN"))

    if breakdown.get("valid_json"):
        overall = breakdown.get("field_accuracy_overall")
        evals.append(Evaluation(
            name="field_accuracy_overall",
            value=float(overall) if overall is not None else 0.0,
            data_type="NUMERIC",
        ))
        if breakdown.get("objects_f1") is not None:
            evals.append(Evaluation(name="objects_f1", value=float(breakdown["objects_f1"]), data_type="NUMERIC"))
    return evals

def run_aggregates(*, item_results, **kwargs):
    from collections import defaultdict
    vals = defaultdict(list)
    for ir in item_results:
        for ev in (getattr(ir, "evaluations", None) or []):
            name = ev.get("name") if isinstance(ev, dict) else getattr(ev, "name", None)
            value = ev.get("value") if isinstance(ev, dict) else getattr(ev, "value", None)
            if name is not None and isinstance(value, (int, float, bool)):
                vals[name].append(float(value))
    return [Evaluation(name=f"mean_{name}", value=sum(xs) / len(xs), data_type="NUMERIC")
            for name, xs in vals.items() if xs]

from functools import lru_cache
from langfuse.openai import openai
import time

PROMPT_NAME = "extraction-system"
TERSE_V1 = "Extract the insurance claim from the text into the structured schema."

@lru_cache(maxsize=16)
def get_system_prompt(version):
    return langfuse.get_prompt(PROMPT_NAME, version=version)

def seed_prompt(version, text, labels, commit_message):
    # Idempotent: prompt versions are immutable, so only create what doesn't exist yet.
    metas = langfuse.api.prompts.list(name=PROMPT_NAME).data
    if any(version in (m.versions or []) for m in metas):
        existing = langfuse.get_prompt(PROMPT_NAME, version=version)
        if existing.prompt.strip() != text.strip():
            raise RuntimeError(
                f"'{PROMPT_NAME}' v{version} already exists with different text. This project "
                f"has an unrelated prompt under that name; change PROMPT_NAME or use a fresh project.")
        print(f"prompt v{version} already exists")
        return
    langfuse.create_prompt(name=PROMPT_NAME, prompt=text, type="text",
                           labels=labels, commit_message=commit_message)
    get_system_prompt.cache_clear()
    print(f"created prompt v{version}")

# Seed v1 (terse baseline). Later we'll add v2/v3 when the results tell us to.
seed_prompt(1, TERSE_V1, ["baseline"], "v1: terse baseline")

def build_messages(system_text, claim_text):
    return [{"role": "system", "content": system_text},
            {"role": "user", "content": claim_text}]

def extract(cfg, claim_text, name):
    system_text = TERSE_V1 if cfg["prompt_version"] is None else get_system_prompt(cfg["prompt_version"]).compile()
    prompt_obj = None if cfg["prompt_version"] is None else get_system_prompt(cfg["prompt_version"])
    t0 = time.perf_counter()
    c = openai.chat.completions.parse(
        model=cfg["model"],
        messages=build_messages(system_text, claim_text),
        response_format=SCHEMAS[cfg["schema"]],
        name=name, langfuse_prompt=prompt_obj, temperature=0)
    parsed = c.choices[0].message.parsed
    return {"prediction": parsed.model_dump(mode="json") if parsed else None,
            "valid_json": parsed is not None, "latency_ms": (time.perf_counter() - t0) * 1000,
            "model": cfg["model"]}

import pandas as pd

def ev_dict(evs):
    return {getattr(e, "name", None): getattr(e, "value", None) for e in (evs or [])}

def _mean(rows, name):
    xs = [float(r[name]) for r in rows if isinstance(r.get(name), (int, float, bool))]
    return sum(xs) / len(xs) if xs else None

def _mean_field_suffix(rows, suffix):
    xs = []
    for r in rows:
        for name, value in r.items():
            if not (isinstance(name, str) and name.startswith("field.")):
                continue
            path = name[len("field."):]
            if path == suffix or path.endswith("." + suffix):
                if isinstance(value, (int, float, bool)):
                    xs.append(float(value))
    return sum(xs) / len(xs) if xs else None

def worst_fields(run, top=10):
    from collections import defaultdict
    hits = defaultdict(lambda: [0, 0])
    for ir in run.item_results:
        for ev in (getattr(ir, "evaluations", None) or []):
            name = ev.get("name") if isinstance(ev, dict) else getattr(ev, "name", None)
            value = ev.get("value") if isinstance(ev, dict) else getattr(ev, "value", None)
            if name and name.startswith("field.") and isinstance(value, (int, float, bool)):
                hits[name][1] += 1
                hits[name][0] += int(bool(value))
    ranked = sorted(hits.items(), key=lambda kv: (kv[1][0] / kv[1][1] if kv[1][1] else 1.0, kv[0]))
    print("lowest field accuracies:")
    for name, (ok, n) in ranked[:top]:
        print(f"  {name}: {ok}/{n} ({ok / n:.2f})")

results, summaries = {}, []

def _round(x):
    return None if x is None else round(x, 3)

def run_round(name, *, model="gpt-4o-mini", schema="descriptive", prompt_version=None):
    cfg = dict(model=model, schema=schema, prompt_version=prompt_version)
    def task(*, item, **kwargs):
        return extract(cfg, item.input, name)
    res = dataset.run_experiment(
        name=name, task=task, evaluators=[extraction_evaluator],
        run_evaluators=[run_aggregates], max_concurrency=8)
    langfuse.flush()
    results[name] = res
    rows = [ev_dict(ir.evaluations) for ir in res.item_results]
    summary = {
        "variant": name,
        "valid_json%": _round(_mean(rows, "valid_json")),
        "field_acc": _round(_mean(rows, "field_accuracy_overall")),
        "object_id": _round(_mean_field_suffix(rows, "object_id")),
        "policyholder": _round(_mean_field_suffix(rows, "policyholder_name")),
        "objects_F1": _round(_mean(rows, "objects_f1")),
    }
    for i, s in enumerate(summaries):
        if s["variant"] == name:
            summaries[i] = summary
            break
    else:
        summaries.append(summary)
    return comparison_table()

def comparison_table():
    cols = ["variant", "valid_json%", "field_acc", "object_id", "policyholder", "objects_F1"]
    return pd.DataFrame([{k: s[k] for k in cols} for s in summaries])

run_round("V1_baseline", model="gpt-4o-mini", schema="bare", prompt_version=1)

worst_fields(results["V1_baseline"])

run_round("V2_schema_spec", model="gpt-4o-mini", schema="descriptive", prompt_version=1)

worst_fields(results["V2_schema_spec"])

RULES_V2 = """\
Extract the insurance claim into the structured schema. Follow these rules:
- IDs: claim_id must match CLM-XXXXXX (6 digits); policy_number must match POL-XXXXXXXXX (9 digits).
  For insured objects use the format implied by object_type: VINxx... for vehicles, PROP-XXXXXX for
  buildings, LIAB-XXXXXX for liability, OBJ-XXXXXX otherwise.
- Dates: output ISO 8601 (YYYY-MM-DD).
- Enums: choose the single closest standardized value; never invent new enum values.
- Do NOT invent information. If policy details are not present in the text, omit policy_details.
  If no insured objects are described, omit insured_objects. Leave optional fields null when absent.
- Extract values verbatim where possible; do not normalize names beyond what the schema requires."""


seed_prompt(2, RULES_V2, ["production"], "v2: explicit rules")

run_round("V3_prompt_rules", model="gpt-4o-mini", schema="descriptive", prompt_version=2)

worst_fields(results["V3_prompt_rules"])

drill_into(results["V3_prompt_rules"], "V3")

RULES_V3 = RULES_V2 + """
- policyholder_name: copy the policyholder's full name exactly as written. NEVER output 'null', 'NA',
  '/', or any placeholder for a field that IS present.
- Insured objects: only emit an object explicitly described as insured or damaged. If none are, return
  an empty list — do NOT invent one to fill an id. When an object exists, its object_id uses the real
  value/format for its type; never copy the literal 'xxxx' placeholder."""


seed_prompt(3, RULES_V3, ["candidate"], "v3: copy name verbatim; stop re-inventing objects")

run_round("V3_fixed", model="gpt-4o-mini", schema="descriptive", prompt_version=3)

worst_fields(results["V3_fixed"])