RFP requirement-extraction pipeline — core extraction code. order of operation 1) segment.py 2) rules.py 3) classify.py

classify.py

"""Stage 2 — batched classification of DEFER units.

Units are sent as contiguous numbered blocks so each unit doubles as context
for its neighbours (no duplicated context tokens). The returned index set is
validated against the sent index set; a mismatch triggers a batch split and
retry rather than silently losing units.
"""
from __future__ import annotations

import re

from .providers import Provider, Usage

SYSTEM_PROMPT = """You classify individual units of text from a Request for Proposal (RFP).
Your ONLY job is to decide whether each unit states a REQUIREMENT the vendor
must respond to. You do not find or summarize requirements — each unit is
given to you already segmented.

DEFINITION
A requirement is anything the vendor could be scored on or held to. Test:
"If we ignored this, could our proposal be marked non-compliant?" If yes, it
is a requirement.

DECISION PROCEDURE — apply in order, stop at the first that decides:
1. WHO is obligated? If the unit obligates the ISSUER (e.g. "the Government
   will provide...") or no one (background, history, scope narrative), label
   not_requirement.
2. IS IT BINDING? shall/must/is required to -> mandatory.
   should/recommended/preferred -> desired.
   may/can/optional/at its discretion -> not_requirement (permission, not
   obligation).
3. IS IT ACTIONABLE? A specific capability, deliverable, standard, or
   constraint -> requirement. A bare definition or reference with no
   obligation -> not_requirement.
4. If none of the above clearly decides, label borderline. Do NOT default to
   mandatory when unsure.

RULES
- Judge each unit as written. Do not infer obligations that aren't stated.
- "will" usually describes the ISSUER's actions — check the subject before
  treating it as a requirement.
- A table row is a requirement ONLY if it states a performance obligation, a
  deliverable, or a threshold the vendor must meet. A row that merely names or
  numbers a contract line item, price element, or cost category is contract
  STRUCTURE, not a requirement -> not_requirement. Examples that are NOT
  requirements: "X006AA | Tool purchases in support of CLIN X006",
  "X001AB | Travel in support of CLIN X001", a CLIN number paired with a task
  title, a pricing/cost cell, or a labor-hour rate line.
- Lines are contiguous; use neighbours for context but judge only the
  numbered unit itself.

OUTPUT FORMAT — one line per unit, pipe-delimited, nothing else. No prose,
no JSON, no markdown fences:
INDEX|LABEL|WHO|TRIGGER

LABEL  = mandatory | desired | not_requirement | borderline
WHO    = vendor | issuer | none
TRIGGER = the exact word or phrase that decided it (or "none")

Return exactly one line for every index you were given."""

LINE_RE = re.compile(r"^\s*(\d+)\s*\|\s*([a-z_]+)\s*\|\s*([a-z]+)\s*\|\s*(.*?)\s*$", re.I)
VALID_LABELS = {"mandatory", "desired", "not_requirement", "borderline"}


def build_batch_prompt(batch: list[dict]) -> str:
    lines = ["Classify each numbered unit. Lines are contiguous.", ""]
    for b in batch:
        lines.append(f"[{b['idx']}] ({b['unit_type']} | {b['reason']}) {b['text']}")
    lines.append("")
    lines.append(f"Return exactly {len(batch)} lines, indices "
                 f"{batch[0]['idx']}-{batch[-1]['idx']}.")
    return "\n".join(lines)


def parse_response(text: str) -> dict[int, tuple[str, str, str]]:
    out: dict[int, tuple[str, str, str]] = {}
    for line in text.splitlines():
        line = line.strip().strip("`")
        if not line:
            continue
        m = LINE_RE.match(line)
        if not m:
            continue
        idx = int(m.group(1))
        label = m.group(2).lower()
        who = m.group(3).lower()
        trigger = m.group(4) or "none"
        if label not in VALID_LABELS:
            continue
        out[idx] = (label, who, trigger)
    return out


