Skip to content

feat: parallel chunk processing for large-document transformations (#990) - #1130

Open
dyzur wants to merge 12 commits into
lfnovo:mainfrom
dyzur:feat/parallel-chunking
Open

feat: parallel chunk processing for large-document transformations (#990)#1130
dyzur wants to merge 12 commits into
lfnovo:mainfrom
dyzur:feat/parallel-chunking

Conversation

@dyzur

@dyzur dyzur commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

Implements automatic parallel chunk processing for large-document transformations.

When a document exceeds the model's context window, the graph automatically:

  1. Falls back from full-content attempt on context-limit error
  2. Splits the content into token-bounded chunks (paragraph -> sentence -> word hierarchical splitting)
  3. Processes chunks in parallel via LangGraph Send
  4. Hierarchically synthesizes results into a single coherent output

For small documents, behavior is identical to the original single-node graph (full content, 8192 output token budget).

Closes #990

Review in cubic

…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.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 issue found across 4 files

Confidence score: 3/5

  • open_notebook/graphs/transformation.py adds 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

Comment thread open_notebook/utils/token_utils.py
Comment thread open_notebook/graphs/transformation.py
Comment thread open_notebook/graphs/transformation.py Outdated
Comment thread open_notebook/utils/token_utils.py Outdated
Comment thread open_notebook/graphs/transformation.py Outdated
# ── Graph nodes ────────────────────────────────────────────────────────


async def try_full_content(state: TransformationState, config: RunnableConfig) -> dict:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>

Comment thread open_notebook/utils/token_utils.py Outdated

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 4 files (changes from recent commits).

Reply with feedback, questions, or to request a fix.

Re-trigger cubic

Comment thread open_notebook/utils/token_utils.py
Comment thread open_notebook/graphs/transformation.py Outdated
Comment thread open_notebook/graphs/transformation.py
Comment thread open_notebook/graphs/transformation.py Outdated
- 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.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread open_notebook/graphs/transformation.py Outdated
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 lfnovo left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 modelsopen_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_hint is dead codetransformation.py:247-252. _build_system_prompt injects section_context, but prompts/transformation/execute.jinja only 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 coveragetransformation.py:387. test_progress_guard_fallback_on_single_item_groups uses ~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 concurrencytransformation.py:216. One Send per chunk with no max_concurrency; N scales with document length → rate-limit storms, compounded across concurrent worker jobs. Issue #990 explicitly asked to mind OPEN_NOTEBOOK_WORKER_MAX_....
  • One chunk failure discards all sibling worktransformation.py:266. process_chunk re-raises on any exception; with parallel Send, 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 insighttransformation.py:~430-450. When no adjacent pair fits the window, texts are joined with literal Part 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.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread open_notebook/graphs/transformation.py Outdated

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread open_notebook/graphs/transformation.py
Comment thread open_notebook/graphs/transformation.py Outdated
dyzur and others added 3 commits July 14, 2026 14:09
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.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread tests/test_graphs.py
Comment on lines +202 to +203
cast(TransformationState, state), {"configurable": {}}
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>
Suggested change
cast(TransformationState, state), {"configurable": {}}
)
cast(TransformationState, state), {"configurable": {}}
)

@dyzur
dyzur requested a review from lfnovo July 14, 2026 21:10

@lfnovo lfnovo left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 if and falls straight to the forced pairwise merge at L516+, which sends Part 1 + Part 2 (= min_combined tokens) to the provider. That branch is only reached because grouping found no adjacent pair fits context_limit - output_budget - 500, i.e. min_combined is 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.pytest_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 = 10 is hard-coded (L52) and the Send fan-out has no max_concurrency; #990 explicitly asked to honor the worker cap. Also, the module-level asyncio.Semaphore is bound to the import-time event loop.
  • Still fail-all on a single chunk error. synthesize_results now 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.

@dyzur
dyzur force-pushed the feat/parallel-chunking branch 2 times, most recently from 6620b19 to d06d45d Compare July 27, 2026 19:20
…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.
@dyzur
dyzur force-pushed the feat/parallel-chunking branch from d06d45d to 64afdd8 Compare July 27, 2026 19:28
@dyzur
dyzur requested a review from lfnovo July 27, 2026 19:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature]: parallel chunk processing for large-document transformations

2 participants