From f9f4629af8d7e462ebbd6de93c3e8649c260128c Mon Sep 17 00:00:00 2001 From: Ander Alvarez Sanz <104446704+aalvsz@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:30:59 +0200 Subject: [PATCH 1/2] Add ProvenanceGuard finance evaluation Signed-off-by: Ander Alvarez Sanz <104446704+aalvsz@users.noreply.github.com> --- nvflow/provenanceguard/README.md | 109 + nvflow/provenanceguard/__init__.py | 52 + nvflow/provenanceguard/decomposer.py | 116 ++ nvflow/provenanceguard/embedder.py | 104 + nvflow/provenanceguard/evaluator.py | 300 +++ nvflow/provenanceguard/nli.py | 160 ++ nvflow/provenanceguard/protected_values.py | 188 ++ nvflow/provenanceguard/protocols.py | 102 + nvflow/provenanceguard/router.py | 143 ++ nvflow/provenanceguard/types.py | 168 ++ .../finance/stages/rl/evaluate_provenance.py | 209 ++ .../finance/utils/rl/provenanceguard.py | 993 ++++++++++ .../recipes/finance/workflows/grpo/base.yaml | 41 + tests/test_provenanceguard.py | 1751 +++++++++++++++++ 14 files changed, 4436 insertions(+) create mode 100644 nvflow/provenanceguard/README.md create mode 100644 nvflow/provenanceguard/__init__.py create mode 100644 nvflow/provenanceguard/decomposer.py create mode 100644 nvflow/provenanceguard/embedder.py create mode 100644 nvflow/provenanceguard/evaluator.py create mode 100644 nvflow/provenanceguard/nli.py create mode 100644 nvflow/provenanceguard/protected_values.py create mode 100644 nvflow/provenanceguard/protocols.py create mode 100644 nvflow/provenanceguard/router.py create mode 100644 nvflow/provenanceguard/types.py create mode 100644 nvflow/recipes/finance/stages/rl/evaluate_provenance.py create mode 100644 nvflow/recipes/finance/utils/rl/provenanceguard.py create mode 100644 tests/test_provenanceguard.py diff --git a/nvflow/provenanceguard/README.md b/nvflow/provenanceguard/README.md new file mode 100644 index 0000000..b93a52a --- /dev/null +++ b/nvflow/provenanceguard/README.md @@ -0,0 +1,109 @@ +# ProvenanceGuard Open v1 + +Uncalibrated, not paper-faithful open approximation of the research +ProvenanceGuard system. Provides claim decomposition, embedding-centroid +source routing, DeBERTa NLI scoring, and protected-value checking to produce +per-rollout-row `allow` / `block` / `unavailable` verdicts. + +## Algorithm + +**routing-nli-v1** — pipeline: decompose, route, NLI, protected-value, decision. + +## Decision Semantics + +- **allow**: every claim is entailed AND every protected value found in evidence. +- **block**: any claim is contradiction, neutral, no_source, or protected_value_mismatch. +- **unavailable**: no evidence, no claims extracted, or any NLI/model error on any claim. If any claim has an error, the row is unavailable even alongside block-worthy verdicts. + +## Evidence Basis + +`retrieval_model_excerpt` — evidence comes from `retrieve_information` tool-call +outputs in the rollout trace, NOT from primary SEC filing verification. + +## Generic vs Finance Layers + +| Layer | Location | Contents | +|---|---|---| +| Generic | `nvflow/provenanceguard/` | Types, protocols, decomposer, embedder, router, NLI, evaluator, protected_values | +| Finance sidecar + CLI | `nvflow/recipes/finance/utils/rl/provenanceguard.py` | Trace extraction, deterministic/atomic sidecar evaluator, CLI | +| Finance stage | `nvflow/recipes/finance/stages/rl/evaluate_provenance.py` | Registered as `finance/grpo/evaluate_provenance` | + +The generic package contains no finance-specific code. Trace extraction, +sidecar evaluation, and CLI live in the finance recipe layer. + +## Enablement + +Uncomment `# - evaluate_provenance` in `pipeline_stages` in +`nvflow/recipes/finance/workflows/grpo/base.yaml`. The stage config +(`stages.evaluate_provenance`) is active by default; only the pipeline_stages +entry is commented out (opt-in). + +## Paths + +- **Input**: `${directories.step-5-collect-rollouts}/{env}/rollout/output-rs.jsonl` +- **Input marker** (required): `output-rs.jsonl.done` +- **Output**: `${directories.provenanceguard-eval}/{env}/provenanceguard-rs.jsonl` +- **Output marker** (created atomically after `os.replace`): `provenanceguard-rs.jsonl.done` +- **Directory**: `${model_output_dir}/provenanceguard-eval` + +Depends only on `collect_rollouts`; no downstream dependencies. + +## Sidecar Protocol + +- Exactly one sidecar row for every physical input line (including blank/malformed/non-object). +- No random UUID; `evaluation_uuid` is a deterministic SHA-256 of raw-line fingerprint + seed + line number + algorithm + config digest. Duplicate identical rows get distinguishable IDs by physical line position. +- Input gate first: both the input file and its sibling `.done` marker must exist before any output state is touched. A missing input gate preserves prior output and output `.done` unchanged. +- Stale output `.done` marker is cleared only after the input gate passes, before mkdir/temp/evaluation. Atomic temp write + `os.replace`; a fresh empty `.done` marker is created only after successful replace. If evaluation fails after the valid input gate, prior output remains intact but the stale `.done` marker stays absent. +- Input file is never modified. + +## Source ID Stability + +Stable canonical IDs (`sec:cik=<10-digit>:accession=<...>:doc=<...>`) are only +produced when a 10-digit zero-padded CIK + accession + document are all +present. Never falls back to storage keys or URLs for the canonical ID. +Multiple keys, any unknown key, or no canonical ID yields `attribution_state: +unavailable`; `source_ids` contains only known canonical candidates. + +## Protected Values + +Full normalized dates/numbers/currency/percent are required for a match. +A bare 4-digit year in the evidence does NOT satisfy a date protected value. + +## NLI Label Validation + +The NLI scorer validates actual model config label names +(entailment/neutral/contradiction, case-insensitive). It never assumes +arbitrary `LABEL_0` ordering — if a label does not match one of the three +canonical classes, it raises `ValueError`. + +## Model IDs and Licenses + +| Model | ID | License (from public model card) | +|---|---|---| +| Routing embedder | `sentence-transformers/all-MiniLM-L6-v2` | Apache-2.0 | +| NLI scorer | `MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli` | MIT | + +Both are lazy-loaded (torch/transformers imports happen inside `_ensure_loaded`, +not at module import time). Production deployments should pin revisions from +Hugging Face Hub metadata; revisions are optional and not hardcoded. + +## Differences from Private/Research Harness + +This open implementation differs from the private/research ProvenanceGuard +harness in several ways. No parity is claimed. + +- **Rule-based decomposition** rather than LLM/Gemma-based decomposition. +- **No token alignment/conflation head** — claims are not aligned to specific + evidence tokens. +- **No RF calibration** — NLI scores are used directly without random forest + calibration. +- Evidence is model-retrieved excerpts, not primary SEC verification. + +## Tests + +```bash +uv run --frozen --offline pytest tests/test_provenanceguard.py -q -o addopts="" +``` + +Tests use deterministic fakes (no model downloads, no network). No real model +execution or download is performed. diff --git a/nvflow/provenanceguard/__init__.py b/nvflow/provenanceguard/__init__.py new file mode 100644 index 0000000..2a2dee1 --- /dev/null +++ b/nvflow/provenanceguard/__init__.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""ProvenanceGuard Open v1 — routing-nli-v1. + +An uncalibrated, not paper-faithful open approximation of the research +ProvenanceGuard system. Provides claim decomposition, embedding-centroid +source routing, DeBERTa NLI scoring, and protected-value checking to +produce per-rollout-row allow / block / unavailable verdicts. + +All Hugging Face model loading is strictly lazy — no model download or +real model execution is performed at import time. +""" + +from nvflow.provenanceguard.protocols import ( + ClaimDecomposer, + Embedder, + NLIScorer, + RoutedClaim, + SourceRouter, +) +from nvflow.provenanceguard.types import ( + AtomicClaim, + ClaimVerdict, + Decision, + EvidenceChunk, + NLIResult, +) + +__all__ = [ + "AtomicClaim", + "ClaimDecomposer", + "ClaimVerdict", + "Decision", + "Embedder", + "EvidenceChunk", + "NLIScorer", + "NLIResult", + "RoutedClaim", + "SourceRouter", +] diff --git a/nvflow/provenanceguard/decomposer.py b/nvflow/provenanceguard/decomposer.py new file mode 100644 index 0000000..9d81012 --- /dev/null +++ b/nvflow/provenanceguard/decomposer.py @@ -0,0 +1,116 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Deterministic rule-based claim decomposer (v1 default). + +**Limitation**: This is a conservative sentence/claim splitter that does +not call any LLM. It splits on sentence boundaries and then further +splits compound claims on conjunctions (``and``, ``;``, ``--``). It may +over-merge compound claims or split mid-clause. The research +ProvenanceGuard uses Gemma-based decomposition, which is not included +in this open implementation. + +Claim IDs are deterministic SHA-256 hashes of the claim text and the +source sentence, so re-evaluation of the same answer yields stable IDs. +""" + +from __future__ import annotations + +import hashlib +import re +from collections.abc import Sequence + +from nvflow.provenanceguard.types import AtomicClaim + +# SEC filing URL pattern — used to extract stated SEC IDs/URLs from claim text. +_SEC_URL_RE = re.compile( + r"https?://(?:www\.)?sec\.gov/[^\s\"'<>]+", + re.IGNORECASE, +) + +# CIK pattern — 10-digit zero-padded numbers. +_CIK_RE = re.compile(r"\b(\d{10})\b") + +# Accession number pattern — e.g. 0001811414-25-000010 +_ACCESSION_RE = re.compile(r"\b(\d{10}-\d{2}-\d{6})\b") + +# Sentence boundary: period, exclamation, question mark followed by space +# or end of string. Also splits on newlines. +_SENTENCE_END_RE = re.compile(r"(?<=[.!?])\s+|\n+") + +# Compound claim splitters within a sentence. +_COMPOUND_SPLIT_RE = re.compile(r"\s*;\s*|\s+--\s+|\s+but also\s+", re.IGNORECASE) + + +def _stable_claim_id(text: str, source_sentence: str) -> str: + """Deterministic 16-char hex hash for a claim.""" + raw = f"{text}|{source_sentence}" + return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16] + + +def _extract_stated_sec_ids(text: str) -> tuple[tuple[str, ...], tuple[str, ...]]: + """Extract SEC CIKs and accession numbers mentioned in claim text.""" + ids: list[str] = [] + for m in _CIK_RE.finditer(text): + ids.append(f"cik:{m.group(1)}") + for m in _ACCESSION_RE.finditer(text): + ids.append(f"accession:{m.group(1)}") + urls = tuple(m.group(0) for m in _SEC_URL_RE.finditer(text)) + return tuple(ids), urls + + +class RuleBasedDecomposer: + """Deterministic sentence/claim splitter. + + Splits the answer into sentences, then further splits compound + claims on semicolons, em-dashes, and ``"but also"`` connectors. + Filters out empty fragments and very short noise (< 3 words). + """ + + def decompose(self, answer: str) -> Sequence[AtomicClaim]: + if not answer or not answer.strip(): + return [] + + # Normalize whitespace but preserve sentence structure. + text = answer.strip() + + # Split into sentences. + sentences = [s.strip() for s in _SENTENCE_END_RE.split(text) if s.strip()] + if not sentences: + sentences = [text] + + claims: list[AtomicClaim] = [] + for sentence in sentences: + # Further split compound claims. + fragments = _COMPOUND_SPLIT_RE.split(sentence) + for frag in fragments: + frag = frag.strip() + if not frag: + continue + # Skip very short fragments (noise). + if len(frag.split()) < 3: + continue + sec_ids, sec_urls = _extract_stated_sec_ids(frag) + claim_id = _stable_claim_id(frag, sentence) + claims.append( + AtomicClaim( + claim_id=claim_id, + text=frag, + source_sentence=sentence, + stated_sec_ids=sec_ids, + stated_sec_urls=sec_urls, + ) + ) + + return claims diff --git a/nvflow/provenanceguard/embedder.py b/nvflow/provenanceguard/embedder.py new file mode 100644 index 0000000..74a4a84 --- /dev/null +++ b/nvflow/provenanceguard/embedder.py @@ -0,0 +1,104 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Lazy Hugging Face embedder for claim/evidence routing. + +Default public model: ``sentence-transformers/all-MiniLM-L6-v2`` +(Apache 2.0). Model construction is strictly lazy — ``torch`` and +``transformers`` imports happen inside ``_ensure_loaded``, not at +module import time. + +Uses the raw Transformers ``AutoTokenizer`` / ``AutoModel`` recipe with +mean pooling over ``last_hidden_state`` with attention mask, L2-normalized, +returning plain float lists — so no ``sentence-transformers`` runtime +dependency is needed. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +DEFAULT_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2" + + +class HFEmbedder: + """Transformers embedder with lazy model loading and mean pooling. + + Args: + model_id: Hugging Face model ID. Defaults to + ``sentence-transformers/all-MiniLM-L6-v2``. + revision: Git revision to pin. ``None`` uses the latest + available revision. Production deployments should pin a + specific revision fetched from Hugging Face Hub metadata + (do not hardcode unverified hashes). + """ + + def __init__( + self, + model_id: str = DEFAULT_EMBEDDING_MODEL, + revision: str | None = None, + ) -> None: + self._model_id = model_id + self._revision = revision + self._model = None + self._tokenizer = None + self._dim: int | None = None + + def _ensure_loaded(self) -> None: + if self._model is not None: + return + from transformers import AutoModel, AutoTokenizer + + kwargs: dict = {} + if self._revision: + kwargs["revision"] = self._revision + + self._tokenizer = AutoTokenizer.from_pretrained(self._model_id, **kwargs) + self._model = AutoModel.from_pretrained(self._model_id, **kwargs) + self._model.eval() + self._dim = self._model.config.hidden_size + + @property + def dimension(self) -> int: + self._ensure_loaded() + assert self._dim is not None + return self._dim + + def embed(self, texts: Sequence[str]) -> Sequence[Sequence[float]]: + import torch + + self._ensure_loaded() + assert self._tokenizer is not None + assert self._model is not None + + encoded = self._tokenizer( + list(texts), + padding=True, + truncation=True, + max_length=512, + return_tensors="pt", + ) + with torch.no_grad(): + outputs = self._model(**encoded) + token_embeddings = outputs.last_hidden_state + attention_mask = encoded["attention_mask"] + mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()) + sum_embeddings = torch.sum(token_embeddings * mask_expanded, dim=1) + sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9) + sentence_embeddings = sum_embeddings / sum_mask + sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1) + return [[float(x) for x in row] for row in sentence_embeddings.detach().cpu().tolist()] + + +__all__ = ["DEFAULT_EMBEDDING_MODEL", "HFEmbedder"] diff --git a/nvflow/provenanceguard/evaluator.py b/nvflow/provenanceguard/evaluator.py new file mode 100644 index 0000000..13ca3c0 --- /dev/null +++ b/nvflow/provenanceguard/evaluator.py @@ -0,0 +1,300 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Core evaluator for ProvenanceGuard Open v1 (routing-nli-v1). + +Pipeline: decompose → route → NLI → protected-value check → decision. + +Decision strictness (conservative): + +- **allow**: every claim is entailed AND every protected value is + found in routed evidence. +- **block**: any claim is contradiction, neutral, no_source, or + protected_value_mismatch. +- **unavailable**: no evidence chunks were provided, no claims were + extracted from the answer, or any NLI/model error prevented at + least one claim from being fully evaluated. If any claim has an + error, the row is unavailable even if another claim is + contradiction or neutral. + +This is an **uncalibrated open approximation** — not paper-faithful. +Evidence is retrieval_model_excerpt only, never primary SEC verification. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from nvflow.provenanceguard.protected_values import check_protected_values +from nvflow.provenanceguard.protocols import ( + ClaimDecomposer, + NLIScorer, + SourceRouter, +) +from nvflow.provenanceguard.types import ( + AtomicClaim, + ClaimVerdict, + Decision, + EvidenceChunk, +) + +ALGORITHM_VERSION = "routing-nli-v1" + + +@dataclass(frozen=True) +class ProvenanceGuardConfig: + """Configuration for the ProvenanceGuard evaluator. + + The fail-closed policy is **fixed and non-configurable**: + contradiction, neutral, no_source, and protected_value_mismatch + always block; any model/trace error yields unavailable. + There are no ``block_on_*`` toggles — the policy cannot be + disabled. + + Attributes: + evidence_excerpt_length: Maximum characters of evidence text + to include in each claim verdict for auditability. + routing_model: Model ID used for embedding/routing (metadata). + nli_model: Model ID used for NLI scoring (metadata). + """ + + evidence_excerpt_length: int = 500 + routing_model: str = "sentence-transformers/all-MiniLM-L6-v2" + nli_model: str = "MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli" + + def to_dict(self) -> dict[str, Any]: + return { + "algorithm": ALGORITHM_VERSION, + "policy": "fixed_fail_closed", + "evidence_excerpt_length": self.evidence_excerpt_length, + "routing_model": self.routing_model, + "nli_model": self.nli_model, + } + + +class ProvenanceGuardEvaluator: + """Orchestrates claim decomposition, routing, NLI, and protected-value checking. + + All model-bearing collaborators (``Embedder``, ``NLIScorer``) are + injected so that tests can use deterministic fakes. No Hugging Face + imports happen during construction or evaluation of this class — + that is the responsibility of ``HFEmbedder`` / ``HFNLI``, which load + lazily. + """ + + def __init__( + self, + *, + decomposer: ClaimDecomposer, + router: SourceRouter, + nli_scorer: NLIScorer, + config: ProvenanceGuardConfig | None = None, + ) -> None: + self._decomposer = decomposer + self._router = router + self._nli = nli_scorer + self._config = config or ProvenanceGuardConfig() + + @property + def config(self) -> ProvenanceGuardConfig: + return self._config + + def evaluate( + self, + answer: str, + evidence: Sequence[EvidenceChunk], + ) -> Decision: + """Evaluate an answer against evidence and return a Decision. + + Args: + answer: The final answer text submitted by the agent. + evidence: Evidence chunks extracted from the tool-call trace. + + Returns: + A :class:`~nvflow.provenanceguard.types.Decision` with status + ``allow``, ``block``, or ``unavailable``. + """ + errors: list[str] = [] + + if not evidence: + return Decision( + status="unavailable", + reason="no_evidence", + errors=("No evidence chunks were provided.",), + ) + + claims = self._decomposer.decompose(answer) + if not claims: + return Decision( + status="unavailable", + reason="no_claims_extracted", + errors=(), + ) + + try: + routed = self._router.route(claims, evidence) + except Exception as exc: + return Decision( + status="unavailable", + reason="routing_error", + errors=(f"Routing failed: {exc!s}",), + ) + + routed_by_claim_id: dict[str, Any] = {} + for rc in routed: + routed_by_claim_id[rc.claim.claim_id] = rc + + verdicts: list[ClaimVerdict] = [] + nli_errors = 0 + + for claim in claims: + verdict = self._evaluate_claim(claim, routed_by_claim_id, errors) + verdicts.append(verdict) + if verdict.errors: + nli_errors += 1 + + if nli_errors > 0: + reason = "all_claims_errored" if nli_errors == len(claims) else "partial_nli_errors" + return Decision( + status="unavailable", + reason=reason, + verdicts=tuple(verdicts), + errors=tuple(errors), + ) + + status, reason = self._aggregate(verdicts) + return Decision( + status=status, + reason=reason, + verdicts=tuple(verdicts), + errors=tuple(errors), + ) + + def _evaluate_claim( + self, + claim: AtomicClaim, + routed_by_claim_id: dict[str, Any], + errors: list[str], + ) -> ClaimVerdict: + """Evaluate a single claim and return its verdict.""" + routed = routed_by_claim_id.get(claim.claim_id) + + if routed is None: + return ClaimVerdict( + claim_id=claim.claim_id, + claim_text=claim.text, + final_label="no_source", + protected_value_outcome="not_applicable", + errors=(), + ) + + chunk = routed.chunk + excerpt = chunk.text[: self._config.evidence_excerpt_length] + + if chunk.attribution_state in ("unavailable", "composite", "unknown"): + return ClaimVerdict( + claim_id=claim.claim_id, + claim_text=claim.text, + routed_source_id=chunk.source_id, + routed_source_ids=chunk.source_ids, + routed_attribution_state=chunk.attribution_state, + routed_chunk_id=chunk.chunk_id, + routing_score=routed.score, + routing_margin=routed.margin, + final_label="no_source", + protected_value_outcome="not_applicable", + evidence_excerpt=excerpt, + errors=(), + ) + + try: + nli_result = self._nli.score(premise=chunk.text, hypothesis=claim.text) + except Exception as exc: + err = f"NLI error for claim {claim.claim_id}: {exc!s}" + errors.append(err) + return ClaimVerdict( + claim_id=claim.claim_id, + claim_text=claim.text, + routed_source_id=chunk.source_id, + routed_source_ids=chunk.source_ids, + routed_attribution_state=chunk.attribution_state, + routed_chunk_id=chunk.chunk_id, + routing_score=routed.score, + routing_margin=routed.margin, + raw_nli_label="neutral", + raw_nli_probabilities=(), + final_label="neutral", + protected_value_outcome="not_applicable", + evidence_excerpt=excerpt, + errors=(err,), + ) + + raw_label = nli_result.label + pv_outcome, _missing = check_protected_values(claim.text, chunk.text) + + final_label = self._apply_policy(raw_label, pv_outcome) + + return ClaimVerdict( + claim_id=claim.claim_id, + claim_text=claim.text, + routed_source_id=chunk.source_id, + routed_source_ids=chunk.source_ids, + routed_attribution_state=chunk.attribution_state, + routed_chunk_id=chunk.chunk_id, + routing_score=routed.score, + routing_margin=routed.margin, + raw_nli_label=raw_label, + raw_nli_probabilities=nli_result.probabilities, + final_label=final_label, + protected_value_outcome=pv_outcome, + evidence_excerpt=excerpt, + errors=(), + ) + + def _apply_policy(self, raw_nli_label: str, pv_outcome: str) -> str: + """Apply fixed fail-closed policy to the raw NLI label.""" + if raw_nli_label == "entailment": + if pv_outcome == "fail": + return "protected_value_mismatch" + return "entailment" + if raw_nli_label == "contradiction": + return "contradiction" + return "neutral" + + def _aggregate(self, verdicts: list[ClaimVerdict]) -> tuple[str, str]: + """Aggregate per-claim verdicts into a top-level decision (fixed fail-closed).""" + labels = [v.final_label for v in verdicts if not v.errors] + if not labels: + return "unavailable", "all_claims_errored" + + if any(lbl == "contradiction" for lbl in labels): + return "block", "contradiction" + if any(lbl == "protected_value_mismatch" for lbl in labels): + return "block", "protected_value_mismatch" + if any(lbl == "neutral" for lbl in labels): + return "block", "neutral" + if any(lbl == "no_source" for lbl in labels): + return "block", "no_source" + if all(lbl == "entailment" for lbl in labels): + return "allow", "all_entailed" + return "block", "unverifiable" + + +__all__ = [ + "ALGORITHM_VERSION", + "ProvenanceGuardConfig", + "ProvenanceGuardEvaluator", +] diff --git a/nvflow/provenanceguard/nli.py b/nvflow/provenanceguard/nli.py new file mode 100644 index 0000000..05391b7 --- /dev/null +++ b/nvflow/provenanceguard/nli.py @@ -0,0 +1,160 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Lazy Hugging Face NLI scorer for claim-evidence entailment checking. + +Default public model: ``MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli`` (MIT). + +Validates the exact unique label set (entailment, neutral, contradiction) +from the model config — fails on duplicates, extras, or missing labels. +Model construction is strictly lazy — the ``transformers`` import happens +inside ``__init__``, not at module import time. + +No lexical/hash runtime fallback. If model loading fails, the caller +must catch the exception and produce an ``unavailable`` decision. +""" + +from __future__ import annotations + +from nvflow.provenanceguard.types import NLIResult + +DEFAULT_NLI_MODEL = "MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli" + +_VALID_NLI_LABELS = frozenset({"entailment", "neutral", "contradiction"}) + + +def _normalize_nli_label(raw_label: str) -> str: + """Normalize a model config label to one of the three canonical NLI classes. + + Validates the *actual* label name from the model config + (``config.id2label``). Never assumes arbitrary ``LABEL_0`` ordering. + + Raises ``ValueError`` if the label does not match entailment, + neutral, or contradiction (case-insensitive). + """ + normalized = raw_label.strip().lower() + if normalized in _VALID_NLI_LABELS: + return normalized + raise ValueError( + f"Unknown NLI label '{raw_label}': expected one of " + f"entailment, neutral, contradiction (case-insensitive). " + "Refusing to assume arbitrary label ordering." + ) + + +class HFNLI: + """DeBERTa sequence-classification NLI scorer with lazy model loading. + + Args: + model_id: Hugging Face model ID. + revision: Git revision to pin. ``None`` uses the latest + available revision. Production deployments should pin a + specific revision fetched from Hugging Face Hub metadata + (do not hardcode unverified hashes). + """ + + def __init__( + self, + model_id: str = DEFAULT_NLI_MODEL, + revision: str | None = None, + ) -> None: + self._model_id = model_id + self._revision = revision + self._model = None + self._tokenizer = None + self._label_map: dict[int, str] | None = None + + def _ensure_loaded(self) -> None: + if self._model is not None: + return + from transformers import AutoModelForSequenceClassification, AutoTokenizer + + kwargs: dict = {} + if self._revision: + kwargs["revision"] = self._revision + + self._tokenizer = AutoTokenizer.from_pretrained(self._model_id, **kwargs) + self._model = AutoModelForSequenceClassification.from_pretrained(self._model_id, **kwargs) + self._model.eval() + + # Build label map from model config — validate actual label names, + # never assume arbitrary LABEL_0 ordering. + id2label = self._model.config.id2label + self._label_map = {} + for idx, label in id2label.items(): + normalized = _normalize_nli_label(str(label)) + self._label_map[int(idx)] = normalized + # Validate exact unique label set: exactly 3 labels, one each + # of entailment / neutral / contradiction, no duplicates, + # extras, or missing. + seen_labels: list[str] = [] + for idx in sorted(self._label_map): + seen_labels.append(self._label_map[idx]) + if len(seen_labels) != 3: + raise ValueError( + f"NLI model has {len(seen_labels)} labels; expected " + f"exactly 3 (entailment, neutral, contradiction)." + ) + label_set = set(seen_labels) + if label_set != _VALID_NLI_LABELS: + missing = _VALID_NLI_LABELS - label_set + extra = label_set - _VALID_NLI_LABELS + parts: list[str] = [] + if missing: + parts.append(f"missing: {sorted(missing)}") + if extra: + parts.append(f"extra: {sorted(extra)}") + if len(label_set) < len(seen_labels): + parts.append("duplicate labels detected") + raise ValueError("NLI label set validation failed: " + "; ".join(parts)) + + def score(self, *, premise: str, hypothesis: str) -> NLIResult: + import torch + + self._ensure_loaded() + assert self._tokenizer is not None + assert self._model is not None + assert self._label_map is not None + + # NLI convention: premise is the evidence (premise), hypothesis + # is the claim to verify. + encoded = self._tokenizer( + premise, + hypothesis, + return_tensors="pt", + truncation=True, + max_length=512, + padding=True, + ) + with torch.no_grad(): + logits = self._model(**encoded).logits + probs = torch.softmax(logits, dim=-1)[0] + + probabilities: list[tuple[str, float]] = [] + for idx in range(len(probs)): + label = self._label_map[int(idx)] + probabilities.append((label, float(probs[idx]))) + + best_idx = int(probs.argmax().item()) + best_label = self._label_map[best_idx] + best_score = float(probs[best_idx].item()) + + return NLIResult( + label=best_label, + score=best_score, + probabilities=tuple(probabilities), + ) + + +__all__ = ["DEFAULT_NLI_MODEL", "HFNLI", "_normalize_nli_label"] diff --git a/nvflow/provenanceguard/protected_values.py b/nvflow/provenanceguard/protected_values.py new file mode 100644 index 0000000..7997ef2 --- /dev/null +++ b/nvflow/provenanceguard/protected_values.py @@ -0,0 +1,188 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Protected value extraction and checking for ProvenanceGuard Open v1. + +Protected values are numeric amounts, currency figures, percentages, +and dates that appear in an entailed claim. If a protected value is +absent from the routed evidence text, the claim cannot be ``allow``ed +— the model may have fabricated the specific number even if the general +claim is entailed. + +This is a conservative string-matching check: it normalizes formats +(billions/millions suffixes, percentage signs, date separators) and +checks for presence in the evidence text. It does NOT parse financial +semantics. + +Full normalized dates/numbers/currency/percent are required for a match. +A bare 4-digit year appearing in the evidence does NOT satisfy a date +protected value — the full normalized date string must be found. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ProtectedValue: + """A protected numeric/date/percentage value extracted from a claim.""" + + raw: str # original text as it appeared in the claim + normalized: str # normalized form for matching + kind: str # "currency" | "percentage" | "date" | "number" + + +# --- Regex patterns ---------------------------------------------------------- + +# Currency: $1.23 billion, $1,234,567, $1.23M, €100, £50 million, etc. +_CURRENCY_RE = re.compile( + r"[$€£¥]\s?\d[\d,]*(?:\.\d+)?\s*(?:billion|million|thousand|trillion|[bmk])?\b", + re.IGNORECASE, +) + +# Percentage: 12.3%, 5 percent, etc. +_PERCENTAGE_RE = re.compile( + r"\d[\d,]*(?:\.\d+)?\s*(?:%|percent\b)", + re.IGNORECASE, +) + +# Date: 2024-01-15, January 15, 2024, Jan 15 2024, Q1 2024, FY2024, 2024-01, etc. +_DATE_RE = re.compile( + r"\b(?:" + r"\d{4}-\d{2}-\d{2}" # 2024-01-15 + r"|\d{4}-\d{2}" # 2024-01 + r"|\d{4}" # 2024 + r"|(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\w*\s+\d{1,2},?\s*\d{4}" # January 15, 2024 + r"|Q[1-4]\s+\d{4}" # Q1 2024 + r"|FY\s*\d{4}" # FY2024 + r"|(?:first|second|third|fourth)\s+quarter\s+\d{4}" # first quarter 2024 + r")\b", + re.IGNORECASE, +) + +# Plain number with magnitude: 1.23 billion, 1,234,567, 12.3M, etc. +# Only matched if it looks like a financial figure (has commas, decimals, +# or magnitude words). Single small integers are not protected. +_NUMBER_RE = re.compile( + r"\b\d{1,3}(?:,\d{3})+(?:\.\d+)?\b" # 1,234,567 + r"|\b\d+(?:\.\d+)?\s*(?:billion|million|thousand|trillion|[bmk])\b", # 1.23 billion + re.IGNORECASE, +) + + +def _normalize_currency(raw: str) -> str: + """Normalize currency for matching: lowercase, remove spaces.""" + return re.sub(r"\s+", "", raw.lower()) + + +def _normalize_percentage(raw: str) -> str: + """Normalize percentage: '12.3 %' -> '12.3%'.""" + return re.sub(r"\s+", "", raw.lower()) + + +def _normalize_date(raw: str) -> str: + """Normalize date: lowercase, collapse spaces.""" + return re.sub(r"\s+", " ", raw.lower()).strip() + + +def _normalize_number(raw: str) -> str: + """Normalize number: lowercase, remove spaces.""" + return re.sub(r"\s+", "", raw.lower()) + + +def extract_protected_values(text: str) -> list[ProtectedValue]: + """Extract all protected values from a claim text. + + Returns a list of :class:`ProtectedValue` with normalized forms. + Order: currency, percentage, date, number (deduplicated by raw text). + """ + if not text: + return [] + + results: list[ProtectedValue] = [] + seen_raw: set[str] = set() + + for match in _CURRENCY_RE.finditer(text): + raw = match.group(0) + if raw not in seen_raw: + results.append( + ProtectedValue(raw=raw, normalized=_normalize_currency(raw), kind="currency") + ) + seen_raw.add(raw) + + for match in _PERCENTAGE_RE.finditer(text): + raw = match.group(0) + if raw not in seen_raw: + results.append( + ProtectedValue(raw=raw, normalized=_normalize_percentage(raw), kind="percentage") + ) + seen_raw.add(raw) + + for match in _DATE_RE.finditer(text): + raw = match.group(0) + if raw not in seen_raw: + results.append(ProtectedValue(raw=raw, normalized=_normalize_date(raw), kind="date")) + seen_raw.add(raw) + + for match in _NUMBER_RE.finditer(text): + raw = match.group(0) + if raw not in seen_raw: + results.append( + ProtectedValue(raw=raw, normalized=_normalize_number(raw), kind="number") + ) + seen_raw.add(raw) + + return results + + +def check_protected_values( + claim_text: str, + evidence_text: str, +) -> tuple[str, list[ProtectedValue]]: + """Check if all protected values in *claim_text* appear in *evidence_text*. + + Returns ``(outcome, missing_values)``: + - ``("pass", [])`` — all protected values found in evidence. + - ``("fail", [...])`` — some protected values missing from evidence. + - ``("not_applicable", [])`` — no protected values in the claim. + """ + protected = extract_protected_values(claim_text) + if not protected: + return "not_applicable", [] + + if not evidence_text: + return "fail", protected + + evidence_lower = evidence_text.lower() + evidence_compact = re.sub(r"\s+", "", evidence_lower) + + missing: list[ProtectedValue] = [] + for pv in protected: + # Check both normalized and compact forms for robustness. + if pv.normalized in evidence_lower or pv.normalized in evidence_compact: + continue + missing.append(pv) + + if missing: + return "fail", missing + return "pass", [] + + +__all__ = [ + "ProtectedValue", + "extract_protected_values", + "check_protected_values", +] diff --git a/nvflow/provenanceguard/protocols.py b/nvflow/provenanceguard/protocols.py new file mode 100644 index 0000000..ee0ba02 --- /dev/null +++ b/nvflow/provenanceguard/protocols.py @@ -0,0 +1,102 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Protocol interfaces for ProvenanceGuard Open v1. + +These protocols enable dependency injection so that unit tests can use +deterministic fake decomposer / embedder / NLI implementations with +``HF_HUB_OFFLINE=1`` and ``TRANSFORMERS_OFFLINE=1``. + +No model download or real model execution is performed during import. +All concrete implementations that touch Hugging Face libraries do so +lazily — only when their methods are called. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Protocol, runtime_checkable + +from nvflow.provenanceguard.types import AtomicClaim, EvidenceChunk, NLIResult + + +@runtime_checkable +class ClaimDecomposer(Protocol): + """Split an assistant answer into atomic claims. + + The v1 default is a deterministic conservative rule-based + sentence/claim splitter. It does NOT call a private LLM. + + Limitation: rule-based splitting may over-merge compound claims or + split mid-clause. The research ProvenanceGuard uses Gemma-based + decomposition, which is not included in this open implementation. + """ + + def decompose(self, answer: str) -> Sequence[AtomicClaim]: ... + + +@runtime_checkable +class Embedder(Protocol): + """Embed text into fixed-dimensional vectors for routing. + + Default public model: ``sentence-transformers/all-MiniLM-L6-v2``. + No lexical/hash runtime fallback — if model loading fails, the + evaluator returns ``unavailable``. + """ + + @property + def dimension(self) -> int: ... + + def embed(self, texts: Sequence[str]) -> Sequence[Sequence[float]]: ... + + +@runtime_checkable +class NLIScorer(Protocol): + """Score a (premise, hypothesis) pair with sequence-classification NLI. + + Default public model: + ``MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli``. + + Returns entailment / neutral / contradiction with probabilities. + No lexical/hash runtime fallback. + """ + + def score(self, *, premise: str, hypothesis: str) -> NLIResult: ... + + +@dataclass(frozen=True) +class RoutedClaim: + """A claim paired with its best evidence chunk and routing scores.""" + + claim: AtomicClaim + chunk: EvidenceChunk + score: float + margin: float + + +@runtime_checkable +class SourceRouter(Protocol): + """Route each claim to its best-matching evidence chunk. + + Uses embedding cosine similarity (centroids per source) to find + the top-1 evidence source, recording the margin to top-2 for + ambiguity detection. + """ + + def route( + self, + claims: Sequence[AtomicClaim], + evidence: Sequence[EvidenceChunk], + ) -> Sequence[RoutedClaim]: ... diff --git a/nvflow/provenanceguard/router.py b/nvflow/provenanceguard/router.py new file mode 100644 index 0000000..3ec3af6 --- /dev/null +++ b/nvflow/provenanceguard/router.py @@ -0,0 +1,143 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Embedding-centroid source router for ProvenanceGuard Open v1. + +For each unique source (identified by ``source_id`` or ``chunk_id``), +computes the mean embedding of its chunks (a centroid). Each atomic +claim is routed to the source whose centroid has the highest cosine +similarity. The margin between top-1 and top-2 is recorded so the +pipeline can flag ambiguous routing. + +Uses the injectable :class:`~nvflow.provenanceguard.protocols.Embedder` +protocol. No lexical/hash runtime fallback — if the embedder fails, +the evaluator must produce an ``unavailable`` decision. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence + +from nvflow.provenanceguard.protocols import Embedder, RoutedClaim +from nvflow.provenanceguard.types import AtomicClaim, EvidenceChunk + + +def _cosine(a: Sequence[float], b: Sequence[float]) -> float: + if not a or not b: + return 0.0 + n = min(len(a), len(b)) + dot = sum(a[i] * b[i] for i in range(n)) + na = math.sqrt(sum(x * x for x in a[:n])) + nb = math.sqrt(sum(x * x for x in b[:n])) + if na == 0.0 or nb == 0.0: + return 0.0 + return dot / (na * nb) + + +def _vec_add_inplace(acc: list[float], v: Sequence[float]) -> None: + if len(acc) != len(v): + if not acc: + acc.extend(float(x) for x in v) + return + raise ValueError(f"embedding dim mismatch: acc={len(acc)} v={len(v)}") + for i, x in enumerate(v): + acc[i] += float(x) + + +def _vec_scale(v: list[float], k: float) -> list[float]: + return [x * k for x in v] + + +class EmbeddingSourceRouter: + """Cosine-on-centroids router. Returns top-1 with margin to top-2.""" + + def __init__(self, embedder: Embedder) -> None: + self._embedder = embedder + + def route( + self, + claims: Sequence[AtomicClaim], + evidence: Sequence[EvidenceChunk], + ) -> Sequence[RoutedClaim]: + if not claims or not evidence: + return [] + + # Group evidence by routing key (source_id, or chunk_id as fallback). + groups: dict[str, list[EvidenceChunk]] = {} + key_order: list[str] = [] + for chunk in evidence: + # Composite/unattributable chunks (attribution_state == "unavailable") + # have source_id=None on purpose. Using chunk_id directly would be + # fine since IDs are unique, but we prefix with "_unattributable:" so + # the group key can never collide with a real source_id and the + # intent is explicit when debugging. + if chunk.attribution_state == "unavailable": + key = f"_unattributable:{chunk.chunk_id}" + else: + key = chunk.source_id or chunk.chunk_id + if key not in groups: + groups[key] = [] + key_order.append(key) + groups[key].append(chunk) + + # Embed evidence + claims in one batch. + all_evidence_texts: list[str] = [] + chunk_keys: list[str] = [] + for key in key_order: + for c in groups[key]: + all_evidence_texts.append(c.text) + chunk_keys.append(key) + claim_texts = [c.text for c in claims] + + embeddings = self._embedder.embed(all_evidence_texts + claim_texts) + if len(embeddings) != len(all_evidence_texts) + len(claim_texts): + raise RuntimeError("embedder returned wrong number of vectors") + + ev_vecs = embeddings[: len(all_evidence_texts)] + cl_vecs = embeddings[len(all_evidence_texts) :] + + # Per-key centroid. + centroids: dict[str, list[float]] = {k: [] for k in key_order} + counts: dict[str, int] = dict.fromkeys(key_order, 0) + for key, vec in zip(chunk_keys, ev_vecs, strict=True): + _vec_add_inplace(centroids[key], vec) + counts[key] += 1 + for key in key_order: + n = counts[key] or 1 + centroids[key] = _vec_scale(centroids[key], 1.0 / n) + + # Representative chunk per key (longest text, deterministic). + rep_chunk: dict[str, EvidenceChunk] = { + key: max(groups[key], key=lambda c: len(c.text)) for key in key_order + } + + routed: list[RoutedClaim] = [] + for claim, cv in zip(claims, cl_vecs, strict=True): + scored = [(key, _cosine(cv, centroids[key])) for key in key_order] + scored.sort(key=lambda kv: kv[1], reverse=True) + top_key, top_score = scored[0] + margin = top_score - (scored[1][1] if len(scored) > 1 else 0.0) + routed.append( + RoutedClaim( + claim=claim, + chunk=rep_chunk[top_key], + score=float(top_score), + margin=float(margin), + ) + ) + return routed + + +__all__ = ["EmbeddingSourceRouter"] diff --git a/nvflow/provenanceguard/types.py b/nvflow/provenanceguard/types.py new file mode 100644 index 0000000..6e970d9 --- /dev/null +++ b/nvflow/provenanceguard/types.py @@ -0,0 +1,168 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Domain data types for ProvenanceGuard Open v1. + +All types are frozen dataclasses so they are hashable and safe to share +across threads. The ``to_dict`` / ``serialize`` helpers produce the JSON +shape written to sidecar files. + +This is an **uncalibrated open approximation** of the research +ProvenanceGuard system. The evidence basis is model-retrieved trace +excerpts, NOT primary SEC filing verification. +""" + +from __future__ import annotations + +import dataclasses +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class EvidenceChunk: + """A single evidence span extracted from a tool-call trace. + + ``content_basis`` is always ``"retrieval_model_excerpt"`` — the text + came from a model retrieval call (``retrieve_information``), not + from primary source verification. + + ``source_id`` is a stable canonical SEC identifier + (``sec:cik=<...>:accession=<...>:doc=<...>``) when the chunk can be + correlated to a specific filing. When multiple storage keys are + combined in one retrieval output and cannot be separated, + ``source_ids`` holds all candidate IDs and ``attribution_state`` is + ``"unavailable"``. + """ + + chunk_id: str + text: str + content_basis: str = "retrieval_model_excerpt" + source_id: str | None = None + source_ids: tuple[str, ...] = () + sec_url: str | None = None + sec_accession: str | None = None + sec_document: str | None = None + sec_cik: str | None = None + storage_keys: tuple[str, ...] = () + char_range: tuple[int, int] | None = None + tool_call_id: str | None = None + tool_result_id: str | None = None + attribution_state: str = "available" + + def to_dict(self) -> dict[str, Any]: + d = dataclasses.asdict(self) + if self.char_range is not None: + d["char_range"] = {"start": self.char_range[0], "end": self.char_range[1]} + d["source_ids"] = list(self.source_ids) + d["storage_keys"] = list(self.storage_keys) + return d + + +@dataclass(frozen=True) +class AtomicClaim: + """A sub-sentence atomic claim extracted from the final answer. + + ``claim_id`` is a deterministic hash of the claim text and source + sentence, so re-evaluation of the same answer yields stable IDs. + + ``stated_sec_ids`` / ``stated_sec_urls`` capture SEC identifiers or + URLs the claim explicitly references in its text. + """ + + claim_id: str + text: str + source_sentence: str + stated_sec_ids: tuple[str, ...] = () + stated_sec_urls: tuple[str, ...] = () + + def to_dict(self) -> dict[str, Any]: + d = dataclasses.asdict(self) + d["stated_sec_ids"] = list(self.stated_sec_ids) + d["stated_sec_urls"] = list(self.stated_sec_urls) + return d + + +@dataclass(frozen=True) +class ClaimVerdict: + """Per-claim outcome with routing, NLI, and protected-value detail. + + ``raw_nli_label`` is the direct output of the NLI model + (``entailment`` / ``neutral`` / ``contradiction``). + + ``final_label`` applies policy adjustments: + - ``no_source`` when no evidence chunk was available for routing. + - ``protected_value_mismatch`` when protected numeric/date/percentage + values in an entailed claim are absent from routed evidence. + """ + + claim_id: str + claim_text: str + routed_source_id: str | None = None + routed_source_ids: tuple[str, ...] = () + routed_attribution_state: str = "available" + routed_chunk_id: str | None = None + routing_score: float = 0.0 + routing_margin: float = 0.0 + raw_nli_label: str = "neutral" + raw_nli_probabilities: tuple[tuple[str, float], ...] = () + final_label: str = "neutral" + protected_value_outcome: str = "not_applicable" + evidence_excerpt: str = "" + errors: tuple[str, ...] = () + + def to_dict(self) -> dict[str, Any]: + return { + "claim_id": self.claim_id, + "claim_text": self.claim_text, + "routed_source_id": self.routed_source_id, + "routed_source_ids": list(self.routed_source_ids), + "routed_attribution_state": self.routed_attribution_state, + "routed_chunk_id": self.routed_chunk_id, + "routing_score": self.routing_score, + "routing_margin": self.routing_margin, + "raw_nli_label": self.raw_nli_label, + "raw_nli_probabilities": dict(self.raw_nli_probabilities), + "final_label": self.final_label, + "protected_value_outcome": self.protected_value_outcome, + "evidence_excerpt": self.evidence_excerpt, + "errors": list(self.errors), + } + + +@dataclass(frozen=True) +class Decision: + """Top-level allow / block / unavailable result for one rollout row.""" + + status: str # allow | block | unavailable + reason: str + verdicts: tuple[ClaimVerdict, ...] = () + errors: tuple[str, ...] = () + + def to_dict(self) -> dict[str, Any]: + return { + "status": self.status, + "reason": self.reason, + "verdicts": [v.to_dict() for v in self.verdicts], + "errors": list(self.errors), + } + + +@dataclass(frozen=True) +class NLIResult: + """Raw NLI model output for a single (premise, hypothesis) pair.""" + + label: str # entailment | neutral | contradiction + score: float + probabilities: tuple[tuple[str, float], ...] = () diff --git a/nvflow/recipes/finance/stages/rl/evaluate_provenance.py b/nvflow/recipes/finance/stages/rl/evaluate_provenance.py new file mode 100644 index 0000000..51496ab --- /dev/null +++ b/nvflow/recipes/finance/stages/rl/evaluate_provenance.py @@ -0,0 +1,209 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""ProvenanceGuard evaluation stage (finance recipe). + +Thin stage wrapper around the finance ProvenanceGuard sidecar CLI +(:mod:`nvflow.recipes.finance.utils.rl.provenanceguard`). Registers as +``finance/grpo/evaluate_provenance`` and submits one CPU Slurm job per +(environment, seed) pair. + +Opt-in stage: commented out in ``pipeline_stages`` in +``nvflow/recipes/finance/workflows/grpo/base.yaml``. To enable, +uncomment the ``# - evaluate_provenance`` line. + +Depends only on ``collect_rollouts``; no downstream dependencies. + +Input path pattern:: + + ${directories.step-5-collect-rollouts}/{env}/rollout/output-rs.jsonl + +Input completion marker (required):: + + ${directories.step-5-collect-rollouts}/{env}/rollout/output-rs.jsonl.done + +Output path pattern:: + + ${directories.provenanceguard-eval}/{env}/provenanceguard-rs.jsonl + +Output completion marker (created atomically after successful replace):: + + ${directories.provenanceguard-eval}/{env}/provenanceguard-rs.jsonl.done + +Does not force ``HF_HUB_OFFLINE`` in production. Does not create remote +directories on the local host. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +from nvflow.core import BaseStage, StageRegistry, console + +_SIDECAR_MODULE = "nvflow.recipes.finance.utils.rl.provenanceguard" + + +def build_evaluate_command( + input_file: str, + output_file: str, + seed: int, + environment: str, + routing_model: str = "sentence-transformers/all-MiniLM-L6-v2", + nli_model: str = "MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli", + routing_revision: str | None = None, + nli_revision: str | None = None, + evidence_excerpt_length: int = 500, +) -> str: + """Build the rendered CLI command string for one ProvenanceGuard job. + + Pure function with no side effects — used by :meth:`execute` and + tested independently. + """ + from nvflow.lib.cli_cmd import build_python_cmd + + flags: dict[str, str | int] = { + "input_file": input_file, + "output_file": output_file, + "seed": seed, + "environment": environment, + "routing_model": routing_model, + "nli_model": nli_model, + "evidence_excerpt_length": evidence_excerpt_length, + } + if routing_revision: + flags["routing_model_revision"] = routing_revision + if nli_revision: + flags["nli_model_revision"] = nli_revision + return build_python_cmd(_SIDECAR_MODULE, **flags) + + +@StageRegistry.register( + recipe="finance", + workflow="grpo", + stage="evaluate_provenance", +) +class EvaluateProvenanceStage(BaseStage): + """CPU stage that runs ProvenanceGuard evaluation on rollout outputs. + + Reads merged rollout files from collect_rollouts output and writes + sidecar provenance verdicts. Depends only on collect_rollouts. + """ + + workflow = "grpo" + + def execute( + self, + config: dict[str, Any], + cluster: str, + expname: str, + run_after: list[str] | None = None, + ) -> None: + """Submit ProvenanceGuard evaluation jobs (one per seed per env).""" + from nemo_skills.pipeline.cli import run_cmd, wrap_arguments + + from nvflow.lib.rl.helpers import resolve_environments + + rollouts_dir = config["rollouts_dir"] + output_dir = config["output_dir"] + + environments = resolve_environments(config) + container = config.get("container", "nemo-skills") + installation_command = config.get("installation_command", "true") + + routing_model = config.get("routing_model", "sentence-transformers/all-MiniLM-L6-v2") + nli_model = config.get("nli_model", "MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli") + routing_revision = config.get("routing_model_revision") + nli_revision = config.get("nli_model_revision") + + evidence_excerpt_length = config.get("evidence_excerpt_length", 500) + + starting_seed = config["starting_seed"] + num_seeds = config["num_random_seeds"] + seeds = config.get( + "seeds", + list(range(starting_seed, starting_seed + num_seeds)), + ) + + for env_name, _env_cfg in environments.items(): + env_rollouts = f"{rollouts_dir}/{env_name}/rollout" + env_output = f"{output_dir}/{env_name}" + + for seed in seeds: + input_file = f"{env_rollouts}/output-rs{seed}.jsonl" + input_done = f"{input_file}.done" + output_file = f"{env_output}/provenanceguard-rs{seed}.jsonl" + + console.status(f"Submitting ProvenanceGuard for {env_name} seed {seed}") + console.detail("Input", input_file) + console.detail("Input marker", input_done) + console.detail("Output", output_file) + + cmd = build_evaluate_command( + input_file=input_file, + output_file=output_file, + seed=seed, + environment=env_name, + routing_model=routing_model, + nli_model=nli_model, + routing_revision=routing_revision, + nli_revision=nli_revision, + evidence_excerpt_length=evidence_excerpt_length, + ) + + run_cmd( + ctx=wrap_arguments(cmd), + cluster=cluster, + log_dir=f"{env_output}/logs", + expname=f"{expname}-{env_name}-seed{seed}", + run_after=run_after, + container=container, + installation_command=installation_command, + num_gpus=config.get("num_gpus", 0), + ) + + console.success( + f"ProvenanceGuard jobs submitted for {len(seeds)} seed(s) " + f"across {len(environments)} env(s)" + ) + + def validate_config(self, config: dict[str, Any]) -> None: + """Validate required paths and no overlap (canonicalized).""" + for field_name in ("rollouts_dir", "output_dir"): + if field_name not in config: + raise ValueError(f"{field_name} is required in evaluate_provenance config") + + rollouts_dir = Path(os.path.expanduser(config["rollouts_dir"])).resolve(strict=False) + output_dir = Path(os.path.expanduser(config["output_dir"])).resolve(strict=False) + + if rollouts_dir == output_dir: + raise ValueError("rollouts_dir and output_dir must not be the same path") + + try: + output_dir.relative_to(rollouts_dir) + except ValueError: + pass + else: + raise ValueError("output_dir must not be inside rollouts_dir") + + try: + rollouts_dir.relative_to(output_dir) + except ValueError: + pass + else: + raise ValueError("rollouts_dir must not be inside output_dir") + + +__all__ = ["EvaluateProvenanceStage", "build_evaluate_command"] diff --git a/nvflow/recipes/finance/utils/rl/provenanceguard.py b/nvflow/recipes/finance/utils/rl/provenanceguard.py new file mode 100644 index 0000000..2e4a190 --- /dev/null +++ b/nvflow/recipes/finance/utils/rl/provenanceguard.py @@ -0,0 +1,993 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Finance ProvenanceGuard sidecar: trace extraction, evaluation, and CLI. + +This module lives in the finance recipe layer (not the generic +``nvflow.provenanceguard`` package). It provides: + +1. **Trace extraction** - parses NeMo-Gym rollout JSONL rows and extracts + structured evidence chunks + the final answer from the tool-call trace. +2. **Deterministic sidecar evaluator** - runs the generic + :class:`~nvflow.provenanceguard.evaluator.ProvenanceGuardEvaluator` + on each row and produces a complete sidecar dict with full metadata. +3. **Atomic file writer** - writes the sidecar JSONL using an atomic + temp-file + ``os.replace`` protocol with strict marker gating. +4. **CLI** - ``python3 -m nvflow.recipes.finance.utils.rl.provenanceguard``. + +Extraction rules +~~~~~~~~~~~~~~~~ + +- **Answer**: last valid ``submit_final_result.final_result``. +- **Evidence**: one ``EvidenceChunk`` per ``retrieve_information`` result, + paired by ``call_id``. +- **Source correlation**: ``sec_filing_search`` metadata + ``parse_html_page`` + storage keys. +- **Stable IDs**: only when 10-digit zero-padded CIK + accession + document + are all present. Never storage-key/URL fallback IDs. +- **Composite chunks**: multiple storage keys, any unknown key, or no + canonical ID -> ``attribution_state="unavailable"``; + ``source_ids`` contains only known canonical candidates. + +Every physical input line (including blank/malformed/non-object) produces +exactly one sidecar row. Output is deterministic: no random UUID. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import sys +import tempfile +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from nvflow.provenanceguard.evaluator import ( + ALGORITHM_VERSION, + ProvenanceGuardConfig, + ProvenanceGuardEvaluator, +) +from nvflow.provenanceguard.types import EvidenceChunk + +SCHEMA_VERSION = "1.0.0" +EVIDENCE_BASIS = "retrieval_model_excerpt" +LIMITATION = ( + "Uncalibrated open approximation of ProvenanceGuard. " + "Evidence is model-retrieved trace excerpts, not primary SEC filing " + "verification. Rule-based claim decomposition may over-merge or " + "split mid-clause. No paper-faithful calibration." +) + +SEC_SEARCH_NAMES = frozenset({"sec_filing_search", "edgar_search"}) +PARSE_HTML_NAME = "parse_html_page" +RETRIEVE_INFO_NAME = "retrieve_information" +SUBMIT_FINAL_NAME = "submit_final_result" + +STORAGE_KEY_RE = re.compile(r"\{\{([^{}]+)\}\}") + +SEC_FILING_URL_RE = re.compile( + r"https?://www\.sec\.gov/Archives/edgar/data/" + r"(?P\d+)/(?P\d+)/" + r"(?P[^\s\"<]+)", + re.IGNORECASE, +) + + +@dataclass(frozen=True) +class FilingMetadata: + """SEC filing metadata extracted from sec_filing_search or URL parsing.""" + + cik: str | None = None + accession: str | None = None + document: str | None = None + url: str | None = None + + def source_id(self) -> str | None: + """Build stable source ID if CIK + accession + document are present. + + CIK must be exactly 10 digits (zero-padded). Never falls back to + storage keys or URLs for the canonical ID. + """ + if self.cik and self.accession and self.document: + cik = self.cik.zfill(10) + acc = str(self.accession) + if len(cik) == 10 and cik.isdigit() and len(acc) == 18 and acc.isdigit(): + return f"sec:cik={cik}:accession={self.accession}:doc={self.document}" + return None + + def to_dict(self) -> dict[str, Any]: + return { + "cik": self.cik, + "accession": self.accession, + "document": self.document, + "url": self.url, + "source_id": self.source_id(), + } + + +@dataclass +class TraceExtraction: + """Result of extracting evidence + answer from one rollout row.""" + + answer: str | None = None + evidence: list[EvidenceChunk] = field(default_factory=list) + extraction_errors: list[str] = field(default_factory=list) + submit_call_id: str | None = None + has_submit: bool = False + + +@dataclass +class SidecarRow: + """One output row for the ProvenanceGuard sidecar JSONL file.""" + + raw_line_fingerprint: str + parse_error: str | None = None + trace: TraceExtraction | None = None + + def to_dict(self) -> dict[str, Any]: + d: dict[str, Any] = { + "raw_line_fingerprint": self.raw_line_fingerprint, + } + if self.parse_error: + d["parse_error"] = self.parse_error + if self.trace: + d["answer"] = self.trace.answer + d["has_submit"] = self.trace.has_submit + d["submit_call_id"] = self.trace.submit_call_id + d["evidence"] = [c.to_dict() for c in self.trace.evidence] + d["extraction_errors"] = list(self.trace.extraction_errors) + return d + + +@dataclass(frozen=True) +class FinanceEvaluatorConfig: + """Configuration for the finance ProvenanceGuard evaluator.""" + + provenance_config: ProvenanceGuardConfig = field(default_factory=ProvenanceGuardConfig) + environment: str = "finance_sec_search" + routing_model_revision: str | None = None + nli_model_revision: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": SCHEMA_VERSION, + "algorithm": ALGORITHM_VERSION, + "evidence_basis": EVIDENCE_BASIS, + "limitation": LIMITATION, + "environment": self.environment, + "thresholds": { + "policy": "fixed_fail_closed", + "evidence_excerpt_length": self.provenance_config.evidence_excerpt_length, + }, + "models": { + "routing_model": self.provenance_config.routing_model, + "routing_model_revision": self.routing_model_revision, + "nli_model": self.provenance_config.nli_model, + "nli_model_revision": self.nli_model_revision, + }, + } + + +def _raw_line_fingerprint(raw_line: str) -> str: + """SHA-256 fingerprint (first 16 hex chars) of the raw JSONL line.""" + return hashlib.sha256(raw_line.encode("utf-8")).hexdigest()[:16] + + +def _iter_conversation_items(row: dict[str, Any]) -> Iterator[dict[str, Any]]: + """Yield conversation items from a rollout row. + + Tries common locations in order: + 1. row["output"] (Responses API output items list) + 2. row["response"]["output"] (nested under response) + 3. row["messages"] (chat completions format) + 4. row["conversation"] (alternative trace field) + """ + for source_key in ("output", "response"): + source = row.get(source_key) + if isinstance(source, dict): + items = source.get("output") + if isinstance(items, list): + yield from (i for i in items if isinstance(i, dict)) + return + if isinstance(source, list): + yield from (i for i in source if isinstance(i, dict)) + return + + for key in ("messages", "conversation"): + items = row.get(key) + if isinstance(items, list): + yield from (i for i in items if isinstance(i, dict)) + return + + +def _extract_filing_metadata(result: Any) -> FilingMetadata: + """Extract CIK, accession, document, and URL from a sec_filing_search result.""" + if not isinstance(result, dict | list): + return FilingMetadata() + + filings: list[dict[str, Any]] = [] + if isinstance(result, list): + filings = [f for f in result if isinstance(f, dict)] + elif isinstance(result, dict): + if "filings" in result and isinstance(result["filings"], list): + filings = [f for f in result["filings"] if isinstance(f, dict)] + else: + filings = [result] + + if not filings: + return FilingMetadata() + + cik = None + accession = None + document = None + url = None + + for filing in filings: + if not isinstance(filing, dict): + continue + if cik is None: + cik = filing.get("cik") or filing.get("CIK") or filing.get("cik_number") + if accession is None: + accession = ( + filing.get("accessionNo") + or filing.get("accession_no") + or filing.get("accession_number") + or filing.get("accession") + or filing.get("AccessionNumber") + ) + if document is None: + document = ( + filing.get("primaryDocument") + or filing.get("document") + or filing.get("file_name") + or filing.get("filename") + ) + if url is None: + url = ( + filing.get("linkToHtml") + or filing.get("url") + or filing.get("filing_url") + or filing.get("link") + ) + if cik and accession and document and url: + break + + if cik: + cik = str(cik).strip() + if len(cik) < 10: + cik = cik.zfill(10) + if accession: + accession = _format_accession(str(accession).strip()) + if document: + document = str(document).strip() + document = document.split("?")[0].split("#")[0] + document = os.path.basename(document) + + return FilingMetadata(cik=cik, accession=accession, document=document, url=url) + + +def _extract_filing_from_url(url: str) -> FilingMetadata: + """Parse CIK, accession, and document from a SEC EDGAR Archives URL. + + The CIK in the URL path is zero-padded to 10 digits. + """ + match = SEC_FILING_URL_RE.search(url) + if not match: + return FilingMetadata(url=url) + + cik = match.group("cik").zfill(10) + accession_raw = match.group("accession") + document = match.group("document") + + accession = _format_accession(accession_raw) + return FilingMetadata(cik=cik, accession=accession, document=document, url=url) + + +def _format_accession(accession_raw: str) -> str | None: + """Strip dashes/non-digits and return the 18-digit canonical form. + + Returns ``None`` unless exactly 18 digits remain after stripping + all non-digit characters. Both shorter and longer values are + rejected (overlong values are not truncated). + """ + digits = re.sub(r"\D", "", accession_raw) + if len(digits) == 18: + return digits + return None + + +def _parse_tool_arguments(raw: Any) -> dict[str, Any]: + """Parse tool call arguments which may be a JSON string or dict.""" + if isinstance(raw, dict): + return raw + if isinstance(raw, str): + try: + parsed = json.loads(raw) + if isinstance(parsed, dict): + return parsed + except (json.JSONDecodeError, TypeError): + pass + return {} + + +def _parse_tool_result(raw: Any) -> Any: + """Parse a tool result output which may be a JSON string or already parsed.""" + if isinstance(raw, dict | list): + return raw + if isinstance(raw, str): + try: + return json.loads(raw) + except (json.JSONDecodeError, TypeError): + return raw + return raw + + +def _is_error_result(parsed: Any) -> bool: + """Check if a parsed tool result is an error response. + + Rejects: + - Dict failures: ``success is False``, ``error``, ``is_error``, + or a ``result`` string containing error indicators. + - Raw failure strings: ``ERROR`` prefix, traceback. + """ + if isinstance(parsed, dict): + if parsed.get("error") or parsed.get("is_error"): + return True + if parsed.get("success") is False: + return True + result_str = parsed.get("result") + if isinstance(result_str, str) and _is_raw_error_string(result_str): + return True + return False + if isinstance(parsed, str): + return _is_raw_error_string(parsed) + return False + + +def _is_raw_error_string(text: str) -> bool: + """Check if a raw string looks like an error, traceback, or failure. + + Rejects strings containing: Error:, Failed, traceback, exception, + unavailable, or similar failure indicators from real finance tools. + Case-insensitive. + """ + stripped = text.strip() + lower = stripped.lower() + if lower.startswith("error"): + return True + if lower.startswith("failed"): + return True + if "traceback" in lower: + return True + if "exception" in lower: + return True + if lower == "unavailable": + return True + return False + + +def _build_key_to_filing_map( + items: list[dict[str, Any]], +) -> dict[str, FilingMetadata]: + """Build a mapping from storage keys to filing metadata. + + Scans for ``sec_filing_search`` / ``edgar_search`` calls (extracts + filing metadata from arguments and results) and ``parse_html_page`` + calls (reads the plain ``key`` argument directly). + """ + key_to_url: dict[str, str] = {} + call_id_to_filing: dict[str, FilingMetadata] = {} + url_to_filing: dict[str, FilingMetadata] = {} + + sec_search_call_ids: set[str] = set() + + for item in items: + item_type = item.get("type", "") + name = item.get("name", "") + call_id = item.get("call_id", item.get("id", "")) + + if item_type == "function_call" and name in SEC_SEARCH_NAMES: + sec_search_call_ids.add(call_id) + args = _parse_tool_arguments(item.get("arguments")) + filing = _extract_filing_metadata(args) + if filing.cik or filing.accession or filing.url: + call_id_to_filing[call_id] = filing + continue + + if item_type == "function_call" and name == PARSE_HTML_NAME: + args = _parse_tool_arguments(item.get("arguments")) + url = args.get("url", "") + key = args.get("key", "") + if key and url: + key_to_url[key] = url + continue + + if item_type in ("function_call_output", "tool_result"): + if not call_id: + continue + if call_id in sec_search_call_ids: + existing = call_id_to_filing.get(call_id, FilingMetadata()) + parsed = _parse_tool_result(item.get("output", item.get("result", ""))) + if isinstance(parsed, dict | list): + enriched = _extract_filing_metadata(parsed) + if enriched.cik or enriched.accession or enriched.url: + filing = FilingMetadata( + cik=enriched.cik or existing.cik, + accession=enriched.accession or existing.accession, + document=enriched.document or existing.document, + url=enriched.url or existing.url, + ) + call_id_to_filing[call_id] = filing + if call_id in call_id_to_filing: + filing = call_id_to_filing[call_id] + if filing.url: + url_to_filing[filing.url] = filing + + key_to_filing: dict[str, FilingMetadata] = {} + for key, url in key_to_url.items(): + filing = url_to_filing.get(url) + if not filing: + filing = _extract_filing_from_url(url) + key_to_filing[key] = filing + + return key_to_filing + + +def _extract_storage_keys_from_text(text: str) -> list[str]: + """Extract ``{{key}}`` storage key references from text.""" + return [m.strip() for m in STORAGE_KEY_RE.findall(text)] + + +def _build_evidence_chunk( + chunk_id: str, + text: str, + keys: list[str], + key_to_filing: dict[str, FilingMetadata], + tool_call_id: str | None, + tool_result_id: str | None, +) -> EvidenceChunk: + """Build an EvidenceChunk from a retrieve_information output. + + Attribution is ``available`` ONLY when the retrieval requested + exactly one key AND that key maps to exactly one canonical + CIK+accession+document source. Any zero keys, multiple keys, + unknown key, or noncanonical mapping yields ``source_id=None``, + ``source_ids`` containing only known canonical IDs, and + ``attribution_state="unavailable"``. + """ + filings: list[FilingMetadata] = [] + for key in keys: + filing = key_to_filing.get(key) + if filing: + filings.append(filing) + + if not filings: + return EvidenceChunk( + chunk_id=chunk_id, + text=text, + source_id=None, + source_ids=(), + tool_call_id=tool_call_id, + tool_result_id=tool_result_id, + attribution_state="unavailable", + ) + + if len(keys) == 1 and len(filings) == 1: + filing = filings[0] + sid = filing.source_id() + if sid: + return EvidenceChunk( + chunk_id=chunk_id, + text=text, + source_id=sid, + source_ids=(), + sec_url=filing.url, + sec_cik=filing.cik, + sec_accession=filing.accession, + sec_document=filing.document, + storage_keys=tuple(keys), + tool_call_id=tool_call_id, + tool_result_id=tool_result_id, + attribution_state="available", + ) + return EvidenceChunk( + chunk_id=chunk_id, + text=text, + source_id=None, + source_ids=(), + sec_url=filing.url, + storage_keys=tuple(keys), + tool_call_id=tool_call_id, + tool_result_id=tool_result_id, + attribution_state="unavailable", + ) + + source_ids = tuple(f.source_id() for f in filings if f.source_id() is not None) + return EvidenceChunk( + chunk_id=chunk_id, + text=text, + source_id=None, + source_ids=source_ids, + sec_url=None, + storage_keys=tuple(keys), + tool_call_id=tool_call_id, + tool_result_id=tool_result_id, + attribution_state="unavailable", + ) + + +def extract_trace(row: dict[str, Any]) -> TraceExtraction: + """Extract answer and evidence from a single rollout row.""" + items = list(_iter_conversation_items(row)) + extraction = TraceExtraction() + + key_to_filing = _build_key_to_filing_map(items) + + call_id_to_result: dict[str, Any] = {} + for item in items: + item_type = item.get("type", "") + if item_type in ("function_call_output", "tool_result"): + call_id = item.get("call_id", item.get("tool_call_id", "")) + output = item.get("output", item.get("result", "")) + if call_id and output is not None: + call_id_to_result[call_id] = output + + last_submit_call_id: str | None = None + last_submit_result: str | None = None + retrieve_calls: list[dict[str, Any]] = [] + + for item in items: + item_type = item.get("type", "") + name = item.get("name", "") + call_id = item.get("call_id", item.get("id", "")) + + if item_type == "function_call" and name == SUBMIT_FINAL_NAME: + args = _parse_tool_arguments(item.get("arguments")) + final_result = args.get("final_result", "") + if final_result and isinstance(final_result, str): + last_submit_call_id = call_id + last_submit_result = final_result + + if item_type == "function_call" and name == RETRIEVE_INFO_NAME: + retrieve_calls.append(item) + + if last_submit_result is not None: + extraction.answer = last_submit_result + extraction.submit_call_id = last_submit_call_id + extraction.has_submit = True + else: + extraction.extraction_errors.append("No valid submit_final_result call found in trace") + + for idx, call in enumerate(retrieve_calls): + call_id = call.get("call_id", call.get("id", f"retrieve_{idx}")) + result_raw = call_id_to_result.get(call_id, "") + if not result_raw: + extraction.extraction_errors.append( + f"retrieve_information call {call_id}: no result output" + ) + continue + + parsed = _parse_tool_result(result_raw) + if _is_error_result(parsed): + extraction.extraction_errors.append( + f"retrieve_information call {call_id}: error result" + ) + continue + + if isinstance(parsed, dict): + text = parsed.get("result", parsed.get("retrieval", "")) + if isinstance(text, dict): + text = json.dumps(text) + text = str(text) if text else "" + elif isinstance(parsed, str): + text = parsed + else: + text = str(parsed) + + if not text.strip(): + extraction.extraction_errors.append( + f"retrieve_information call {call_id}: empty result" + ) + continue + + args = _parse_tool_arguments(call.get("arguments")) + prompt = args.get("prompt", "") + keys = _extract_storage_keys_from_text(prompt) + + chunk_id = f"ev:{call_id}" + chunk = _build_evidence_chunk( + chunk_id=chunk_id, + text=text, + keys=keys, + key_to_filing=key_to_filing, + tool_call_id=call_id, + tool_result_id=call_id, + ) + extraction.evidence.append(chunk) + + return extraction + + +def parse_rollout_line(raw_line: str) -> SidecarRow: + """Parse a single JSONL line and return a SidecarRow. + + Every valid or malformed JSONL line yields exactly one SidecarRow. + """ + fingerprint = _raw_line_fingerprint(raw_line) + stripped = raw_line.strip() + if not stripped: + return SidecarRow( + raw_line_fingerprint=fingerprint, + parse_error="empty line", + ) + + try: + row = json.loads(stripped) + except json.JSONDecodeError as exc: + return SidecarRow( + raw_line_fingerprint=fingerprint, + parse_error=f"JSONDecodeError: {exc!s}", + ) + + if not isinstance(row, dict): + return SidecarRow( + raw_line_fingerprint=fingerprint, + parse_error=f"expected JSON object, got {type(row).__name__}", + ) + + trace = extract_trace(row) + return SidecarRow( + raw_line_fingerprint=fingerprint, + trace=trace, + ) + + +def _response_id(row: dict[str, Any]) -> str | None: + """Extract response ID from a rollout row.""" + for key in ("response_id", "id", "response_id_str"): + val = row.get(key) + if isinstance(val, str) and val: + return val + resp = row.get("response") + if isinstance(resp, dict): + rid = resp.get("id") + if isinstance(rid, str) and rid: + return rid + return None + + +def _row_uuid(row: dict[str, Any]) -> str | None: + """Extract UUID from a rollout row.""" + for key in ("uuid", "task_uuid", "sample_id"): + val = row.get(key) + if isinstance(val, str) and val: + return val + return None + + +def _row_indices(row: dict[str, Any]) -> tuple[int | None, int | None]: + """Extract (_ng_task_index, _ng_rollout_index) from a rollout row.""" + task_idx = row.get("_ng_task_index") + rollout_idx = row.get("_ng_rollout_index") + return ( + int(task_idx) if isinstance(task_idx, int | float) else None, + int(rollout_idx) if isinstance(rollout_idx, int | float) else None, + ) + + +def _row_fingerprint(row: dict[str, Any]) -> str: + """Deterministic fingerprint of the rollout row content.""" + raw = json.dumps(row, sort_keys=True, default=str) + return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16] + + +def _evaluation_uuid( + fingerprint: str, + seed: int, + line_number: int, + config: FinanceEvaluatorConfig, +) -> str: + """Deterministic evaluation UUID from raw-line fingerprint + seed + line number + algorithm/config. + + Uses the raw-line fingerprint (SHA-256 of the raw JSONL line) rather + than the parsed dict fingerprint, avoiding collisions on malformed + lines that all parse to ``{}``. The physical line number is included + so duplicate identical rows still get distinguishable IDs by their + position in the input file. The config digest captures the + algorithm version and policy settings so the same row evaluated under + different configurations yields a different UUID. + """ + config_digest = hashlib.sha256( + json.dumps(config.to_dict(), sort_keys=True, default=str).encode("utf-8") + ).hexdigest()[:16] + raw = f"eval:{fingerprint}:{seed}:{line_number}:{ALGORITHM_VERSION}:{config_digest}" + return hashlib.sha256(raw.encode("utf-8")).hexdigest() + + +def evaluate_row( + raw_line: str, + row: dict[str, Any], + seed: int, + evaluator: ProvenanceGuardEvaluator, + config: FinanceEvaluatorConfig, + line_number: int = 0, +) -> dict[str, Any]: + """Evaluate a single rollout row and produce a complete sidecar dict.""" + sidecar = parse_rollout_line(raw_line) + trace = sidecar.trace + + task_idx, rollout_idx = _row_indices(row) + response_id = _response_id(row) + row_uuid = _row_uuid(row) + fingerprint = _row_fingerprint(row) + + output: dict[str, Any] = { + **config.to_dict(), + "seed": seed, + "line_number": line_number, + "uuid": row_uuid, + "_ng_task_index": task_idx, + "_ng_rollout_index": rollout_idx, + "response_id": response_id, + "fingerprint": fingerprint, + "raw_line_fingerprint": sidecar.raw_line_fingerprint, + "evaluation_uuid": _evaluation_uuid( + sidecar.raw_line_fingerprint, seed, line_number, config + ), + } + + if sidecar.parse_error: + output["parse_error"] = sidecar.parse_error + output["verdict"] = { + "status": "unavailable", + "reason": "parse_error", + "verdicts": [], + "errors": [sidecar.parse_error], + } + output["evidence"] = [] + output["extraction_errors"] = [sidecar.parse_error] + return output + + if trace is None: + output["parse_error"] = "trace extraction returned None" + output["verdict"] = { + "status": "unavailable", + "reason": "trace_extraction_failed", + "verdicts": [], + "errors": ["trace extraction returned None"], + } + output["evidence"] = [] + output["extraction_errors"] = ["trace extraction returned None"] + return output + + evidence = trace.evidence + extraction_errors = list(trace.extraction_errors) + + output["answer"] = trace.answer + output["has_submit"] = trace.has_submit + output["submit_call_id"] = trace.submit_call_id + output["evidence"] = [c.to_dict() for c in evidence] + output["extraction_errors"] = extraction_errors + + if not trace.has_submit or trace.answer is None: + decision = { + "status": "unavailable", + "reason": "no_submit_final_result", + "verdicts": [], + "errors": extraction_errors, + } + else: + decision_obj = evaluator.evaluate(trace.answer, evidence) + decision = decision_obj.to_dict() + if extraction_errors: + decision = { + "status": "unavailable", + "reason": "trace_extraction_errors", + "verdicts": decision.get("verdicts", []), + "errors": list(extraction_errors) + list(decision.get("errors", [])), + } + + output["verdict"] = decision + return output + + +def evaluate_seed( + input_path: str, + output_path: str, + seed: int, + evaluator: ProvenanceGuardEvaluator, + config: FinanceEvaluatorConfig, +) -> int: + """Evaluate one seed rollout file and write the sidecar JSONL atomically. + + Protocol: + 1. Require BOTH input_path and its sibling .done marker to exist + before touching any output state. A missing input gate + preserves prior output and output .done unchanged. + 2. After both input gates pass, remove any stale output .done + marker before mkdir/temp/evaluation so a failed rerun cannot + leave a stale completion marker. + 3. Write all rows to a temp file in the output file directory. + 4. os.replace temp -> final (atomic on same filesystem). + 5. Create a fresh empty .done marker only after successful replace. + + Returns the number of rows written. + """ + done_path = output_path + ".done" + + input_done = input_path + ".done" + if not os.path.isfile(input_path): + raise FileNotFoundError(f"Input file does not exist: {input_path}") + if not os.path.isfile(input_done): + raise RuntimeError(f"Input completion marker not found: {input_done}") + + if os.path.exists(done_path): + os.remove(done_path) + + out_dir = os.path.dirname(output_path) + os.makedirs(out_dir or ".", exist_ok=True) + + count = 0 + fd, tmp_path = tempfile.mkstemp( + dir=out_dir or ".", + prefix=".pg_tmp_", + suffix=".jsonl", + ) + try: + with ( + os.fdopen(fd, "w", encoding="utf-8") as out_f, + open(input_path, encoding="utf-8") as in_f, + ): + for raw_line in in_f: + stripped = raw_line.strip() + if not stripped: + row: dict[str, Any] = {} + else: + try: + row = json.loads(stripped) + if not isinstance(row, dict): + row = {} + except json.JSONDecodeError: + row = {} + + result = evaluate_row(raw_line, row, seed, evaluator, config, count) + out_f.write(json.dumps(result, ensure_ascii=False, default=str) + "\n") + count += 1 + + os.replace(tmp_path, output_path) + except Exception: + if os.path.exists(tmp_path): + os.remove(tmp_path) + raise + + Path(done_path).touch() + return count + + +def sidecar_filename(seed: int) -> str: + """Return the sidecar output filename for a given seed.""" + return f"provenanceguard-rs{seed}.jsonl" + + +def merged_filename(seed: int) -> str: + """Return the merged rollout filename for a given seed.""" + return f"output-rs{seed}.jsonl" + + +def done_marker(seed: int) -> str: + """Return the completion marker filename for a given seed.""" + return f"{merged_filename(seed)}.done" + + +def build_evaluator_from_config( + config: FinanceEvaluatorConfig, +) -> ProvenanceGuardEvaluator: + """Build a ProvenanceGuardEvaluator with lazy HF model collaborators. + + Model loading is deferred; this function constructs the objects but + does not download or load any model weights. + """ + from nvflow.provenanceguard.decomposer import RuleBasedDecomposer + from nvflow.provenanceguard.embedder import HFEmbedder + from nvflow.provenanceguard.nli import HFNLI + from nvflow.provenanceguard.router import EmbeddingSourceRouter + + pg_config = config.provenance_config + embedder = HFEmbedder( + model_id=pg_config.routing_model, + revision=config.routing_model_revision, + ) + nli_scorer = HFNLI( + model_id=pg_config.nli_model, + revision=config.nli_model_revision, + ) + return ProvenanceGuardEvaluator( + decomposer=RuleBasedDecomposer(), + router=EmbeddingSourceRouter(embedder), + nli_scorer=nli_scorer, + config=pg_config, + ) + + +def main(argv: Sequence[str] | None = None) -> int: + """CLI entry point for ProvenanceGuard finance evaluation.""" + parser = argparse.ArgumentParser( + description="Run ProvenanceGuard evaluation on a finance rollout file." + ) + parser.add_argument("--input_file", required=True) + parser.add_argument("--output_file", required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--environment", default="finance_sec_search") + parser.add_argument( + "--routing_model", + default="sentence-transformers/all-MiniLM-L6-v2", + ) + parser.add_argument( + "--nli_model", + default="MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli", + ) + parser.add_argument("--routing_model_revision", default=None) + parser.add_argument("--nli_model_revision", default=None) + parser.add_argument("--evidence_excerpt_length", type=int, default=500) + + args = parser.parse_args(argv) + + pg_config = ProvenanceGuardConfig( + evidence_excerpt_length=args.evidence_excerpt_length, + routing_model=args.routing_model, + nli_model=args.nli_model, + ) + finance_config = FinanceEvaluatorConfig( + provenance_config=pg_config, + environment=args.environment, + routing_model_revision=args.routing_model_revision, + nli_model_revision=args.nli_model_revision, + ) + + evaluator = build_evaluator_from_config(finance_config) + + count = evaluate_seed( + input_path=args.input_file, + output_path=args.output_file, + seed=args.seed, + evaluator=evaluator, + config=finance_config, + ) + print(f"ProvenanceGuard: {count} rows written to {args.output_file}") + return 0 + + +__all__ = [ + "EVIDENCE_BASIS", + "LIMITATION", + "SCHEMA_VERSION", + "FinanceEvaluatorConfig", + "build_evaluator_from_config", + "done_marker", + "evaluate_row", + "evaluate_seed", + "extract_trace", + "merged_filename", + "parse_rollout_line", + "sidecar_filename", +] + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/nvflow/recipes/finance/workflows/grpo/base.yaml b/nvflow/recipes/finance/workflows/grpo/base.yaml index cdf0e87..30bbed9 100644 --- a/nvflow/recipes/finance/workflows/grpo/base.yaml +++ b/nvflow/recipes/finance/workflows/grpo/base.yaml @@ -98,6 +98,7 @@ pipeline_stages: - prepare_data # Step 4: Run ng_prepare_data (CPU) - prefetch_cache # Pre-warm SEC metadata cache (CPU, optional, no step-N) - collect_rollouts # Step 5: Rollout collection + reward profiling + filter + # - evaluate_provenance # ProvenanceGuard evaluation (opt-in; depends on collect_rollouts) # - compute_rewards # Step 6: Re-judge rollouts with different judge (optional) - train_validation_split # Step 7: Final train/val split on reward-filtered data - training # Step 8: GRPO training with NeMo-Gym environment @@ -126,6 +127,7 @@ directories: cache-finance-sec-search: ${base_output_dir}/cache/finance_sec_search # Model-specific stages (rollouts, training, eval). step-5-collect-rollouts: ${model_output_dir}/step-5-collect-rollouts + provenanceguard-eval: ${model_output_dir}/provenanceguard-eval step-6-compute-rewards: ${model_output_dir}/step-6-compute-rewards step-7-train-validation-split: ${model_output_dir}/step-7-train-validation-split step-8-training: ${model_output_dir}/step-8-training @@ -532,6 +534,45 @@ stages: filter: min_reward_std: 1e-6 # remove questions with zero reward variance (no GRPO gradient) + # -------------------------------------------------------------------------- + # ProvenanceGuard Evaluation [OPT-IN] + # -------------------------------------------------------------------------- + # Runs ProvenanceGuard Open v1 provenance evaluation on the merged rollout + # files produced by collect_rollouts. Produces sidecar JSONL files + # (provenanceguard-rs.jsonl) with per-row allow/block/unavailable + # verdicts, evidence metadata, and extraction errors. + # + # CPU-only: uses all-MiniLM-L6-v2 (routing) + DeBERTa NLI (lazy-loaded). + # Depends only on collect_rollouts. No downstream dependencies. + # + # Input: ${directories.step-5-collect-rollouts}/{env}/rollout/output-rs.jsonl + # Marker: ${directories.step-5-collect-rollouts}/{env}/rollout/output-rs.jsonl.done + # Output: ${directories.provenanceguard-eval}/{env}/provenanceguard-rs.jsonl + # Marker: ${directories.provenanceguard-eval}/{env}/provenanceguard-rs.jsonl.done + # + # To enable, uncomment "# - evaluate_provenance" in pipeline_stages above. + evaluate_provenance: + output_dir: ${directories.provenanceguard-eval} + rollouts_dir: ${directories.step-5-collect-rollouts} + environments: ${environments} + container: "nemo-skills" + installation_command: "true" + num_gpus: 0 + # ProvenanceGuard model IDs (lazy-loaded, CPU-only). + routing_model: "sentence-transformers/all-MiniLM-L6-v2" + nli_model: "MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli" + routing_model_revision: null + nli_model_revision: null + # Fixed fail-closed policy: contradiction, neutral, no_source, + # protected_value_mismatch always block; model/trace errors => unavailable. + # No configurable block_on_* toggles — the policy is fixed by design. + evidence_excerpt_length: 500 + # Seeds dynamically inherited from collect_rollouts rollout settings. + starting_seed: ${stages.collect_rollouts.rollout.starting_seed} + num_random_seeds: ${stages.collect_rollouts.rollout.num_random_seeds} + dependencies: + - collect_rollouts + # -------------------------------------------------------------------------- # Stage 6: Compute Rewards (re-judge rollouts) [OPTIONAL] # -------------------------------------------------------------------------- diff --git a/tests/test_provenanceguard.py b/tests/test_provenanceguard.py new file mode 100644 index 0000000..e5f7480 --- /dev/null +++ b/tests/test_provenanceguard.py @@ -0,0 +1,1751 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Offline tests for ProvenanceGuard Open v1. + +All tests use deterministic fakes — no model downloads, no network. +Covers: evaluator decision semantics, NLI label validation, finance trace +extraction, atomic sidecar output, YAML parse, registry discovery, opt-in +ordering, exact paths, marker gating, environment/model/revision/policy +propagation, every-line output, malformed/non-object, unchanged input, +repeat-byte determinism, marker atomicity, mixed error+block unavailable, +known+unknown key, zero-pad URL CIK, real nested response.output trace. +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +import os +from collections.abc import Sequence +from pathlib import Path + +import pytest +import yaml + +from nvflow.provenanceguard.decomposer import RuleBasedDecomposer +from nvflow.provenanceguard.embedder import DEFAULT_EMBEDDING_MODEL +from nvflow.provenanceguard.evaluator import ( + ALGORITHM_VERSION, + ProvenanceGuardConfig, + ProvenanceGuardEvaluator, +) +from nvflow.provenanceguard.nli import DEFAULT_NLI_MODEL, _normalize_nli_label +from nvflow.provenanceguard.protected_values import ( + check_protected_values, + extract_protected_values, +) +from nvflow.provenanceguard.router import EmbeddingSourceRouter +from nvflow.provenanceguard.types import EvidenceChunk, NLIResult +from nvflow.recipes.finance.utils.rl.provenanceguard import ( + EVIDENCE_BASIS, + LIMITATION, + SCHEMA_VERSION, + FinanceEvaluatorConfig, + done_marker, + evaluate_row, + evaluate_seed, + merged_filename, + parse_rollout_line, + sidecar_filename, +) + +# Cache for the directly-loaded evaluate_provenance stage module. +# Loading via spec_from_file_location avoids the package __init__ +# auto-importing every finance stage under the test interpreter. +_eval_stage_mod = None + + +def _load_eval_stage_module(): + """Load evaluate_provenance.py directly, bypassing __init__ auto-import. + + Cached so the @StageRegistry.register decorator runs only once, + avoiding duplicate registration errors. + """ + global _eval_stage_mod + if _eval_stage_mod is not None: + return _eval_stage_mod + import sys + + existing = sys.modules.get("nvflow.recipes.finance.stages.rl.evaluate_provenance") + if existing is not None: + _eval_stage_mod = existing + return _eval_stage_mod + spec = importlib.util.spec_from_file_location( + "_pg_eval_stage_private", + Path("nvflow/recipes/finance/stages/rl/evaluate_provenance.py"), + ) + _eval_stage_mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(_eval_stage_mod) + return _eval_stage_mod + + +# --------------------------------------------------------------------------- +# Deterministic fakes +# --------------------------------------------------------------------------- + + +class FakeEmbedder: + """Deterministic embedder using simple character-level hashing.""" + + @property + def dimension(self) -> int: + return 16 + + def embed(self, texts: Sequence[str]) -> Sequence[Sequence[float]]: + vectors = [] + for text in texts: + vec = [0.0] * 16 + for i, ch in enumerate(text[:16]): + vec[i] = ord(ch) / 128.0 + vectors.append(vec) + return vectors + + +class FakeNLIScorer: + """NLI scorer that returns a configurable label, with optional per-claim error.""" + + def __init__(self, label: str = "entailment") -> None: + self._label = label + self._should_error = False + self._error_on_claim: str | None = None + + def set_error(self, should_error: bool = True) -> None: + self._should_error = should_error + + def set_error_on_claim(self, claim_text: str) -> None: + self._error_on_claim = claim_text + + def score(self, *, premise: str, hypothesis: str) -> NLIResult: + if self._should_error: + raise RuntimeError("Fake NLI model error") + if self._error_on_claim and self._error_on_claim in hypothesis: + raise RuntimeError(f"Fake NLI error for claim: {hypothesis}") + return NLIResult( + label=self._label, + score=0.99, + probabilities=( + ("entailment", 0.99 if self._label == "entailment" else 0.01), + ("neutral", 0.99 if self._label == "neutral" else 0.005), + ("contradiction", 0.99 if self._label == "contradiction" else 0.005), + ), + ) + + +def _make_evaluator( + nli_label: str = "entailment", + config: ProvenanceGuardConfig | None = None, +) -> tuple[ProvenanceGuardEvaluator, FakeNLIScorer]: + nli = FakeNLIScorer(label=nli_label) + evaluator = ProvenanceGuardEvaluator( + decomposer=RuleBasedDecomposer(), + router=EmbeddingSourceRouter(FakeEmbedder()), + nli_scorer=nli, + config=config or ProvenanceGuardConfig(), + ) + return evaluator, nli + + +def _make_evidence( + text: str = "Revenue was $1.23 billion in 2024.", + source_id: str | None = "sec:cik=0001811414:accession=0001811414-25-000010:doc=10-K", + attribution_state: str = "available", + source_ids: tuple[str, ...] = (), +) -> EvidenceChunk: + return EvidenceChunk( + chunk_id="ev:test", + text=text, + source_id=source_id, + source_ids=source_ids, + attribution_state=attribution_state, + ) + + +def _make_finance_evaluator( + nli_label: str = "entailment", + config: ProvenanceGuardConfig | None = None, +) -> tuple[ProvenanceGuardEvaluator, FakeNLIScorer, FinanceEvaluatorConfig]: + evaluator, nli = _make_evaluator(nli_label=nli_label, config=config) + finance_config = FinanceEvaluatorConfig( + provenance_config=evaluator.config, + environment="finance_sec_search", + routing_model_revision="abc123", + nli_model_revision="def456", + ) + return evaluator, nli, finance_config + + +def _make_rollout_row( + answer: str = "Revenue was $1.23 billion in 2024.", + evidence_text: str = "Revenue was $1.23 billion in 2024.", +) -> dict: + return { + "uuid": "test-uuid-001", + "response_id": "resp-001", + "_ng_task_index": 0, + "_ng_rollout_index": 0, + "output": [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_1", + "arguments": json.dumps({"prompt": "What was the revenue?"}), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_1", + "output": json.dumps({"result": evidence_text}), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_1", + "arguments": json.dumps({"final_result": answer}), + }, + ], + } + + +# --------------------------------------------------------------------------- +# YAML parse tests +# --------------------------------------------------------------------------- + + +class TestYamlParse: + def test_base_yaml_parses(self): + path = Path("nvflow/recipes/finance/workflows/grpo/base.yaml") + with open(path) as f: + data = yaml.safe_load(f) + assert data is not None + assert data["recipe"] == "finance" + assert data["workflow"]["name"] == "grpo" + + def test_provenanceguard_eval_directory_present(self): + path = Path("nvflow/recipes/finance/workflows/grpo/base.yaml") + with open(path) as f: + data = yaml.safe_load(f) + dirs = data["directories"] + assert "provenanceguard-eval" in dirs + assert dirs["provenanceguard-eval"] == "${model_output_dir}/provenanceguard-eval" + + def test_evaluate_provenance_stage_config_present(self): + path = Path("nvflow/recipes/finance/workflows/grpo/base.yaml") + with open(path) as f: + data = yaml.safe_load(f) + stages = data["stages"] + assert "evaluate_provenance" in stages + ep = stages["evaluate_provenance"] + assert ep["dependencies"] == ["collect_rollouts"] + assert ep["container"] == "nemo-skills" + assert ep["routing_model"] == DEFAULT_EMBEDDING_MODEL + assert ep["nli_model"] == DEFAULT_NLI_MODEL + + def test_evaluate_provenance_commented_in_pipeline_stages(self): + path = Path("nvflow/recipes/finance/workflows/grpo/base.yaml") + text = path.read_text() + assert "# - evaluate_provenance" in text + raw_lines = text.splitlines() + collect_idx = next( + i for i, line in enumerate(raw_lines) if line.strip().startswith("- collect_rollouts") + ) + eval_idx = next(i for i, line in enumerate(raw_lines) if "# - evaluate_provenance" in line) + assert eval_idx == collect_idx + 1 + + +# --------------------------------------------------------------------------- +# Registry discovery tests +# --------------------------------------------------------------------------- + + +class TestRegistryDiscovery: + def test_stage_registered(self): + from nvflow.core import StageRegistry + + stage_cls = StageRegistry.get( + recipe="finance", + workflow="grpo", + stage="evaluate_provenance", + ) + assert stage_cls is not None + assert stage_cls.workflow == "grpo" + + def test_stage_not_in_default_pipeline(self): + path = Path("nvflow/recipes/finance/workflows/grpo/base.yaml") + with open(path) as f: + data = yaml.safe_load(f) + stages = data["pipeline_stages"] + assert "evaluate_provenance" not in stages + + +# --------------------------------------------------------------------------- +# Opt-in ordering tests +# --------------------------------------------------------------------------- + + +class TestOptInOrdering: + def test_evaluate_provenance_after_collect_rollouts(self): + path = Path("nvflow/recipes/finance/workflows/grpo/base.yaml") + text = path.read_text() + lines = text.splitlines() + collect_idx = next( + i for i, line in enumerate(lines) if line.strip().startswith("- collect_rollouts") + ) + eval_idx = next(i for i, line in enumerate(lines) if "# - evaluate_provenance" in line) + assert eval_idx == collect_idx + 1 + + def test_no_downstream_dependency_on_evaluate_provenance(self): + path = Path("nvflow/recipes/finance/workflows/grpo/base.yaml") + with open(path) as f: + data = yaml.safe_load(f) + for stage_name, stage_cfg in data["stages"].items(): + deps = stage_cfg.get("dependencies", []) + if deps: + assert "evaluate_provenance" not in deps, ( + f"Stage '{stage_name}' should not depend on evaluate_provenance" + ) + + +# --------------------------------------------------------------------------- +# Exact paths tests +# --------------------------------------------------------------------------- + + +class TestExactPaths: + def test_input_path_pattern(self): + path = Path("nvflow/recipes/finance/workflows/grpo/base.yaml") + text = path.read_text() + assert "step-5-collect-rollouts}/{env}/rollout/output-rs" in text + assert "rollout/output-rs.jsonl.done" in text + + def test_output_path_pattern(self): + path = Path("nvflow/recipes/finance/workflows/grpo/base.yaml") + text = path.read_text() + assert "provenanceguard-eval}/{env}/provenanceguard-rs" in text + assert "provenanceguard-rs.jsonl.done" in text + + def test_sidecar_filename(self): + assert sidecar_filename(0) == "provenanceguard-rs0.jsonl" + assert sidecar_filename(3) == "provenanceguard-rs3.jsonl" + + def test_merged_filename(self): + assert merged_filename(0) == "output-rs0.jsonl" + assert merged_filename(7) == "output-rs7.jsonl" + + def test_done_marker(self): + assert done_marker(0) == "output-rs0.jsonl.done" + assert done_marker(3) == "output-rs3.jsonl.done" + + +# --------------------------------------------------------------------------- +# Marker gating tests +# --------------------------------------------------------------------------- + + +class TestMarkerGating: + def test_marker_created_after_successful_write(self, tmp_path): + evaluator, _, config = _make_finance_evaluator() + input_file = tmp_path / "rollout" / "output-rs0.jsonl" + input_file.parent.mkdir(parents=True) + input_file.write_text(json.dumps(_make_rollout_row()) + "\n") + Path(str(input_file) + ".done").touch() + output_file = str(tmp_path / "out" / "provenanceguard-rs0.jsonl") + count = evaluate_seed(str(input_file), output_file, 0, evaluator, config) + assert count == 1 + assert os.path.exists(output_file) + assert os.path.exists(output_file + ".done") + + def test_stale_marker_replaced_with_fresh_empty(self, tmp_path): + evaluator, _, config = _make_finance_evaluator() + input_file = tmp_path / "rollout" / "output-rs0.jsonl" + input_file.parent.mkdir(parents=True) + input_file.write_text(json.dumps(_make_rollout_row()) + "\n") + Path(str(input_file) + ".done").touch() + output_file = str(tmp_path / "out" / "provenanceguard-rs0.jsonl") + stale_marker = output_file + ".done" + Path(stale_marker).parent.mkdir(parents=True, exist_ok=True) + Path(stale_marker).write_text("stale") + assert os.path.exists(stale_marker) + evaluate_seed(str(input_file), output_file, 0, evaluator, config) + assert os.path.exists(stale_marker) + assert Path(stale_marker).read_text() == "" + assert os.path.exists(output_file) + + def test_failing_rerun_preserves_output_removes_marker(self, tmp_path): + """If evaluation fails after valid input gate, old output stays + but stale .done marker is absent.""" + evaluator, _, config = _make_finance_evaluator() + input_file = tmp_path / "input.jsonl" + input_file.write_text(json.dumps(_make_rollout_row()) + "\n") + Path(str(input_file) + ".done").touch() + output_file = str(tmp_path / "output.jsonl") + output_file_path = Path(output_file) + output_file_path.write_text("prior output content") + done_path = output_file + ".done" + Path(done_path).write_text("stale marker") + + # Force evaluation failure by using a broken evaluator. + class _BrokenEvaluator: + def evaluate(self, answer, evidence): + raise RuntimeError("evaluation crashed") + + with pytest.raises(RuntimeError, match="evaluation crashed"): + evaluate_seed(str(input_file), output_file, 0, _BrokenEvaluator(), config) + + assert output_file_path.read_text() == "prior output content" + assert not os.path.exists(done_path) + + def test_missing_input_marker_preserves_output_and_marker(self, tmp_path): + """Missing input .done preserves old output and marker bytes.""" + evaluator, _, config = _make_finance_evaluator() + input_file = tmp_path / "input.jsonl" + input_file.write_text(json.dumps(_make_rollout_row()) + "\n") + # NOTE: no .done marker for input + output_file = tmp_path / "output.jsonl" + output_file.write_text("prior output") + done_path = str(output_file) + ".done" + Path(done_path).write_text("prior done bytes") + with pytest.raises(RuntimeError, match="Input completion marker"): + evaluate_seed(str(input_file), str(output_file), 0, evaluator, config) + assert output_file.read_text() == "prior output" + assert Path(done_path).read_text() == "prior done bytes" + + def test_successful_rerun_replaces_stale_marker_fresh_empty(self, tmp_path): + """Successful valid rerun replaces stale marker with fresh empty marker.""" + evaluator, _, config = _make_finance_evaluator() + input_file = tmp_path / "input.jsonl" + input_file.write_text(json.dumps(_make_rollout_row()) + "\n") + Path(str(input_file) + ".done").touch() + output_file = str(tmp_path / "output.jsonl") + done_path = output_file + ".done" + Path(done_path).write_text("stale content") + evaluate_seed(str(input_file), output_file, 0, evaluator, config) + assert os.path.exists(output_file) + assert os.path.exists(done_path) + assert Path(done_path).read_text() == "" + + +# --------------------------------------------------------------------------- +# Environment/model/revision/policy propagation tests +# --------------------------------------------------------------------------- + + +class TestPropagation: + def test_environment_propagated(self): + row = _make_rollout_row() + raw_line = json.dumps(row) + evaluator, _, config = _make_finance_evaluator() + config = FinanceEvaluatorConfig( + provenance_config=evaluator.config, + environment="my_custom_env", + ) + result = evaluate_row(raw_line, row, 0, evaluator, config) + assert result["environment"] == "my_custom_env" + + def test_model_ids_propagated(self): + row = _make_rollout_row() + raw_line = json.dumps(row) + evaluator, _, config = _make_finance_evaluator() + result = evaluate_row(raw_line, row, 0, evaluator, config) + assert result["models"]["routing_model"] == DEFAULT_EMBEDDING_MODEL + assert result["models"]["nli_model"] == DEFAULT_NLI_MODEL + + def test_revisions_propagated(self): + row = _make_rollout_row() + raw_line = json.dumps(row) + evaluator, _, config = _make_finance_evaluator() + result = evaluate_row(raw_line, row, 0, evaluator, config) + assert result["models"]["routing_model_revision"] == "abc123" + assert result["models"]["nli_model_revision"] == "def456" + + def test_nullable_revisions(self): + row = _make_rollout_row() + raw_line = json.dumps(row) + evaluator, _, _ = _make_finance_evaluator() + config = FinanceEvaluatorConfig( + provenance_config=evaluator.config, + routing_model_revision=None, + nli_model_revision=None, + ) + result = evaluate_row(raw_line, row, 0, evaluator, config) + assert result["models"]["routing_model_revision"] is None + assert result["models"]["nli_model_revision"] is None + + def test_policy_config_propagated(self): + row = _make_rollout_row() + raw_line = json.dumps(row) + pg_config = ProvenanceGuardConfig( + evidence_excerpt_length=200, + ) + evaluator, _, config = _make_finance_evaluator(config=pg_config) + result = evaluate_row(raw_line, row, 0, evaluator, config) + assert result["thresholds"]["policy"] == "fixed_fail_closed" + assert "block_on_contradiction" not in result["thresholds"] + assert "block_on_neutral" not in result["thresholds"] + assert "block_on_no_source" not in result["thresholds"] + assert "block_on_protected_value_mismatch" not in result["thresholds"] + assert result["thresholds"]["evidence_excerpt_length"] == 200 + + +# --------------------------------------------------------------------------- +# Every line / malformed / unchanged input tests +# --------------------------------------------------------------------------- + + +class TestEveryLine: + def test_every_line_produces_sidecar_row(self, tmp_path): + evaluator, _, config = _make_finance_evaluator() + input_file = tmp_path / "input.jsonl" + input_file.write_text( + json.dumps(_make_rollout_row()) + "\n" + "\n" + "not valid json\n" + "[1, 2, 3]\n" + ) + Path(str(input_file) + ".done").touch() + output_file = str(tmp_path / "output.jsonl") + count = evaluate_seed(str(input_file), output_file, 0, evaluator, config) + assert count == 4 + + def test_blank_line_produces_sidecar_row(self): + sidecar = parse_rollout_line("") + assert sidecar.parse_error is not None + assert sidecar.parse_error == "empty line" + assert sidecar.raw_line_fingerprint + + def test_malformed_json_produces_sidecar_row(self): + raw = '{"broken": json' + sidecar = parse_rollout_line(raw) + assert sidecar.parse_error is not None + assert "JSONDecodeError" in sidecar.parse_error + assert sidecar.raw_line_fingerprint + + def test_non_object_json_produces_sidecar_row(self): + raw = "[1, 2, 3]" + sidecar = parse_rollout_line(raw) + assert sidecar.parse_error is not None + assert "expected JSON object" in sidecar.parse_error + + +class TestUnchangedInput: + def test_input_file_not_modified(self, tmp_path): + evaluator, _, config = _make_finance_evaluator() + row = _make_rollout_row() + raw_line = json.dumps(row) + "\n" + input_file = tmp_path / "input.jsonl" + input_file.write_text(raw_line) + original = input_file.read_bytes() + Path(str(input_file) + ".done").touch() + output_file = str(tmp_path / "output.jsonl") + evaluate_seed(str(input_file), output_file, 0, evaluator, config) + assert input_file.read_bytes() == original + + +# --------------------------------------------------------------------------- +# Determinism tests +# --------------------------------------------------------------------------- + + +class TestDeterminism: + def test_repeat_byte_determinism(self, tmp_path): + evaluator, _, config = _make_finance_evaluator() + row = _make_rollout_row() + raw_line = json.dumps(row) + "\n" + input_file = tmp_path / "input.jsonl" + input_file.write_text(raw_line * 3) + Path(str(input_file) + ".done").touch() + out1 = str(tmp_path / "out1.jsonl") + out2 = str(tmp_path / "out2.jsonl") + evaluate_seed(str(input_file), out1, 0, evaluator, config) + evaluate_seed(str(input_file), out2, 0, evaluator, config) + assert Path(out1).read_bytes() == Path(out2).read_bytes() + + def test_no_random_uuid(self): + row = _make_rollout_row() + raw_line = json.dumps(row) + evaluator, _, config = _make_finance_evaluator() + result1 = evaluate_row(raw_line, row, 0, evaluator, config, line_number=0) + result2 = evaluate_row(raw_line, row, 0, evaluator, config, line_number=0) + assert result1["evaluation_uuid"] == result2["evaluation_uuid"] + config_digest = hashlib.sha256( + json.dumps(config.to_dict(), sort_keys=True, default=str).encode("utf-8") + ).hexdigest()[:16] + expected = hashlib.sha256( + f"eval:{result1['raw_line_fingerprint']}:0:0:{ALGORITHM_VERSION}:{config_digest}".encode() + ).hexdigest() + assert result1["evaluation_uuid"] == expected + + +# --------------------------------------------------------------------------- +# Marker atomicity tests +# --------------------------------------------------------------------------- + + +class TestMarkerAtomicity: + def test_no_marker_on_error(self, tmp_path): + evaluator, _, config = _make_finance_evaluator() + input_file = tmp_path / "nonexistent.jsonl" + output_file = str(tmp_path / "out.jsonl") + with pytest.raises(FileNotFoundError): + evaluate_seed(str(input_file), output_file, 0, evaluator, config) + assert not os.path.exists(output_file) + assert not os.path.exists(output_file + ".done") + + def test_temp_file_cleaned_on_error(self, tmp_path): + evaluator, _, config = _make_finance_evaluator() + output_file = str(tmp_path / "out.jsonl") + done_path = output_file + ".done" + Path(done_path).write_text("stale") + with pytest.raises(FileNotFoundError): + evaluate_seed(str(tmp_path / "nonexistent.jsonl"), output_file, 0, evaluator, config) + temp_files = list(tmp_path.glob(".pg_tmp_*")) + assert len(temp_files) == 0 + + +# --------------------------------------------------------------------------- +# Mixed contradiction-or-neutral + error => unavailable +# --------------------------------------------------------------------------- + + +class TestMixedVerdicts: + def test_contradiction_plus_error_is_unavailable(self): + evidence = [ + _make_evidence(text="Revenue was $1.23 billion in 2024."), + _make_evidence(text="The sky is blue."), + ] + evaluator, nli = _make_evaluator(nli_label="entailment") + answer = "Revenue was $1.23 billion in 2024. The sky is blue today." + nli.set_error_on_claim("The sky is blue today") + decision = evaluator.evaluate(answer, evidence) + assert decision.status == "unavailable" + assert decision.reason == "partial_nli_errors" + + def test_neutral_plus_error_is_unavailable(self): + evidence = [_make_evidence(text="Some unrelated text.")] + evaluator, nli = _make_evaluator(nli_label="neutral") + answer = "Revenue was $1.23 billion in 2024." + nli.set_error_on_claim("Revenue") + decision = evaluator.evaluate(answer, evidence) + assert decision.status == "unavailable" + + +# --------------------------------------------------------------------------- +# Known + unknown key tests +# --------------------------------------------------------------------------- + + +class TestKnownUnknownKey: + def test_multiple_keys_yields_unavailable(self): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_1", + "arguments": json.dumps( + { + "url": "https://www.sec.gov/Archives/edgar/data/1811414/000181141425000010/10-K.htm", + "key": "filing_10k", + } + ), + }, + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_2", + "arguments": json.dumps( + { + "url": "https://www.sec.gov/Archives/edgar/data/1811414/000181141425000020/10-K.htm", + "key": "filing_10k_2", + } + ), + }, + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_1", + "arguments": json.dumps({"prompt": "Search {{filing_10k}} and {{filing_10k_2}}"}), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_1", + "output": json.dumps({"result": "Some text"}), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_1", + "arguments": json.dumps({"final_result": "Some answer."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 1 + chunk = sidecar.trace.evidence[0] + assert chunk.source_id is None + assert chunk.attribution_state == "unavailable" + assert len(chunk.source_ids) == 2 + + def test_unknown_key_yields_unavailable(self): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_1", + "arguments": json.dumps({"prompt": "Search {{unknown_key}}"}), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_1", + "output": json.dumps({"result": "Some text"}), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_1", + "arguments": json.dumps({"final_result": "Some answer."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 1 + chunk = sidecar.trace.evidence[0] + assert chunk.source_id is None + assert chunk.attribution_state == "unavailable" + + +# --------------------------------------------------------------------------- +# Zero-pad URL CIK tests +# --------------------------------------------------------------------------- + + +class TestZeroPadCik: + def test_url_cik_zero_padded(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _extract_filing_from_url, + ) + + url = "https://www.sec.gov/Archives/edgar/data/1811414/000181141425000010/10-K.htm" + filing = _extract_filing_from_url(url) + assert filing.cik == "0001811414" + assert len(filing.cik) == 10 + sid = filing.source_id() + assert sid is not None + assert "cik=0001811414" in sid + + def test_already_padded_cik_preserved(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _extract_filing_from_url, + ) + + url = "https://www.sec.gov/Archives/edgar/data/0001811414/000181141425000010/10-K.htm" + filing = _extract_filing_from_url(url) + assert filing.cik == "0001811414" + assert len(filing.cik) == 10 + + +# --------------------------------------------------------------------------- +# Real nested response.output trace tests +# --------------------------------------------------------------------------- + + +class TestNestedResponseOutput: + def test_nested_response_output_trace(self): + row = { + "uuid": "test-nested-001", + "response": { + "id": "resp-nested-001", + "output": [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_nested", + "arguments": json.dumps({"prompt": "What was the revenue?"}), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_nested", + "output": json.dumps({"result": "Revenue was $1.23 billion in 2024."}), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_nested", + "arguments": json.dumps( + {"final_result": "Revenue was $1.23 billion in 2024."} + ), + }, + ], + }, + } + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert sidecar.trace.has_submit is True + assert sidecar.trace.answer == "Revenue was $1.23 billion in 2024." + assert len(sidecar.trace.evidence) == 1 + assert sidecar.trace.evidence[0].text == "Revenue was $1.23 billion in 2024." + + def test_flat_output_trace(self): + row = _make_rollout_row() + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert sidecar.trace.has_submit is True + assert sidecar.trace.answer == "Revenue was $1.23 billion in 2024." + assert len(sidecar.trace.evidence) == 1 + + +# --------------------------------------------------------------------------- +# NLI label validation tests +# --------------------------------------------------------------------------- + + +class TestNLILabelValidation: + def test_entailment_label_accepted(self): + assert _normalize_nli_label("entailment") == "entailment" + assert _normalize_nli_label("ENTAILMENT") == "entailment" + + def test_neutral_label_accepted(self): + assert _normalize_nli_label("neutral") == "neutral" + assert _normalize_nli_label("Neutral") == "neutral" + + def test_contradiction_label_accepted(self): + assert _normalize_nli_label("contradiction") == "contradiction" + + def test_label_0_rejected(self): + with pytest.raises(ValueError, match="Refusing to assume"): + _normalize_nli_label("label_0") + + def test_label_1_rejected(self): + with pytest.raises(ValueError, match="Refusing to assume"): + _normalize_nli_label("label_1") + + def test_label_2_rejected(self): + with pytest.raises(ValueError, match="Refusing to assume"): + _normalize_nli_label("label_2") + + def test_arbitrary_label_rejected(self): + with pytest.raises(ValueError, match="Unknown NLI label"): + _normalize_nli_label("some_random_label") + + +# --------------------------------------------------------------------------- +# Evaluator decision semantics tests +# --------------------------------------------------------------------------- + + +class TestEvaluatorSemantics: + def test_no_evidence_returns_unavailable(self): + evaluator, _ = _make_evaluator() + decision = evaluator.evaluate("Some answer.", []) + assert decision.status == "unavailable" + assert decision.reason == "no_evidence" + + def test_no_claims_returns_unavailable(self): + evaluator, _ = _make_evaluator() + decision = evaluator.evaluate("", [_make_evidence()]) + assert decision.status == "unavailable" + assert decision.reason == "no_claims_extracted" + + def test_all_entailed_returns_allow(self): + evaluator, _ = _make_evaluator(nli_label="entailment") + evidence = [_make_evidence(text="Revenue was $1.23 billion in 2024.")] + decision = evaluator.evaluate("Revenue was $1.23 billion in 2024.", evidence) + assert decision.status == "allow" + assert decision.reason == "all_entailed" + + def test_contradiction_returns_block(self): + evaluator, _ = _make_evaluator(nli_label="contradiction") + evidence = [_make_evidence(text="Revenue was $2.00 billion in 2024.")] + decision = evaluator.evaluate("Revenue was $1.23 billion in 2024.", evidence) + assert decision.status == "block" + assert decision.reason == "contradiction" + + def test_neutral_returns_block(self): + evaluator, _ = _make_evaluator(nli_label="neutral") + evidence = [_make_evidence(text="Some unrelated text here.")] + decision = evaluator.evaluate("Revenue was $1.23 billion in 2024.", evidence) + assert decision.status == "block" + assert decision.reason == "neutral" + + def test_protected_value_mismatch_returns_block(self): + evaluator, _ = _make_evaluator(nli_label="entailment") + evidence = [_make_evidence(text="Revenue was $2.00 billion in 2024.")] + decision = evaluator.evaluate("Revenue was $1.23 billion in 2024.", evidence) + assert decision.status == "block" + assert decision.reason == "protected_value_mismatch" + + def test_nli_error_returns_unavailable(self): + evaluator, nli = _make_evaluator() + nli.set_error(True) + evidence = [_make_evidence(text="Some text.")] + decision = evaluator.evaluate("Revenue was $1.23 billion.", evidence) + assert decision.status == "unavailable" + + +# --------------------------------------------------------------------------- +# Protected values: no year-only date match +# --------------------------------------------------------------------------- + + +class TestProtectedValuesNoYearOnly: + def test_year_only_does_not_match(self): + outcome, missing = check_protected_values( + "The fiscal year ended December 31, 2024.", + "In 2024 the company grew.", + ) + assert outcome == "fail" + assert len(missing) > 0 + + def test_full_date_matches(self): + outcome, _ = check_protected_values( + "The fiscal year ended December 31, 2024.", + "The fiscal year ended December 31, 2024.", + ) + assert outcome == "pass" + + +# --------------------------------------------------------------------------- +# Sidecar output schema tests +# --------------------------------------------------------------------------- + + +class TestSidecarSchema: + def test_output_contains_all_required_fields(self): + row = _make_rollout_row() + raw_line = json.dumps(row) + evaluator, _, config = _make_finance_evaluator() + result = evaluate_row(raw_line, row, 0, evaluator, config) + required_top = [ + "schema_version", + "algorithm", + "evidence_basis", + "limitation", + "environment", + "seed", + "line_number", + "uuid", + "_ng_task_index", + "_ng_rollout_index", + "response_id", + "fingerprint", + "raw_line_fingerprint", + "evaluation_uuid", + "thresholds", + "models", + "answer", + "has_submit", + "submit_call_id", + "evidence", + "extraction_errors", + "verdict", + ] + for key in required_top: + assert key in result, f"Missing required field: {key}" + + def test_output_schema_and_algorithm(self): + row = _make_rollout_row() + raw_line = json.dumps(row) + evaluator, _, config = _make_finance_evaluator() + result = evaluate_row(raw_line, row, 0, evaluator, config) + assert result["schema_version"] == SCHEMA_VERSION + assert result["algorithm"] == ALGORITHM_VERSION + assert result["evidence_basis"] == EVIDENCE_BASIS + assert LIMITATION in result["limitation"] + + def test_parse_error_row_has_unavailable_verdict(self): + raw = '{"broken": json' + evaluator, _, config = _make_finance_evaluator() + result = evaluate_row(raw, {}, 0, evaluator, config) + assert result["verdict"]["status"] == "unavailable" + assert result["parse_error"] is not None + assert result["evidence"] == [] + + def test_no_submit_row_has_unavailable_verdict(self): + row = _make_rollout_row() + row["output"] = [i for i in row["output"] if i.get("name") != "submit_final_result"] + raw_line = json.dumps(row) + evaluator, _, config = _make_finance_evaluator() + result = evaluate_row(raw_line, row, 0, evaluator, config) + assert result["verdict"]["status"] == "unavailable" + assert result["verdict"]["reason"] == "no_submit_final_result" + + def test_one_chunk_per_retrieve_result(self): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r1", + "arguments": json.dumps({"prompt": "query1"}), + }, + { + "type": "function_call_output", + "call_id": "call_r1", + "output": json.dumps({"result": "Evidence one."}), + }, + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r2", + "arguments": json.dumps({"prompt": "query2"}), + }, + { + "type": "function_call_output", + "call_id": "call_r2", + "output": json.dumps({"result": "Evidence two."}), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_s1", + "arguments": json.dumps({"final_result": "Some answer."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 2 + assert sidecar.trace.evidence[0].text == "Evidence one." + assert sidecar.trace.evidence[1].text == "Evidence two." + + def test_last_submit_final_result_used(self): + row = _make_rollout_row() + row["output"].append( + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_2", + "arguments": json.dumps({"final_result": "The later answer."}), + } + ) + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert sidecar.trace.answer == "The later answer." + assert sidecar.trace.submit_call_id == "call_submit_2" + + +# --------------------------------------------------------------------------- +# Realistic nested response.output trace tests +# --------------------------------------------------------------------------- + +_FILING_URL = "https://www.sec.gov/Archives/edgar/data/1811414/000181141425000010/10-K.htm" + + +def _make_realistic_trace() -> dict: + """A realistic nested response.output trace with all 4 finance tool types.""" + return { + "uuid": "real-trace-001", + "response_id": "resp-real-001", + "_ng_task_index": 0, + "_ng_rollout_index": 0, + "output": [ + { + "type": "function_call", + "name": "sec_filing_search", + "call_id": "call_sec_1", + "arguments": json.dumps({"query": "Apple 10-K", "form_type": "10-K"}), + }, + { + "type": "function_call_output", + "call_id": "call_sec_1", + "output": json.dumps( + { + "filings": [ + { + "cik": "1811414", + "accessionNo": "0001811414-25-000010", + "primaryDocument": "10-K.htm", + "linkToHtml": _FILING_URL, + } + ] + } + ), + }, + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_1", + "arguments": json.dumps({"url": _FILING_URL, "key": "filing_10k"}), + }, + { + "type": "function_call_output", + "call_id": "call_parse_1", + "output": json.dumps({"result": "Page stored successfully"}), + }, + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_1", + "arguments": json.dumps({"prompt": "What was the revenue? {{filing_10k}}"}), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_1", + "output": json.dumps({"result": "Revenue was $1.23 billion in 2024."}), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_1", + "arguments": json.dumps({"final_result": "Revenue was $1.23 billion in 2024."}), + }, + ], + } + + +class TestRealisticTrace: + def test_realistic_trace_extraction(self): + row = _make_realistic_trace() + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert sidecar.trace.has_submit is True + assert sidecar.trace.answer == "Revenue was $1.23 billion in 2024." + assert len(sidecar.trace.evidence) == 1 + chunk = sidecar.trace.evidence[0] + assert chunk.text == "Revenue was $1.23 billion in 2024." + assert chunk.source_id is not None + assert chunk.attribution_state == "available" + assert "cik=0001811414" in chunk.source_id + assert "accession=000181141425000010" in chunk.source_id + assert "doc=10-K.htm" in chunk.source_id + + def test_realistic_trace_evaluates_allow(self): + row = _make_realistic_trace() + raw_line = json.dumps(row) + evaluator, _, config = _make_finance_evaluator(nli_label="entailment") + result = evaluate_row(raw_line, row, 0, evaluator, config) + assert result["verdict"]["status"] == "allow" + assert result["verdict"]["reason"] == "all_entailed" + + +# --------------------------------------------------------------------------- +# parse_html_page plain key tests +# --------------------------------------------------------------------------- + + +class TestParsePlainKey: + def test_parse_html_page_reads_plain_key(self): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_1", + "arguments": json.dumps({"url": _FILING_URL, "key": "filing_10k"}), + }, + { + "type": "function_call_output", + "call_id": "call_parse_1", + "output": json.dumps({"result": "Page stored"}), + }, + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_1", + "arguments": json.dumps({"prompt": "Revenue {{filing_10k}}"}), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_1", + "output": json.dumps({"result": "Revenue was $1.23 billion in 2024."}), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_1", + "arguments": json.dumps({"final_result": "Revenue was $1.23 billion in 2024."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 1 + chunk = sidecar.trace.evidence[0] + assert chunk.source_id is not None + assert chunk.attribution_state == "available" + + +# --------------------------------------------------------------------------- +# Error exclusion tests +# --------------------------------------------------------------------------- + + +class TestErrorExclusion: + def test_raw_error_string_excluded(self): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r1", + "arguments": json.dumps({"prompt": "query"}), + }, + { + "type": "function_call_output", + "call_id": "call_r1", + "output": "Error: Failed to retrieve data", + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_s1", + "arguments": json.dumps({"final_result": "Some answer."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 0 + assert any("error" in e.lower() for e in sidecar.trace.extraction_errors) + + def test_traceback_excluded(self): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r1", + "arguments": json.dumps({"prompt": "query"}), + }, + { + "type": "function_call_output", + "call_id": "call_r1", + "output": "Traceback (most recent call last):\n File ...", + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_s1", + "arguments": json.dumps({"final_result": "Some answer."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 0 + + def test_unavailable_string_excluded(self): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r1", + "arguments": json.dumps({"prompt": "query"}), + }, + { + "type": "function_call_output", + "call_id": "call_r1", + "output": "unavailable", + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_s1", + "arguments": json.dumps({"final_result": "Some answer."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 0 + + def test_dict_error_result_excluded(self): + """A dict-shaped failed retrieval (success=False, error) must be excluded. + + This test does not skip. It proves that structured tool results + are preserved as Any (not coerced to str) so error dicts with + ``success`` false / ``error`` cannot become evidence. + """ + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r1", + "arguments": json.dumps({"prompt": "query"}), + }, + { + "type": "function_call_output", + "call_id": "call_r1", + "output": {"success": False, "error": "retrieval failed"}, + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_s1", + "arguments": json.dumps({"final_result": "Some answer."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 0 + assert any("error" in e.lower() for e in sidecar.trace.extraction_errors) + + raw_line = json.dumps(row) + evaluator, _, config = _make_finance_evaluator(nli_label="entailment") + result = evaluate_row(raw_line, row, 0, evaluator, config) + assert result["verdict"]["status"] == "unavailable" + assert result["verdict"]["reason"] == "trace_extraction_errors" + assert len(result["evidence"]) == 0 + + +# --------------------------------------------------------------------------- +# Normalized accession tests +# --------------------------------------------------------------------------- + + +class TestNormalizedAccession: + def test_accession_dashes_stripped(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _format_accession, + ) + + assert _format_accession("0001811414-25-000010") == "000181141425000010" + + def test_accession_already_digits(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _format_accession, + ) + + assert _format_accession("000181141425000010") == "000181141425000010" + + def test_accession_short_returns_none(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _format_accession, + ) + + assert _format_accession("123") is None + + def test_accession_overlong_returns_none(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _format_accession, + ) + + assert _format_accession("000181141425000010999") is None + assert _format_accession("0001811414-25-000010-999") is None + + def test_source_id_uses_digit_only_accession(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _extract_filing_from_url, + ) + + filing = _extract_filing_from_url(_FILING_URL) + assert filing.accession == "000181141425000010" + sid = filing.source_id() + assert sid is not None + assert "accession=000181141425000010" in sid + + +# --------------------------------------------------------------------------- +# Percentage extraction tests +# --------------------------------------------------------------------------- + + +class TestPercentageExtraction: + def test_percentage_at_end_of_string(self): + result = extract_protected_values("Revenue grew 12.3%") + assert len(result) == 1 + assert result[0].kind == "percentage" + assert "12.3%" in result[0].normalized + + def test_percentage_before_word(self): + result = extract_protected_values("Growth was 12.3% YoY") + assert len(result) == 1 + assert result[0].kind == "percentage" + + def test_percentage_protected_value_check_pass(self): + outcome, _ = check_protected_values("Margin was 12.3%", "Margin was 12.3%") + assert outcome == "pass" + + def test_percentage_protected_value_check_fail(self): + outcome, missing = check_protected_values("Margin was 12.3%", "Margin was 15.0%") + assert outcome == "fail" + assert len(missing) == 1 + + +# --------------------------------------------------------------------------- +# Unattributable chunk blocks no_source without NLI +# --------------------------------------------------------------------------- + + +class TestUnattributableChunk: + def test_unavailable_chunk_blocks_no_source_without_nli(self): + evidence = [ + _make_evidence( + text="Some text about revenue.", + source_id=None, + attribution_state="unavailable", + ), + ] + evaluator, nli = _make_evaluator(nli_label="entailment") + nli.set_error(True) + decision = evaluator.evaluate("Revenue was $1.23 billion.", evidence) + assert decision.status == "block" + assert decision.reason == "no_source" + assert len(decision.verdicts) == 1 + assert decision.verdicts[0].final_label == "no_source" + assert decision.verdicts[0].errors == () + + def test_composite_chunk_blocks_no_source(self): + evidence = [ + _make_evidence( + text="Some text.", + source_id=None, + attribution_state="composite", + ), + ] + evaluator, nli = _make_evaluator(nli_label="entailment") + nli.set_error(True) + decision = evaluator.evaluate("Revenue was $1.23 billion.", evidence) + assert decision.status == "block" + assert decision.reason == "no_source" + + +# --------------------------------------------------------------------------- +# Top trace failure: extraction errors force unavailable preserving verdicts +# --------------------------------------------------------------------------- + + +class TestTopTraceFailure: + def test_extraction_errors_force_unavailable_preserving_verdicts(self): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r1", + "arguments": json.dumps({"prompt": "query1"}), + }, + { + "type": "function_call_output", + "call_id": "call_r1", + "output": json.dumps({"result": "Revenue was $1.23 billion in 2024."}), + }, + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r2", + "arguments": json.dumps({"prompt": "query2"}), + }, + { + "type": "function_call_output", + "call_id": "call_r2", + "output": "Error: Failed to retrieve data", + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_s1", + "arguments": json.dumps({"final_result": "Revenue was $1.23 billion in 2024."}), + }, + ] + raw_line = json.dumps(row) + evaluator, _, config = _make_finance_evaluator(nli_label="entailment") + result = evaluate_row(raw_line, row, 0, evaluator, config) + assert result["verdict"]["status"] == "unavailable" + assert result["verdict"]["reason"] == "trace_extraction_errors" + assert len(result["verdict"]["verdicts"]) > 0 + + +# --------------------------------------------------------------------------- +# Model adapter tests +# --------------------------------------------------------------------------- + + +class TestModelAdapters: + def test_embedder_lazy_loading(self): + from nvflow.provenanceguard.embedder import HFEmbedder + + emb = HFEmbedder() + assert emb._model is None + assert emb._tokenizer is None + + def test_nli_lazy_loading(self): + from nvflow.provenanceguard.nli import HFNLI + + nli = HFNLI() + assert nli._model is None + assert nli._tokenizer is None + assert nli._label_map is None + + def test_no_model_imports_at_import_time(self): + import importlib + + mod = importlib.import_module("nvflow.provenanceguard.embedder") + assert mod is not None + mod2 = importlib.import_module("nvflow.provenanceguard.nli") + assert mod2 is not None + + def test_embed_returns_plain_floats(self): + """HFEmbedder.embed must return list[list[float]], not tensor scalars. + + Uses a testable subclass with fake tokenizer/model to exercise the + full embed() pipeline (mean pooling, L2 norm, return) without + loading any real model. The assertion rejects scalar tensors + (which the old list(v) approach produced) and proves every + element is a plain Python float. + """ + import torch + + from nvflow.provenanceguard.embedder import HFEmbedder + + class _FakeTokenizer: + def __call__(self, texts, **kwargs): + batch = len(texts) + seq_len = 3 + return { + "input_ids": torch.zeros(batch, seq_len, dtype=torch.long), + "attention_mask": torch.ones(batch, seq_len, dtype=torch.long), + } + + class _FakeModel: + class Config: + hidden_size = 4 + + config = Config() + + def eval(self): + return self + + def __call__(self, **kwargs): + class _Output: + last_hidden_state = torch.randn(2, 3, 4) + + return _Output() + + class _TestableEmbedder(HFEmbedder): + def _ensure_loaded(self): + if self._model is not None: + return + self._dim = 4 + self._tokenizer = _FakeTokenizer() + self._model = _FakeModel() + + emb = _TestableEmbedder() + result = emb.embed(["hello", "world"]) + + assert isinstance(result, list) + assert len(result) == 2 + for vec in result: + assert isinstance(vec, list) + for val in vec: + assert isinstance(val, float), f"Expected Python float, got {type(val).__name__}" + assert not hasattr(val, "item"), ( + "Value is a tensor scalar, not a plain Python float" + ) + + # Sanity: the old list(v) approach would produce tensor scalars. + old_result = [list(v) for v in torch.randn(2, 3)] + assert hasattr(old_result[0][0], "item"), ( + "Sanity check: old list(v) approach should produce tensor scalars" + ) + + +# --------------------------------------------------------------------------- +# Seed interpolation tests +# --------------------------------------------------------------------------- + + +class TestSeedInterpolation: + def test_omegaconf_seed_interpolation(self): + p = Path("nvflow/recipes/finance/workflows/grpo/base.yaml") + with open(p) as f: + data = yaml.safe_load(f) + eval_stage = data["stages"]["evaluate_provenance"] + assert eval_stage["starting_seed"] == "${stages.collect_rollouts.rollout.starting_seed}" + assert ( + eval_stage["num_random_seeds"] == "${stages.collect_rollouts.rollout.num_random_seeds}" + ) + + def test_evaluate_provenance_has_num_gpus_zero(self): + p = Path("nvflow/recipes/finance/workflows/grpo/base.yaml") + with open(p) as f: + data = yaml.safe_load(f) + assert data["stages"]["evaluate_provenance"]["num_gpus"] == 0 + + +# --------------------------------------------------------------------------- +# Build command tests +# --------------------------------------------------------------------------- + + +class TestBuildCommand: + def test_build_command_contains_module_and_flags(self): + mod = _load_eval_stage_module() + build_evaluate_command = mod.build_evaluate_command + cmd = build_evaluate_command( + input_file="/input/r.jsonl", + output_file="/output/p.jsonl", + seed=3, + environment="finance_sec_search", + ) + assert "python3 -m" in cmd + assert "nvflow.recipes.finance.utils.rl.provenanceguard" in cmd + assert "--input_file" in cmd + assert "--output_file" in cmd + assert "--seed" in cmd + assert "--environment" in cmd + + def test_build_command_is_pure(self): + mod = _load_eval_stage_module() + build_evaluate_command = mod.build_evaluate_command + cmd1 = build_evaluate_command( + input_file="/in.jsonl", + output_file="/out.jsonl", + seed=0, + environment="test", + ) + cmd2 = build_evaluate_command( + input_file="/in.jsonl", + output_file="/out.jsonl", + seed=0, + environment="test", + ) + assert cmd1 == cmd2 + + +# --------------------------------------------------------------------------- +# Path overlap tests +# --------------------------------------------------------------------------- + + +class TestPathOverlap: + def _stage(self): + mod = _load_eval_stage_module() + return mod.EvaluateProvenanceStage() + + def test_equal_paths_rejected(self, tmp_path): + stage = self._stage() + p = str(tmp_path / "same") + with pytest.raises(ValueError, match="same path"): + stage.validate_config({"rollouts_dir": p, "output_dir": p}) + + def test_output_inside_rollouts_rejected(self, tmp_path): + stage = self._stage() + with pytest.raises(ValueError, match="inside rollouts_dir"): + stage.validate_config( + { + "rollouts_dir": str(tmp_path), + "output_dir": str(tmp_path / "sub"), + } + ) + + def test_rollouts_inside_output_rejected(self, tmp_path): + stage = self._stage() + with pytest.raises(ValueError, match="inside output_dir"): + stage.validate_config( + { + "rollouts_dir": str(tmp_path / "sub"), + "output_dir": str(tmp_path), + } + ) + + def test_non_overlapping_paths_accepted(self, tmp_path): + stage = self._stage() + stage.validate_config( + { + "rollouts_dir": str(tmp_path / "rollouts"), + "output_dir": str(tmp_path / "output"), + } + ) + + +# --------------------------------------------------------------------------- +# Missing input marker tests +# --------------------------------------------------------------------------- + + +class TestMissingInputMarker: + def test_missing_input_marker_preserves_output(self, tmp_path): + evaluator, _, config = _make_finance_evaluator() + input_file = tmp_path / "input.jsonl" + input_file.write_text(json.dumps(_make_rollout_row()) + "\n") + output_file = tmp_path / "output.jsonl" + output_file.write_text("prior output") + done_path = str(output_file) + ".done" + Path(done_path).write_text("prior done") + with pytest.raises(RuntimeError, match="Input completion marker"): + evaluate_seed(str(input_file), str(output_file), 0, evaluator, config) + assert output_file.read_text() == "prior output" + assert Path(done_path).read_text() == "prior done" + + def test_missing_input_file_raises(self, tmp_path): + evaluator, _, config = _make_finance_evaluator() + output_file = str(tmp_path / "output.jsonl") + with pytest.raises(FileNotFoundError, match="Input file does not exist"): + evaluate_seed( + str(tmp_path / "nonexistent.jsonl"), + output_file, + 0, + evaluator, + config, + ) + + +# --------------------------------------------------------------------------- +# Fixed policy tests +# --------------------------------------------------------------------------- + + +class TestFixedPolicy: + def test_config_has_no_block_on_fields(self): + cfg = ProvenanceGuardConfig() + assert not hasattr(cfg, "block_on_contradiction") + assert not hasattr(cfg, "block_on_neutral") + assert not hasattr(cfg, "block_on_no_source") + assert not hasattr(cfg, "block_on_protected_value_mismatch") + + def test_config_rejects_block_on_kwargs(self): + with pytest.raises(TypeError): + ProvenanceGuardConfig(block_on_contradiction=False) # type: ignore[call-arg] + + def test_to_dict_has_no_block_on(self): + result = FinanceEvaluatorConfig().to_dict() + thresholds = result["thresholds"] + assert "block_on_contradiction" not in thresholds + assert "block_on_neutral" not in thresholds + assert "block_on_no_source" not in thresholds + assert "block_on_protected_value_mismatch" not in thresholds + assert thresholds["policy"] == "fixed_fail_closed" + + +# --------------------------------------------------------------------------- +# Deterministic per-line IDs for duplicate/invalid lines +# --------------------------------------------------------------------------- + + +class TestDuplicateLineIDs: + def test_duplicate_lines_get_different_uuids(self, tmp_path): + evaluator, _, config = _make_finance_evaluator() + row = _make_rollout_row() + raw_line = json.dumps(row) + input_file = tmp_path / "input.jsonl" + input_file.write_text(raw_line + "\n" + raw_line + "\n") + Path(str(input_file) + ".done").touch() + output_file = str(tmp_path / "output.jsonl") + evaluate_seed(str(input_file), output_file, 0, evaluator, config) + lines = Path(output_file).read_text().strip().split("\n") + r1 = json.loads(lines[0]) + r2 = json.loads(lines[1]) + assert r1["evaluation_uuid"] != r2["evaluation_uuid"] + assert r1["raw_line_fingerprint"] == r2["raw_line_fingerprint"] + assert r1["line_number"] == 0 + assert r2["line_number"] == 1 + + def test_duplicate_invalid_lines_get_different_uuids(self, tmp_path): + evaluator, _, config = _make_finance_evaluator() + input_file = tmp_path / "input.jsonl" + input_file.write_text("not valid json\nnot valid json\n") + Path(str(input_file) + ".done").touch() + output_file = str(tmp_path / "output.jsonl") + evaluate_seed(str(input_file), output_file, 0, evaluator, config) + lines = Path(output_file).read_text().strip().split("\n") + r1 = json.loads(lines[0]) + r2 = json.loads(lines[1]) + assert r1["evaluation_uuid"] != r2["evaluation_uuid"] + assert r1["raw_line_fingerprint"] == r2["raw_line_fingerprint"] From d645a09e86e651f9f9f46d104578619772a26c0d Mon Sep 17 00:00:00 2001 From: Ander Alvarez Sanz <104446704+aalvsz@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:01:33 +0200 Subject: [PATCH 2/2] Harden ProvenanceGuard finance trace parsing Signed-off-by: Ander Alvarez Sanz <104446704+aalvsz@users.noreply.github.com> --- nvflow/provenanceguard/nli.py | 26 +- .../finance/utils/rl/provenanceguard.py | 217 ++-- tests/test_provenanceguard.py | 1013 ++++++++++++++++- 3 files changed, 1160 insertions(+), 96 deletions(-) diff --git a/nvflow/provenanceguard/nli.py b/nvflow/provenanceguard/nli.py index 05391b7..bf2e5f6 100644 --- a/nvflow/provenanceguard/nli.py +++ b/nvflow/provenanceguard/nli.py @@ -84,23 +84,28 @@ def _ensure_loaded(self) -> None: if self._revision: kwargs["revision"] = self._revision - self._tokenizer = AutoTokenizer.from_pretrained(self._model_id, **kwargs) - self._model = AutoModelForSequenceClassification.from_pretrained(self._model_id, **kwargs) - self._model.eval() + # Load into locals first — do not publish _model / _tokenizer / + # _label_map until ALL validation succeeds. This ensures that + # a failed duplicate, missing, or incomplete id2label does not + # poison state: subsequent evaluations re-attempt loading and + # fail again rather than reusing partial/cached state. + tokenizer = AutoTokenizer.from_pretrained(self._model_id, **kwargs) + model = AutoModelForSequenceClassification.from_pretrained(self._model_id, **kwargs) + model.eval() # Build label map from model config — validate actual label names, # never assume arbitrary LABEL_0 ordering. - id2label = self._model.config.id2label - self._label_map = {} + id2label = model.config.id2label + label_map: dict[int, str] = {} for idx, label in id2label.items(): normalized = _normalize_nli_label(str(label)) - self._label_map[int(idx)] = normalized + label_map[int(idx)] = normalized # Validate exact unique label set: exactly 3 labels, one each # of entailment / neutral / contradiction, no duplicates, # extras, or missing. seen_labels: list[str] = [] - for idx in sorted(self._label_map): - seen_labels.append(self._label_map[idx]) + for idx in sorted(label_map): + seen_labels.append(label_map[idx]) if len(seen_labels) != 3: raise ValueError( f"NLI model has {len(seen_labels)} labels; expected " @@ -119,6 +124,11 @@ def _ensure_loaded(self) -> None: parts.append("duplicate labels detected") raise ValueError("NLI label set validation failed: " + "; ".join(parts)) + # Publish only after all validation succeeds. + self._tokenizer = tokenizer + self._model = model + self._label_map = label_map + def score(self, *, premise: str, hypothesis: str) -> NLIResult: import torch diff --git a/nvflow/recipes/finance/utils/rl/provenanceguard.py b/nvflow/recipes/finance/utils/rl/provenanceguard.py index 2e4a190..8608774 100644 --- a/nvflow/recipes/finance/utils/rl/provenanceguard.py +++ b/nvflow/recipes/finance/utils/rl/provenanceguard.py @@ -89,6 +89,7 @@ ) + @dataclass(frozen=True) class FilingMetadata: """SEC filing metadata extracted from sec_filing_search or URL parsing.""" @@ -216,6 +217,21 @@ def _iter_conversation_items(row: dict[str, Any]) -> Iterator[dict[str, Any]]: return +def _canonicalize_document(document: str | None) -> str | None: + """Normalize a document name: strip query/fragment, take basename. + + Ensures URL-derived and metadata-derived document names are identical: + ``10-K.htm``, ``10-K.htm?output=1``, and ``10-K.htm#part1`` all + yield ``10-K.htm``. + """ + if not document: + return None + doc = str(document).strip() + doc = doc.split("?")[0].split("#")[0] + doc = os.path.basename(doc) + return doc if doc else None + + def _extract_filing_metadata(result: Any) -> FilingMetadata: """Extract CIK, accession, document, and URL from a sec_filing_search result.""" if not isinstance(result, dict | list): @@ -275,9 +291,13 @@ def _extract_filing_metadata(result: Any) -> FilingMetadata: if accession: accession = _format_accession(str(accession).strip()) if document: - document = str(document).strip() - document = document.split("?")[0].split("#")[0] - document = os.path.basename(document) + document = _canonicalize_document(document) + + if url: + url_filing = _extract_filing_from_url(url) + cik = cik or url_filing.cik + accession = accession or url_filing.accession + document = document or url_filing.document return FilingMetadata(cik=cik, accession=accession, document=document, url=url) @@ -285,7 +305,9 @@ def _extract_filing_metadata(result: Any) -> FilingMetadata: def _extract_filing_from_url(url: str) -> FilingMetadata: """Parse CIK, accession, and document from a SEC EDGAR Archives URL. - The CIK in the URL path is zero-padded to 10 digits. + The CIK in the URL path is zero-padded to 10 digits. The document + name is canonicalized (query/fragment stripped, basename taken) so + that URL-derived and metadata-derived document names are identical. """ match = SEC_FILING_URL_RE.search(url) if not match: @@ -293,7 +315,7 @@ def _extract_filing_from_url(url: str) -> FilingMetadata: cik = match.group("cik").zfill(10) accession_raw = match.group("accession") - document = match.group("document") + document = _canonicalize_document(match.group("document")) accession = _format_accession(accession_raw) return FilingMetadata(cik=cik, accession=accession, document=document, url=url) @@ -338,47 +360,97 @@ def _parse_tool_result(raw: Any) -> Any: return raw -def _is_error_result(parsed: Any) -> bool: - """Check if a parsed tool result is an error response. +def _try_parse_json(text: str) -> Any: + """Attempt to parse a string as JSON; return None on failure.""" + try: + return json.loads(text) + except (json.JSONDecodeError, TypeError): + return None + - Rejects: - - Dict failures: ``success is False``, ``error``, ``is_error``, - or a ``result`` string containing error indicators. - - Raw failure strings: ``ERROR`` prefix, traceback. +def _decode_tool_envelope(raw: Any) -> tuple[str | None, str | None]: + """Strictly decode a function_call_output envelope for pinned Gym. + + Returns ``(payload, error_reason)``. On success, ``payload`` is the + unwrapped content string and ``error_reason`` is ``None``. On failure, + ``payload`` is ``None`` and ``error_reason`` explains why. + + Accepts: + - Pinned Gym envelope: ``{"results": }`` with no top-level + ``error`` key, a string payload that does not begin ``ERROR:`` and + is not a nested JSON ``{"error": ...}`` payload. + - Legacy producer shape: ``{"success": true, "result": }`` + — accepted only when ``success is True`` plus a string ``result``. + + Rejects (returning ``(None, reason)``): + - Top-level agent ``{"error": ...}`` (timeout / exception envelope). + - Missing or unknown envelope keys. + - Non-string ``results`` / ``result`` payloads. + - ``results`` payload beginning ``ERROR:``. + - Nested JSON ``{"error": ...}`` inside ``results`` (time-budget / + no-company error payloads from the resource server). + - Legacy ``success is False`` or missing ``success``. + - All raw strings and other non-dict outputs (unstructured). """ - if isinstance(parsed, dict): - if parsed.get("error") or parsed.get("is_error"): - return True - if parsed.get("success") is False: - return True - result_str = parsed.get("result") - if isinstance(result_str, str) and _is_raw_error_string(result_str): - return True - return False - if isinstance(parsed, str): - return _is_raw_error_string(parsed) - return False + parsed = _parse_tool_result(raw) + + if not isinstance(parsed, dict): + return None, "unstructured output" + if "error" in parsed: + return None, "agent error envelope" -def _is_raw_error_string(text: str) -> bool: - """Check if a raw string looks like an error, traceback, or failure. + if "results" in parsed: + results = parsed["results"] + if not isinstance(results, str): + return None, "non-string results payload" - Rejects strings containing: Error:, Failed, traceback, exception, - unavailable, or similar failure indicators from real finance tools. - Case-insensitive. + stripped = results.strip() + if stripped.upper().startswith("ERROR:"): + return None, "ERROR: payload" + + nested = _try_parse_json(results) + if isinstance(nested, dict) and "error" in nested: + return None, "nested JSON error payload" + + return results, None + + if "success" in parsed: + if parsed["success"] is not True: + return None, "legacy shape with success=False" + result = parsed.get("result") + if not isinstance(result, str): + return None, "legacy shape with non-string result" + return result, None + + return None, "unknown envelope" + + +def _is_parse_success(payload: str, key: str) -> bool: + """Check whether a parse_html_page results payload indicates success. + + In pinned Gym, a successful ``parse_html_page`` returns a results + string whose lines include the exact marker (the + ``_save_tool_output`` line):: + + SUCCESS: The result has been saved to the data storage under the key: {key}. + + Failed parses return arbitrary ``str(e)`` text without the marker. + This function checks for the exact marker on any line, bound to the + expected ``key``, so that the ``WARNING:`` overwrite case (which + still contains the exact line) is accepted while abbreviated, + unrelated, or wrong-key ``SUCCESS:`` strings are rejected. """ - stripped = text.strip() - lower = stripped.lower() - if lower.startswith("error"): - return True - if lower.startswith("failed"): - return True - if "traceback" in lower: - return True - if "exception" in lower: - return True - if lower == "unavailable": - return True + stripped = payload.strip() + if not stripped: + return False + expected = ( + "SUCCESS: The result has been saved to the data storage " + f"under the key: {key}." + ) + for line in stripped.splitlines(): + if line.strip() == expected: + return True return False @@ -389,13 +461,17 @@ def _build_key_to_filing_map( Scans for ``sec_filing_search`` / ``edgar_search`` calls (extracts filing metadata from arguments and results) and ``parse_html_page`` - calls (reads the plain ``key`` argument directly). + calls (reads the plain ``key`` argument, paired by ``call_id`` to + its output). The key-to-URL mapping is updated only after a strict + successful pinned parse output (``SUCCESS:`` marker); failed retries + preserve the last successful mapping. """ key_to_url: dict[str, str] = {} call_id_to_filing: dict[str, FilingMetadata] = {} url_to_filing: dict[str, FilingMetadata] = {} sec_search_call_ids: set[str] = set() + parse_call_ids: dict[str, tuple[str, str]] = {} for item in items: item_type = item.get("type", "") @@ -414,8 +490,8 @@ def _build_key_to_filing_map( args = _parse_tool_arguments(item.get("arguments")) url = args.get("url", "") key = args.get("key", "") - if key and url: - key_to_url[key] = url + if key and url and call_id: + parse_call_ids[call_id] = (key, url) continue if item_type in ("function_call_output", "tool_result"): @@ -423,27 +499,44 @@ def _build_key_to_filing_map( continue if call_id in sec_search_call_ids: existing = call_id_to_filing.get(call_id, FilingMetadata()) - parsed = _parse_tool_result(item.get("output", item.get("result", ""))) - if isinstance(parsed, dict | list): - enriched = _extract_filing_metadata(parsed) - if enriched.cik or enriched.accession or enriched.url: - filing = FilingMetadata( - cik=enriched.cik or existing.cik, - accession=enriched.accession or existing.accession, - document=enriched.document or existing.document, - url=enriched.url or existing.url, - ) - call_id_to_filing[call_id] = filing + raw_output = item.get("output", item.get("result", "")) + payload, _error = _decode_tool_envelope(raw_output) + if payload is not None: + decoded = _try_parse_json(payload) + if isinstance(decoded, (dict, list)): + enriched = _extract_filing_metadata(decoded) + if enriched.cik or enriched.accession or enriched.url: + filing = FilingMetadata( + cik=enriched.cik or existing.cik, + accession=enriched.accession or existing.accession, + document=enriched.document or existing.document, + url=enriched.url or existing.url, + ) + call_id_to_filing[call_id] = filing if call_id in call_id_to_filing: filing = call_id_to_filing[call_id] if filing.url: url_to_filing[filing.url] = filing + elif call_id in parse_call_ids: + key, url = parse_call_ids[call_id] + raw_output = item.get("output", item.get("result", "")) + payload, _error = _decode_tool_envelope(raw_output) + if payload is not None and _is_parse_success(payload, key): + key_to_url[key] = url key_to_filing: dict[str, FilingMetadata] = {} for key, url in key_to_url.items(): filing = url_to_filing.get(url) if not filing: filing = _extract_filing_from_url(url) + elif filing.url: + url_filing = _extract_filing_from_url(filing.url) + filing = FilingMetadata( + cik=filing.cik or url_filing.cik, + accession=filing.accession or url_filing.accession, + document=filing.document or url_filing.document, + url=filing.url, + ) key_to_filing[key] = filing return key_to_filing @@ -583,23 +676,13 @@ def extract_trace(row: dict[str, Any]) -> TraceExtraction: ) continue - parsed = _parse_tool_result(result_raw) - if _is_error_result(parsed): + text, error_reason = _decode_tool_envelope(result_raw) + if text is None: extraction.extraction_errors.append( - f"retrieve_information call {call_id}: error result" + f"retrieve_information call {call_id}: {error_reason}" ) continue - if isinstance(parsed, dict): - text = parsed.get("result", parsed.get("retrieval", "")) - if isinstance(text, dict): - text = json.dumps(text) - text = str(text) if text else "" - elif isinstance(parsed, str): - text = parsed - else: - text = str(parsed) - if not text.strip(): extraction.extraction_errors.append( f"retrieve_information call {call_id}: empty result" diff --git a/tests/test_provenanceguard.py b/tests/test_provenanceguard.py index e5f7480..ed3f955 100644 --- a/tests/test_provenanceguard.py +++ b/tests/test_provenanceguard.py @@ -206,7 +206,7 @@ def _make_rollout_row( { "type": "function_call_output", "call_id": "call_retrieve_1", - "output": json.dumps({"result": evidence_text}), + "output": json.dumps({"results": evidence_text}), }, { "type": "function_call", @@ -273,6 +273,17 @@ class TestRegistryDiscovery: def test_stage_registered(self): from nvflow.core import StageRegistry + # Prior tests (e.g. test_core.py) may call StageRegistry.clear() + # after collection, wiping entries added by @StageRegistry.register + # during module import. Re-register via the public API if missing. + if not StageRegistry.has( + recipe="finance", workflow="grpo", stage="evaluate_provenance" + ): + mod = _load_eval_stage_module() + StageRegistry.register( + recipe="finance", workflow="grpo", stage="evaluate_provenance", + )(mod.EvaluateProvenanceStage) + stage_cls = StageRegistry.get( recipe="finance", workflow="grpo", @@ -657,6 +668,13 @@ def test_multiple_keys_yields_unavailable(self): } ), }, + { + "type": "function_call_output", + "call_id": "call_parse_1", + "output": json.dumps( + {"results": "SUCCESS: The result has been saved to the data storage under the key: filing_10k."} + ), + }, { "type": "function_call", "name": "parse_html_page", @@ -668,6 +686,13 @@ def test_multiple_keys_yields_unavailable(self): } ), }, + { + "type": "function_call_output", + "call_id": "call_parse_2", + "output": json.dumps( + {"results": "SUCCESS: The result has been saved to the data storage under the key: filing_10k_2."} + ), + }, { "type": "function_call", "name": "retrieve_information", @@ -677,7 +702,7 @@ def test_multiple_keys_yields_unavailable(self): { "type": "function_call_output", "call_id": "call_retrieve_1", - "output": json.dumps({"result": "Some text"}), + "output": json.dumps({"results": "Some text"}), }, { "type": "function_call", @@ -706,7 +731,7 @@ def test_unknown_key_yields_unavailable(self): { "type": "function_call_output", "call_id": "call_retrieve_1", - "output": json.dumps({"result": "Some text"}), + "output": json.dumps({"results": "Some text"}), }, { "type": "function_call", @@ -774,7 +799,7 @@ def test_nested_response_output_trace(self): { "type": "function_call_output", "call_id": "call_retrieve_nested", - "output": json.dumps({"result": "Revenue was $1.23 billion in 2024."}), + "output": json.dumps({"results": "Revenue was $1.23 billion in 2024."}), }, { "type": "function_call", @@ -990,7 +1015,7 @@ def test_one_chunk_per_retrieve_result(self): { "type": "function_call_output", "call_id": "call_r1", - "output": json.dumps({"result": "Evidence one."}), + "output": json.dumps({"results": "Evidence one."}), }, { "type": "function_call", @@ -1001,7 +1026,7 @@ def test_one_chunk_per_retrieve_result(self): { "type": "function_call_output", "call_id": "call_r2", - "output": json.dumps({"result": "Evidence two."}), + "output": json.dumps({"results": "Evidence two."}), }, { "type": "function_call", @@ -1051,21 +1076,27 @@ def _make_realistic_trace() -> dict: "type": "function_call", "name": "sec_filing_search", "call_id": "call_sec_1", - "arguments": json.dumps({"query": "Apple 10-K", "form_type": "10-K"}), + "arguments": json.dumps({"ticker": "AAPL", "form_types": ["10-K"]}), }, { "type": "function_call_output", "call_id": "call_sec_1", "output": json.dumps( { - "filings": [ - { - "cik": "1811414", - "accessionNo": "0001811414-25-000010", - "primaryDocument": "10-K.htm", - "linkToHtml": _FILING_URL, - } - ] + "results": json.dumps( + [ + { + "ticker": "AAPL", + "company_name": "Apple Inc.", + "form": "10-K", + "filing_date": "2025-01-01", + "report_date": "2024-12-31", + "accession_number": "0001811414-25-000010", + "filing_url": _FILING_URL, + } + ], + indent=2, + ) } ), }, @@ -1078,7 +1109,9 @@ def _make_realistic_trace() -> dict: { "type": "function_call_output", "call_id": "call_parse_1", - "output": json.dumps({"result": "Page stored successfully"}), + "output": json.dumps( + {"results": "SUCCESS: The result has been saved to the data storage under the key: filing_10k."} + ), }, { "type": "function_call", @@ -1089,7 +1122,7 @@ def _make_realistic_trace() -> dict: { "type": "function_call_output", "call_id": "call_retrieve_1", - "output": json.dumps({"result": "Revenue was $1.23 billion in 2024."}), + "output": json.dumps({"results": "Revenue was $1.23 billion in 2024."}), }, { "type": "function_call", @@ -1144,7 +1177,9 @@ def test_parse_html_page_reads_plain_key(self): { "type": "function_call_output", "call_id": "call_parse_1", - "output": json.dumps({"result": "Page stored"}), + "output": json.dumps( + {"results": "SUCCESS: The result has been saved to the data storage under the key: filing_10k."} + ), }, { "type": "function_call", @@ -1155,7 +1190,7 @@ def test_parse_html_page_reads_plain_key(self): { "type": "function_call_output", "call_id": "call_retrieve_1", - "output": json.dumps({"result": "Revenue was $1.23 billion in 2024."}), + "output": json.dumps({"results": "Revenue was $1.23 billion in 2024."}), }, { "type": "function_call", @@ -1202,7 +1237,7 @@ def test_raw_error_string_excluded(self): sidecar = parse_rollout_line(json.dumps(row)) assert sidecar.trace is not None assert len(sidecar.trace.evidence) == 0 - assert any("error" in e.lower() for e in sidecar.trace.extraction_errors) + assert any("unstructured" in e.lower() for e in sidecar.trace.extraction_errors) def test_traceback_excluded(self): row = _make_rollout_row() @@ -1424,7 +1459,7 @@ def test_extraction_errors_force_unavailable_preserving_verdicts(self): { "type": "function_call_output", "call_id": "call_r1", - "output": json.dumps({"result": "Revenue was $1.23 billion in 2024."}), + "output": json.dumps({"results": "Revenue was $1.23 billion in 2024."}), }, { "type": "function_call", @@ -1749,3 +1784,939 @@ def test_duplicate_invalid_lines_get_different_uuids(self, tmp_path): r2 = json.loads(lines[1]) assert r1["evaluation_uuid"] != r2["evaluation_uuid"] assert r1["raw_line_fingerprint"] == r2["raw_line_fingerprint"] + + +# --------------------------------------------------------------------------- +# Strict envelope decoder tests (pinned Gym interface) +# --------------------------------------------------------------------------- + + +class TestStrictEnvelopeDecoder: + """Strict envelope decoder for pinned Gym function_call_output items.""" + + def test_gym_results_envelope_accepted(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _decode_tool_envelope, + ) + + payload, error = _decode_tool_envelope( + {"results": "Revenue was $1.23 billion."} + ) + assert payload == "Revenue was $1.23 billion." + assert error is None + + def test_gym_results_envelope_json_string_accepted(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _decode_tool_envelope, + ) + + payload, error = _decode_tool_envelope( + {"results": json.dumps({"key": "value"})} + ) + assert payload is not None + assert error is None + + def test_agent_error_envelope_rejected(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _decode_tool_envelope, + ) + + payload, error = _decode_tool_envelope( + {"error": "Tool call timed out after 30s."} + ) + assert payload is None + assert "agent error" in error + + def test_results_error_prefix_rejected(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _decode_tool_envelope, + ) + + payload, error = _decode_tool_envelope( + {"results": "ERROR: Retrieval LLM call failed: timeout"} + ) + assert payload is None + assert "ERROR" in error + + def test_nested_json_error_rejected(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _decode_tool_envelope, + ) + + payload, error = _decode_tool_envelope( + {"results": json.dumps({"error": "Time budget exhausted."})} + ) + assert payload is None + assert "nested" in error + + @pytest.mark.parametrize( + "raw_output", + [ + "remote end hung up unexpectedly", + "upstream disconnected", + "Connection reset by peer", + "Request timed out", + ], + ) + def test_arbitrary_raw_string_rejected(self, raw_output): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _decode_tool_envelope, + ) + + payload, error = _decode_tool_envelope(raw_output) + assert payload is None + assert "unstructured" in error + + def test_legacy_success_accepted(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _decode_tool_envelope, + ) + + payload, error = _decode_tool_envelope( + {"success": True, "result": "Revenue was $1.23 billion."} + ) + assert payload == "Revenue was $1.23 billion." + assert error is None + + def test_legacy_failure_rejected(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _decode_tool_envelope, + ) + + payload, error = _decode_tool_envelope( + {"success": False, "result": "error"} + ) + assert payload is None + assert "success" in error.lower() + + def test_legacy_missing_success_rejected(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _decode_tool_envelope, + ) + + payload, error = _decode_tool_envelope( + {"result": "Revenue was $1.23 billion."} + ) + assert payload is None + assert "unknown" in error + + def test_non_string_results_rejected(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _decode_tool_envelope, + ) + + payload, error = _decode_tool_envelope({"results": 42}) + assert payload is None + assert "non-string" in error + + def test_unknown_envelope_rejected(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _decode_tool_envelope, + ) + + payload, error = _decode_tool_envelope({"foo": "bar"}) + assert payload is None + assert "unknown" in error + + +# --------------------------------------------------------------------------- +# Pinned Gym producer-to-sidecar integration test +# --------------------------------------------------------------------------- + + +class TestGymPinnedIntegration: + """End-to-sidecar using response.output items matching the pinned example.""" + + def test_full_pinned_gym_trace(self): + filing_url = ( + "https://www.sec.gov/Archives/edgar/data/320193/" + "000032019322000108/aapl-20220924.htm" + ) + row = { + "uuid": "gym-pinned-001", + "response": { + "id": "resp-gym-001", + "output": [ + { + "type": "function_call", + "name": "sec_filing_search", + "call_id": "call_sec_1", + "arguments": json.dumps( + {"ticker": "AAPL", "form_types": ["10-K"]} + ), + }, + { + "type": "function_call_output", + "call_id": "call_sec_1", + "output": json.dumps( + { + "results": json.dumps( + [ + { + "ticker": "AAPL", + "company_name": "Apple Inc.", + "form": "10-K", + "filing_date": "2022-10-28", + "report_date": "2022-09-24", + "accession_number": "0000320193-22-000108", + "filing_url": filing_url, + } + ], + indent=2, + ) + } + ), + }, + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_1", + "arguments": json.dumps( + {"url": filing_url, "key": "aapl_10k_2022"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_parse_1", + "output": json.dumps( + { + "results": ( + "SUCCESS: The result has been saved to the data " + "storage under the key: aapl_10k_2022." + ) + } + ), + }, + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_1", + "arguments": json.dumps( + {"prompt": "Find FTE employees {{aapl_10k_2022}}"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_1", + "output": json.dumps( + { + "results": ( + "As of September 24, 2022, the Company had " + "approximately 164,000 full-time equivalent " + "employees." + ) + } + ), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_1", + "arguments": json.dumps( + {"final_result": "Apple had approximately 164,000 FTE."} + ), + }, + ], + }, + } + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert sidecar.trace.has_submit is True + assert sidecar.trace.answer == "Apple had approximately 164,000 FTE." + assert len(sidecar.trace.evidence) == 1 + chunk = sidecar.trace.evidence[0] + assert "164,000 full-time equivalent" in chunk.text + assert chunk.source_id is not None + assert chunk.attribution_state == "available" + assert "cik=0000320193" in chunk.source_id + assert "accession=000032019322000108" in chunk.source_id + assert "doc=aapl-20220924.htm" in chunk.source_id + + +# --------------------------------------------------------------------------- +# Pinned Gym failure envelopes +# --------------------------------------------------------------------------- + + +class TestGymPinnedFailures: + """Pinned Gym failure envelopes must produce no evidence.""" + + def test_agent_error_envelope_no_evidence(self): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r1", + "arguments": json.dumps({"prompt": "query"}), + }, + { + "type": "function_call_output", + "call_id": "call_r1", + "output": json.dumps( + {"error": "Tool call timed out after 30s."} + ), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_s1", + "arguments": json.dumps({"final_result": "Some answer."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert len(sidecar.trace.evidence) == 0 + assert any("agent error" in e for e in sidecar.trace.extraction_errors) + + def test_retrieve_error_prefix_no_evidence(self): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r1", + "arguments": json.dumps({"prompt": "query"}), + }, + { + "type": "function_call_output", + "call_id": "call_r1", + "output": json.dumps( + {"results": "ERROR: Retrieval LLM call failed: timeout"} + ), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_s1", + "arguments": json.dumps({"final_result": "Some answer."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert len(sidecar.trace.evidence) == 0 + assert any("ERROR" in e for e in sidecar.trace.extraction_errors) + + def test_nested_json_error_no_evidence(self): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r1", + "arguments": json.dumps({"prompt": "query"}), + }, + { + "type": "function_call_output", + "call_id": "call_r1", + "output": json.dumps( + {"results": json.dumps({"error": "Time budget exhausted."})} + ), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_s1", + "arguments": json.dumps({"final_result": "Some answer."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert len(sidecar.trace.evidence) == 0 + assert any("nested" in e for e in sidecar.trace.extraction_errors) + + @pytest.mark.parametrize( + "raw_output", + [ + "remote end hung up unexpectedly", + "upstream disconnected", + ], + ) + def test_arbitrary_raw_string_no_evidence(self, raw_output): + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_r1", + "arguments": json.dumps({"prompt": "query"}), + }, + { + "type": "function_call_output", + "call_id": "call_r1", + "output": raw_output, + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_s1", + "arguments": json.dumps({"final_result": "Some answer."}), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert len(sidecar.trace.evidence) == 0 + assert any("unstructured" in e for e in sidecar.trace.extraction_errors) + + +# --------------------------------------------------------------------------- +# Fix 2: Transactional NLI label configuration +# --------------------------------------------------------------------------- + + +class TestNLITransactionalLoading: + """Failed label validation must not poison _model/_label_map state.""" + + @staticmethod + def _make_fake_transformers(id2label: dict): + import types + + class _FakeModel: + def __init__(self): + class _Config: + pass + + self.config = _Config() + self.config.id2label = id2label + + def eval(self): + return self + + class _FakeAutoModel: + @staticmethod + def from_pretrained(*args, **kwargs): + return _FakeModel() + + class _FakeTokenizer: + @staticmethod + def from_pretrained(*args, **kwargs): + return _FakeTokenizer() + + mod = types.ModuleType("transformers") + mod.AutoModelForSequenceClassification = _FakeAutoModel + mod.AutoTokenizer = _FakeTokenizer + return mod + + def test_duplicate_id2label_fails_and_does_not_poison(self): + import sys + + from nvflow.provenanceguard.nli import HFNLI + + fake_mod = self._make_fake_transformers( + {0: "entailment", 1: "entailment", 2: "contradiction"} + ) + original = sys.modules.get("transformers") + sys.modules["transformers"] = fake_mod + try: + nli = HFNLI() + with pytest.raises(ValueError, match="duplicate labels"): + nli._ensure_loaded() + assert nli._model is None + assert nli._tokenizer is None + assert nli._label_map is None + with pytest.raises(ValueError, match="duplicate labels"): + nli._ensure_loaded() + assert nli._model is None + assert nli._label_map is None + finally: + if original is not None: + sys.modules["transformers"] = original + else: + sys.modules.pop("transformers", None) + + def test_incomplete_id2label_fails_and_does_not_poison(self): + import sys + + from nvflow.provenanceguard.nli import HFNLI + + fake_mod = self._make_fake_transformers( + {0: "entailment", 1: "contradiction"} + ) + original = sys.modules.get("transformers") + sys.modules["transformers"] = fake_mod + try: + nli = HFNLI() + with pytest.raises(ValueError, match="2 labels"): + nli._ensure_loaded() + assert nli._model is None + assert nli._tokenizer is None + assert nli._label_map is None + with pytest.raises(ValueError, match="2 labels"): + nli._ensure_loaded() + assert nli._model is None + assert nli._label_map is None + finally: + if original is not None: + sys.modules["transformers"] = original + else: + sys.modules.pop("transformers", None) + + def test_valid_id2label_publishes_state(self): + import sys + + from nvflow.provenanceguard.nli import HFNLI + + fake_mod = self._make_fake_transformers( + {0: "entailment", 1: "neutral", 2: "contradiction"} + ) + original = sys.modules.get("transformers") + sys.modules["transformers"] = fake_mod + try: + nli = HFNLI() + nli._ensure_loaded() + assert nli._model is not None + assert nli._tokenizer is not None + assert nli._label_map == {0: "entailment", 1: "neutral", 2: "contradiction"} + finally: + if original is not None: + sys.modules["transformers"] = original + else: + sys.modules.pop("transformers", None) + + +# --------------------------------------------------------------------------- +# Fix 3: Canonicalize URL-derived SEC document IDs +# --------------------------------------------------------------------------- + + +class TestCanonicalDocumentIds: + """URL-derived document names must match metadata-derived document names.""" + + def test_url_document_strips_query_string(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _extract_filing_from_url, + ) + + base = ( + "https://www.sec.gov/Archives/edgar/data/1811414/" + "000181141425000010/10-K.htm" + ) + filing_plain = _extract_filing_from_url(base) + filing_query = _extract_filing_from_url(base + "?output=1") + assert filing_plain.document == "10-K.htm" + assert filing_query.document == "10-K.htm" + assert filing_plain.source_id() == filing_query.source_id() + + def test_url_document_strips_fragment(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _extract_filing_from_url, + ) + + base = ( + "https://www.sec.gov/Archives/edgar/data/1811414/" + "000181141425000010/10-K.htm" + ) + filing_plain = _extract_filing_from_url(base) + filing_frag = _extract_filing_from_url(base + "#part1") + assert filing_plain.document == "10-K.htm" + assert filing_frag.document == "10-K.htm" + assert filing_plain.source_id() == filing_frag.source_id() + + def test_url_query_and_fragment_equivalent(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _extract_filing_from_url, + ) + + base = ( + "https://www.sec.gov/Archives/edgar/data/1811414/" + "000181141425000010/10-K.htm" + ) + filing_q = _extract_filing_from_url(base + "?output=1") + filing_f = _extract_filing_from_url(base + "#part1") + assert filing_q.document == filing_f.document + assert filing_q.source_id() == filing_f.source_id() + + def test_canonicalize_document_helper(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _canonicalize_document, + ) + + assert _canonicalize_document("10-K.htm") == "10-K.htm" + assert _canonicalize_document("10-K.htm?output=1") == "10-K.htm" + assert _canonicalize_document("10-K.htm#part1") == "10-K.htm" + assert _canonicalize_document("path/to/10-K.htm?x=1") == "10-K.htm" + assert _canonicalize_document(None) is None + assert _canonicalize_document("") is None + + def test_url_and_metadata_document_equivalence(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _extract_filing_from_url, + _extract_filing_metadata, + ) + + url = ( + "https://www.sec.gov/Archives/edgar/data/1811414/" + "000181141425000010/10-K.htm" + ) + url_filing = _extract_filing_from_url(url + "?output=1") + meta_filing = _extract_filing_metadata( + { + "filings": [ + { + "cik": "1811414", + "accessionNo": "0001811414-25-000010", + "primaryDocument": "10-K.htm?output=1", + "linkToHtml": url, + } + ] + } + ) + assert url_filing.document == meta_filing.document + assert url_filing.source_id() == meta_filing.source_id() + + +# --------------------------------------------------------------------------- +# Call-id-paired parse tracking: failed retry preserves last successful key +# --------------------------------------------------------------------------- + + +class TestParseCallIdPairedTracking: + """parse_html_page key-to-URL must update only on strict SUCCESS output. + + In pinned Gym, a successful parse returns ``{results: "SUCCESS: ..."}`` + and a failure returns ``{results: str(e)}``. A failed retry of the same + key with a different URL must preserve the previously stored filing, not + fabricate the new URL. + """ + + _URL_A = ( + "https://www.sec.gov/Archives/edgar/data/1811414/" + "000181141425000010/10-K.htm" + ) + _URL_B = ( + "https://www.sec.gov/Archives/edgar/data/1811414/" + "000181141425000020/10-K.htm" + ) + + def test_failed_retry_preserves_successful_source(self): + """Success A then failed B same key: retrieve must map to source A.""" + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_A", + "arguments": json.dumps( + {"url": self._URL_A, "key": "filing_k"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_parse_A", + "output": json.dumps( + { + "results": ( + "SUCCESS: The result has been saved to the data " + "storage under the key: filing_k.\n" + "The data_storage currently contains the following " + "keys:\nfiling_k\n" + ) + } + ), + }, + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_B", + "arguments": json.dumps( + {"url": self._URL_B, "key": "filing_k"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_parse_B", + "output": json.dumps( + {"results": "HTTPSConnectionPool(host='www.sec.gov'): Read timed out."} + ), + }, + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_1", + "arguments": json.dumps( + {"prompt": "Revenue {{filing_k}}"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_1", + "output": json.dumps( + {"results": "Revenue was $1.23 billion in 2024."} + ), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_1", + "arguments": json.dumps( + {"final_result": "Revenue was $1.23 billion in 2024."} + ), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 1 + chunk = sidecar.trace.evidence[0] + assert chunk.source_id is not None + assert chunk.attribution_state == "available" + assert "accession=000181141425000010" in chunk.source_id + assert chunk.sec_url == self._URL_A + + def test_ordinary_success_still_maps(self): + """A single successful parse still maps the key to the filing.""" + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_1", + "arguments": json.dumps( + {"url": self._URL_A, "key": "filing_k"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_parse_1", + "output": json.dumps( + { + "results": ( + "SUCCESS: The result has been saved to the data " + "storage under the key: filing_k." + ) + } + ), + }, + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_1", + "arguments": json.dumps( + {"prompt": "Revenue {{filing_k}}"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_1", + "output": json.dumps( + {"results": "Revenue was $1.23 billion in 2024."} + ), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_1", + "arguments": json.dumps( + {"final_result": "Revenue was $1.23 billion in 2024."} + ), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 1 + chunk = sidecar.trace.evidence[0] + assert chunk.source_id is not None + assert chunk.attribution_state == "available" + assert "accession=000181141425000010" in chunk.source_id + assert chunk.sec_url == self._URL_A + + def test_no_success_output_no_mapping(self): + """parse_html_page with no paired output must not map the key.""" + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_1", + "arguments": json.dumps( + {"url": self._URL_A, "key": "filing_k"} + ), + }, + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_1", + "arguments": json.dumps( + {"prompt": "Revenue {{filing_k}}"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_1", + "output": json.dumps( + {"results": "Revenue was $1.23 billion in 2024."} + ), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_1", + "arguments": json.dumps( + {"final_result": "Revenue was $1.23 billion in 2024."} + ), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 1 + chunk = sidecar.trace.evidence[0] + assert chunk.source_id is None + assert chunk.attribution_state == "unavailable" + + def test_error_envelope_does_not_map(self): + """parse_html_page output with agent error must not map the key.""" + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_1", + "arguments": json.dumps( + {"url": self._URL_A, "key": "filing_k"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_parse_1", + "output": json.dumps( + {"error": "Tool call timed out after 30s."} + ), + }, + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_1", + "arguments": json.dumps( + {"prompt": "Revenue {{filing_k}}"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_1", + "output": json.dumps( + {"results": "Revenue was $1.23 billion in 2024."} + ), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_1", + "arguments": json.dumps( + {"final_result": "Revenue was $1.23 billion in 2024."} + ), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 1 + chunk = sidecar.trace.evidence[0] + assert chunk.source_id is None + assert chunk.attribution_state == "unavailable" + + def test_success_then_success_overwrites(self): + """Two successful parses of the same key: last success wins.""" + row = _make_rollout_row() + row["output"] = [ + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_A", + "arguments": json.dumps( + {"url": self._URL_A, "key": "filing_k"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_parse_A", + "output": json.dumps( + { + "results": ( + "SUCCESS: The result has been saved to the data " + "storage under the key: filing_k." + ) + } + ), + }, + { + "type": "function_call", + "name": "parse_html_page", + "call_id": "call_parse_B", + "arguments": json.dumps( + {"url": self._URL_B, "key": "filing_k"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_parse_B", + "output": json.dumps( + { + "results": ( + "SUCCESS: The result has been saved to the data " + "storage under the key: filing_k." + ) + } + ), + }, + { + "type": "function_call", + "name": "retrieve_information", + "call_id": "call_retrieve_1", + "arguments": json.dumps( + {"prompt": "Revenue {{filing_k}}"} + ), + }, + { + "type": "function_call_output", + "call_id": "call_retrieve_1", + "output": json.dumps( + {"results": "Revenue was $1.23 billion in 2024."} + ), + }, + { + "type": "function_call", + "name": "submit_final_result", + "call_id": "call_submit_1", + "arguments": json.dumps( + {"final_result": "Revenue was $1.23 billion in 2024."} + ), + }, + ] + sidecar = parse_rollout_line(json.dumps(row)) + assert sidecar.trace is not None + assert len(sidecar.trace.evidence) == 1 + chunk = sidecar.trace.evidence[0] + assert chunk.source_id is not None + assert chunk.attribution_state == "available" + assert "accession=000181141425000020" in chunk.source_id + assert chunk.sec_url == self._URL_B + + def test_is_parse_success_helper(self): + from nvflow.recipes.finance.utils.rl.provenanceguard import ( + _is_parse_success, + ) + + assert _is_parse_success( + "SUCCESS: The result has been saved to the data " + "storage under the key: filing_k.", + "filing_k", + ) is True + assert ( + _is_parse_success( + "WARNING: key exists.\n" + "SUCCESS: The result has been saved to the data " + "storage under the key: filing_k.\n" + "The data_storage currently contains the following keys:\nfiling_k\n", + "filing_k", + ) + is True + ) + assert _is_parse_success( + "SUCCESS: saved under key: filing_k.", "filing_k" + ) is False + assert _is_parse_success( + "SUCCESS: operation completed.", "filing_k" + ) is False + assert _is_parse_success( + "SUCCESS: The result has been saved to the data " + "storage under the key: other_k.", + "filing_k", + ) is False + assert _is_parse_success("Connection timed out", "filing_k") is False + assert _is_parse_success("", "filing_k") is False + assert _is_parse_success(" ", "filing_k") is False + assert _is_parse_success("ERROR: url is required.", "filing_k") is False