def classify_units(units: list[dict], provider: Provider, batch_size: int = 30,
                   cache: dict | None = None, max_depth: int = 2):
    """Classify DEFER units. Returns (results_by_unit_id, usage, stats)."""
    cache = cache if cache is not None else {}
    results: dict[str, dict] = {}
    total = Usage()
    stats = {"batches": 0, "retries": 0, "cache_hits": 0,
             "unparsed_units": 0, "validation_failures": 0, "errors": []}

    pending = []
    for u in units:
        ck = f"{provider.model_id}:{u['text_hash']}"
        if ck in cache:
            results[u["unit_id"]] = dict(cache[ck])
            stats["cache_hits"] += 1
        else:
            pending.append(u)

    for start in range(0, len(pending), batch_size):
        chunk = pending[start:start + batch_size]
        _do_batch(chunk, provider, results, total, stats, cache, depth=0,
                  max_depth=max_depth)

    return results, total, stats


def _do_batch(chunk, provider, results, total, stats, cache, depth, max_depth):
    batch = [{"idx": i + 1, "text": u["text"], "unit_type": u["unit_type"],
              "reason": u["reason"]} for i, u in enumerate(chunk)]
    prompt = build_batch_prompt(batch)

    try:
        text, usage = provider.complete(SYSTEM_PROMPT, prompt)
    except Exception as exc:  # noqa: BLE001
        u = Usage(errors=1)
        total += u
        stats["validation_failures"] += 1
        msg = f"{type(exc).__name__}: {exc}"
        if msg not in stats["errors"]:
            stats["errors"].append(msg)
            # surface it once, loudly — a swallowed 400 is undebuggable
            print(f"\n  [API ERROR] {msg}\n")
        for unit in chunk:
            results[unit["unit_id"]] = {"label": "borderline", "who": "none",
                                        "trigger": f"api_error:{type(exc).__name__}",
                                        "status": "error"}
        return

    # accumulate usage
    for k in ("input_tokens", "output_tokens", "cache_read_tokens",
              "cache_write_tokens", "calls", "errors"):
        setattr(total, k, getattr(total, k) + getattr(usage, k))
    total.latency_s = round(total.latency_s + usage.latency_s, 3)
    stats["batches"] += 1

    parsed = parse_response(text)
    sent = {b["idx"] for b in batch}
    got = set(parsed)

    # ---- hard validation: index sets must match exactly ----
    if got != sent:
        stats["validation_failures"] += 1
        if depth < max_depth and len(chunk) > 1:
            stats["retries"] += 1
            mid = len(chunk) // 2
            _do_batch(chunk[:mid], provider, results, total, stats, cache,
                      depth + 1, max_depth)
            _do_batch(chunk[mid:], provider, results, total, stats, cache,
                      depth + 1, max_depth)
            return

    for i, unit in enumerate(chunk, start=1):
        if i in parsed:
            label, who, trigger = parsed[i]
            rec = {"label": label, "who": who, "trigger": trigger, "status": "ok"}
        else:
            stats["unparsed_units"] += 1
            rec = {"label": "borderline", "who": "none", "trigger": "unparsed",
                   "status": "unparsed"}
        results[unit["unit_id"]] = rec
        cache[f"{provider.model_id}:{unit['text_hash']}"] = dict(rec)

rules.py

"""Stage 1 — rules router, and Stage 4 — reconciliation.

The router is NOT a classifier. It peels off cases it is *sure* about and
defers everything else to the model. Being unsure is a legal move: it just
means DEFER. Every decision carries a reason string for the audit trail.
"""
from __future__ import annotations

import re

MANDATORY = re.compile(
    r"\b(shall|must|is\s+required\s+to|are\s+required\s+to|will\s+be\s+required)\b", re.I)
SOFT = re.compile(
    r"\b(should|will|may|can|expected\s+to|responsible\s+for|recommended|preferred)\b", re.I)
