feat: parallel chunk processing for large-document transformations (#990) - #1130
feat: parallel chunk processing for large-document transformations (#990)#1130dyzur wants to merge 12 commits into
Conversation
…fnovo#990) Automatically splits large documents into token-bounded chunks when the content exceeds the model's context window, processes chunks in parallel via LangGraph Send, and hierarchically synthesizes results. For small documents, behavior is identical to the original single-node graph (full content, 8192 output token budget). The chunking path is triggered only when the LLM reports a context-limit error.
There was a problem hiding this comment.
1 issue found across 4 files
Confidence score: 3/5
open_notebook/graphs/transformation.pyadds a complex parallel chunk-processing pipeline (regex error parsing, hierarchical splitting, fan-out, and merge reduction) without matching test coverage, which raises regression risk around incorrect transforms or hard-to-diagnose failures in production; add focused unit/integration tests for chunk boundaries, error parsing, and merge correctness before merging.
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="open_notebook/graphs/transformation.py">
<violation number="1" location="open_notebook/graphs/transformation.py:149">
P2: The parallel chunk processing feature introduces significant complexity (regex-based error parsing, hierarchical text splitting, parallel Send fan-out, and hierarchical LLM merge reduction) without corresponding test coverage. While `try_full_content` is tested for both success and failure paths, the entire chunking/synthesis pathway is uncovered — meaning a regression in `parse_context_limit_error`, `chunk_text_by_tokens`, `process_chunk`, or `synthesize_results` would not be caught by CI. Consider adding tests for at least the core utilities (`parse_context_limit_error` with representative error strings from each provider, `chunk_text_by_tokens` edge cases) and one integration test that exercises the full chunk → parallel process → synthesis flow.</violation>
</file>
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
| # ── Graph nodes ──────────────────────────────────────────────────────── | ||
|
|
||
|
|
||
| async def try_full_content(state: TransformationState, config: RunnableConfig) -> dict: |
There was a problem hiding this comment.
P2: The parallel chunk processing feature introduces significant complexity (regex-based error parsing, hierarchical text splitting, parallel Send fan-out, and hierarchical LLM merge reduction) without corresponding test coverage. While try_full_content is tested for both success and failure paths, the entire chunking/synthesis pathway is uncovered — meaning a regression in parse_context_limit_error, chunk_text_by_tokens, process_chunk, or synthesize_results would not be caught by CI. Consider adding tests for at least the core utilities (parse_context_limit_error with representative error strings from each provider, chunk_text_by_tokens edge cases) and one integration test that exercises the full chunk → parallel process → synthesis flow.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At open_notebook/graphs/transformation.py, line 149:
<comment>The parallel chunk processing feature introduces significant complexity (regex-based error parsing, hierarchical text splitting, parallel Send fan-out, and hierarchical LLM merge reduction) without corresponding test coverage. While `try_full_content` is tested for both success and failure paths, the entire chunking/synthesis pathway is uncovered — meaning a regression in `parse_context_limit_error`, `chunk_text_by_tokens`, `process_chunk`, or `synthesize_results` would not be caught by CI. Consider adding tests for at least the core utilities (`parse_context_limit_error` with representative error strings from each provider, `chunk_text_by_tokens` edge cases) and one integration test that exercises the full chunk → parallel process → synthesis flow.</comment>
<file context>
@@ -1,77 +1,403 @@
+# ── Graph nodes ────────────────────────────────────────────────────────
+
+
+async def try_full_content(state: TransformationState, config: RunnableConfig) -> dict:
+ """
+ Optimistically process the full content in one call.
</file context>
There was a problem hiding this comment.
All reported issues were addressed across 4 files (changes from recent commits).
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
- Fix [-0:] returning entire previous chunk (P1): use for-else to cleanly fall through to continuation-marker-only when all overlap sizes overshoot. - Fix stalled synthesis reordering document sections (P1): merge the smallest adjacent pair instead of sorting by token count, preserving source order. - Guard against oversized pairwise merge (P1): check combined token count before invoking LLM; concatenate all remaining texts when even the smallest adjacent pair exceeds the context window. - Remove dead max_iterations guard (P3): both normal and stalled paths always shrink texts, so the guard was unreachable.
There was a problem hiding this comment.
All reported issues were addressed across 2 files (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
When no adjacent pair fits the context window during synthesis, re-chunk the oversized texts at a smaller token budget using chunk_text_by_tokens and continue the merge loop, rather than returning concatenated raw partial outputs.
lfnovo
left a comment
There was a problem hiding this comment.
Full review run on top of the cubic pass. The five P1s cubic raised are all genuinely addressed in f617006 — verified against the current diff, not re-litigating them.
One new P1 blocks merge, plus some P2s worth a look.
P1 (blocker)
max_tokens mismatch defeats the feature on sub-~82k-context models — open_notebook/graphs/transformation.py:137-140 vs :285 (and merges at :395, :490).
Chunks are sized reserving output_budget = min(FULL_CONTENT_MAX_TOKENS, int(context_limit * 0.10)), so for any effective context below ~82k, output_budget < 8192 and max_chunk_tokens = context - system - output_budget. But the actual chunk/merge LLM calls then request max_tokens=FULL_CONTENT_MAX_TOKENS (8192, hard-coded). That makes chunk_input + 8192 + system > context_limit, so providers that count requested max_tokens against the window reject the call — the chunk re-fails with the same context-limit error, process_chunk raises, and the whole transformation aborts.
This breaks the feature precisely for the small/mid-context models it exists to serve (e.g. a 32k model on a 40k-token doc). Fix: have the chunk and merge calls request max_tokens=output_budget instead of the constant.
P2 (worth addressing, discussable if here or follow-up)
section_hintis dead code —transformation.py:247-252._build_system_promptinjectssection_context, butprompts/transformation/execute.jinjaonly renders{{ instructions }}, so the "section X of Y" hint never reaches the model. Chunks are processed with no partial-document awareness.- Progress-guard has no real test coverage —
transformation.py:387.test_progress_guard_fallback_on_single_item_groupsuses ~2000-token texts against a ~6873-token grouping threshold, so groups always hold multiple items and the fallback branch (the point of the anti-hang guard) never executes. - Unbounded fan-out concurrency —
transformation.py:216. OneSendper chunk with nomax_concurrency; N scales with document length → rate-limit storms, compounded across concurrent worker jobs. Issue #990 explicitly asked to mindOPEN_NOTEBOOK_WORKER_MAX_.... - One chunk failure discards all sibling work —
transformation.py:266.process_chunkre-raises on any exception; with parallelSend, a single transient error (e.g. a rate limit) aborts the entire transformation with no retry or partial-result path. - Concatenation fallback leaks scaffolding into the stored insight —
transformation.py:~430-450. When no adjacent pair fits the window, texts are joined with literalPart N:/---markers and returned as the final output with no LLM merge — meta-structure leaks into the persisted insight for non-summary transformations (the one place issue #990's "meta-hints shouldn't leak" requirement is actually violated).
Great work overall — the chunking hierarchy and hierarchical synthesis are solid, and the cubic fixes landed cleanly. Just the P1 to sort before merge.
There was a problem hiding this comment.
All reported issues were addressed across 1 file (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
…chunk failures, test coverage
There was a problem hiding this comment.
All reported issues were addressed across 3 files (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Tip: cubic used a learning from your PR history. Let your coding agent read cubic learnings directly with the cubic MCP.
Re-trigger cubic
Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>
…swallowing A failed chunk was silently producing a partial transformation with a plain-text error placeholder — transient errors (rate limits, network blips) bypassed the command/API retry mechanism and an incomplete insight could be saved. Now process_chunk stores the classified error class and message in ChunkResult, letting sibling chunks finish. synthesize_results checks for any chunk errors after the merge loop and raises an aggregated typed exception before add_insight is called.
…r pair merge The max(512, ...) floor could produce pieces too large to pair when the available context (context_limit - overhead) is small. Now verifies 2 * re_chunk_budget + overhead <= context_limit and raises ConfigurationError if not, instead of looping indefinitely.
There was a problem hiding this comment.
1 issue found across 1 file (changes from recent commits).
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="open_notebook/graphs/transformation.py">
<violation number="1" location="open_notebook/graphs/transformation.py:484">
P0: Oversized adjacent results now bypass re-chunking and are sent to the forced merge. Keep the `min_combined + overhead > context_limit` guard around the re-chunk path, with this new two-piece feasibility check nested inside it; otherwise the context-limit failure recurs for ordinary model windows.</violation>
</file>
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
| # Verify that two re-chunked pieces can actually be merged. | ||
| # If not, the transformation cannot proceed — fail with a | ||
| # typed error so the caller can retry or reconfigure. | ||
| if 2 * re_chunk_budget + overhead > context_limit: |
There was a problem hiding this comment.
P0: Oversized adjacent results now bypass re-chunking and are sent to the forced merge. Keep the min_combined + overhead > context_limit guard around the re-chunk path, with this new two-piece feasibility check nested inside it; otherwise the context-limit failure recurs for ordinary model windows.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At open_notebook/graphs/transformation.py, line 484:
<comment>Oversized adjacent results now bypass re-chunking and are sent to the forced merge. Keep the `min_combined + overhead > context_limit` guard around the re-chunk path, with this new two-piece feasibility check nested inside it; otherwise the context-limit failure recurs for ordinary model windows.</comment>
<file context>
@@ -472,10 +472,22 @@ async def synthesize_results(
+ # Verify that two re-chunked pieces can actually be merged.
+ # If not, the transformation cannot proceed — fail with a
+ # typed error so the caller can retry or reconfigure.
+ if 2 * re_chunk_budget + overhead > context_limit:
+ raise ConfigurationError(
+ "The model's context window is too small to "
</file context>
Makes context_limit, output_budget, chunks, chunk_results, total_chunks NotRequired in TransformationState (they're only set when chunking is triggered) and error_class/error_message NotRequired in ChunkResult. Adds cast() calls in tests where plain dicts are passed as typed state objects, matching the existing codebase pattern.
There was a problem hiding this comment.
1 issue found across 4 files (changes from recent commits).
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="tests/test_graphs.py">
<violation number="1" location="tests/test_graphs.py:202">
P3: Continuation-line indentation is inconsistent with PEP 8 and with the other test methods in the same file that use this same cast pattern correctly. The argument line and closing paren in the `with`-block tests don't align with their opening statement. Fix indentation to match the pattern used in `test_no_chunking_passthrough` and `test_empty_chunk_results` for consistency.</violation>
</file>
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
| cast(TransformationState, state), {"configurable": {}} | ||
| ) |
There was a problem hiding this comment.
P3: Continuation-line indentation is inconsistent with PEP 8 and with the other test methods in the same file that use this same cast pattern correctly. The argument line and closing paren in the with-block tests don't align with their opening statement. Fix indentation to match the pattern used in test_no_chunking_passthrough and test_empty_chunk_results for consistency.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At tests/test_graphs.py, line 202:
<comment>Continuation-line indentation is inconsistent with PEP 8 and with the other test methods in the same file that use this same cast pattern correctly. The argument line and closing paren in the `with`-block tests don't align with their opening statement. Fix indentation to match the pattern used in `test_no_chunking_passthrough` and `test_empty_chunk_results` for consistency.</comment>
<file context>
@@ -198,7 +198,9 @@ async def test_no_chunking_passthrough(self):
state = {"output": "direct output"}
- result = await synthesize_results(state, {"configurable": {}})
+ result = await synthesize_results(
+ cast(TransformationState, state), {"configurable": {}}
+ )
assert result == {"output": "direct output"}
</file context>
| cast(TransformationState, state), {"configurable": {}} | |
| ) | |
| cast(TransformationState, state), {"configurable": {}} | |
| ) |
lfnovo
left a comment
There was a problem hiding this comment.
Re-reviewed against the current branch. The max_tokens blocker from my earlier pass is genuinely fixed (chunk/merge calls now derive from output_budget at L294/425/532), and section_hint reaches the model and the scaffolding no longer leaks into the stored insight. But one P0 remains — and the tests that claim to cover it don't.
P0 — the oversized-result re-chunk branch is dead code
open_notebook/graphs/transformation.py:485. The raise ConfigurationError(...) inside the if 2 * re_chunk_budget + overhead > context_limit: block is unconditional, so everything after it (the logger.warning, the re_chunked build loop, texts = re_chunked, and continue at L492–515) is unreachable. Consequences:
- When the guard is false (the normal case), control skips the whole
ifand falls straight to the forced pairwise merge at L516+, which sendsPart 1 + Part 2(=min_combinedtokens) to the provider. That branch is only reached because grouping found no adjacent pair fitscontext_limit - output_budget - 500, i.e.min_combinedis large — so the oversized pair is shipped with no guard and the context-limit rejection recurs on ordinary sub-~82k windows. This is exactly the failure the feature exists to prevent. - The re-chunk remediation the "re-chunk oversized results" / "prevent hang" commits describe never actually executes.
Fix: make the raise conditional on genuine infeasibility and let the re-chunk block run (and continue) when re-chunking is feasible — i.e. the re-chunk loop needs to sit on the reachable path, not after the raise.
Tests don't cover the re-chunk path
tests/test_graphs.py — test_stalled_synthesis_rechunk_fallback and test_progress_guard_fallback_on_single_item_groups are tautological: with the default context_limit=8192 and their inputs, the L484 guard is always false, so the re-chunk path (dead code anyway) never runs; both simply loop over a mock that returns a fixed string and would pass whether or not re-chunking works. Please add a test that drives the re-chunk/guard path with a small context_limit and asserts re-chunking actually occurred (and that no oversized pair is forwarded to the merge). process_chunk and the full try_full_content → context-error → chunk → synthesize flow are also currently untested.
P2 (discussable — here or follow-up)
- Concurrency ignores
OPEN_NOTEBOOK_WORKER_MAX_TASKS.MAX_CONCURRENT_CHUNKS = 10is hard-coded (L52) and theSendfan-out has nomax_concurrency; #990 explicitly asked to honor the worker cap. Also, the module-levelasyncio.Semaphoreis bound to the import-time event loop. - Still fail-all on a single chunk error.
synthesize_resultsnow aggregates a classified error instead of aborting immediately (an improvement), but no partial insight is ever saved and there's no per-chunk retry. Whether a partial-result path is in scope is your call.
The chunking hierarchy and hierarchical synthesis are solid — it's the P0 (dead re-chunk branch) plus a real test for it that block merge. The two P2s are design calls I'm happy to defer to a follow-up if you'd rather land the core first.
6620b19 to
d06d45d
Compare
…al raise The guard in synthesize_results had the re-chunk remediation code inside the if block after an unconditional raise, making it unreachable. When the guard is TRUE (infeasible), raise ConfigurationError. When FALSE (feasible), run the re-chunk loop to split oversized texts so subsequent merge rounds can handle them — instead of falling through to the forced pairwise merge which ships Part 1+Part 2 without any size guard. Tests added/improved: - test_progress_guard_fallback_on_single_item_groups: added context_limit to state + assert provision_langchain_model was called. - test_stalled_synthesis_rechunk_fallback: same improvements. - test_rechunk_guard_raises_on_infeasible_context: verifies ConfigurationError is raised when context_limit is 600 (guard TRUE). - test_rechunk_actually_splits_texts: verifies multiple merge rounds occur (proving re-chunking happened) with context_limit=4000. - test_process_chunk_success: unit test for process_chunk node. - test_process_chunk_error_classified: verifies errors are classified and stored (not re-raised) in the ChunkResult.
d06d45d to
64afdd8
Compare
Implements automatic parallel chunk processing for large-document transformations.
When a document exceeds the model's context window, the graph automatically:
For small documents, behavior is identical to the original single-node graph (full content, 8192 output token budget).
Closes #990