Skip to content

fix(dataprep): make Contextual Retrieval max_tokens env-configurable

Jérôme Revillard requested to merge fix/contextual-max-tokens-env into main

Problem

Contextual Retrieval silently degraded to raw-chunk fallback on large documents. The potato guide 1782829218963_ee591cfd (38 chunks, 67K chars) on el-salvador failed all 3 doc-level context attempts with JSONDecodeError: Unterminated string starting at line 1 column 13 (char 12) → 0 contexts generated.

Root cause

_context_doc_call (and the per-chunk calls) hardcoded max_tokens=200. The model writes ~196 tokens for a large doc (67K-char input ≈ 16K prompt tokens) — right at the 200 cap. Under concurrent vLLM load the JSON {"context":"..."} was truncated mid-string → unterminated string at char 12 (the value-quote position) → JSONDecodeError. No env knob existed.

Live probe (el-salvador vLLM, exact prod params): completion_tokens=196 at the 200 cap; with max_tokens=512 the model produces the same 196 with finish_reason=stop. 200 left zero headroom.

Fix

Add DATAPREP_CONTEXTUAL_MAX_TOKENS (default 512) and wire it into all three context-generation calls:

  • _context_doc_call (doc-level)
  • per-chunk batched (200*len(batch) → CONTEXTUAL_MAX_TOKENS*len(batch))
  • per-chunk single

512 gives comfortable margin at no extra cost (the model stops early at ~196). Consistent with the labeling path's env-tunable caps.

Rollout

Config + docs: env, docker-compose.yaml, deploy/ansible/templates/env.j2, CLAUDE.md, GENIE.AI-Installation-Configuration-Guide.md, GENIE.AI-Data-Labelling-Strategy.md §7.

Verification

  • TDD: pytest asserting the env-driven cap reaches the vLLM call (would fail at hardcoded 200).
  • Full tests/test_dataprep.py: 73 passed, ruff clean.
  • Based on current main (includes the labeling fix from !216 (merged)); rebase clean.
  • Pending: redeploy el-salvador + re-ingest the potato doc → confirm doc-level contexts generate (no Unterminated string).

Merge request reports