IMPERATIVE_VERBS = {
    "provide", "submit", "ensure", "describe", "include", "demonstrate",
    "comply", "deliver", "maintain", "support", "perform", "develop",
    "prepare", "coordinate", "conduct", "establish", "implement", "identify",
    "report", "furnish", "attend",
}
PAGE_FURNITURE = re.compile(r"^(page\s+\d+|\d+\s+of\s+\d+|rev\.?\s*\d+)\b", re.I)

# Subject detection (used by both the router and reconciliation)
ISSUER_SUBJECT = re.compile(
    r"^(the\s+)?(government|cor|contracting officer|gsa|client)\b", re.I)
VENDOR_TOKEN = re.compile(r"\b(contractor|offeror|vendor|bidder)\b", re.I)

# External references genuinely point outside this document — always worth a
# flag. Case-sensitive on the keyword: a real reference is "Appendix F", while
# "exhibit a genuine concern" is prose, not a cross-reference.
EXTERNAL_XREF = re.compile(
    r"\b(Attachment|Appendix|Exhibit)\s+[A-Z0-9]{1,3}\b"
    r"|(?i:see the attached|incorporated by reference)")
# Internal section references ("Section F.6") only matter if the section is
# absent from the document — otherwise they're just navigation.
SECTION_REF = re.compile(r"\bsection\s+([A-Z]\.\d+(?:\.\d+)*)", re.I)

AUTO_YES, DEFER, AUTO_NO = "AUTO_YES", "DEFER", "AUTO_NO"


def route(unit) -> tuple[str, str]:
    """Return (route, reason). unit is a Unit or dict-like."""
    text = unit["text"] if isinstance(unit, dict) else unit.text
    utype = unit["unit_type"] if isinstance(unit, dict) else unit.unit_type
    words = text.split()

    # 1. structural noise
    if PAGE_FURNITURE.match(text):
        return AUTO_NO, "page_furniture"
    if utype == "sentence" and len(words) < 4:
        return AUTO_NO, "too_short"

    # 2. high-confidence requirement: binding modal (a legal term of art).
    #    BUT if the sentence's subject is the issuer ("The Government shall..."),
    #    the shall binds them, not the vendor — send it to the model to judge
    #    direction rather than auto-accepting it as a vendor requirement.
    m = MANDATORY.search(text)
    if m:
        if ISSUER_SUBJECT.match(text):
            return DEFER, f"issuer_subject_shall:{m.group().lower()}"
        return AUTO_YES, f"mandatory_modal:{m.group().lower()}"

    # 3. ambiguous middle -> the model decides
    if utype == "table_row":
        return DEFER, "table_row"
    s = SOFT.search(text)
    if s:
        return DEFER, f"soft_modal:{s.group().lower()}"
    if words and words[0].lower().strip(".,:;") in IMPERATIVE_VERBS:
        return DEFER, "imperative"

    # 4. no obligation signal at all
    return AUTO_NO, "no_obligation_signal"


def classify_xref(text: str, known_sections: set[str] | None = None) -> str | None:
    """Return 'external', 'unresolved_section', or None.

    An internal 'Section X.Y' reference whose target exists in the document is
    navigation, not a gap, so it returns None (no flag).
    """
    if EXTERNAL_XREF.search(text):
        return "external"
    m = SECTION_REF.search(text)
    if m:
        ref = m.group(1).upper()
        if known_sections is None:
            return None  # can't verify -> don't cry wolf
        # resolved if the referenced section (or a parent) exists in the doc
        if ref in known_sections or any(
                ref.startswith(s) or s.startswith(ref) for s in known_sections):
            return None
        return "unresolved_section"
    return None


def has_xref(text: str) -> bool:
    """Back-compat: any reference at all (used where sections aren't known)."""
    return bool(EXTERNAL_XREF.search(text) or SECTION_REF.search(text))


# Signals that CONTRADICT a mandatory label: permission / optionality. A real
# model returns a free-form trigger phrase, so we must NOT flag merely because
# the trigger isn't a modal verb — only when it actively contradicts the label.
_CONTRADICTS_MANDATORY = re.compile(
    r"\b(may|optional|at (its|the) discretion|should|recommended|preferred"
    r"|encouraged|if desired|as needed)\b", re.I)


