feat: add parallel chunk processing for large documents in transformations - #529
feat: add parallel chunk processing for large documents in transformations#529kevincolten wants to merge 9 commits into
Conversation
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
2 issues found across 6 files
Prompt for AI agents (all issues)
Check if these issues are valid — if so, understand the root cause of each and fix them.
<file name="open_notebook/utils/token_utils.py">
<violation number="1" location="open_notebook/utils/token_utils.py:153">
P2: `is_context_limit_error` matches generic substrings like "limit" and "exceeded", so common rate-limit errors (e.g., "Rate limit exceeded") will be treated as context-length errors. In transformation, that triggers parallel chunk retries, likely worsening rate limiting and causing retry spikes.</violation>
<violation number="2" location="open_notebook/utils/token_utils.py:268">
P2: Sentence-level splitting can still create chunks that exceed max_tokens when a single sentence is longer than the limit, violating the function’s contract and potentially re-triggering context-limit errors.</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
There was a problem hiding this comment.
1 issue found across 2 files (changes from recent commits).
Prompt for AI agents (all issues)
Check if these issues are valid — if so, understand the root cause of each and fix them.
<file name="open_notebook/utils/token_utils.py">
<violation number="1" location="open_notebook/utils/token_utils.py:310">
P2: `current_chunk` is set to a list of words, but remaining chunks are joined with `"\n\n"`, so an oversized sentence fragment at the end will be emitted with double newlines between every word instead of spaces.</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
… synthesis When a transformation's full content exceeds the model's context window, split it into token-sized chunks, process them in parallel, and synthesize the partial results back into one output. The synthesis is reduced hierarchically — chunk results are batched to fit a token budget and combined in rounds — so merging many (or large) results never overflows the context window either. - token_utils: add context-limit error parsing (OpenAI/Anthropic/Google), is_context_limit_error, token-aware text chunking, and output-buffer helpers. - transformation: try_full_content -> fan_out_chunks -> process_chunk -> synthesize_results (with budgeted reduce). Preserves the single-call fast path and upstream's classify_error behavior. - tests: chunking helpers, fan-out routing, and a regression test asserting synthesis batches instead of overflowing. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
3c76e16 to
19448a1
Compare
lfnovo
left a comment
lfnovo
left a comment
There was a problem hiding this comment.
Thanks @kevincolten — this is a genuinely nice piece of work: the optimistic-full → chunk-on-context-error → parallel Send → hierarchical synthesis flow is well structured, it keeps Prompter intact, doesn't disturb the token_count fix, and it's well tested (the 145K-token real run is a great touch). Two blocking items before it can merge, plus a few smaller notes.
Blocking
1. Output cap silently drops from 8192 → 4096 on the normal (non-chunked) path.
try_full_content sets output_buffer = DEFAULT_OUTPUT_TOKENS (4096) and calls provision_langchain_model(..., max_tokens=output_buffer). main currently uses max_tokens=8192. Since try_full runs for every transformation (not just large docs), any transformation whose output exceeds 4096 tokens would now be truncated — a silent regression on the common path. Please preserve the current 8192 default for the full-content attempt (or make the output budget configurable and default it to 8192).
2. Module-level asyncio.Semaphore bound at import.
_chunk_semaphore = asyncio.Semaphore(_CHUNK_CONCURRENCY_LIMIT) # module scopeThere are no other module-level asyncio primitives in the codebase, and this one is a footgun: an asyncio.Semaphore binds to the first event loop it's awaited from. This graph is imported by the worker (run_transformation_command) and can also be exercised from the API's loop; awaiting the same module-level semaphore from two loops raises RuntimeError: bound to a different event loop. Please create it inside the function (or lazily per-invocation) rather than at import.
Non-blocking (worth addressing)
-
Context-limit detection parses provider error strings (
is_context_limit_error/get_context_limit_from_error). That's inherently provider-format-dependent and will silently fall back toDEFAULT_CONTEXT_LIMIT(8192) when a wording doesn't match, which can mis-size chunks. A short comment on the supported formats + how the fallback behaves would help future maintenance. -
Chunk meta-prefix leaks into content. Each chunk is sent as
"[Processing section X of N from a larger document]\n\n{chunk}". For summarization that's fine, but for extraction-style transformations that instruction text can bleed into or skew the output. Consider putting the "section X of N" hint in the system prompt instead of the user content. -
_CHUNK_CONCURRENCY_LIMIT = 3is a second concurrency layer on top of the worker's own limit (OPEN_NOTEBOOK_WORKER_MAX_TASKS, see #893). Worth a comment noting the interaction, or deriving it from the same config.
Process note
There's no linked issue — features go through an approved issue first (CONTRIBUTING). The feature itself is well-aligned, so this is easy to formalize; I'll get an issue opened to track it. Once (1) and (2) are addressed I'm happy to re-review. (Heads up: maintainerCanModify is off on this PR, so these need to come from your side.)
|
Opened #990 to track this feature (with the review notes captured as design/acceptance criteria) — this PR is the implementation for it. Once the two blocking items above are addressed, happy to re-review. 🙏 |
- Restore the 8192 output cap on the full-content path: bump DEFAULT_OUTPUT_TOKENS from 4096 to 8192 so the optimistic attempt matches the pre-chunking max_tokens and never truncates outputs on the common path. - Create chunk semaphores lazily per event loop instead of at module import: an asyncio.Semaphore binds to the loop it is first awaited from, and the graph runs from both the worker's and the API's loops. - Move the "section X of N" hint from the user content into the system prompt so it can't bleed into extraction-style outputs. - Document the supported provider error formats and the DEFAULT_CONTEXT_LIMIT fallback behavior in token_utils. - Note the interaction between _CHUNK_CONCURRENCY_LIMIT and the worker's task limit (OPEN_NOTEBOOK_WORKER_MAX_TASKS). - Add tests for the 8192 cap, verbatim chunk content, and per-loop semaphores. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
1 issue found across 3 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="open_notebook/graphs/transformation.py">
<violation number="1" location="open_notebook/graphs/transformation.py:248">
P1: Chunk and synthesis execution paths lack the exception classification used in `try_full_content`, which means raw provider exceptions can propagate with unsanitized messages instead of the user-friendly, truncated errors the rest of the codebase guarantees. `process_chunk` and `_synthesize_once` should wrap their LLM calls in a `try/except Exception` block and route errors through `classify_error()`, matching the pattern in `try_full_content` and the other graphs. Without this, failures during parallel chunk or synthesis processing leak internal details and produce inconsistent user-facing behavior.</violation>
</file>
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
| SystemMessage(content=chunk_system_prompt), | ||
| HumanMessage(content=state["chunk"]), | ||
| ] | ||
| async with _get_chunk_semaphore(): |
There was a problem hiding this comment.
P1: Chunk and synthesis execution paths lack the exception classification used in try_full_content, which means raw provider exceptions can propagate with unsanitized messages instead of the user-friendly, truncated errors the rest of the codebase guarantees. process_chunk and _synthesize_once should wrap their LLM calls in a try/except Exception block and route errors through classify_error(), matching the pattern in try_full_content and the other graphs. Without this, failures during parallel chunk or synthesis processing leak internal details and produce inconsistent user-facing behavior.
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 248:
<comment>Chunk and synthesis execution paths lack the exception classification used in `try_full_content`, which means raw provider exceptions can propagate with unsanitized messages instead of the user-friendly, truncated errors the rest of the codebase guarantees. `process_chunk` and `_synthesize_once` should wrap their LLM calls in a `try/except Exception` block and route errors through `classify_error()`, matching the pattern in `try_full_content` and the other graphs. Without this, failures during parallel chunk or synthesis processing leak internal details and produce inconsistent user-facing behavior.</comment>
<file context>
@@ -214,16 +234,18 @@ async def process_chunk(state: ChunkState, config: RunnableConfig) -> dict:
+ HumanMessage(content=state["chunk"]),
]
- async with _chunk_semaphore:
+ async with _get_chunk_semaphore():
chain = await provision_langchain_model(
str(payload),
</file context>
|
Thanks for the thorough review @lfnovo! All items addressed in 69e5156: Blocking:
Non-blocking:
Also linked the PR to #990. All 204 tests pass, ruff clean. Ready for re-review 🙏 |
…chunking # Conflicts: # open_notebook/graphs/transformation.py
The chunking work split run_transformation into try_full_content -> process_chunk -> synthesize_results. try_full_content owns the non-chunking path and its add_insight() call, so it is where upstream's propagation contract now lives.
- parse_context_limit_error: widen to Optional[Tuple[Optional[int], int]]. The docstring already documented that tokens_sent may be None when only the limit is parseable, and its sole caller (get_context_limit_from_error) already returns that wider type. - test_graphs: narrow await_args before attribute access, and type the process_chunk state as ChunkState.
There was a problem hiding this comment.
1 issue found across 1 file (changes from recent commits).
Confidence score: 5/5
- In
tests/test_add_insight_failure_propagation.py, the test name/scope currently suggests compiled-graph propagation coverage while it directly callstry_full_content(), which could create a false sense of end-to-end failure-path protection and let graph-wiring regressions slip through unnoticed—either rename it to reflect the node-level contract or invoke the compiled graph path explicitly.
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_add_insight_failure_propagation.py">
<violation number="1" location="tests/test_add_insight_failure_propagation.py:114">
P3: The test now exercises `try_full_content()` directly, but its `run_transformation` name and graph-level class imply that compiled-graph propagation is covered. Renaming it to the node-level contract or invoking the compiled graph would keep the test's stated coverage accurate.</violation>
</file>
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
|
|
||
| with pytest.raises(DatabaseOperationError): | ||
| await run_transformation(state, config={"configurable": {}}) | ||
| await try_full_content(state, config={"configurable": {}}) |
There was a problem hiding this comment.
P3: The test now exercises try_full_content() directly, but its run_transformation name and graph-level class imply that compiled-graph propagation is covered. Renaming it to the node-level contract or invoking the compiled graph would keep the test's stated coverage accurate.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At tests/test_add_insight_failure_propagation.py, line 114:
<comment>The test now exercises `try_full_content()` directly, but its `run_transformation` name and graph-level class imply that compiled-graph propagation is covered. Renaming it to the node-level contract or invoking the compiled graph would keep the test's stated coverage accurate.</comment>
<file context>
@@ -103,13 +111,13 @@ async def test_add_insight_failure_propagates_out_of_run_transformation(self):
with pytest.raises(DatabaseOperationError):
- await run_transformation(state, config={"configurable": {}})
+ await try_full_content(state, config={"configurable": {}})
mock_add_insight.assert_awaited_once()
</file context>
2 participants
Description
Enable smaller context models to process large documents by automatically chunking content that exceeds context limits and processing chunks in parallel using LangGraph's Send API.
How it works:
try_full_content- Attempts to process entire document optimisticallyfan_out_chunks- Creates parallelSend()calls for each chunkprocess_chunk- Processes chunks concurrently via LangGraph Send APIsynthesize_results- Merges chunk results into unified outputRelated Issue
Fixes #990
Type of Change
How Has This Been Tested?
uv run pytest)Test Details:
Design Alignment
Which design principles does this PR support? (See DESIGN_PRINCIPLES.md)
Explanation:
Checklist
Code Quality
Testing
make rufforruff check . --fixDocumentation
Screenshots (if applicable)
N/A - Backend changes only
Additional Context
Files Modified:
open_notebook/graphs/transformation.py- Full restructure to use Send APItests/test_graphs.py- Updated imports for new function namesNew Components:
ChunkResult/ChunkState- TypedDicts for parallel processingtry_full_content()- Optimistic processing with error-based fallbackfan_out_chunks()- Conditional edge creating Send objectsprocess_chunk()- Individual chunk processorsynthesize_results()- Result aggregator usingAnnotated[list, operator.add]Pre-Submission Verification
Before submitting, please verify: