diff --git a/open_notebook/graphs/transformation.py b/open_notebook/graphs/transformation.py index 1eff8c6c8..1acdf8611 100644 --- a/open_notebook/graphs/transformation.py +++ b/open_notebook/graphs/transformation.py @@ -1,7 +1,27 @@ +"""Transformation graph with large-document support. + +Runs a transformation over a source/text. It first tries the full content +optimistically; if the model's context window is exceeded, it splits the content +into token-sized chunks, processes them in parallel, and synthesizes the partial +results back into a single output. + +The synthesis is itself reduced **hierarchically**: combining many chunk results +into one call can exceed the context window too (each result is up to +``output_buffer`` tokens, and there can be many chunks), so results are batched +to fit a token budget and synthesized in rounds until one result remains. +""" + +import asyncio +import operator +import weakref +from typing import Annotated, List, Optional, Union + from ai_prompter import Prompter from langchain_core.messages import HumanMessage, SystemMessage from langchain_core.runnables import RunnableConfig from langgraph.graph import END, START, StateGraph +from langgraph.types import Send +from loguru import logger from typing_extensions import TypedDict from open_notebook.ai.provision import provision_langchain_model @@ -11,67 +31,367 @@ from open_notebook.utils import clean_thinking_content from open_notebook.utils.error_classifier import classify_error from open_notebook.utils.text_utils import extract_text_content +from open_notebook.utils.token_utils import ( + DEFAULT_CONTEXT_LIMIT, + DEFAULT_OUTPUT_TOKENS, + SAFETY_BUFFER, + calculate_output_buffer, + chunk_text_by_tokens, + get_context_limit_from_error, + is_context_limit_error, + token_count, +) + +# Bound how many chunk/synthesis LLM calls run at once so a large document (many +# chunks) doesn't fan out into a burst of provider requests. Note this is a +# second concurrency layer on top of the worker's own task limit +# (OPEN_NOTEBOOK_WORKER_MAX_TASKS, see #893): the worst-case number of +# concurrent provider calls is roughly worker tasks x this limit. +_CHUNK_CONCURRENCY_LIMIT = 3 + +# Semaphores are created lazily per event loop: an asyncio.Semaphore binds to +# the loop it is first awaited from, and this graph runs from both the worker's +# and the API's loops. All Send() nodes of one invocation share a loop, so a +# per-loop semaphore still bounds the fan-out of a single transformation. +_chunk_semaphores: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Semaphore]" = ( + weakref.WeakKeyDictionary() +) + + +def _get_chunk_semaphore() -> asyncio.Semaphore: + loop = asyncio.get_running_loop() + sem = _chunk_semaphores.get(loop) + if sem is None: + sem = asyncio.Semaphore(_CHUNK_CONCURRENCY_LIMIT) + _chunk_semaphores[loop] = sem + return sem + + +class ChunkResult(TypedDict): + """Result from processing a single chunk.""" + + idx: int + result: str -class TransformationState(TypedDict): +class TransformationState(TypedDict, total=False): + """State for the transformation graph. + + ``input_text``/``source``/``transformation``/``output`` are the public + fields; the rest are populated only when the content exceeds the context + window and the graph falls back to chunking. + """ + input_text: str source: Source transformation: Transformation output: str + # Chunking fields (populated when full content exceeds the context limit) + system_prompt: Optional[str] + model_id: Optional[str] + output_buffer: int + context_limit: Optional[int] + chunks: Optional[List[str]] + chunk_results: Annotated[list, operator.add] # collects parallel results + needs_chunking: bool + title: Optional[str] + + +class ChunkState(TypedDict): + """State for processing a single chunk (used by Send).""" + + system_prompt: str + model_id: Optional[str] + output_buffer: int + title: str + chunk: str + chunk_idx: int + total_chunks: int -async def run_transformation(state: dict, config: RunnableConfig) -> dict: +def _get_source(state: dict) -> Optional[Source]: source_obj = state.get("source") - source: Source = source_obj if isinstance(source_obj, Source) else None # type: ignore[assignment] + return source_obj if isinstance(source_obj, Source) else None + + +def _extract_response_content(response) -> str: + return clean_thinking_content(extract_text_content(response.content)) + + +def _build_system_prompt(state: dict) -> str: + transformation: Transformation = state["transformation"] + # transformation.prompt is user-controlled free text. Never compile it as + # Jinja template *source* (Prompter(template_text=...)) - pass it as a + # plain render variable into a fixed, developer-authored template instead. + # See docs/7-DEVELOPMENT/security.md (GHSA-f35w-wx37-26q7). + instructions = transformation.prompt + default_prompts: DefaultPrompts = DefaultPrompts(transformation_instructions=None) + if default_prompts.transformation_instructions: + instructions = ( + f"{default_prompts.transformation_instructions}\n\n{instructions}" + ) + return Prompter(prompt_template="transformation/execute").render( + data={**state, "instructions": instructions} + ) + + +def _get_content(state: dict) -> str: content = state.get("input_text") + if not content: + source = _get_source(state) + if source: + content = source.full_text + return str(content) if content else "" + + +async def try_full_content(state: dict, config: RunnableConfig) -> dict: + """Try processing the full content; on a context-limit error, fall back to + chunking by returning chunking parameters and ``needs_chunking=True``.""" + source = _get_source(state) + content = _get_content(state) assert source or content, "No content to transform" + transformation: Transformation = state["transformation"] + title = transformation.title + system_prompt = _build_system_prompt(state) + output_buffer = DEFAULT_OUTPUT_TOKENS + model_id = config.get("configurable", {}).get("model_id") + + payload = [SystemMessage(content=system_prompt), HumanMessage(content=content)] try: - if not content: - content = source.full_text - # transformation.prompt is user-controlled free text. Never compile it as - # Jinja template *source* (Prompter(template_text=...)) - pass it as a - # plain render variable into a fixed, developer-authored template instead. - # See docs/7-DEVELOPMENT/security.md (GHSA-f35w-wx37-26q7). - instructions = transformation.prompt - default_prompts: DefaultPrompts = DefaultPrompts(transformation_instructions=None) - if default_prompts.transformation_instructions: - instructions = f"{default_prompts.transformation_instructions}\n\n{instructions}" - - system_prompt = Prompter(prompt_template="transformation/execute").render( - data={**state, "instructions": instructions} + chain = await provision_langchain_model( + str(payload), model_id, "transformation", max_tokens=output_buffer ) - content_str = str(content) if content else "" - payload = [SystemMessage(content=system_prompt), HumanMessage(content=content_str)] + response = await chain.ainvoke(payload) + cleaned_content = _extract_response_content(response) + + if source: + await source.add_insight(title, cleaned_content) + + return {"output": cleaned_content, "needs_chunking": False} + + except OpenNotebookError: + raise + except Exception as e: + if not is_context_limit_error(e): + error_class, user_message = classify_error(e) + raise error_class(user_message) from e + + tokens_sent, context_limit = get_context_limit_from_error( + e, DEFAULT_CONTEXT_LIMIT + ) + output_buffer = calculate_output_buffer(context_limit) + logger.info( + f"Transformation '{title}' exceeded context " + f"({tokens_sent or '?'} tokens > {context_limit}); chunking." + ) + + system_overhead = token_count(system_prompt) + available = int(context_limit * SAFETY_BUFFER) - system_overhead - output_buffer + chunk_size = max(available, 500) + chunks = chunk_text_by_tokens(content, chunk_size) + logger.info(f"Split content into {len(chunks)} chunks for '{title}'") + + return { + "needs_chunking": True, + "system_prompt": system_prompt, + "model_id": model_id, + "output_buffer": output_buffer, + "context_limit": context_limit, + "chunks": chunks, + "chunk_results": [], + "title": title, + } + + +def fan_out_chunks(state: dict) -> Union[List[Send], str]: + """Conditional edge: fan out chunks to parallel ``process_chunk`` nodes, or + route straight to ``synthesize`` when no chunking is needed.""" + if not state.get("needs_chunking", False): + return "synthesize" + chunks = state.get("chunks", []) + if not chunks: + return "synthesize" + + total_chunks = len(chunks) + logger.info(f"Fanning out {total_chunks} chunks for parallel processing") + return [ + Send( + "process_chunk", + { + "system_prompt": state["system_prompt"], + "model_id": state.get("model_id"), + "output_buffer": state["output_buffer"], + "title": state.get("title", ""), + "chunk": chunk, + "chunk_idx": idx, + "total_chunks": total_chunks, + }, + ) + for idx, chunk in enumerate(chunks) + ] + + +async def process_chunk(state: ChunkState, config: RunnableConfig) -> dict: + """Process a single chunk (runs in parallel via LangGraph's Send API).""" + idx = state["chunk_idx"] + total = state["total_chunks"] + title = state.get("title", "transformation") + + logger.info(f"Processing chunk {idx + 1}/{total} for '{title}'") + # The section hint lives in the system prompt so it can't bleed into the + # output of extraction-style transformations that echo the user content. + chunk_system_prompt = ( + f"{state['system_prompt']}\n\n" + f"[Note: The input is section {idx + 1} of {total} of a larger " + f"document. Apply the instructions to this section only.]" + ) + payload = [ + SystemMessage(content=chunk_system_prompt), + HumanMessage(content=state["chunk"]), + ] + async with _get_chunk_semaphore(): chain = await provision_langchain_model( str(payload), - config.get("configurable", {}).get("model_id"), + state.get("model_id"), "transformation", - max_tokens=8192, + max_tokens=state["output_buffer"], ) + response = await chain.ainvoke(payload) + + logger.info(f"Chunk {idx + 1}/{total} for '{title}' completed") + return {"chunk_results": [{"idx": idx, "result": _extract_response_content(response)}]} + +def _batch_results_by_tokens(results: List[str], budget: int) -> List[List[str]]: + """Greedily group consecutive results so each group fits within ``budget`` + tokens. A lone result that exceeds the budget gets its own group.""" + batches: List[List[str]] = [] + current: List[str] = [] + current_tokens = 0 + for r in results: + t = token_count(r) + if current and current_tokens + t > budget: + batches.append(current) + current, current_tokens = [], 0 + current.append(r) + current_tokens += t + if current: + batches.append(current) + return batches + + +async def _synthesize_once( + results: List[str], state: dict, synthesis_prompt: str +) -> str: + """One synthesis LLM call merging a list of partial results.""" + combined_text = "\n\n---\n\n".join( + f"## Result from Part {i + 1}:\n{r}" for i, r in enumerate(results) + ) + payload = [ + SystemMessage(content=synthesis_prompt), + HumanMessage(content=combined_text), + ] + async with _get_chunk_semaphore(): + chain = await provision_langchain_model( + str(payload), + state.get("model_id"), + "transformation", + max_tokens=state.get("output_buffer", DEFAULT_OUTPUT_TOKENS), + ) response = await chain.ainvoke(payload) + return _extract_response_content(response) + + +async def _reduce_results( + results: List[str], state: dict, synthesis_prompt: str, budget: int +) -> str: + """Hierarchically reduce many chunk results into one, in rounds, so the + combined synthesis input never exceeds the model's context window.""" + round_num = 0 + while len(results) > 1: + round_num += 1 + batches = _batch_results_by_tokens(results, budget) + # If nothing grouped (each result alone exceeds budget), force pairs so + # the reduction always makes progress. + if len(batches) == len(results): + batches = [results[i : i + 2] for i in range(0, len(results), 2)] + logger.info( + f"Synthesis reduce round {round_num}: {len(results)} results -> " + f"{len(batches)} batch(es)" + ) + next_results: List[str] = [] + for batch in batches: + if len(batch) == 1: + next_results.append(batch[0]) + else: + next_results.append( + await _synthesize_once(batch, state, synthesis_prompt) + ) + results = next_results + return results[0] - # Clean thinking content from the response - response_content = extract_text_content(response.content) - cleaned_content = clean_thinking_content(response_content) +async def synthesize_results(state: dict, config: RunnableConfig) -> dict: + """Synthesize chunk results into the final output (no-op if full content + already succeeded).""" + source = _get_source(state) + title = state.get("title") or state["transformation"].title + + if state.get("output") and not state.get("needs_chunking", False): + return {} + + chunk_results: List[ChunkResult] = state.get("chunk_results", []) + if not chunk_results: + logger.warning(f"No chunk results to synthesize for '{title}'") + return {"output": ""} + + sorted_results = sorted(chunk_results, key=lambda x: x["idx"]) + + if len(sorted_results) == 1: + result = sorted_results[0]["result"] if source: - await source.add_insight(transformation.title, cleaned_content) + await source.add_insight(title, result) + return {"output": result} - return { - "output": cleaned_content, - } - except OpenNotebookError: - raise - except Exception as e: - error_class, user_message = classify_error(e) - raise error_class(user_message) from e + logger.info(f"Synthesizing {len(sorted_results)} chunk results for '{title}'") + synthesis_prompt = f"""You previously processed a large document in {len(sorted_results)} parts using the following instructions: + +{state.get("system_prompt", "")} + +Below are the results from each part. Your task is to synthesize these into a single, coherent output that combines the key information from all parts. Remove any redundancy and create a unified result. + +Do NOT simply concatenate - intelligently merge and synthesize the information.""" + + # Budget the synthesis input to the context window so combining many results + # (or large results) never overflows. + context_limit = state.get("context_limit") or DEFAULT_CONTEXT_LIMIT + output_buffer = state.get("output_buffer", DEFAULT_OUTPUT_TOKENS) + budget = max( + int(context_limit * SAFETY_BUFFER) + - token_count(synthesis_prompt) + - output_buffer, + 1000, + ) + + result = await _reduce_results( + [r["result"] for r in sorted_results], state, synthesis_prompt, budget + ) + + logger.info(f"Completed transformation '{title}' with parallel chunking") + if source: + await source.add_insight(title, result) + return {"output": result} +# Build the graph agent_state = StateGraph(TransformationState) -agent_state.add_node("agent", run_transformation) # type: ignore[type-var] -agent_state.add_edge(START, "agent") -agent_state.add_edge("agent", END) +agent_state.add_node("try_full", try_full_content) # type: ignore[type-var] +agent_state.add_node("process_chunk", process_chunk) # type: ignore[type-var] +agent_state.add_node("synthesize", synthesize_results) # type: ignore[type-var] +agent_state.add_edge(START, "try_full") +agent_state.add_conditional_edges("try_full", fan_out_chunks, ["process_chunk", "synthesize"]) +agent_state.add_edge("process_chunk", "synthesize") +agent_state.add_edge("synthesize", END) graph = agent_state.compile() diff --git a/open_notebook/utils/token_utils.py b/open_notebook/utils/token_utils.py index 224f769c2..8693aa8c2 100644 --- a/open_notebook/utils/token_utils.py +++ b/open_notebook/utils/token_utils.py @@ -1,9 +1,12 @@ """ Token utilities for Open Notebook. -Handles token counting and cost calculations for language models. +Handles token counting, cost calculations, context-limit error parsing, and +text chunking for working within model context windows. """ import os +import re +from typing import List, Optional, Tuple from open_notebook.config import TIKTOKEN_CACHE_DIR @@ -11,6 +14,21 @@ # tokenizer encodings are cached persistently in the data folder os.environ["TIKTOKEN_CACHE_DIR"] = TIKTOKEN_CACHE_DIR +# Safety buffer: use 90% of the context limit to leave headroom for output and +# for tokenizer disagreement between our estimate and the provider's. +SAFETY_BUFFER = 0.90 + +# Conservative fallback context limit (tokens) when an error can't be parsed. +DEFAULT_CONTEXT_LIMIT = 8192 + +# Initial output buffer (tokens). Matches the max_tokens the transformation +# graph has always used for the full-content attempt; lowering it would +# silently truncate outputs on the common (non-chunked) path. +DEFAULT_OUTPUT_TOKENS = 8192 + +# Fraction of the context window reserved for the model's output. +OUTPUT_RATIO = 0.10 + def token_count(input_string: str) -> int: """ @@ -55,3 +73,241 @@ def token_cost(token_count: int, cost_per_million: float = 0.150) -> float: float: The calculated cost for the given token count. """ return cost_per_million * (token_count / 1_000_000) + + +def parse_context_limit_error(error: Exception) -> Optional[Tuple[Optional[int], int]]: + """Extract token counts from a context-limit error message. + + Context-limit detection is inherently provider-format-dependent: providers + report the limit only in free-text error messages, so this parses the known + wordings and will return ``None`` for formats it doesn't recognize (callers + then fall back to ``DEFAULT_CONTEXT_LIMIT`` — see + ``get_context_limit_from_error``). Supported formats: + + - OpenAI: ``"maximum context length is 8192 tokens... 10000 tokens"`` + - Anthropic: ``"prompt is too long: 10000 tokens > 8192 maximum"`` + - Google: ``"input token count (10000) exceeds the maximum (8192)"`` + - Generic variants: ``"10000 tokens > 8192"``, ``"tokens (10000) exceeded + ... limit 8192"``, or a lone ``"max/maximum ... "`` + + Returns: + ``(tokens_sent, context_limit)`` if parseable, else ``None``. + ``tokens_sent`` may be ``None`` if only the limit could be found. + """ + error_str = str(error) if error else "" + if not error_str: + return None + + match = re.search( + r"tokens?\s*\(?(\d+)\)?\s*(?:exceeded|>).*?(?:limit|max|maximum)[^\d]*(\d+)", + error_str, + re.IGNORECASE, + ) + if match: + return int(match.group(1)), int(match.group(2)) + + match = re.search( + r"(\d+)\s*tokens?\s*(?:>|exceeded)\s*(\d+)\s*(?:max|limit)", + error_str, + re.IGNORECASE, + ) + if match: + return int(match.group(1)), int(match.group(2)) + + # OpenAI: "maximum context length is 8192 tokens...10000 tokens" + match = re.search( + r"maximum context length is (\d+) tokens.*?(\d+) tokens", + error_str, + re.DOTALL, + ) + if match: + return int(match.group(2)), int(match.group(1)) + + # Anthropic: "prompt is too long: 10000 tokens > 8192 maximum" + match = re.search( + r"prompt is too long: (\d+) tokens? > (\d+) maximum", error_str + ) + if match: + return int(match.group(1)), int(match.group(2)) + + # Google: "input token count (10000) exceeds the maximum (8192)" + match = re.search( + r"input token count \((\d+)\) exceeds the maximum \((\d+)\)", error_str + ) + if match: + return int(match.group(1)), int(match.group(2)) + + # Generic: "10000 tokens > 8192" + match = re.search(r"(\d+)\s*tokens?\s*>\s*(\d+)", error_str) + if match: + return int(match.group(1)), int(match.group(2)) + + # Only the limit is known. + match = re.search( + r"(?:max|maximum)\s*(?:context)?\s*(?:length|limit|window)?\s*(?:is|of|:)?\s*(\d+)", + error_str, + re.IGNORECASE, + ) + if match: + return None, int(match.group(1)) + + match = re.search( + r"(?:context|limit|max|window)[^\d]*(\d{4,})", error_str, re.IGNORECASE + ) + if match: + return None, int(match.group(1)) + + return None + + +def is_context_limit_error(error: Exception) -> bool: + """Return True if the exception is a context/token-limit error (and not a + rate limit, auth, quota, network, etc.). + + Like ``parse_context_limit_error``, this matches provider wordings by + keyword: non-context errors are excluded first, then context keywords are + required. Unknown wordings return False, so callers treat the error as a + regular failure rather than chunking on it.""" + error_msg = str(error).lower() + + non_context_keywords = [ + "rate limit", + "rate_limit", + "ratelimit", + "too many requests", + "quota exceeded", + "unauthorized", + "authentication", + "forbidden", + "not found", + "invalid api key", + "billing", + "insufficient_quota", + "server error", + "internal error", + "timeout", + "connection", + ] + for keyword in non_context_keywords: + if keyword in error_msg: + return False + + context_keywords = [ + "maximum context length", + "token limit", + "too many tokens", + "prompt is too long", + "exceeds the maximum", + "context window", + "max_tokens", + "token count", + "context length", + "input too long", + "request too large", + "payload size exceeds", + "you requested", + ] + for keyword in context_keywords: + if keyword in error_msg: + return True + + return False + + +def get_context_limit_from_error( + error: Exception, default_limit: int = DEFAULT_CONTEXT_LIMIT +) -> Tuple[Optional[int], int]: + """Parse a context-limit error, falling back to ``default_limit`` when the + limit can't be parsed. Returns ``(tokens_sent, context_limit)``. + + The fallback is deliberately conservative: an unrecognized wording yields + ``DEFAULT_CONTEXT_LIMIT`` (8192), which may over-chunk a large-context + model (more, smaller chunks) but never produces chunks that are too big to + process.""" + parsed = parse_context_limit_error(error) + if parsed: + return parsed + return None, default_limit + + +def calculate_output_buffer(context_limit: int) -> int: + """Reserve a slice of the context window for the model's output.""" + return int(context_limit * OUTPUT_RATIO) + + +def chunk_text_by_tokens(text: str, max_tokens: int) -> List[str]: + """Split text into chunks that each fit within ``max_tokens``. + + Hierarchical: paragraphs, then sentences, then words, then a hard character + split as a last resort. + """ + if not text: + return [] + if token_count(text) <= max_tokens: + return [text] + + paragraphs = re.split(r"\n\s*\n", text) + if len(paragraphs) > 1: + chunks = _merge_splits(paragraphs, max_tokens) + if chunks: + return chunks + + sentences = re.split(r"(?<=[.!?])\s+", text) + if len(sentences) > 1: + chunks = _merge_splits(sentences, max_tokens) + if chunks: + return chunks + + chunks = _merge_splits(text.split(), max_tokens, separator=" ") + if chunks: + return chunks + + # Last resort: hard split (rough estimate 1 token ~= 4 chars). + char_limit = max_tokens * 4 + return [text[i : i + char_limit] for i in range(0, len(text), char_limit)] + + +def _merge_splits( + parts: List[str], max_tokens: int, separator: str = "\n\n" +) -> List[str]: + """Merge split parts into chunks respecting the token limit, recursing into + finer splits for any single part that's still too large.""" + chunks: List[str] = [] + current_chunk = "" + + for part in parts: + part = part.strip() + if not part: + continue + + if token_count(part) > max_tokens: + if current_chunk: + chunks.append(current_chunk) + current_chunk = "" + if separator == "\n\n": + sub_chunks = _merge_splits( + re.split(r"(?<=[.!?])\s+", part), max_tokens, " " + ) + elif separator == " ": + char_limit = max_tokens * 4 + sub_chunks = [ + part[i : i + char_limit] + for i in range(0, len(part), char_limit) + ] + else: + sub_chunks = [part] + chunks.extend(sub_chunks) + continue + + candidate = f"{current_chunk}{separator}{part}" if current_chunk else part + if token_count(candidate) <= max_tokens: + current_chunk = candidate + else: + if current_chunk: + chunks.append(current_chunk) + current_chunk = part + + if current_chunk: + chunks.append(current_chunk) + + return chunks diff --git a/tests/test_add_insight_failure_propagation.py b/tests/test_add_insight_failure_propagation.py index 1e9575048..c3a01da02 100644 --- a/tests/test_add_insight_failure_propagation.py +++ b/tests/test_add_insight_failure_propagation.py @@ -61,11 +61,19 @@ async def test_still_raises_invalid_input_for_empty_type(self): class TestTransformationGraphPropagatesFailure: - """open_notebook/graphs/transformation.py: run_transformation().""" + """open_notebook/graphs/transformation.py: try_full_content(). + + Upstream tested this contract on run_transformation(), the single node the + transformation graph used to have. Chunking split that node into + try_full_content -> (fan_out_chunks -> process_chunk) -> synthesize_results; + try_full_content is the direct analogue — it owns the non-chunking path and + its add_insight() call. The contract is unchanged: a failed insight + submission must propagate instead of being reported as success. + """ @pytest.mark.asyncio async def test_add_insight_failure_propagates_out_of_run_transformation(self): - from open_notebook.graphs.transformation import run_transformation + from open_notebook.graphs.transformation import try_full_content source = make_source() transformation = MagicMock(title="Summary", prompt="Summarize this") @@ -103,13 +111,13 @@ async def test_add_insight_failure_propagates_out_of_run_transformation(self): source.full_text = "full text of the source" with pytest.raises(DatabaseOperationError): - await run_transformation(state, config={"configurable": {}}) + await try_full_content(state, config={"configurable": {}}) mock_add_insight.assert_awaited_once() @pytest.mark.asyncio async def test_successful_add_insight_returns_output_normally(self): - from open_notebook.graphs.transformation import run_transformation + from open_notebook.graphs.transformation import try_full_content source = make_source() transformation = MagicMock(title="Summary", prompt="Summarize this") @@ -144,10 +152,14 @@ async def test_successful_add_insight_returns_output_normally(self): mock_prompter_cls.return_value.render.return_value = "rendered prompt" source.full_text = "full text of the source" - result = await run_transformation(state, config={"configurable": {}}) + result = await try_full_content(state, config={"configurable": {}}) mock_add_insight.assert_awaited_once() - assert result == {"output": "the transformation output"} + # needs_chunking=False: the full content fit, so the graph skips chunking. + assert result == { + "output": "the transformation output", + "needs_chunking": False, + } class TestSourceGraphTransformContentPropagatesFailure: diff --git a/tests/test_graphs.py b/tests/test_graphs.py index f894a5b01..74fabbb9e 100644 --- a/tests/test_graphs.py +++ b/tests/test_graphs.py @@ -17,7 +17,10 @@ from open_notebook.graphs.tools import get_current_timestamp from open_notebook.graphs.transformation import ( TransformationState, - run_transformation, + _batch_results_by_tokens, + fan_out_chunks, + synthesize_results, + try_full_content, ) from open_notebook.graphs.transformation import ( graph as transformation_graph, @@ -129,8 +132,8 @@ def test_transformation_state_structure(self): assert state["output"] == "" @pytest.mark.asyncio - async def test_run_transformation_assertion_no_content(self): - """Test transformation raises assertion with no content.""" + async def test_try_full_content_assertion_no_content(self): + """try_full_content raises an assertion when there's no content.""" from unittest.mock import MagicMock from open_notebook.domain.transformation import Transformation @@ -146,7 +149,7 @@ async def test_run_transformation_assertion_no_content(self): config: RunnableConfig = {"configurable": {"model_id": None}} with pytest.raises(AssertionError, match="No content to transform"): - await run_transformation(state, config) + await try_full_content(state, config) def test_transformation_graph_compilation(self): """Test that transformation graph compiles correctly.""" @@ -154,6 +157,171 @@ def test_transformation_graph_compilation(self): assert hasattr(transformation_graph, "invoke") assert hasattr(transformation_graph, "ainvoke") + def test_fan_out_chunks_routes_to_synthesize_without_chunking(self): + """fan_out_chunks goes straight to synthesize when no chunking needed.""" + assert fan_out_chunks({"needs_chunking": False}) == "synthesize" + assert fan_out_chunks({"needs_chunking": True, "chunks": []}) == "synthesize" + + +class TestTransformationFullContentPath: + """Tests for the optimistic full-content attempt.""" + + @pytest.mark.asyncio + async def test_full_content_uses_8192_output_cap(self): + """The full-content attempt must keep the pre-chunking 8192 output cap; + a lower cap would silently truncate outputs on the common path.""" + from open_notebook.domain.transformation import Transformation + + mock_transformation = MagicMock(spec=Transformation) + mock_transformation.prompt = "Summarize the document." + mock_transformation.title = "Summary" + + state = { + "input_text": "Some short content.", + "transformation": mock_transformation, + "source": None, + } + + resp = MagicMock() + resp.content = "summary" + fake_chain = MagicMock() + fake_chain.ainvoke = AsyncMock(return_value=resp) + provision = AsyncMock(return_value=fake_chain) + + with patch( + "open_notebook.graphs.transformation.provision_langchain_model", + new=provision, + ): + result = await try_full_content(state, {"configurable": {}}) + + assert result == {"output": "summary", "needs_chunking": False} + assert provision.await_args is not None + assert provision.await_args.kwargs["max_tokens"] == 8192 + + +class TestProcessChunk: + """Tests for the parallel chunk processor.""" + + @pytest.mark.asyncio + async def test_chunk_sent_verbatim_with_hint_in_system_prompt(self): + """The section hint must live in the system prompt, not the user + content, so it can't bleed into extraction-style outputs.""" + from open_notebook.graphs.transformation import ChunkState, process_chunk + + state: ChunkState = { + "system_prompt": "Extract all names.", + "model_id": None, + "output_buffer": 800, + "title": "Names", + "chunk": "Alice met Bob.", + "chunk_idx": 1, + "total_chunks": 3, + } + + resp = MagicMock() + resp.content = "Alice, Bob" + fake_chain = MagicMock() + fake_chain.ainvoke = AsyncMock(return_value=resp) + + with patch( + "open_notebook.graphs.transformation.provision_langchain_model", + new=AsyncMock(return_value=fake_chain), + ): + result = await process_chunk(state, {"configurable": {}}) + + payload = fake_chain.ainvoke.await_args.args[0] + system_message, human_message = payload + assert human_message.content == "Alice met Bob." + assert system_message.content.startswith("Extract all names.") + assert "section 2 of 3" in system_message.content + assert result == {"chunk_results": [{"idx": 1, "result": "Alice, Bob"}]} + + def test_chunk_semaphore_is_per_event_loop(self): + """Each event loop gets its own semaphore (a shared one would raise + 'bound to a different event loop' across worker/API loops).""" + import asyncio + + from open_notebook.graphs.transformation import _get_chunk_semaphore + + async def grab_twice(): + return _get_chunk_semaphore(), _get_chunk_semaphore() + + loop1 = asyncio.new_event_loop() + try: + sem_a, sem_b = loop1.run_until_complete(grab_twice()) + finally: + loop1.close() + + loop2 = asyncio.new_event_loop() + try: + sem_c, _ = loop2.run_until_complete(grab_twice()) + finally: + loop2.close() + + assert sem_a is sem_b # same loop reuses its semaphore + assert sem_a is not sem_c # different loop gets a fresh one + + +class TestTransformationChunkingReduce: + """Tests for the large-document chunking + hierarchical synthesis reduce.""" + + def test_batch_results_by_tokens_respects_budget(self): + from open_notebook.graphs.transformation import token_count + + results = ["word " * 100 for _ in range(10)] # ~100 tokens each + budget = 250 + batches = _batch_results_by_tokens(results, budget) + + assert sum(len(b) for b in batches) == 10 # nothing dropped + for b in batches: + assert len(b) == 1 or sum(token_count(r) for r in b) <= budget + + @pytest.mark.asyncio + async def test_synthesize_batches_instead_of_overflowing(self): + """Many/large chunk results must be reduced in context-sized batches, + not concatenated into one oversized synthesis call.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from open_notebook.graphs.transformation import token_count + + chunk_results = [{"idx": i, "result": "word " * 1000} for i in range(12)] + state = { + "output": None, + "needs_chunking": True, + "chunk_results": chunk_results, + "title": "Dense Summary", + "system_prompt": "Summarize the document.", + "output_buffer": 1000, + "context_limit": 8000, + "model_id": None, + "source": None, + "transformation": MagicMock(title="Dense Summary"), + } + + seen_call_tokens: list[int] = [] + + async def fake_ainvoke(payload): + seen_call_tokens.append(token_count(payload[-1].content)) + resp = MagicMock() + resp.content = "merged" # small result so the reduction converges + return resp + + fake_chain = MagicMock() + fake_chain.ainvoke = fake_ainvoke + + budget = int(8000 * 0.90) # generous upper bound for any single call + with patch( + "open_notebook.graphs.transformation.provision_langchain_model", + new=AsyncMock(return_value=fake_chain), + ): + result = await synthesize_results(state, {"configurable": {}}) + + assert result["output"] == "merged" + assert len(seen_call_tokens) > 1, "should batch into multiple calls" + assert all( + n <= budget for n in seen_call_tokens + ), f"a synthesis call exceeded budget: {seen_call_tokens}" + # ============================================================================ # TEST SUITE 4: Source Graph - Title Preservation