def reconcile(record: dict, known_sections: set[str] | None = None) -> list[str]:
    """Return a list of human-review flags for one classified record."""
    flags: list[str] = []
    text = record["text"]
    label = record.get("label", "")
    head = " ".join(text.split()[:6])

    # model says requirement, but the subject is the issuer (safety net — the
    # router now sends most of these to the model, but catch any that slip by)
    if label in {"mandatory", "desired"} and ISSUER_SUBJECT.match(text) \
            and not VENDOR_TOKEN.search(head):
        flags.append("issuer_subject_but_labeled_requirement")

    # model says not-a-requirement, but binding modal present + vendor subject.
    # Ignore bare stub lines ("The Contractor shall:") that only introduce a list.
    if label == "not_requirement" and MANDATORY.search(text) \
            and VENDOR_TOKEN.search(text) and not text.rstrip().endswith(":"):
        flags.append("binding_modal_but_labeled_non_requirement")

    # model labeled it mandatory, but the trigger it cited grants permission or
    # is optional — a genuine contradiction worth a human's eye.
    trig = (record.get("trigger") or "").lower()
    if label == "mandatory" and _CONTRADICTS_MANDATORY.search(trig) \
            and not MANDATORY.search(trig):
        flags.append("trigger_contradicts_mandatory")

    xref = classify_xref(text, known_sections)
    if xref == "external":
        flags.append("unresolved_cross_reference")
    elif xref == "unresolved_section":
        flags.append("unresolved_section_reference")

    return flags

runner.py

"""Orchestrates a single pipeline run and writes an immutable run artifact."""
from __future__ import annotations

import dataclasses
import hashlib
import json
import platform
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

import yaml

from . import rules
from .classify import SYSTEM_PROMPT, classify_units
from .providers import Usage, build_provider
from .segment import segment_document, with_context

RUNS_DIR = Path(__file__).resolve().parent.parent / "runs"


def load_models(cfg_path: str | Path) -> dict:
    with open(cfg_path) as f:
        return yaml.safe_load(f)["models"]


def prompt_version() -> str:
    return hashlib.sha1(SYSTEM_PROMPT.encode()).hexdigest()[:8]


