diff --git a/docs/plugins/data-designer-retrieval-sdg/multimodal-reference.md b/docs/plugins/data-designer-retrieval-sdg/multimodal-reference.md index 5afa89f..55581dc 100644 --- a/docs/plugins/data-designer-retrieval-sdg/multimodal-reference.md +++ b/docs/plugins/data-designer-retrieval-sdg/multimodal-reference.md @@ -20,6 +20,7 @@ image paths inside source records resolve relative to their JSONL file. | `concurrency` | `8` | Model-call concurrency, from 1 to 128. | | `batch_size` | `30` | Request batch size, from 1 to 128. | | `missing_response_attempts` | `3` | Total attempts for missing or invalid rows in normally completed batches, from 1 to 3. | +| `failed_record_policy` | `error` | Use `skip` to record exhausted model-validation failures and continue with usable records. | | `seed` | `42` | Seed for instruction sampling, selection, and query-group splitting. Model output can still vary. | | `resume` | `false` | Reuse a compatible run and cached responses. | | `ratios` | `{train: 0.8, validation: 0.0, evaluation: 0.2}` | Query-group split fractions, each from 0 to 1 and summing to 1. | @@ -41,6 +42,58 @@ provider retry behavior. Batch size and retry settings do not define a total API-call, cost, or elapsed-time budget. Refer to [Resume an Interrupted Run](multimodal.md#resume-an-interrupted-run) for recovery. +### Continue After Failed Records + +For a new run, add this to the existing YAML configuration: + +```yaml +failed_record_policy: skip +``` + +The default `error` policy stops when required structured responses remain +unavailable. With `skip`, the structured column records a diagnostic after Data +Designer's native validation retries are exhausted. This adds no retry loop and +does not change the native retry budget. `missing_response_attempts` keeps its +separate meaning; a recorded failure is terminal and is not retried by that loop +or by resume. Unexplained missing rows, duplicate identities, provider errors, +and integrity errors are never converted into skipped records. + +| Failed stage | Disposition | +| --- | --- | +| TOC confirmation | Preserve the source and omit unverified TOC projections. | +| Visual or document enrichment | Preserve original text/images and omit the failed enrichment. | +| Summary or summary combination | Omit the unavailable summary from generation inputs. | +| Summary judgment | Exclude the summary; a failed judgment cannot become a pass. | +| Query generation | Exclude the context's requested query slots. | +| Query metadata, self-sufficiency, relevance, or localization | Exclude the affected query; successful responses remain cached. | + +The original corpus and quality thresholds are preserved. If no queries pass, +export still fails. Inspect `generation_failures.json` in the run directory and +portable bundle for stage, request ID, source/context/query identity where +applicable, validation detail, native attempt limit, and the last invalid +response when available. The portable ledger is included in the manifest's +checksums. `report.json.generation_failures` counts failures by stage, failed +requested slots, and excluded query IDs. Failed query-generation slots contribute +to `requested_slots`, separately from model abstentions and quality rejections. +Earlier planning failures have no known query-slot count. +Queries that fail a later judgment or localization still count in +`funnel.generated` and the `retention_fraction` denominator, once per query. +For example, three generated queries with one localization failure and two +accepted queries report three generated and a retention fraction of `2/3`. + +The same policy is supported by `RetrievalStructuredColumnConfig`. In skip mode, +non-null output cells are JSON strings containing either a validated response +or a `_retrieval_sdg_failure` diagnostic. The multimodal workflow decodes these +before its typed consumers. Direct column users must handle the diagnostic +before treating a cell as a model response. Failed records have `None` in any +configured trace and reasoning-content columns, including when every record in +a batch fails, so downstream expressions can still reference those columns. + +Resume preserves both cached successes and recorded failures. Changing the +policy or upgrading the implementation changes the frozen run identity; use a +new output directory. Failure rates and exclusions can affect sample composition, +so include them when comparing generation yield. + ## Source Records Each line in `sources_file` is one JSON object describing an independently diff --git a/docs/plugins/data-designer-retrieval-sdg/multimodal.md b/docs/plugins/data-designer-retrieval-sdg/multimodal.md index cb16e3b..14bba8a 100644 --- a/docs/plugins/data-designer-retrieval-sdg/multimodal.md +++ b/docs/plugins/data-designer-retrieval-sdg/multimodal.md @@ -264,6 +264,11 @@ reviewed relevance labels when comparing retrieval models. ## Resume an Interrupted Run +For runs that should continue past persistently malformed model responses, +set `failed_record_policy: skip` when creating a new run. See +[Continue After Failed Records](multimodal-reference.md#continue-after-failed-records) +for stage behavior and failure diagnostics. The default remains `error`. + Preserve the run directory and inspect the error and saved artifacts. After resolving an external issue, set `resume: true` in the same configuration and run the command again. Keep the inputs, remaining configuration, package code, and diff --git a/plugins/data-designer-retrieval-sdg/docs/multimodal-reference.md b/plugins/data-designer-retrieval-sdg/docs/multimodal-reference.md index 5afa89f..55581dc 100644 --- a/plugins/data-designer-retrieval-sdg/docs/multimodal-reference.md +++ b/plugins/data-designer-retrieval-sdg/docs/multimodal-reference.md @@ -20,6 +20,7 @@ image paths inside source records resolve relative to their JSONL file. | `concurrency` | `8` | Model-call concurrency, from 1 to 128. | | `batch_size` | `30` | Request batch size, from 1 to 128. | | `missing_response_attempts` | `3` | Total attempts for missing or invalid rows in normally completed batches, from 1 to 3. | +| `failed_record_policy` | `error` | Use `skip` to record exhausted model-validation failures and continue with usable records. | | `seed` | `42` | Seed for instruction sampling, selection, and query-group splitting. Model output can still vary. | | `resume` | `false` | Reuse a compatible run and cached responses. | | `ratios` | `{train: 0.8, validation: 0.0, evaluation: 0.2}` | Query-group split fractions, each from 0 to 1 and summing to 1. | @@ -41,6 +42,58 @@ provider retry behavior. Batch size and retry settings do not define a total API-call, cost, or elapsed-time budget. Refer to [Resume an Interrupted Run](multimodal.md#resume-an-interrupted-run) for recovery. +### Continue After Failed Records + +For a new run, add this to the existing YAML configuration: + +```yaml +failed_record_policy: skip +``` + +The default `error` policy stops when required structured responses remain +unavailable. With `skip`, the structured column records a diagnostic after Data +Designer's native validation retries are exhausted. This adds no retry loop and +does not change the native retry budget. `missing_response_attempts` keeps its +separate meaning; a recorded failure is terminal and is not retried by that loop +or by resume. Unexplained missing rows, duplicate identities, provider errors, +and integrity errors are never converted into skipped records. + +| Failed stage | Disposition | +| --- | --- | +| TOC confirmation | Preserve the source and omit unverified TOC projections. | +| Visual or document enrichment | Preserve original text/images and omit the failed enrichment. | +| Summary or summary combination | Omit the unavailable summary from generation inputs. | +| Summary judgment | Exclude the summary; a failed judgment cannot become a pass. | +| Query generation | Exclude the context's requested query slots. | +| Query metadata, self-sufficiency, relevance, or localization | Exclude the affected query; successful responses remain cached. | + +The original corpus and quality thresholds are preserved. If no queries pass, +export still fails. Inspect `generation_failures.json` in the run directory and +portable bundle for stage, request ID, source/context/query identity where +applicable, validation detail, native attempt limit, and the last invalid +response when available. The portable ledger is included in the manifest's +checksums. `report.json.generation_failures` counts failures by stage, failed +requested slots, and excluded query IDs. Failed query-generation slots contribute +to `requested_slots`, separately from model abstentions and quality rejections. +Earlier planning failures have no known query-slot count. +Queries that fail a later judgment or localization still count in +`funnel.generated` and the `retention_fraction` denominator, once per query. +For example, three generated queries with one localization failure and two +accepted queries report three generated and a retention fraction of `2/3`. + +The same policy is supported by `RetrievalStructuredColumnConfig`. In skip mode, +non-null output cells are JSON strings containing either a validated response +or a `_retrieval_sdg_failure` diagnostic. The multimodal workflow decodes these +before its typed consumers. Direct column users must handle the diagnostic +before treating a cell as a model response. Failed records have `None` in any +configured trace and reasoning-content columns, including when every record in +a batch fails, so downstream expressions can still reference those columns. + +Resume preserves both cached successes and recorded failures. Changing the +policy or upgrading the implementation changes the frozen run identity; use a +new output directory. Failure rates and exclusions can affect sample composition, +so include them when comparing generation yield. + ## Source Records Each line in `sources_file` is one JSON object describing an independently diff --git a/plugins/data-designer-retrieval-sdg/docs/multimodal.md b/plugins/data-designer-retrieval-sdg/docs/multimodal.md index cb16e3b..14bba8a 100644 --- a/plugins/data-designer-retrieval-sdg/docs/multimodal.md +++ b/plugins/data-designer-retrieval-sdg/docs/multimodal.md @@ -264,6 +264,11 @@ reviewed relevance labels when comparing retrieval models. ## Resume an Interrupted Run +For runs that should continue past persistently malformed model responses, +set `failed_record_policy: skip` when creating a new run. See +[Continue After Failed Records](multimodal-reference.md#continue-after-failed-records) +for stage behavior and failure diagnostics. The default remains `error`. + Preserve the run directory and inspect the error and saved artifacts. After resolving an external issue, set `resume: true` in the same configuration and run the command again. Keep the inputs, remaining configuration, package code, and diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/config.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/config.py index 26321ff..d62b702 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/config.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/config.py @@ -12,7 +12,9 @@ class RetrievalStructuredColumnConfig(LLMStructuredColumnConfig): """Generate schema-validated retrieval JSON, with or without a Markdown fence. All model, image, prompt, skip, and trace settings retain their native - Data Designer meanings. Only response decoding uses the plugin parser. + Data Designer meanings. The optional skip policy records exhausted model + validation failures as tagged diagnostics instead of dropping the row. """ column_type: Literal["retrieval-structured"] = "retrieval-structured" + failed_record_policy: Literal["error", "skip"] = "error" diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/export.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/export.py index 93c78fb..c77f994 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/export.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/export.py @@ -6,6 +6,7 @@ from __future__ import annotations import errno +import json import tempfile from collections import Counter from pathlib import Path @@ -308,7 +309,24 @@ def write_multimodal_bundle( } from data_designer_retrieval_sdg.multimodal.reporting import generation_report - counts.update(generation_report(sources, candidates, outcomes)) + failed_query_ids = set() + if config.failed_record_policy == "skip": + failures = json.loads((config.output_dir / "generation_failures.json").read_text()) + write_json(root / "generation_failures.json", failures) + failed_query_ids = {row["query_id"] for row in failures["records"] if "query_id" in row} + failed_slots = sum( + len(row["instructions"]) for row in outcomes if row.get("selection_reason") == "generation_failed" + ) + counts["requested_slots"] += failed_slots + counts["generation_failures"] = { + "policy": "skip", + "failed_records": len(failures["records"]), + "failed_requested_slots": failed_slots, + "excluded_query_ids": len(failed_query_ids), + "by_stage": dict(Counter(row["stage"] for row in failures["records"])), + "diagnostics": "generation_failures.json", + } + counts.update(generation_report(sources, candidates, outcomes, failed_query_ids)) write_json(root / "report.json", counts) write_json( root / "run_manifest.json", diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/inference.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/inference.py index c5dcef2..75237e5 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/inference.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/inference.py @@ -22,9 +22,10 @@ from data_designer_retrieval_sdg.multimodal.bounds import check_request from data_designer_retrieval_sdg.multimodal.config import RetrievalStructuredColumnConfig -from data_designer_retrieval_sdg.multimodal.models import GenerationContext, MultimodalSDGConfig +from data_designer_retrieval_sdg.multimodal.models import FailedGeneration, GenerationContext, MultimodalSDGConfig from data_designer_retrieval_sdg.multimodal.source_models import RetrievalSource from data_designer_retrieval_sdg.multimodal.storage import digest, fingerprint, write_json +from data_designer_retrieval_sdg.multimodal.structured import FAILURE_KEY @dataclass(frozen=True) @@ -60,6 +61,73 @@ def source_request( ) +def request_identity(context: GenerationContext, **details) -> dict: + """Attach source membership and optional query/slot details to a failure record. + + Args: + context: Original context owning the request. + **details: Additional query, slot, or instruction provenance. + + Returns: + JSON-compatible identity retained with any generation failure. + """ + return {"context_id": context.context_id, "source_unit_ids": context.unit_ids, **details} + + +def generate_records( + inference: DataDesignerInference, requests: list[Request], stage: str, records: list[dict] +) -> list[BaseModel]: + """Generate aligned responses and retain contextual diagnostics for failed ones. + + Args: + inference: Run-owned response cache and native Data Designer transport. + requests: Ordered requests, including any identical requests shared by records. + stage: Workflow stage responsible for these requests. + records: One context/source/query identity per request. + + Returns: + Validated responses or explicit failures in the original record order. + Consumers decide the stage-specific disposition without shifting alignment. + + Raises: + ValueError: Request, identity, or response counts disagree. + """ + if len(requests) != len(records): + raise ValueError("Every generation request requires a record identity") + responses = inference.generate(requests) + for request, record, response in zip(requests, records, responses, strict=True): + if isinstance(response, FailedGeneration): + key = inference.request_key(request) + identity = fingerprint([stage, key, record]) + inference.failures[identity] = { + **record, + "stage": stage, + "request_id": key, + "failure": response.model_dump(mode="json"), + } + return responses + + +def parse_failure(value: object, policy: str) -> FailedGeneration | None: + """Accept only a well-formed diagnostic under the explicit skip policy. + + Args: + value: Decoded native response or cached payload. + policy: Configured error/skip policy. + + Returns: + The typed failure, or None for an ordinary response. + + Raises: + ValueError: A failure marker is malformed or the skip policy is disabled. + """ + if not isinstance(value, dict) or FAILURE_KEY not in value: + return None + if policy != "skip" or set(value) != {FAILURE_KEY}: + raise ValueError("Saved failure requires the skip policy and a valid diagnostic") + return FailedGeneration.model_validate(value[FAILURE_KEY]) + + class DataDesignerInference: """Run schema-homogeneous batches with immutable per-response caches. @@ -71,6 +139,7 @@ class DataDesignerInference: def __init__(self, root: Path, config: MultimodalSDGConfig): self.root = root self.config = config + self.failures: dict[str, dict] = {} def request_key(self, request: Request) -> str: """Bind responses to exact prompts, schema, pixels and model settings.""" @@ -97,7 +166,8 @@ def cached(self, request: Request) -> BaseModel | None: value = json.loads(path.read_text()) if value["request_id"] != key or fingerprint(value["response"]) != value["response_sha256"]: raise ValueError(f"Response cache integrity failure: {key}") - return request.schema.model_validate(value["response"]) + failure = parse_failure(value["response"], self.config.failed_record_policy) + return failure if failure is not None else request.schema.model_validate(value["response"]) def generate(self, requests: list[Request]) -> list[BaseModel]: """Retry only missing rows within the configured finite attempt budget. @@ -105,7 +175,9 @@ def generate(self, requests: list[Request]) -> list[BaseModel]: Missing outputs fail the stage after preserving successful responses and attempt evidence. Exceptions stop immediately; only normally completed batches with missing/invalid rows get another attempt. Native DD's own - schema corrections remain bounded inside each attempt. + schema corrections remain bounded inside each attempt. In skip mode, + explicit exhausted validation failures are cached terminal results; + unexplained missing rows still fail the stage. """ groups: dict[str, dict[str, Request]] = {} for request in requests: @@ -169,6 +241,7 @@ def run_batch(self, items: list[tuple[str, Request]]) -> None: model_alias="retrieval", prompt="{{ request_text }}", output_format=request.schema, + failed_record_policy=self.config.failed_record_policy, multi_modal_context=[dd.ImageContext(column_name="request_images")] if request.images else None, ) ) @@ -226,6 +299,7 @@ def collect_responses(self, items: list[tuple[str, Request]], dataset: Path, att raise ValueError("Unexpected or duplicate Data Designer response identity") seen.add(key) found: set[str] = set() + failed: list[str] = [] invalid: list[str] = [] for batch in batches: for row in pq.read_table(batch).to_pylist(): @@ -235,15 +309,20 @@ def collect_responses(self, items: list[tuple[str, Request]], dataset: Path, att raw = row["response"] request = expected[key] try: - response = ( - request.schema.model_validate_json(raw) - if isinstance(raw, str) - else request.schema.model_validate(raw) - ) + value = json.loads(raw) if isinstance(raw, str) else raw except ValueError: invalid.append(key) continue - payload = response.model_dump(mode="json") + failure = parse_failure(value, self.config.failed_record_policy) + if failure is not None: + payload = {FAILURE_KEY: failure.model_dump(mode="json")} + failed.append(key) + else: + try: + payload = request.schema.model_validate(value).model_dump(mode="json") + except ValueError: + invalid.append(key) + continue write_json( self.root / "responses" / f"{key}.json", { @@ -257,8 +336,9 @@ def collect_responses(self, items: list[tuple[str, Request]], dataset: Path, att attempt / "outcome.json", { "requested": len(items), - "returned": len(found), + "returned": len(found) - len(failed), "missing": sorted(expected.keys() - found), "invalid": invalid, + **({"failed": failed} if self.config.failed_record_policy == "skip" else {}), }, ) diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/models.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/models.py index 9430968..1d4059f 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/models.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/models.py @@ -44,6 +44,15 @@ def explicit_budget(self): return self +class FailedGeneration(StrictModel): + """An exhausted validation failure, never a generated answer or judgment.""" + + kind: Literal["model_validation_exhausted"] + detail: str = Field(min_length=1) + max_model_attempts: int = Field(ge=1) + last_invalid_response: str | None = None + + class QueryInstruction(StrictModel): """A requested information need; matching its style is diagnostic, not a quality gate.""" @@ -110,6 +119,7 @@ class MultimodalSDGConfig(StrictModel): summary_near_duplicate_threshold: float | None = Field(default=None, gt=0, le=1) instructions_per_context: int | None = Field(default=None, gt=0) missing_response_attempts: int = Field(default=3, ge=1, le=3) + failed_record_policy: Literal["error", "skip"] = "error" require_verbatim_quotes: bool = False instructions: list[QueryInstruction] = Field(default_factory=lambda: list(DEFAULT_INSTRUCTIONS), min_length=1) relevance_threshold: int = Field(default=4, ge=1, le=5) diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/planning.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/planning.py index 8698861..3174e3b 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/planning.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/planning.py @@ -162,6 +162,8 @@ def related_contexts( def summary_passes(row: dict, config: MultimodalSDGConfig) -> bool: """Apply the configured summary gate without inventing missing judgments.""" + if row.get("summary_judgment_failed"): + return False judgment = row["summary_judgment"] return judgment is None or min(value["grade"] for value in judgment.values()) >= config.summary_quality_threshold @@ -216,7 +218,9 @@ def select_summaries(rows: list[dict], config: MultimodalSDGConfig) -> list[dict for row in sorted(rows, key=summary_rank): row["selection_reason"] = "selected" if not summary_passes(row, config): - row["selection_reason"] = "summary_quality" + row["selection_reason"] = ( + "summary_judgment_failed" if row.get("summary_judgment_failed") else "summary_quality" + ) continue duplicate = next((other for other in kept if duplicate_summary(row, other, config)), None) if duplicate is not None: diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/reporting.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/reporting.py index 38f01b1..c4a49ed 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/reporting.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/reporting.py @@ -9,13 +9,19 @@ from data_designer_retrieval_sdg.multimodal.source_models import RetrievalSource -def generation_report(sources: list[RetrievalSource], candidates: list[Candidate], outcomes: list[dict]) -> dict: +def generation_report( + sources: list[RetrievalSource], + candidates: list[Candidate], + outcomes: list[dict], + failed_query_ids: set[str] | None = None, +) -> dict: """Count the source coverage, sequential funnel and independent diagnostic labels. Args: sources: Full canonical corpus, including distractors. candidates: Accepted and rejected records. outcomes: All selected, filtered and budgeted summaries. + failed_query_ids: Generated queries excluded by later model-validation failures. Returns: JSON-compatible diagnostic counts. Requested labels never gate acceptance. @@ -31,7 +37,8 @@ def generation_report(sources: list[RetrievalSource], candidates: list[Candidate accepted = [c for c in candidates if c.accepted] positives = {s.unit_id for c in accepted for s in c.localization.supports} survivors = candidates - funnel = {"generated": len(candidates)} + generated_count = len({c.query_id for c in candidates} | (failed_query_ids or set())) + funnel = {"generated": generated_count} for name in ("insufficient_standalone_query", "query_contains_answer", "insufficient_relevance"): survivors = [c for c in survivors if name not in c.rejection_reasons] funnel["after_" + name] = len(survivors) @@ -47,7 +54,7 @@ def generation_report(sources: list[RetrievalSource], candidates: list[Candidate for stage, ids in {"planned": planned, "generated": generated, "accepted_positive": positives}.items() }, "funnel": funnel, - "retention_fraction": len(accepted) / len(candidates) if candidates else 0.0, + "retention_fraction": len(accepted) / generated_count if generated_count else 0.0, "image_dependent_fraction": sum(any(s.modality != "text" for s in c.localization.supports) for c in accepted) / len(accepted) if accepted diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/sections.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/sections.py index e78b886..2f447b8 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/sections.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/sections.py @@ -11,8 +11,13 @@ from data_designer_retrieval_sdg.multimodal import prompts from data_designer_retrieval_sdg.multimodal.bounds import check_request -from data_designer_retrieval_sdg.multimodal.inference import Request -from data_designer_retrieval_sdg.multimodal.models import GenerationContext, MultimodalSDGConfig, TOCConfirmation +from data_designer_retrieval_sdg.multimodal.inference import Request, generate_records +from data_designer_retrieval_sdg.multimodal.models import ( + FailedGeneration, + GenerationContext, + MultimodalSDGConfig, + TOCConfirmation, +) from data_designer_retrieval_sdg.multimodal.planning import automatic_contexts, context_chars from data_designer_retrieval_sdg.multimodal.source_models import RetrievalSource from data_designer_retrieval_sdg.multimodal.storage import fingerprint, write_json @@ -93,16 +98,22 @@ def confirm_tocs(sources, config, inference) -> dict[str, dict]: config, ) responses = ( - inference.generate( + generate_records( + inference, [ Request(prompts.render("check_toc", page=s.text, language=s.language), TOCConfirmation, "judge") for s in bounded - ] + ], + "toc-confirmation", + [{"source_unit_ids": [s.unit_id]} for s in bounded], ) if bounded else [] ) - result = {s.unit_id: response.model_dump() for s, response in zip(bounded, responses, strict=True)} + result = { + s.unit_id: {"generation_failed": True} if isinstance(response, FailedGeneration) else response.model_dump() + for s, response in zip(bounded, responses, strict=True) + } for source in candidates: if source.unit_id not in result: result[source.unit_id] = { @@ -316,7 +327,7 @@ def section_contexts( Without inference, only headings contribute structural evidence. Explicit caller contexts bypass this automatic planner in the workflow. Confirmation - failures propagate through the existing retry policy, never become approvals. + failures follow the configured error/skip policy and never become approvals. """ confirmations = confirm_tocs(sources, config, inference) if inference is not None else {} contexts, audit = build_sections(sources, config, confirmations) diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/structured.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/structured.py index 3ba8b7f..c2209d8 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/structured.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/structured.py @@ -12,10 +12,14 @@ from typing import Any from data_designer.engine.column_generators.generators.llm_completion import LLMStructuredCellGenerator +from data_designer.engine.models.errors import ModelGenerationValidationFailureError from data_designer.engine.models.parsers.errors import ParserException from data_designer.engine.models.recipes.response_recipes import StructuredResponseRecipe +from data_designer_retrieval_sdg.multimodal.models import FailedGeneration + _JSON_FENCE = re.compile(r"```json[ \t]*\r?\n(?P.*?)\r?\n```", re.DOTALL | re.IGNORECASE) +FAILURE_KEY = "_retrieval_sdg_failure" def _reject_constant(value: str) -> None: @@ -81,3 +85,82 @@ class RetrievalStructuredCellGenerator(LLMStructuredCellGenerator): def response_recipe(self) -> RetrievalStructuredResponseRecipe: """Return the plugin parser with native schema validation and no pruning.""" return RetrievalStructuredResponseRecipe(json_schema=self.config.output_format, pruning=False) + + def record_result(self, data: dict) -> dict: + """Persist uniform JSON strings in skip mode to avoid mixed Arrow structs. + + Args: + data: Native row containing the validated response or failure diagnostic. + + Returns: + The row with a JSON-string cell in skip mode, unchanged in error mode. + """ + if self.config.failed_record_policy == "skip": + data[self.config.name] = json.dumps(data[self.config.name], ensure_ascii=False, allow_nan=False) + return data + + def record_failure(self, data: dict, error: ModelGenerationValidationFailureError) -> dict: + """Retain a validation diagnostic under the original native row identity. + + Args: + data: Original row, including successful sibling columns. + error: Validation failure after Data Designer's native retry budget. + + Returns: + The row with a tagged diagnostic in the failed column. + + Raises: + ModelGenerationValidationFailureError: Under the default error policy. + """ + if self.config.failed_record_policy != "skip": + raise error + cause, response, seen = error, None, set() + while cause is not None and id(cause) not in seen: + seen.add(id(cause)) + if isinstance(cause, ParserException) and cause.source is not None: + response = cause.source + cause = cause.__cause__ or cause.__context__ + failure = FailedGeneration( + kind="model_validation_exhausted", + detail=error.detail or str(error), + max_model_attempts=(self.max_conversation_correction_steps + 1) * (self.max_conversation_restarts + 1), + last_invalid_response=response, + ) + data[self.config.name] = {FAILURE_KEY: failure.model_dump()} + for column in self.config.side_effect_columns: + data[column] = None + return self.record_result(data) + + def generate(self, data: dict) -> dict: + """Apply the opt-in policy after native validation retries are exhausted. + + Args: + data: Native input row. + + Returns: + The row with a validated response or an explicitly enabled diagnostic. + + Raises: + ModelGenerationValidationFailureError: Under the default error policy. + """ + try: + return self.record_result(super().generate(data)) + except ModelGenerationValidationFailureError as error: + return self.record_failure(data, error) + + async def agenerate(self, data: dict) -> dict: + """Apply the same policy while retaining native asynchronous generation. + + Args: + data: Native input row. + + Returns: + The row with a validated response or an explicitly enabled diagnostic. + + Raises: + ModelGenerationValidationFailureError: Under the default error policy. + """ + try: + return self.record_result(await super().agenerate(data)) + except ModelGenerationValidationFailureError as error: + return self.record_failure(data, error) diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/summary_stages.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/summary_stages.py index 673311c..b520b40 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/summary_stages.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/summary_stages.py @@ -12,10 +12,11 @@ from data_designer_retrieval_sdg.multimodal import prompts from data_designer_retrieval_sdg.multimodal.bounds import fit_contexts from data_designer_retrieval_sdg.multimodal.combinations import semantic_combinations -from data_designer_retrieval_sdg.multimodal.inference import Request, source_request +from data_designer_retrieval_sdg.multimodal.inference import Request, generate_records, request_identity, source_request from data_designer_retrieval_sdg.multimodal.models import ( ContextSummary, Description, + FailedGeneration, GenerationContext, SummaryJudgment, VisualDescription, @@ -84,7 +85,9 @@ def summary_requests(context, sources, visual, descriptions=None): section=enriched_text(context, sources, visual), language=context.language, document_description="\n".join( - dict.fromkeys(descriptions[(sources[k].document_id, sources[k].language)] for k in context.unit_ids) + dict.fromkeys( + descriptions.get((sources[k].document_id, sources[k].language), "") for k in context.unit_ids + ) ), ) schema = ContextSummary @@ -105,8 +108,14 @@ def visual_enrichment(sources, inference) -> dict: ) for s in selected ] - responses = inference.generate(requests) - return {s.unit_id: r.model_dump() for s, r in zip(selected, responses, strict=True)} + responses = generate_records( + inference, requests, "visual-enrichment", [{"source_unit_ids": [s.unit_id]} for s in selected] + ) + return { + s.unit_id: r.model_dump() + for s, r in zip(selected, responses, strict=True) + if not isinstance(r, FailedGeneration) + } def describe_documents(sources, visual, config, inference) -> dict: @@ -116,9 +125,13 @@ def describe_documents(sources, visual, config, inference) -> dict: chunks = bounded_enriched_contexts(documents, sources, visual, config) builder = partial(summary_requests, sources=by_id, visual=visual) chunks = fit_contexts(chunks, builder, config) - responses = inference.generate([builder(c)[0] for c in chunks]) + responses = generate_records( + inference, [builder(c)[0] for c in chunks], "document-description", [request_identity(c) for c in chunks] + ) descriptions = defaultdict(list) for context, response in zip(chunks, responses, strict=True): + if isinstance(response, FailedGeneration): + continue descriptions[(by_id[context.unit_ids[0]].document_id, context.language)].append(response.description) return {key: "\n".join(values) for key, values in descriptions.items()} @@ -181,8 +194,14 @@ def plan_summaries(sources, contexts, config, inference) -> list[dict]: contexts = fit_contexts(contexts, builder, config) if config.contexts_file is None: verify_section_coverage(contexts, sources, 2 * config.section_size) - summaries = inference.generate([builder(c)[0] for c in contexts]) - rows = [summary_row(c, s, by_id, visual) for c, s in zip(contexts, summaries, strict=True)] + summaries = generate_records( + inference, [builder(c)[0] for c in contexts], "summary", [request_identity(c) for c in contexts] + ) + rows = [ + summary_row(c, s, by_id, visual) + for c, s in zip(contexts, summaries, strict=True) + if not isinstance(s, FailedGeneration) + ] combinations = semantic_combinations(rows, config) combination_builder = partial(combination_requests, rows=rows) combination_contexts = [ @@ -204,8 +223,21 @@ def plan_summaries(sources, contexts, config, inference) -> list[dict]: "members": [[rows[i]["context"]["context_id"] for i in group] for group in combinations], }, ) - combined = inference.generate([combination_builder(c)[0] for c in combination_contexts]) + combined = generate_records( + inference, + [combination_builder(c)[0] for c in combination_contexts], + "summary-combinations", + [ + { + "summary_context_ids": [rows[i]["context"]["context_id"] for i in group], + "source_unit_ids": list(dict.fromkeys(key for i in group for key in rows[i]["context"]["unit_ids"])), + } + for group in combinations + ], + ) for group, summary in zip(combinations, combined, strict=True): + if isinstance(summary, FailedGeneration): + continue ids = list(dict.fromkeys(key for i in group for key in rows[i]["context"]["unit_ids"])) context = GenerationContext( context_id="combined_" + fingerprint([rows[i]["context"]["context_id"] for i in group])[:32], @@ -214,7 +246,8 @@ def plan_summaries(sources, contexts, config, inference) -> list[dict]: ) rows.append(summary_row(context, summary, by_id, visual, combined=True)) if config.judge_summaries: - judgments = inference.generate( + judgments = generate_records( + inference, [ Request( prompts.render( @@ -227,9 +260,14 @@ def plan_summaries(sources, contexts, config, inference) -> list[dict]: "judge", ) for row in rows - ] + ], + "summary-judgment", + [request_identity(GenerationContext.model_validate(row["context"])) for row in rows], ) for row, judgment in zip(rows, judgments, strict=True): - row["summary_judgment"] = judgment.model_dump() + if isinstance(judgment, FailedGeneration): + row["summary_judgment_failed"] = True + else: + row["summary_judgment"] = judgment.model_dump() write_json(root / "summaries.json", rows) return rows diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/workflow.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/workflow.py index c5f0947..52d6a93 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/workflow.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/multimodal/workflow.py @@ -21,10 +21,17 @@ split_context, validate_deployment, ) -from data_designer_retrieval_sdg.multimodal.inference import DataDesignerInference, Request, source_request +from data_designer_retrieval_sdg.multimodal.inference import ( + DataDesignerInference, + Request, + generate_records, + request_identity, + source_request, +) from data_designer_retrieval_sdg.multimodal.models import ( Candidate, ContextSummary, + FailedGeneration, GenerationContext, Localization, MultimodalSDGConfig, @@ -164,9 +171,15 @@ def summarize_contexts(contexts, sources, config, inference) -> list[dict]: """ builder = partial(evidence_summary_requests, sources=sources) contexts = fit_contexts(contexts, builder, config) - summaries = inference.generate([builder(c)[0] for c in contexts]) + summaries = generate_records( + inference, [builder(c)[0] for c in contexts], "summary", [request_identity(c) for c in contexts] + ) + successful = [(c, s) for c, s in zip(contexts, summaries, strict=True) if not isinstance(s, FailedGeneration)] + contexts = [c for c, _ in successful] + summaries = [s for _, s in successful] judgments = ( - inference.generate( + generate_records( + inference, [ Request( prompts.render("judgment", summary=s.summary, persona=config.persona, language=c.language), @@ -174,7 +187,9 @@ def summarize_contexts(contexts, sources, config, inference) -> list[dict]: "judge", ) for c, s in zip(contexts, summaries, strict=True) - ] + ], + "summary-judgment", + [request_identity(c) for c in contexts], ) if config.judge_summaries else [None] * len(contexts) @@ -183,7 +198,8 @@ def summarize_contexts(contexts, sources, config, inference) -> list[dict]: { "context": c.model_dump(), "summary": s.model_dump(), - "summary_judgment": j.model_dump() if j else None, + "summary_judgment": j.model_dump() if j is not None and not isinstance(j, FailedGeneration) else None, + **({"summary_judgment_failed": True} if isinstance(j, FailedGeneration) else {}), "slots": {"queries": []}, } for c, s, j in zip(contexts, summaries, judgments, strict=True) @@ -283,9 +299,18 @@ def bounded_query_batches(selected, outcomes, sources, config, inference): row["instructions"] = [ item.model_dump() for item in context_instructions(config, row["context"]["context_id"]) ] - responses = inference.generate([builder(GenerationContext.model_validate(row["context"]))[0] for row in rows]) + contexts = [GenerationContext.model_validate(row["context"]) for row in rows] + responses = generate_records( + inference, + [builder(context)[0] for context in contexts], + "query-generation", + [request_identity(c, instructions=row["instructions"]) for c, row in zip(contexts, rows, strict=True)], + ) pending = [] for row, batch in zip(rows, responses, strict=True): + if isinstance(batch, FailedGeneration): + row["selection_reason"] = "generation_failed" + continue context = GenerationContext.model_validate(row["context"]) try: for slot in batch.queries: @@ -358,8 +383,27 @@ def generate_candidates( raise ValueError("Generated query must be nonempty and declare evidence") pending.append((context, slot)) judgments = [judgment_requests(c, slot, by_id) for c, slot in pending] - metadata = inference.generate([requests[0] for requests in judgments]) - sufficiency = inference.generate([requests[1] for requests in judgments]) + records = [ + request_identity( + c, + query_id="q_" + fingerprint([c.context_id, slot.slot, slot.query])[:32], + query=slot.query, + slot=slot.slot, + instructions=[instructions[c.context_id][slot.slot].model_dump()], + ) + for c, slot in pending + ] + metadata = generate_records(inference, [requests[0] for requests in judgments], "query-metadata", records) + sufficiency = generate_records(inference, [requests[1] for requests in judgments], "self-sufficiency", records) + relevance = generate_records(inference, [requests[2] for requests in judgments], "relevance", records) + valid = [ + i + for i, values in enumerate(zip(metadata, sufficiency, relevance, strict=True)) + if not any(isinstance(value, FailedGeneration) for value in values) + ] + pending, judgments, records, metadata, sufficiency, relevance = ( + [items[i] for i in valid] for items in (pending, judgments, records, metadata, sufficiency, relevance) + ) query_judgments = [ QueryJudgment( self_sufficiency=s.self_sufficiency, @@ -370,10 +414,14 @@ def generate_candidates( ) for s, m in zip(sufficiency, metadata, strict=True) ] - relevance = inference.generate([requests[2] for requests in judgments]) reasons_by_index = [quality_reasons(q, r, config) for q, r in zip(query_judgments, relevance, strict=True)] localized_indexes = [i for i, reasons in enumerate(reasons_by_index) if not reasons] - localizations = inference.generate([judgments[i][3] for i in localized_indexes]) + localizations = generate_records( + inference, + [judgments[i][3] for i in localized_indexes], + "localization", + [records[i] for i in localized_indexes], + ) by_index = dict(zip(localized_indexes, localizations, strict=True)) candidates = [] for i, ((context, slot), query_judge, relevance_judge) in enumerate( @@ -381,6 +429,8 @@ def generate_candidates( ): reasons = reasons_by_index[i] localization = by_index.get(i) + if isinstance(localization, FailedGeneration): + continue if localization is not None: localization = normalize_localization(localization, context) reasons.extend(localization_reasons(localization, context, by_id, config.require_verbatim_quotes)) @@ -507,11 +557,25 @@ def run_multimodal_sdg(config: MultimodalSDGConfig) -> Path: frozen.append(source.model_copy(update={"images": images})) inference = DataDesignerInference(root / "inference", config) candidates, outcomes = generate_candidates(config, frozen, contexts, inference) + if config.failed_record_policy == "skip": + write_json( + root / "generation_failures.json", + {"policy": "skip", "records": [inference.failures[key] for key in sorted(inference.failures)]}, + ) write_json(root / "candidates.json", [c.model_dump(mode="json") for c in candidates]) write_json(root / "context_outcomes.json", outcomes) handoff = export_multimodal_bundle(root / "bundle", frozen, candidates, outcomes, config) write_json( root / "complete.json", - {"artifacts": [artifact(p, root) for p in sorted((root / "bundle").rglob("*")) if p.is_file()]}, + { + "artifacts": [ + artifact(p, root) + for p in [ + *sorted((root / "bundle").rglob("*")), + *([root / "generation_failures.json"] if config.failed_record_policy == "skip" else []), + ] + if p.is_file() + ] + }, ) return handoff diff --git a/plugins/data-designer-retrieval-sdg/tests/multimodal/test_failed_records.py b/plugins/data-designer-retrieval-sdg/tests/multimodal/test_failed_records.py new file mode 100644 index 0000000..ceaaca0 --- /dev/null +++ b/plugins/data-designer-retrieval-sdg/tests/multimodal/test_failed_records.py @@ -0,0 +1,313 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Failed records use the real DD parser, retries, persistence, workflow and export.""" + +from __future__ import annotations + +import json +import re +from collections import Counter +from pathlib import Path +from unittest.mock import AsyncMock + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from data_designer.engine.models.facade import ModelFacade +from data_designer.engine.testing import make_stub_completion_response +from data_designer.interface.errors import DataDesignerGenerationError +from test_multimodal_sdg import fixture_config, scripted_response + +from data_designer_retrieval_sdg.multimodal import models, run_multimodal_sdg +from data_designer_retrieval_sdg.multimodal.inference import DataDesignerInference, Request +from data_designer_retrieval_sdg.multimodal.storage import digest +from data_designer_retrieval_sdg.multimodal.structured import FAILURE_KEY + + +class ControlledTransport: + """Fail one schema/source at the transport boundary without replacing generation.""" + + def __init__(self, schema: str | tuple[str, ...], unit: str | None = "unit-0"): + self.schemas = {schema} if isinstance(schema, str) else set(schema) + self.unit = unit + self.provider_unit: str | None = None + self.prompt_contains: str | None = None + self.calls = Counter() + + async def complete(self, messages, **kwargs): + """Read the real rendered schema and return controlled model text.""" + content = messages[0].to_dict()["content"] + text = ( + "\n".join(item["text"] for item in content if item["type"] == "text") + if isinstance(content, list) + else content + ) + schema = json.loads(text.split("\n")[1].split("\n")[0]) + title = schema["title"] + model = ( + models.QueryBatch.for_slot_count(schema["properties"]["queries"]["minItems"]) + if title.startswith("QueryBatch") + else getattr(models, title) + ) + prompt = text.split("\n\n* Your response must be in JSON format.")[0] + match = re.search(r"unit-\d+", prompt) + unit = match.group() if match else "unselected" + name = "QueryBatch" if issubclass(model, models.QueryBatch) else title + self.calls[name, unit] += 1 + if name == "QueryBatch" and unit == self.provider_unit: + raise PermissionError("Controlled provider permission failure") + if ( + name in self.schemas + and (self.unit is None or unit == self.unit) + and (self.prompt_contains is None or self.prompt_contains in prompt) + ): + return make_stub_completion_response(content="not JSON") + if model is models.ContextSummary: + response = models.ContextSummary(summary=f"Valve specification for {unit}", visual_evidence="") + elif model is models.VisualDescription: + response = models.VisualDescription(description="A pressure diagram", has_visual_content=True) + elif model is models.Description: + response = models.Description(description="Valve documentation") + else: + response = scripted_response(Request(prompt, model, "generator")) + return make_stub_completion_response(content=response.model_dump_json()) + + +def install_transport(monkeypatch, transport: ControlledTransport) -> AsyncMock: + """Replace only the hosted model call; every parser and workflow consumer is real.""" + monkeypatch.setenv("TEST_SDG_KEY", "mocked-credential") + completion = AsyncMock(side_effect=transport.complete) + monkeypatch.setattr(ModelFacade, "acompletion", completion) + return completion + + +def skip_config(tmp_path: Path, strategy: str = "sections") -> models.MultimodalSDGConfig: + """Keep three contexts plus a corpus-only distractor in a small real workflow.""" + config = fixture_config(tmp_path) + sources = [json.loads(line) for line in config.sources_file.read_text().splitlines()] + config.sources_file.write_text("".join(json.dumps(row) + "\n" for row in [*sources[:3], sources[-1]])) + config.contexts_file.write_text( + "".join(json.dumps({"context_id": f"context-{i}", "unit_ids": [f"unit-{i}"]}) + "\n" for i in range(3)) + ) + return config.model_copy(update={"context_strategy": strategy, "failed_record_policy": "skip", "concurrency": 1}) + + +@pytest.mark.parametrize( + "schema,stage,accepted", + [ + ("VisualDescription", "visual-enrichment", 3), + ("Description", "document-description", 3), + ("ContextSummary", "summary", 2), + ("SummaryJudgment", "summary-judgment", 2), + ("QueryBatch", "query-generation", 2), + ("QueryMetadata", "query-metadata", 2), + ("SelfSufficiencyJudgment", "self-sufficiency", 2), + ("RelevanceJudgment", "relevance", 2), + ("Localization", "localization", 2), + ], +) +def test_failed_record_preserves_other_queries_and_completed_resume(tmp_path, monkeypatch, schema, stage, accepted): + config = skip_config(tmp_path) + original_sources = config.sources_file.read_bytes() + transport = ControlledTransport(schema) + completion = install_transport(monkeypatch, transport) + handoff = run_multimodal_sdg(config) + report = json.loads((handoff.parent / "report.json").read_text()) + ledger = json.loads((handoff.parent / "generation_failures.json").read_text()) + assert ledger == json.loads((config.output_dir / "generation_failures.json").read_text()) + assert report["accepted_query_ids"] == accepted + assert report["corpus_units"] == 4 + assert config.sources_file.read_bytes() == original_sources + assert len(ledger["records"]) == 1 + record = ledger["records"][0] + generated = accepted + int("query_id" in record) + assert report["funnel"]["generated"] == generated + assert report["retention_fraction"] == pytest.approx(accepted / generated) + assert record["stage"] == stage + assert "unit-0" in record["source_unit_ids"] + assert record["failure"]["kind"] == "model_validation_exhausted" + assert transport.calls[schema, "unit-0"] == record["failure"]["max_model_attempts"] + assert report["generation_failures"]["by_stage"] == {stage: 1} + if schema == "QueryBatch": + assert report["requested_slots"] == 9 + assert report["generation_failures"]["failed_requested_slots"] == 3 + assert report["abstentions"] == 4 + if schema == "SummaryJudgment": + outcomes = json.loads((handoff.parent / "context_outcomes.json").read_text()) + failed = next(row for row in outcomes if row["context"]["context_id"] == "context-0") + assert failed["selection_reason"] == "summary_judgment_failed" + assert failed["summary_judgment"] is None and not failed["slots"]["queries"] + if "query_id" in record: + candidates = [json.loads(line) for line in (handoff.parent / "source_records.jsonl").read_text().splitlines()] + assert record["query_id"] not in {row["query_id"] for row in candidates} + assert report["generation_failures"]["excluded_query_ids"] == 1 + files = {p: digest(p) for p in config.output_dir.rglob("*") if p.is_file()} + calls = completion.await_count + assert run_multimodal_sdg(config.model_copy(update={"resume": True})) == handoff + assert completion.await_count == calls + assert all(digest(p) == checksum for p, checksum in files.items()) + + +def test_multiple_judgment_failures_count_one_generated_query(tmp_path, monkeypatch): + config = skip_config(tmp_path) + install_transport(monkeypatch, ControlledTransport(("QueryMetadata", "RelevanceJudgment"))) + handoff = run_multimodal_sdg(config) + report = json.loads((handoff.parent / "report.json").read_text()) + assert report["generation_failures"]["failed_records"] == 2 + assert report["generation_failures"]["excluded_query_ids"] == 1 + assert report["candidate_records"] == 2 + assert report["funnel"]["generated"] == 3 + assert report["funnel"]["accepted_localized"] == 2 + assert report["retention_fraction"] == pytest.approx(2 / 3) + + +@pytest.mark.parametrize("schema", ["ContextSummary", "SummaryJudgment"]) +def test_unit_contexts_exclude_failed_summaries(tmp_path, monkeypatch, schema): + config = skip_config(tmp_path, "unit") + install_transport(monkeypatch, ControlledTransport(schema)) + handoff = run_multimodal_sdg(config) + assert json.loads((handoff.parent / "report.json").read_text())["accepted_query_ids"] == 2 + + +@pytest.mark.parametrize("schema", ["VisualDescription", "Description"]) +def test_all_enrichment_failures_keep_sources_and_allow_generation(tmp_path, monkeypatch, schema): + config = skip_config(tmp_path) + install_transport(monkeypatch, ControlledTransport(schema, unit=None)) + handoff = run_multimodal_sdg(config) + report = json.loads((handoff.parent / "report.json").read_text()) + assert report["corpus_units"] == 4 and report["accepted_query_ids"] == 3 + assert report["generation_failures"]["failed_records"] > 0 + + +def test_toc_failure_keeps_sources_without_unverified_projection(tmp_path, monkeypatch): + config = skip_config(tmp_path).model_copy(update={"contexts_file": None, "section_size": 1}) + rows = [json.loads(line) for line in config.sources_file.read_text().splitlines()] + rows[0]["text"] = "Table of contents\nAlpha ... 1\nBeta ... 2\nGamma ... 3" + config.sources_file.write_text("".join(json.dumps(row) + "\n" for row in rows)) + install_transport(monkeypatch, ControlledTransport("TOCConfirmation", unit=None)) + handoff = run_multimodal_sdg(config) + report = json.loads((handoff.parent / "report.json").read_text()) + assert report["corpus_units"] == 4 + assert report["generation_failures"]["by_stage"] == {"toc-confirmation": 1} + audit = json.loads((config.output_dir / "planning/section_boundaries.json").read_text()) + assert audit["coverage"]["exact_ordered_coverage"] + assert audit["toc_confirmations"]["unit-0"] == {"generation_failed": True} + assert not audit["documents"][0]["accepted"] + + +def one_combination(rows, config): + """Choose a deterministic pair without loading an embedding model.""" + return [(0, 1)] + + +def test_failed_combination_omits_only_unavailable_summary(tmp_path, monkeypatch): + config = skip_config(tmp_path) + transport = ControlledTransport("ContextSummary") + transport.prompt_contains = "Write a combined summary" + install_transport(monkeypatch, transport) + monkeypatch.setattr("data_designer_retrieval_sdg.multimodal.summary_stages.semantic_combinations", one_combination) + handoff = run_multimodal_sdg(config) + report = json.loads((handoff.parent / "report.json").read_text()) + assert report["accepted_query_ids"] == 3 + records = json.loads((handoff.parent / "generation_failures.json").read_text())["records"] + assert len(records) == 1 and records[0]["stage"] == "summary-combinations" + assert records[0]["source_unit_ids"] == ["unit-0", "unit-1"] + assert records[0]["summary_context_ids"] == ["context-0", "context-1"] + + +def test_failed_related_summary_keeps_original_contexts(tmp_path, monkeypatch): + config = skip_config(tmp_path, "unit").model_copy(update={"related_contexts_per_context": 1}) + sources = [json.loads(line) for line in config.sources_file.read_text().splitlines()] + sources[1]["document_id"] = "another-document" + config.sources_file.write_text("".join(json.dumps(row) + "\n" for row in sources)) + transport = ControlledTransport("ContextSummary") + transport.prompt_contains = '"unit_id": "unit-1"' + install_transport(monkeypatch, transport) + handoff = run_multimodal_sdg(config) + records = json.loads((handoff.parent / "generation_failures.json").read_text())["records"] + assert records and all(record["stage"] == "summary" and len(record["source_unit_ids"]) > 1 for record in records) + candidates = json.loads((config.output_dir / "candidates.json").read_text()) + assert {"context-0", "context-1", "context-2"} <= {candidate["context_id"] for candidate in candidates} + + +@pytest.mark.parametrize("schema", ["ContextSummary", "QueryBatch", "SummaryJudgment", "Localization"]) +def test_all_required_responses_fail_without_fabricating_export(tmp_path, monkeypatch, schema): + config = skip_config(tmp_path) + install_transport(monkeypatch, ControlledTransport(schema, unit=None)) + with pytest.raises(ValueError, match="No candidates passed"): + run_multimodal_sdg(config) + assert not (config.output_dir / "complete.json").exists() + assert not (config.output_dir / "bundle").exists() + assert json.loads((config.output_dir / "generation_failures.json").read_text())["records"] + + +def test_interrupted_resume_reuses_successes_and_terminal_failures(tmp_path, monkeypatch): + config = skip_config(tmp_path) + transport = ControlledTransport("QueryBatch") + transport.provider_unit = "unit-2" + install_transport(monkeypatch, transport) + with pytest.raises((RuntimeError, DataDesignerGenerationError)): + run_multimodal_sdg(config) + assert not (config.output_dir / "complete.json").exists() + saved = {p: digest(p) for p in config.output_dir.rglob("*.parquet")} + before = dict(transport.calls) + transport.provider_unit = None + handoff = run_multimodal_sdg(config.model_copy(update={"resume": True})) + assert all(digest(p) == checksum for p, checksum in saved.items()) + assert transport.calls["QueryBatch", "unit-0"] == before["QueryBatch", "unit-0"] + assert transport.calls["QueryBatch", "unit-1"] == before["QueryBatch", "unit-1"] + assert transport.calls["QueryBatch", "unit-2"] == before["QueryBatch", "unit-2"] + 1 + report = json.loads((handoff.parent / "report.json").read_text()) + assert report["accepted_query_ids"] == 2 + assert report["generation_failures"]["failed_records"] == 1 + + +@pytest.mark.parametrize("artifact", ["generation_failures.json", "bundle/generation_failures.json"]) +def test_completed_resume_checks_failure_ledger(tmp_path, monkeypatch, artifact): + config = skip_config(tmp_path) + install_transport(monkeypatch, ControlledTransport("QueryBatch")) + run_multimodal_sdg(config) + (config.output_dir / artifact).write_text("{}") + with pytest.raises(ValueError, match="integrity"): + run_multimodal_sdg(config.model_copy(update={"resume": True})) + + +def test_default_policy_still_stops_persistent_validation_failure(tmp_path, monkeypatch): + config = skip_config(tmp_path).model_copy(update={"failed_record_policy": "error"}) + install_transport(monkeypatch, ControlledTransport("QueryBatch")) + with pytest.raises((RuntimeError, DataDesignerGenerationError)): + run_multimodal_sdg(config) + assert not (config.output_dir / "complete.json").exists() + assert not (config.output_dir / "generation_failures.json").exists() + + +@pytest.mark.parametrize("keys", [["unexpected"], ["expected", "expected"]]) +def test_skip_never_relaxes_native_row_identity_checks(tmp_path, keys): + config = skip_config(tmp_path) + inference = DataDesignerInference(tmp_path / "inference", config) + request = Request("Summarize", models.ContextSummary, "generator") + dataset = tmp_path / "native" + dataset.mkdir() + pq.write_table( + pa.Table.from_pylist([{"request_id": key, "response": None} for key in keys]), dataset / "batch_000.parquet" + ) + with pytest.raises(ValueError, match="Unexpected or duplicate"): + inference.collect_responses([("expected", request)], dataset, tmp_path / "attempt") + + +def test_cached_failure_integrity_and_policy_are_checked(tmp_path, monkeypatch): + config = skip_config(tmp_path) + install_transport(monkeypatch, ControlledTransport("ContextSummary", unit=None)) + inference = DataDesignerInference(tmp_path / "inference", config) + request = Request("Summarize", models.ContextSummary, "generator") + assert isinstance(inference.generate([request])[0], models.FailedGeneration) + strict = DataDesignerInference(inference.root, config.model_copy(update={"failed_record_policy": "error"})) + with pytest.raises(ValueError, match="requires the skip policy"): + strict.generate([request]) + path = next((inference.root / "responses").glob("*.json")) + value = json.loads(path.read_text()) + value["response"][FAILURE_KEY]["detail"] = "changed" + path.write_text(json.dumps(value)) + with pytest.raises(ValueError, match="integrity"): + inference.generate([request]) diff --git a/plugins/data-designer-retrieval-sdg/tests/multimodal/test_multimodal_sdg.py b/plugins/data-designer-retrieval-sdg/tests/multimodal/test_multimodal_sdg.py index a6acadd..0a75af7 100644 --- a/plugins/data-designer-retrieval-sdg/tests/multimodal/test_multimodal_sdg.py +++ b/plugins/data-designer-retrieval-sdg/tests/multimodal/test_multimodal_sdg.py @@ -375,8 +375,10 @@ def run_batch(self, items): self.calls = getattr(self, "calls", 0) + 1 -def test_missing_response_has_finite_attempts_without_fabrication(tmp_path): - inference = MissingInference(tmp_path / "inference", fixture_config(tmp_path)) +@pytest.mark.parametrize("policy", ["error", "skip"]) +def test_missing_response_has_finite_attempts_without_fabrication(tmp_path, policy): + config = fixture_config(tmp_path).model_copy(update={"failed_record_policy": policy}) + inference = MissingInference(tmp_path / "inference", config) with pytest.raises(RuntimeError, match="Missing structured"): inference.generate([Request("one", ContextSummary, "judge")]) assert inference.calls == 3 diff --git a/plugins/data-designer-retrieval-sdg/tests/multimodal/test_structured.py b/plugins/data-designer-retrieval-sdg/tests/multimodal/test_structured.py index 04155fe..d35406d 100644 --- a/plugins/data-designer-retrieval-sdg/tests/multimodal/test_structured.py +++ b/plugins/data-designer-retrieval-sdg/tests/multimodal/test_structured.py @@ -17,6 +17,7 @@ import pandas as pd import pytest from data_designer.engine.column_generators.generators.llm_completion import LLMStructuredCellGenerator +from data_designer.engine.errors import DataDesignerError from data_designer.engine.models.errors import ModelGenerationValidationFailureError from data_designer.engine.models.facade import ModelFacade from data_designer.engine.models.parsers.errors import ParserException @@ -28,6 +29,7 @@ from data_designer_retrieval_sdg.multimodal.config import RetrievalStructuredColumnConfig from data_designer_retrieval_sdg.multimodal.plugins import retrieval_structured_plugin from data_designer_retrieval_sdg.multimodal.structured import ( + FAILURE_KEY, RetrievalStructuredCellGenerator, RetrievalStructuredResponseRecipe, ) @@ -43,9 +45,11 @@ def response_recipe() -> RetrievalStructuredResponseRecipe: return RetrievalStructuredResponseRecipe(Decision.model_json_schema(), pruning=False) -def test_generation_keeps_native_sync_and_async_implementations() -> None: - assert RetrievalStructuredCellGenerator.generate is LLMStructuredCellGenerator.generate - assert RetrievalStructuredCellGenerator.agenerate is LLMStructuredCellGenerator.agenerate +def test_failed_record_policy_defaults_to_error() -> None: + config = RetrievalStructuredColumnConfig( + name="decision", model_alias="judge", prompt="Judge", output_format=Decision + ) + assert config.failed_record_policy == "error" @pytest.mark.parametrize("fenced", [False, True]) @@ -107,6 +111,88 @@ def model_facade() -> ModelFacade: ) +def structured_generator( + facade: ModelFacade, + policy: str, + with_trace: dd.TraceType = dd.TraceType.NONE, + extract_reasoning_content: bool = False, +) -> LLMStructuredCellGenerator: + """Use native model generation with two validation attempts and explicit policy.""" + resources = MagicMock() + resources.run_config = dd.RunConfig(max_conversation_correction_steps=1, max_conversation_restarts=0) + resources.model_registry.get_model.return_value = facade + return RetrievalStructuredCellGenerator( + RetrievalStructuredColumnConfig( + name="decision", + model_alias="judge", + prompt="Judge", + output_format=Decision, + failed_record_policy=policy, + with_trace=with_trace, + extract_reasoning_content=extract_reasoning_content, + ), + resources, + ) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("policy", ["error", "skip"]) +@pytest.mark.parametrize( + "with_trace,extract_reasoning_content", + [ + (dd.TraceType.NONE, False), + (dd.TraceType.ALL_MESSAGES, True), + (dd.TraceType.LAST_MESSAGE, False), + (dd.TraceType.NONE, True), + ], +) +def test_structured_column_policy_uses_native_validation_budget( + monkeypatch, asynchronous, policy, with_trace, extract_reasoning_content +): + facade = model_facade() + monkeypatch.setattr(facade, "_get_mcp_facade", lambda _: None) + completion = (AsyncMock if asynchronous else MagicMock)( + return_value=make_stub_completion_response(content="not JSON") + ) + monkeypatch.setattr(facade, "acompletion" if asynchronous else "completion", completion) + generator = structured_generator(facade, policy, with_trace, extract_reasoning_content) + if policy == "error": + with pytest.raises(ModelGenerationValidationFailureError): + asyncio.run(generator.agenerate({})) if asynchronous else generator.generate({}) + else: + result = ( + asyncio.run(generator.agenerate({"id": "same-row"})) + if asynchronous + else generator.generate({"id": "same-row"}) + ) + diagnostic = json.loads(result["decision"])[FAILURE_KEY] + assert diagnostic["max_model_attempts"] == 2 + assert result["id"] == "same-row" + assert diagnostic["last_invalid_response"] == "not JSON" + for column, enabled in ( + ("decision__trace", with_trace != dd.TraceType.NONE), + ("decision__reasoning_content", extract_reasoning_content), + ): + if enabled: + assert result[column] is None + else: + assert column not in result + assert completion.call_count == 2 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +def test_structured_column_never_converts_provider_errors(monkeypatch, asynchronous): + facade = model_facade() + monkeypatch.setattr(facade, "_get_mcp_facade", lambda _: None) + completion = (AsyncMock if asynchronous else MagicMock)(side_effect=PermissionError("Controlled provider failure")) + monkeypatch.setattr(facade, "acompletion" if asynchronous else "completion", completion) + generator = structured_generator(facade, "skip") + with pytest.raises(DataDesignerError) as raised: + asyncio.run(generator.agenerate({})) if asynchronous else generator.generate({}) + assert not isinstance(raised.value, ModelGenerationValidationFailureError) + assert completion.call_count == 1 + + def test_native_corrections_are_bounded_and_keep_a_valid_rejection(monkeypatch: pytest.MonkeyPatch) -> None: """Mock only transport completion; run the host's real correction loop.""" facade = model_facade() @@ -190,6 +276,48 @@ def test_registered_column_runs_through_preview(tmp_path: Path, monkeypatch: pyt assert completion.await_count == 1 +def test_failed_batch_preserves_trace_columns_for_downstream_expressions(tmp_path, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "test-key-not-a-real-credential") + completion = AsyncMock(return_value=make_stub_completion_response(content="not JSON")) + monkeypatch.setattr(ModelFacade, "acompletion", completion) + builder = dd.DataDesignerConfigBuilder( + model_configs=[ + dd.ModelConfig(alias="judge", model="nvidia/test-model", provider="nvidia", skip_health_check=True) + ] + ) + builder.with_seed_dataset(dd.DataFrameSeedSource(df=pd.DataFrame([{"source": "First"}, {"source": "Second"}]))) + builder.add_column( + RetrievalStructuredColumnConfig( + name="decision", + model_alias="judge", + prompt="Judge {{ source }}.", + output_format=Decision, + failed_record_policy="skip", + with_trace=dd.TraceType.ALL_MESSAGES, + extract_reasoning_content=True, + ) + ) + builder.add_column( + dd.ExpressionColumnConfig(name="trace_missing", expr="{{ decision__trace is none }}", dtype="bool") + ) + builder.add_column( + dd.ExpressionColumnConfig( + name="reasoning_missing", expr="{{ decision__reasoning_content is none }}", dtype="bool" + ) + ) + designer = ddi.DataDesigner(artifact_path=tmp_path) + designer.set_run_config( + dd.RunConfig(max_conversation_correction_steps=1, max_conversation_restarts=0, otel_metrics_port=None) + ) + result = designer.preview(builder, num_records=2).dataset + assert len(result) == 2 + assert result["decision__trace"].tolist() == [None, None] + assert result["decision__reasoning_content"].tolist() == [None, None] + assert result["trace_missing"].all() and result["reasoning_missing"].all() + assert all(FAILURE_KEY in json.loads(value) for value in result["decision"]) + assert completion.await_count == 4 + + def test_retrieval_structured_plugin_valid() -> None: assert_valid_plugin(retrieval_structured_plugin)