Skip to content

fix(dataprep): recover partial batch context responses + clarify log

Jérôme Revillard requested to merge fix/dataprep-partial-batch-recovery into main

Problem

_context_batch_call returned early when the LLM produced a PARTIAL response (some chunk contexts present, some missing). Missing chunk indices silently degraded to raw text — the batch 'succeeded' (non-empty dict) so no fallback fired.

Observed live on el-salvador per_chunk ingest: 47/48 contextualized (1 silent fallback) at batch=8. Trace spans proved batches returned 6-7/8 contexts with completion_tokens 300-400 — well under max_tokens=4096, so NOT truncation. Granite (and most LLMs) omit entries nondeterministically in batched structured-JSON responses, amplified by large input (100k doc-context budget).

Fix

After a successful batch parse, compute missing = valid_ids - out.keys(). Recover each missing chunk via _context_single_call (existing per-chunk path, 3 retries each). Emits dataprep.partial_recovery_count span attribute + INFO log so recovery is visible, not silent. Every chunk now either gets a batched context OR a recovered per-chunk context — raw fallback only on hard failure (exception or empty parse).

Also clarifies the 'X/Y chunks contextualized' log: now states '(Z using raw text after batch + per-chunk recovery failed)' so operators don't misread the post-recovery count as silent fallbacks.

Tests

6 new tests in TestContextualRetrieval:

  • full success → all returned, no recovery, 1 batch call
  • partial response (1/3 missing) → recovered via per-chunk
  • near-empty (3/4 missing) → all missing recovered
  • empty parse → full per-chunk fallback + WARN logged
  • batch exception → full per-chunk fallback + WARN logged
  • single-chunk batch → direct per-chunk, no batch API call

Full dataprep suite: 95 passed, no regressions. npm run lint:py + format:py:check clean.

Live validation

On el-salvador per_chunk ingestion of the tomato doc: 6/48 chunks recovered on a batch=8 run (proven via spans). Batch=4 runs recovered 0 (smaller batches omit less) — recovery mechanism stays as the safety net regardless.

Independent of MR !230 (merged) (instrumentation) — this is a self-contained dataprep fix.

Merge request reports