def run_once(doc_path, model_key, models_cfg, *, batch_size=30, context_window=0,
             repeat_idx=0, cache=None, cache_enabled=False, run_group=None,
             runs_dir=RUNS_DIR) -> dict:
    t_start = time.time()
    spec = models_cfg[model_key]
    if "params" in spec and "seed" in spec["params"]:
        spec = json.loads(json.dumps(spec))  # deep copy
        spec["params"]["seed"] = spec["params"]["seed"] + repeat_idx
    else:
        spec = json.loads(json.dumps(spec))
        spec.setdefault("params", {})["seed"] = repeat_idx

    # ---- Stage 0: segment (deterministic) ----
    doc_path = Path(doc_path)
    source_bytes = doc_path.read_bytes()
    units = [dataclasses.asdict(u) for u in segment_document(doc_path)]
    unit_objs = segment_document(doc_path)

    # ---- Stage 1: rules router ----
    counts = {"AUTO_YES": 0, "DEFER": 0, "AUTO_NO": 0}
    for i, u in enumerate(units):
        r, reason = rules.route(u)
        u["route"], u["reason"] = r, reason
        u["xref"] = rules.has_xref(u["text"])
        counts[r] += 1
        if context_window:
            b, a = with_context(unit_objs, i, context_window)
            u["context_before"], u["context_after"] = b, a

    # ---- Stage 2: classify DEFER ----
    provider = build_provider(spec)
    deferred = [u for u in units if u["route"] == "DEFER"]
    cls, usage, cstats = classify_units(deferred, provider,
                                        batch_size=batch_size, cache=cache)

    for u in units:
        if u["route"] == "AUTO_YES":
            u.update(label="mandatory", who="vendor", trigger=u["reason"],
                     status="rules")
        elif u["route"] == "AUTO_NO":
            u.update(label="not_requirement", who="none", trigger=u["reason"],
                     status="rules")
        else:
            u.update(cls.get(u["unit_id"],
                             {"label": "borderline", "who": "none",
                              "trigger": "missing", "status": "missing"}))

    # ---- Stage 4: reconcile ----
    # Build the set of section identifiers that actually exist in this document,
    # so internal "Section X.Y" references resolve instead of flagging.
    import re as _re
    known_sections = set()
    for u in units:
        m = _re.match(r"([A-Z]\.\d+(?:\.\d+)*)", u["section_path"])
        if m:
            known_sections.add(m.group(1).upper())
    for u in units:
        u["flags"] = rules.reconcile(u, known_sections)

    reqs = [u for u in units if u["label"] in ("mandatory", "desired")]
    pricing = spec.get("pricing", {})

    run_id = f"{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%S')}-{model_key}-r{repeat_idx}-{uuid.uuid4().hex[:6]}"
    record = {
        "run_id": run_id,
        "run_group": run_group,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "doc": str(doc_path),
        "doc_id": units[0]["doc_id"] if units else "empty",
        "source": {
            "extension": doc_path.suffix.lower(),
            "bytes": len(source_bytes),
            "sha256": hashlib.sha256(source_bytes).hexdigest(),
        },
        "model_key": model_key,
        "provider": spec["provider"],
        "model_id": spec["model_id"],
        "params": spec.get("params", {}),
        "batch_size": batch_size,
        "context_window": context_window,
        "cache_enabled": cache_enabled,
        "repeat_idx": repeat_idx,
        "prompt_version": prompt_version(),
        "python": platform.python_version(),
        "stage_counts": counts,
        "totals": {
            "units": len(units),
            "requirements": len(reqs),
            "mandatory": sum(u["label"] == "mandatory" for u in units),
            "desired": sum(u["label"] == "desired" for u in units),
            "borderline": sum(u["label"] == "borderline" for u in units),
            "not_requirement": sum(u["label"] == "not_requirement" for u in units),
            "flagged": sum(bool(u["flags"]) for u in units),
        },
        "usage": dataclasses.asdict(usage),
        "cost_usd": usage.cost(pricing),
        "classify_stats": cstats,
        "wall_time_s": round(time.time() - t_start, 2),
    }

    runs_dir = Path(runs_dir)
    runs_dir.mkdir(parents=True, exist_ok=True)
    out = runs_dir / run_id
    out.mkdir(exist_ok=True)
    with open(out / "units.jsonl", "w") as f:
        for u in units:
            f.write(json.dumps(u) + "\n")
    with open(out / "run.json", "w") as f:
        json.dump(record, f, indent=2)

    record["_units"] = units
    record["_dir"] = str(out)
    return record


def load_run(run_dir: str | Path) -> dict:
    run_dir = Path(run_dir)
    rec = json.loads((run_dir / "run.json").read_text())
    rec["_units"] = [json.loads(l) for l in open(run_dir / "units.jsonl")]
    rec["_dir"] = str(run_dir)
    return rec

segment.py

"""Stage 0 — deterministic document loading and segmentation.

No LLM here. Same bytes in => byte-identical units.jsonl out. This is the
determinism floor of the whole pipeline: unit count and unit IDs are fixed
by code, never by a model.
"""
from __future__ import annotations

import hashlib
import json
import re
from dataclasses import dataclass, asdict
from pathlib import Path

from .docio import load_document  # noqa: F401  (re-exported)


# --------------------------------------------------------------------------
# Segmentation
# --------------------------------------------------------------------------

# Recognizes common RFP heading formats, anchored at line start:
#   SECTION C ...            (federal)
#   C.1  /  H.10.4           (letter-dotted)
#   1.0  /  3.3.1            (number-dotted, incl. trailing .0)
#   S-1  /  N-1  /  Y-1a     (lettered question items)
#   ARTICLE 5 / PART 2       (occasional variants)
SECTION_HDR = re.compile(
    r"^(?:\*+)?\s*("
    r"SECTION\s+[A-Z0-9]\b[^*\n]*"          # SECTION C, SECTION 4
    r"|(?:ARTICLE|PART)\s+[A-Z0-9]+\b[^*\n]*"
    r"|[A-Z]\.\d+(?:\.\d+)*"                # C.1, H.10.4
    r"|\d+\.\d+(?:\.\d+)*"                  # 1.1, 3.3.1
    r"|\d+\.0*(?=\s|$)"                     # 1.0, 2.0 (section roots)
    r"|[A-Z]-\d+[a-z]?(?=[\s.:)]|$)"        # S-1, N-1, Y-1a
    r")", re.I)
BULLET = re.compile(r"^\s*[-•*]\s+")
SENT_SPLIT = re.compile(r"(?<=[.:;])\s+(?=[A-Z(])")
ABBREV = re.compile(r"\b([A-Z])\.\s")


@dataclass
class Unit:
    unit_id: str
    doc_id: str
    section_path: str
    unit_type: str          # sentence | table_row
    text: str
    text_hash: str
    ordinal: int


def _clean(s: str) -> str:
    s = s.replace("**", "").replace("\t", " ")
    return re.sub(r"\s+", " ", s).strip()


def _is_heading(raw: str) -> bool:
    t = _clean(raw)
    if not t:
        return True
    if SECTION_HDR.match(t) or t.startswith("#"):
        return True
    if t.startswith("(END OF") or t.startswith("NOTE:"):
        return True
    # ALL CAPS short line => title
    if re.match(r"^[A-Z][A-Z0-9 \-/&()',.]{5,}$", t) and len(t.split()) <= 9:
        return True
    return False


def _split_sentences(text: str) -> list[str]:
    protected = ABBREV.sub(r"\1<DOT> ", text)
    return [p.replace("<DOT>", ".").strip()
            for p in SENT_SPLIT.split(protected) if p.strip()]


def segment(text: str, doc_id: str, min_words: int = 3) -> list[Unit]:
    units: list[Unit] = []
    section_path = "Preamble"
    n = 0

    for raw in text.splitlines():
        line = raw.rstrip()
        cleaned = _clean(line)
        if not cleaned:
            continue

        if SECTION_HDR.match(cleaned) or cleaned.startswith("#"):
            section_path = cleaned.lstrip("# ").strip()
            continue

        # --- table rows: one row = one unit, headers preserved in text ---
        if line.lstrip().startswith("|"):
            cells = [_clean(c) for c in line.split("|")[1:-1]]
            cells = [c for c in cells if c]
            if not cells or all(set(c) <= {"-", ":"} for c in cells):
                continue
            n += 1
            body = " | ".join(cells)
            units.append(Unit(f"u{n:05d}", doc_id, section_path, "table_row",
                              body, _hash(body), n))
            continue

        if _is_heading(line):
            continue

        body = BULLET.sub("", cleaned)
        for sent in _split_sentences(body):
            if len(sent.split()) < min_words:
                continue
            n += 1
            units.append(Unit(f"u{n:05d}", doc_id, section_path, "sentence",
                              sent, _hash(sent), n))

    return units


def _hash(s: str) -> str:
    return hashlib.sha1(s.encode()).hexdigest()[:12]


def write_units(units: list[Unit], path: str | Path) -> None:
    with open(path, "w") as f:
        for u in units:
            f.write(json.dumps(asdict(u)) + "\n")


def segment_document(path: str | Path, doc_id: str | None = None) -> list[Unit]:
    path = Path(path)
    doc_id = doc_id or path.stem
    return segment(load_document(path), doc_id)


def with_context(units: list[Unit], i: int, window: int = 2) -> tuple[str, str]:
    """Neighbouring text for anaphora resolution ('this', 'the following')."""
    before = " ".join(u.text for u in units[max(0, i - window):i])
    after = " ".join(u.text for u in units[i + 1:i + 1 + window])
    return before, after