diff --git a/open_notebook/graphs/transformation.py b/open_notebook/graphs/transformation.py index 1eff8c6c87..a1969c3435 100644 --- a/open_notebook/graphs/transformation.py +++ b/open_notebook/graphs/transformation.py @@ -1,77 +1,623 @@ +""" +Transformation graph — applies LLM-driven transformations to source content. + +For small documents the entire content is processed in one LLM call (same as +the original single-node graph). For large documents that exceed the model's +context window, the graph transparently splits the content into token-bounded +chunks, processes them in parallel via LangGraph Send, and hierarchically +synthesizes the results — all without any user configuration. +""" + +import asyncio +import operator +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 typing_extensions import TypedDict +from langgraph.types import Send +from loguru import logger +from typing_extensions import NotRequired, TypedDict from open_notebook.ai.provision import provision_langchain_model from open_notebook.domain.notebook import Source from open_notebook.domain.transformation import DefaultPrompts, Transformation -from open_notebook.exceptions import OpenNotebookError +from open_notebook.exceptions import ( + AuthenticationError, + ConfigurationError, + ExternalServiceError, + NetworkError, + RateLimitError, +) 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, + SAFETY_BUFFER, + chunk_text_by_tokens, + parse_context_limit_error, + token_count, +) + +# ── Constants ────────────────────────────────────────────────────────── + +# Token budget for the full-content attempt (matching the original graph). +FULL_CONTENT_MAX_TOKENS = 8192 + +# Max concurrent chunk LLM calls. Prevents rate-limit storms when a +# document is split into many chunks, especially across concurrent +# worker jobs. +MAX_CONCURRENT_CHUNKS = 10 + +# Module-level semaphore shared by all process_chunk invocations. +_chunk_semaphore = asyncio.Semaphore(MAX_CONCURRENT_CHUNKS) + +# ── State types ──────────────────────────────────────────────────────── + + +class ChunkResult(TypedDict): + """Result from processing a single chunk.""" + + index: int + output: str + error_class: NotRequired[str] # set when chunk processing failed + error_message: NotRequired[str] # classified user-facing message class TransformationState(TypedDict): + """Shared graph state for the entire transformation pipeline.""" + input_text: str - source: Source + source: Optional[Source] transformation: Transformation output: str + # Chunking fields — populated only when the full-content path fails + chunks: NotRequired[List[str]] + chunk_results: NotRequired[Annotated[List[ChunkResult], operator.add]] + total_chunks: NotRequired[int] + # Parsed model context window (set when chunking is triggered) + context_limit: NotRequired[int] + # Output token budget derived from context_limit (set alongside it) + output_budget: NotRequired[int] + +class ChunkState(TypedDict): + """Per-chunk state dispatched via Send.""" -async def run_transformation(state: dict, config: RunnableConfig) -> dict: - source_obj = state.get("source") - source: Source = source_obj if isinstance(source_obj, Source) else None # type: ignore[assignment] + input_text: str + source: Optional[Source] + transformation: Transformation + chunk_index: int + chunk_text: str + total_chunks: int + # Output token budget for this model (derived from context_limit) + output_budget: int + + +# ── Helpers ──────────────────────────────────────────────────────────── + + +def _get_source(state) -> Optional[Source]: + source = state.get("source") + return source if isinstance(source, Source) else None + + +def _get_content(state) -> str: content = state.get("input_text") + if content: + return content + source = _get_source(state) + if source: + return source.full_text or "" + return "" + + +def _build_system_prompt(state, instructions: str, section_hint: str = "") -> str: + """Render the system prompt, optionally with a section context hint.""" + prompt_data = {**state, "instructions": instructions} + if section_hint: + prompt_data["section_context"] = section_hint + return Prompter(prompt_template="transformation/execute").render(data=prompt_data) + + +def _extract_response_content(response) -> str: + """Extract and clean text from an LLM response.""" + content = extract_text_content(response.content) + return clean_thinking_content(content) + + +async def _handle_llm_error(exc: Exception, content: str, system_prompt: str) -> dict: + """ + Handle an LLM error: context-limit errors trigger chunking, all others + are re-raised with appropriate classification. + """ + error_str = str(exc) + parsed = parse_context_limit_error(error_str) + + if not parsed: + # Not a context-limit error — re-raise classified + exc_class, message = classify_error(exc) + raise exc_class(message) from exc + + context_limit, _tokens_sent = parsed + logger.info(f"Content exceeds context window ({context_limit}). Chunking.") + + # Calculate a safe chunk size: leave room for system prompt + output + system_prompt_tokens = token_count(system_prompt) + output_budget = min(FULL_CONTENT_MAX_TOKENS, int(context_limit * 0.10)) + max_chunk_tokens = context_limit - system_prompt_tokens - output_budget + max_chunk_tokens = int(max_chunk_tokens * SAFETY_BUFFER) + max_chunk_tokens = max(max_chunk_tokens, 512) + + chunks = chunk_text_by_tokens(content, max_chunk_tokens) + + if len(chunks) == 1: + # Single chunk means even the split content is still too large + # for this model's context window — propagate original error + exc_class, message = classify_error(exc) + raise exc_class(message) from exc + + logger.info(f"Split content into {len(chunks)} chunks") + return { + "output": "", + "chunks": chunks, + "total_chunks": len(chunks), + "context_limit": context_limit, + "output_budget": output_budget, + } + + +# ── Graph nodes ──────────────────────────────────────────────────────── + + +async def try_full_content(state: TransformationState, config: RunnableConfig) -> dict: + """ + Optimistically process the full content in one call. + + On success, sets ``output`` and returns empty chunks. On a context-limit + error, chunks the content for parallel processing. + """ + source = _get_source(state) + content = _get_content(state) assert source or content, "No content to transform" - transformation: Transformation = state["transformation"] - 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}" + transformation: Transformation = state["transformation"] + instructions = transformation.prompt - system_prompt = Prompter(prompt_template="transformation/execute").render( - data={**state, "instructions": instructions} + default_prompts: DefaultPrompts = DefaultPrompts(transformation_instructions=None) + if default_prompts.transformation_instructions: + instructions = ( + f"{default_prompts.transformation_instructions}\n\n{instructions}" ) - content_str = str(content) if content else "" - payload = [SystemMessage(content=system_prompt), HumanMessage(content=content_str)] + + system_prompt = _build_system_prompt(state, instructions) + content_str = str(content) if content else "" + payload = [ + SystemMessage(content=system_prompt), + HumanMessage(content=content_str), + ] + + # ── Optimistic full-content attempt ────────────────────────────── + try: chain = await provision_langchain_model( str(payload), config.get("configurable", {}).get("model_id"), "transformation", - max_tokens=8192, + max_tokens=FULL_CONTENT_MAX_TOKENS, ) - response = await chain.ainvoke(payload) + output = _extract_response_content(response) + except Exception as e: + # Check if this is a context-limit error we can recover from + return await _handle_llm_error(e, content, system_prompt) + # ── End of LLM call section ─────────────────────────────────────── - # Clean thinking content from the response - response_content = extract_text_content(response.content) - cleaned_content = clean_thinking_content(response_content) + if source: + await source.add_insight(transformation.title, output) + + return {"output": output, "chunks": []} + + +def fan_out_or_synthesize( + state: TransformationState, +) -> Union[List[Send], str]: + """ + Route to chunk processing or directly to synthesis. + + Returns ``List[Send]`` for parallel chunk processing when chunks are + present, or the string ``"synthesize"`` to skip directly to output. + """ + chunks = state.get("chunks") + if chunks: + total = state.get("total_chunks", len(chunks)) + output_budget = state.get( + "output_budget", FULL_CONTENT_MAX_TOKENS + ) + return [ + Send( + "process_chunk", + { + "input_text": state.get("input_text", ""), + "source": _get_source(state), + "transformation": state["transformation"], + "chunk_index": i, + "chunk_text": chunk_text, + "total_chunks": total, + "output_budget": output_budget, + }, + ) + for i, chunk_text in enumerate(chunks) + ] + # No chunking needed — output is already set by try_full_content + return "synthesize" - if source: - await source.add_insight(transformation.title, cleaned_content) +async def process_chunk(state: ChunkState, config: RunnableConfig) -> dict: + """Process a single chunk and return its partial result.""" + try: + chunk_text = state["chunk_text"] + transformation: Transformation = state["transformation"] + chunk_index = state["chunk_index"] + total_chunks = state["total_chunks"] + + instructions = transformation.prompt + + default_prompts: DefaultPrompts = DefaultPrompts(transformation_instructions=None) + if default_prompts.transformation_instructions: + instructions = ( + f"{default_prompts.transformation_instructions}\n\n{instructions}" + ) + + # Provide a section hint so the LLM knows this is part of a larger doc + section_hint = ( + f"This is section {chunk_index + 1} of {total_chunks} of the source " + f"document. Process this section according to the instructions." + ) + + system_prompt = _build_system_prompt(state, instructions, section_hint=section_hint) + content_str = str(chunk_text) if chunk_text else "" + payload = [ + SystemMessage(content=system_prompt), + HumanMessage(content=content_str), + ] + + chunk_output_budget = state.get( + "output_budget", FULL_CONTENT_MAX_TOKENS + ) + async with _chunk_semaphore: + chain = await provision_langchain_model( + str(payload), + config.get("configurable", {}).get("model_id"), + "transformation", + max_tokens=chunk_output_budget, + ) + response = await chain.ainvoke(payload) + cleaned = _extract_response_content(response) + + logger.debug(f"Processed chunk {chunk_index + 1}/{total_chunks}") + return {"chunk_results": [{"index": chunk_index, "output": cleaned}]} + except Exception as e: + logger.opt(exception=True).error( + "Failed to process chunk {}: {}", chunk_index + 1, e + ) + # Classify the error so the retry mechanism gets the right + # exception type. Don't re-raise immediately (let sibling + # chunks finish), but store the typed error so + # synthesize_results can surface an aggregated failure + # before saving an incomplete insight. + exc_class, exc_msg = classify_error(e) return { - "output": cleaned_content, + "chunk_results": [ + { + "index": chunk_index, + "output": "", + "error_class": exc_class.__name__, + "error_message": exc_msg, + } + ] } - except OpenNotebookError: - raise + + +async def synthesize_results( + state: TransformationState, config: RunnableConfig +) -> dict: + """ + Synthesize partial chunk results into a final output. + + Handles three cases: + 1. No chunking happened — pass through the output from try_full_content. + 2. Single chunk result — use it directly. + 3. Multiple results — hierarchically reduce via LLM merging rounds. + """ + output = state.get("output") + if output: + # Direct path — no chunking occurred + return {"output": output} + + chunk_results: List[ChunkResult] = state.get("chunk_results", []) + transformation: Transformation = state["transformation"] + + if len(chunk_results) == 0: + return {"output": ""} + + if len(chunk_results) == 1: + error_class = chunk_results[0].get("error_class") + if error_class: + _ERROR_CLASSES = { + "AuthenticationError": AuthenticationError, + "ConfigurationError": ConfigurationError, + "ExternalServiceError": ExternalServiceError, + "NetworkError": NetworkError, + "RateLimitError": RateLimitError, + } + exc_cls = _ERROR_CLASSES.get( + error_class, ExternalServiceError + ) + raise exc_cls(chunk_results[0]["error_message"]) + try: + result = chunk_results[0]["output"] + source = _get_source(state) + if source: + await source.add_insight(transformation.title, result) + return {"output": result} + except Exception as e: + exc_class, message = classify_error(e) + raise exc_class(message) from e + + instructions = transformation.prompt + + default_prompts: DefaultPrompts = DefaultPrompts(transformation_instructions=None) + if default_prompts.transformation_instructions: + instructions = ( + f"{default_prompts.transformation_instructions}\n\n{instructions}" + ) + + # Sort by index to maintain document order + results = sorted(chunk_results, key=lambda r: r["index"]) + texts = [r["output"] for r in results] + + # Use the model's actual context window and output budget (parsed + # from the error that triggered chunking), with safe defaults. + context_limit = state.get("context_limit", DEFAULT_CONTEXT_LIMIT) + output_budget = state.get( + "output_budget", int(context_limit * 0.10) + ) + + while len(texts) > 1: + # Group texts within token budget + groups: List[List[str]] = [] + current_group: List[str] = [] + current_tokens = 0 + + for text in texts: + t = token_count(text) + if ( + current_tokens + t > context_limit - output_budget - 500 + and current_group + ): + groups.append(current_group) + current_group = [] + current_tokens = 0 + current_group.append(text) + current_tokens += t + + if current_group: + groups.append(current_group) + + # Merge each group via LLM + merged: List[str] = [] + made_progress = False + + for group in groups: + if len(group) == 1: + merged.append(group[0]) + continue + + made_progress = True + + try: + chain = await provision_langchain_model( + f"Merge {len(group)} results", + config.get("configurable", {}).get("model_id"), + "transformation", + max_tokens=output_budget, + ) + merge_content = "\n\n---\n\n".join( + f"Section {i + 1}:\n{text}" for i, text in enumerate(group) + ) + merge_prompt = ( + f"Below are {len(group)} partial results from processing " + f"different sections of a document. Merge them into a single " + f"coherent output following this instruction: " + f"{instructions}\n\n{merge_content}" + ) + response = await chain.ainvoke( + [ + SystemMessage( + content="You are merging partial transformation " + "results into a final output." + ), + HumanMessage(content=merge_prompt), + ] + ) + merged.append(_extract_response_content(response)) + except Exception as e: + exc_class, message = classify_error(e) + raise exc_class(message) from e + + if not made_progress: + # All groups are single-item — the loop would never shrink. + # Fall back to merging the smallest adjacent pair (preserving + # document order so downstream transformations remain coherent). + logger.warning( + "Synthesis stalled: all {} groups are single items within " + "the token budget (limit={}). " + "Forcing pairwise merge of the smallest adjacent pair.", + len(groups), + context_limit, + ) + + # Find the adjacent pair with the smallest combined size + min_combined = float("inf") + merge_idx = 0 + for i in range(len(texts) - 1): + combined = token_count(texts[i]) + token_count(texts[i + 1]) + if combined < min_combined: + min_combined = combined + merge_idx = i + + # If even the smallest adjacent pair exceeds the context + # window, re-chunk oversized results so pieces fit in + # future merge rounds, instead of returning raw partials. + overhead = token_count(instructions) + 300 + available = context_limit - overhead + + # Calculate re-chunk budget: 40% of available so ~2.5 pieces + # fit within the context window for merging. + re_chunk_budget = max(200, int(available * 0.4)) + + # 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 " + "synthesize chunk results given the current " + f"instructions (limit={context_limit}). Use a " + "model with a larger context window or provide " + "shorter transformation instructions." + ) + + # Re-chunking is feasible — split oversized texts so they + # can be merged in subsequent rounds. + logger.warning( + "No adjacent pair fits the context window (limit={}, " + "smallest pair={} tkn). Re-chunking {} texts at " + "{} tokens each for further synthesis.", + context_limit, + min_combined, + len(texts), + re_chunk_budget, + ) + re_chunked: List[str] = [] + for chunk_text in texts: + if token_count(chunk_text) > re_chunk_budget: + re_chunked.extend( + chunk_text_by_tokens( + chunk_text, + max_chunk_tokens=re_chunk_budget, + overlap_chars=0, + ) + ) + else: + re_chunked.append(chunk_text) + texts = re_chunked + # Continue the loop to merge the re-chunked pieces + continue + + left = texts[merge_idx] + right = texts[merge_idx + 1] + force_content = ( + f"Part 1:\n{left}\n\n---\n\nPart 2:\n{right}" + ) + merge_prompt = ( + f"Merge the following two partial results into one " + f"coherent output following this instruction: " + f"{instructions}\n\n{force_content}" + ) + try: + chain = await provision_langchain_model( + merge_prompt, + config.get("configurable", {}).get("model_id"), + "transformation", + max_tokens=output_budget, + ) + response = await chain.ainvoke( + [ + SystemMessage( + content="You are merging partial transformation " + "results into a final output." + ), + HumanMessage(content=merge_prompt), + ] + ) + merged_result = _extract_response_content(response) + # Replace the merged pair with the result, preserving order + texts = ( + texts[:merge_idx] + + [merged_result] + + texts[merge_idx + 2 :] + ) + except Exception as e: + exc_class, message = classify_error(e) + raise exc_class(message) from e + continue + + texts = merged + + final_output = texts[0] if texts else "" + + # If any chunk failed, surface an aggregated classified error + # before saving an incomplete insight. This preserves the retry + # mechanism (transient errors won't be swallowed) while still + # allowing sibling chunks to complete. + chunk_errors = [ + r for r in chunk_results if r.get("error_class") + ] + if chunk_errors: + failed_count = len(chunk_errors) + first = chunk_errors[0] + _ERROR_CLASSES = { + "AuthenticationError": AuthenticationError, + "ConfigurationError": ConfigurationError, + "ExternalServiceError": ExternalServiceError, + "NetworkError": NetworkError, + "RateLimitError": RateLimitError, + } + cls_name = first["error_class"] or "ExternalServiceError" + exc_cls = _ERROR_CLASSES.get(cls_name, ExternalServiceError) + msg = ( + f"Partial transformation failure: {failed_count} of " + f"{len(chunk_results)} section(s) could not be processed. " + f"{first['error_message']}" + ) + raise exc_cls(msg) + + try: + source = _get_source(state) + if source: + await source.add_insight(transformation.title, final_output) except Exception as e: - error_class, user_message = classify_error(e) - raise error_class(user_message) from e + exc_class, message = classify_error(e) + raise exc_class(message) from e + + return {"output": final_output} +# ── Graph construction ───────────────────────────────────────────────── + 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) +agent_state.add_node("process_chunk", process_chunk) +agent_state.add_node("synthesize", synthesize_results) + +agent_state.add_edge(START, "try_full") + +# try_full routes to process_chunk (via Send) when chunking, or to +# synthesize directly when the full content succeeded. +agent_state.add_conditional_edges( + "try_full", + fan_out_or_synthesize, + { + "process_chunk": "process_chunk", + "synthesize": "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 224f769c21..2c78e7dee0 100644 --- a/open_notebook/utils/token_utils.py +++ b/open_notebook/utils/token_utils.py @@ -1,9 +1,11 @@ """ Token utilities for Open Notebook. -Handles token counting and cost calculations for language models. +Handles token counting, cost calculations, and text chunking. """ import os +import re +from typing import List, Optional, Tuple from open_notebook.config import TIKTOKEN_CACHE_DIR @@ -11,6 +13,17 @@ # tokenizer encodings are cached persistently in the data folder os.environ["TIKTOKEN_CACHE_DIR"] = TIKTOKEN_CACHE_DIR +# Default context limit when it cannot be parsed from an error message. +# Used as the chunk size so each chunk fits within the default budget. +DEFAULT_CONTEXT_LIMIT = 8192 + +# Safety buffer: keep chunk content below the raw limit to leave room +# for the system prompt and output tokens. +SAFETY_BUFFER = 0.85 + +# Overlap between adjacent chunks in characters (≈1–2 sentences). +CHUNK_OVERLAP_CHARS = 200 + def token_count(input_string: str) -> int: """ @@ -55,3 +68,199 @@ 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_str: str) -> Optional[Tuple[int, int]]: + """ + Parse an LLM error string to extract the context limit and tokens sent. + + Supports Anthropic, OpenAI, Google/Gemini, and generic error formats. + + Each pattern explicitly identifies which captured group is the context + limit vs tokens sent, rather than relying on a size heuristic (which + is incorrect when the sent tokens outnumber the model's limit, as in + Gemini's "X exceeded the limit of Y" format). + + Args: + error_str: The raw error string from the LLM provider. + + Returns: + (context_limit, tokens_sent) if both values can be parsed, else None. + """ + # Each entry: (regex, limit_group_num, sent_group_num) + # 1 — Anthropic: "your request was X tokens ... model's maximum is Y" + # 2 — OpenAI: "maximum context length is Y tokens ... you requested X" + # 3 — Google: "The number of tokens (X) exceeded the limit (Y)" + # 4 — Generic: "context_length_exceeded: X > Y" + patterns: List[Tuple[str, int, int]] = [ + ( + r"(?:you have|sent|used|have|requested|your request was).*?" + r"(\d+).*?(?:token|character).*?(?:maximum|limit|context).*?(\d+)", + 2, + 1, + ), + ( + r"(?:maximum|limit|context).*?(\d+).*?(?:token|character).*?" + r"requested.*?(\d+)", + 1, + 2, + ), + ( + r"number of tokens.*?(\d+).*?exceeded.*?(\d+)", + 2, + 1, + ), + ( + r"context[_\s]length[_\s]exceeded.*?(\d+).*?(\d+)", + 2, + 1, + ), + ] + for pattern, limit_group, sent_group in patterns: + match = re.search(pattern, error_str, re.IGNORECASE) + if match: + return ( + int(match.group(limit_group)), + int(match.group(sent_group)), + ) + return None + + +def chunk_text_by_tokens( + text: str, + max_chunk_tokens: int = int(DEFAULT_CONTEXT_LIMIT * SAFETY_BUFFER), + overlap_chars: int = CHUNK_OVERLAP_CHARS, +) -> List[str]: + """ + Split text into overlapping token-bounded chunks. + + Splits hierarchically: paragraphs first, then sentences, then word + boundaries. Adds overlap_chars of trailing text from the previous + chunk so context carries across chunk boundaries. + + Args: + text: The text to split. + max_chunk_tokens: Maximum tokens per chunk (default ~6963). + overlap_chars: Characters of overlap from the previous chunk. + + Returns: + A list of text chunks, each within the token budget. + Returns [text] unchanged if it fits in a single chunk. + """ + if token_count(text) <= max_chunk_tokens: + return [text] + + # Split on paragraph boundaries first + paragraphs = text.split("\n\n") + raw_chunks: List[str] = [] + current: List[str] = [] + current_tokens = 0 + + def _flush() -> None: + nonlocal current, current_tokens + if current: + raw_chunks.append("\n\n".join(current)) + current = [] + current_tokens = 0 + + for para in paragraphs: + para_tokens = token_count(para) + if current_tokens + para_tokens > max_chunk_tokens and current: + _flush() + current.append(para) + current_tokens += para_tokens + _flush() + + # If any single paragraph exceeds max_chunk_tokens, split it further + # by sentences, then by brute force. + final_chunks: List[str] = [] + for chunk in raw_chunks: + if token_count(chunk) <= max_chunk_tokens: + final_chunks.append(chunk) + else: + # Split oversized chunk by sentence boundaries + sentences = re.split(r"(?<=[.!?])\s+", chunk) + sentence_chunks: List[str] = [] + cur_sents: List[str] = [] + cur_sent_tokens = 0 + for sent in sentences: + sent_tokens = token_count(sent) + if cur_sent_tokens + sent_tokens > max_chunk_tokens and cur_sents: + sentence_chunks.append(" ".join(cur_sents)) + cur_sents = [] + cur_sent_tokens = 0 + cur_sents.append(sent) + cur_sent_tokens += sent_tokens + if cur_sents: + sentence_chunks.append(" ".join(cur_sents)) + + # Word-level fallback: if any sentence chunk still exceeds the + # token budget (e.g. a run-on with no sentence boundaries), split + # by whitespace word boundaries. This is the final brute-force + # level after paragraph → sentence → word. + word_chunks: List[str] = [] + for sc in sentence_chunks: + if token_count(sc) <= max_chunk_tokens: + word_chunks.append(sc) + else: + words = sc.split() + cur_words: List[str] = [] + cur_word_tokens = 0 + for w in words: + w_tokens = token_count(w) + if cur_word_tokens + w_tokens > max_chunk_tokens and cur_words: + word_chunks.append(" ".join(cur_words)) + cur_words = [] + cur_word_tokens = 0 + + # If a single word/string exceeds the token budget, + # split it at character boundaries as a last resort. + # This handles whitespace-free content (e.g. long + # base64/MD5 hashes, minified code, or CJK text + # without spaces) that no sentence/word split can + # break further. + if w_tokens > max_chunk_tokens: + if cur_words: + word_chunks.append(" ".join(cur_words)) + cur_words = [] + cur_word_tokens = 0 + # Proportional character-level split. + # o200k ~0.75 tokens/char for English; we + # overshoot slightly and let the next stage + # of hierarchical splitting clean up. + n_chunks = max(1, w_tokens // max_chunk_tokens + 1) + chunk_len = max(1, len(w) // n_chunks) + for start in range(0, len(w), chunk_len): + piece = w[start : start + chunk_len] + word_chunks.append(piece) + continue + + cur_words.append(w) + cur_word_tokens += w_tokens + if cur_words: + word_chunks.append(" ".join(cur_words)) + final_chunks.extend(word_chunks) + + # Add overlap from previous chunk and re-check token budget. + # The overlap + continuation marker can push a chunk over + # max_chunk_tokens; gracefully reduce the overlap when that happens. + CONTINUATION_MARKER = "\n\n[... continuation from previous section ...]\n\n" + + result: List[str] = [] + for i, chunk in enumerate(final_chunks): + if i > 0 and overlap_chars > 0: + base_chunk = chunk # preserve for retry below + # Try progressively smaller overlap sizes + chunk_upper_bound = max_chunk_tokens * 1.15 + for step_chars in [overlap_chars, overlap_chars // 2, overlap_chars // 4]: + overlap = final_chunks[i - 1][-step_chars:] + chunk = overlap + CONTINUATION_MARKER + base_chunk + if token_count(chunk) <= chunk_upper_bound: + break + # Continue with smaller overlap on next iteration + else: + # All overlap sizes overshoot — use just the continuation marker + chunk = CONTINUATION_MARKER + base_chunk + result.append(chunk) + + return result diff --git a/prompts/transformation/execute.jinja b/prompts/transformation/execute.jinja index 4b18195e1f..5f63e5c9d4 100644 --- a/prompts/transformation/execute.jinja +++ b/prompts/transformation/execute.jinja @@ -1,4 +1,6 @@ -{{ instructions }} +{% if section_context %}{{ section_context }} + +{% endif %}{{ instructions }} # MATH FORMATTING diff --git a/tests/test_add_insight_failure_propagation.py b/tests/test_add_insight_failure_propagation.py index 1e95750482..ccd9ba9fc8 100644 --- a/tests/test_add_insight_failure_propagation.py +++ b/tests/test_add_insight_failure_propagation.py @@ -61,11 +61,14 @@ 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().""" @pytest.mark.asyncio - async def test_add_insight_failure_propagates_out_of_run_transformation(self): - from open_notebook.graphs.transformation import run_transformation + async def test_add_insight_failure_propagates_out_of_try_full_content(self): + from open_notebook.graphs.transformation import ( + TransformationState, + try_full_content, + ) source = make_source() transformation = MagicMock(title="Summary", prompt="Summarize this") @@ -86,9 +89,7 @@ async def test_add_insight_failure_propagates_out_of_run_transformation(self): "open_notebook.graphs.transformation.DefaultPrompts", return_value=MagicMock(transformation_instructions=None), ), - patch( - "open_notebook.graphs.transformation.Prompter" - ) as mock_prompter_cls, + patch("open_notebook.graphs.transformation.Prompter") as mock_prompter_cls, patch( "open_notebook.graphs.transformation.provision_langchain_model", new=AsyncMock(return_value=fake_chain), @@ -103,13 +104,19 @@ 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( + cast(TransformationState, 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 ( + TransformationState, + try_full_content, + ) source = make_source() transformation = MagicMock(title="Summary", prompt="Summarize this") @@ -130,9 +137,7 @@ async def test_successful_add_insight_returns_output_normally(self): "open_notebook.graphs.transformation.DefaultPrompts", return_value=MagicMock(transformation_instructions=None), ), - patch( - "open_notebook.graphs.transformation.Prompter" - ) as mock_prompter_cls, + patch("open_notebook.graphs.transformation.Prompter") as mock_prompter_cls, patch( "open_notebook.graphs.transformation.provision_langchain_model", new=AsyncMock(return_value=fake_chain), @@ -144,10 +149,13 @@ 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( + cast(TransformationState, state), config={"configurable": {}} + ) mock_add_insight.assert_awaited_once() - assert result == {"output": "the transformation output"} + assert result["output"] == "the transformation output" + assert result.get("chunks") == [] class TestSourceGraphTransformContentPropagatesFailure: diff --git a/tests/test_graphs.py b/tests/test_graphs.py index f894a5b014..ce937e115b 100644 --- a/tests/test_graphs.py +++ b/tests/test_graphs.py @@ -16,8 +16,9 @@ from open_notebook.graphs.prompt import PatternChainState, graph from open_notebook.graphs.tools import get_current_timestamp from open_notebook.graphs.transformation import ( + ChunkState, TransformationState, - run_transformation, + try_full_content, ) from open_notebook.graphs.transformation import ( graph as transformation_graph, @@ -121,6 +122,9 @@ def test_transformation_state_structure(self): source=mock_source, transformation=mock_transformation, output="", + chunks=[], + chunk_results=[], + total_chunks=0, ) assert state["input_text"] == "Test text" @@ -129,8 +133,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): + """Test try_full_content raises assertion with no content.""" from unittest.mock import MagicMock from open_notebook.domain.transformation import Transformation @@ -146,7 +150,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(cast(TransformationState, state), config) def test_transformation_graph_compilation(self): """Test that transformation graph compiles correctly.""" @@ -154,9 +158,398 @@ def test_transformation_graph_compilation(self): assert hasattr(transformation_graph, "invoke") assert hasattr(transformation_graph, "ainvoke") + def test_transformation_state_with_chunking_fields(self): + """Test TransformationState accepts the new chunking fields.""" + from unittest.mock import MagicMock + + from open_notebook.domain.notebook import Source + from open_notebook.domain.transformation import Transformation + + mock_source = MagicMock(spec=Source) + mock_transformation = MagicMock(spec=Transformation) + + state = TransformationState( + input_text="Test text", + source=mock_source, + transformation=mock_transformation, + output="", + chunks=["chunk1", "chunk2"], + chunk_results=[{"index": 0, "output": "result1"}], + total_chunks=2, + context_limit=8192, + ) + + assert state["chunks"] == ["chunk1", "chunk2"] + assert state["chunk_results"] == [{"index": 0, "output": "result1"}] + assert state["total_chunks"] == 2 + assert state["context_limit"] == 8192 + + +# ============================================================================ +# TEST SUITE 4: Synthesize Results (parallel chunking) +# ============================================================================ + + +class TestSynthesizeResults: + """Test suite for synthesize_results progress guard and edge cases.""" + + @pytest.mark.asyncio + async def test_no_chunking_passthrough(self): + """When output is already set (no chunking), pass through.""" + from open_notebook.graphs.transformation import synthesize_results + + state = {"output": "direct output"} + result = await synthesize_results( + cast(TransformationState, state), {"configurable": {}} + ) + assert result == {"output": "direct output"} + + @pytest.mark.asyncio + async def test_empty_chunk_results(self): + """No chunk results returns empty output.""" + from unittest.mock import MagicMock + + from open_notebook.graphs.transformation import synthesize_results + + state = { + "output": "", + "chunk_results": [], + "transformation": MagicMock(title="Test"), + } + result = await synthesize_results( + cast(TransformationState, state), {"configurable": {}} + ) + assert result == {"output": ""} + + @pytest.mark.asyncio + async def test_single_chunk_result(self): + """Single chunk result returns its output directly.""" + from unittest.mock import MagicMock + + from open_notebook.graphs.transformation import synthesize_results + + mock_source = MagicMock() + mock_source.add_insight = AsyncMock() + + state = { + "output": "", + "chunk_results": [{"index": 0, "output": "single result"}], + "source": mock_source, + "transformation": MagicMock(title="Test"), + } + + result = await synthesize_results( + cast(TransformationState, state), {"configurable": {}} + ) + assert result == {"output": "single result"} + + @pytest.mark.asyncio + async def test_progress_guard_fallback_on_single_item_groups(self): + """ + When every merge round produces only single-item groups (because + each chunk result exceeds the grouping threshold), the fallback + re-chunks oversized texts and continues merging instead of hanging. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from open_notebook.graphs.transformation import synthesize_results + + # Each text must exceed the single-group budget (~6873 tokens for + # default context_limit=8192) so every group has exactly one item. + big_text = "Large content. " * 3000 # ~9000 tokens each + chunk_results = [ + {"index": i, "output": big_text} for i in range(5) + ] + + fake_response = MagicMock() + fake_response.content = "merged output" + fake_chain = AsyncMock() + fake_chain.ainvoke = AsyncMock(return_value=fake_response) + + state = { + "output": "", + "chunk_results": chunk_results, + "source": None, + "transformation": MagicMock(title="Test", prompt="Summarize"), + "context_limit": 8192, + } + + with ( + patch( + "open_notebook.graphs.transformation.DefaultPrompts", + return_value=MagicMock(transformation_instructions=None), + ), + patch( + "open_notebook.graphs.transformation.provision_langchain_model", + new=AsyncMock(return_value=fake_chain), + ) as mock_provision, + ): + result = await synthesize_results( + cast(TransformationState, state), {"configurable": {}} + ) + + # Should return a single merged output without hanging + assert result["output"] + assert isinstance(result["output"], str) + # Verify the LLM was called at least once (merge happened) + assert mock_provision.await_count >= 1 + + @pytest.mark.asyncio + async def test_stalled_synthesis_rechunk_fallback(self): + """ + When even the smallest adjacent pair exceeds the context window, + the fallback re-chunks oversized texts and continues merging + instead of returning raw partial outputs. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from open_notebook.graphs.transformation import ( + synthesize_results, + ) + + # Create texts that are each larger than the group budget and + # whose combined size also exceeds the context window. + extra_large = "Huge content for testing rechunk fallback path. " * 2000 + chunk_results = [ + {"index": i, "output": extra_large} for i in range(3) + ] + + fake_response = MagicMock() + fake_response.content = "synthesized output" + fake_chain = AsyncMock() + fake_chain.ainvoke = AsyncMock(return_value=fake_response) + + state = { + "output": "", + "chunk_results": chunk_results, + "source": None, + "transformation": MagicMock(title="Test", prompt="Summarize"), + "context_limit": 8192, + } + + with ( + patch( + "open_notebook.graphs.transformation.DefaultPrompts", + return_value=MagicMock(transformation_instructions=None), + ), + patch( + "open_notebook.graphs.transformation.provision_langchain_model", + new=AsyncMock(return_value=fake_chain), + ) as mock_provision, + ): + result = await synthesize_results( + cast(TransformationState, state), {"configurable": {}} + ) + + assert result["output"] + assert isinstance(result["output"], str) + # Verify the LLM was called — re-chunked texts merge in subsequent + # rounds so at least one merge call should happen + assert mock_provision.await_count >= 1 + + @pytest.mark.asyncio + async def test_rechunk_guard_raises_on_infeasible_context(self): + """ + When the context window is too small for even re-chunked texts to be + merged (2 * re_chunk_budget + overhead > context_limit), the guard + raises ConfigurationError instead of silently building oversized pairs. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from open_notebook.exceptions import ConfigurationError + from open_notebook.graphs.transformation import synthesize_results + + # Create oversized texts that will stall in single-item groups. + big_text = "X" * 20000 # ~5000 tokens + chunk_results = [ + {"index": i, "output": big_text} for i in range(3) + ] + + fake_response = MagicMock() + fake_response.content = "merged output" + fake_chain = AsyncMock() + fake_chain.ainvoke = AsyncMock(return_value=fake_response) + + # Use a tiny context_limit so re-chunking is infeasible: + # overhead = token_count("Summarize") + 300 ~= 304 + # available = context_limit - overhead = 600 - 304 = 296 + # re_chunk_budget = max(200, int(296 * 0.4)) = max(200, 118) = 200 + # 2 * 200 + 304 = 704 > 600 -> guard fires -> ConfigurationError + state = { + "output": "", + "chunk_results": chunk_results, + "source": None, + "transformation": MagicMock(title="Test", prompt="Summarize"), + "context_limit": 600, + } + + with ( + patch( + "open_notebook.graphs.transformation.DefaultPrompts", + return_value=MagicMock(transformation_instructions=None), + ), + patch( + "open_notebook.graphs.transformation.provision_langchain_model", + new=AsyncMock(return_value=fake_chain), + ), + ): + with pytest.raises( + ConfigurationError, + match="context window is too small", + ): + await synthesize_results( + cast(TransformationState, state), + {"configurable": {}}, + ) + + @pytest.mark.asyncio + async def test_rechunk_actually_splits_texts(self): + """ + When the smallest adjacent pair exceeds the context window and + re-chunking is feasible, verify that the LLM is called enough + times to indicate re-chunking + multiple merge rounds occurred. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from open_notebook.graphs.transformation import ( + synthesize_results, + ) + + # Extra-large texts that stall in single-item groups and whose + # adjacent pairs exceed context_limit = 4000. + extra_large = "Huge chunk content for re-chunk verification. " * 1000 + chunk_results = [ + {"index": i, "output": extra_large} for i in range(3) + ] + + fake_response = MagicMock() + fake_response.content = "merged output" + fake_chain = AsyncMock() + fake_chain.ainvoke = AsyncMock(return_value=fake_response) + + state = { + "output": "", + "chunk_results": chunk_results, + "source": None, + "transformation": MagicMock(title="Test", prompt="Summarize"), + "context_limit": 4000, + } + + with ( + patch( + "open_notebook.graphs.transformation.DefaultPrompts", + return_value=MagicMock(transformation_instructions=None), + ), + patch( + "open_notebook.graphs.transformation.provision_langchain_model", + new=AsyncMock(return_value=fake_chain), + ) as mock_provision, + ): + result = await synthesize_results( + cast(TransformationState, state), + {"configurable": {}}, + ) + + assert result["output"] + assert isinstance(result["output"], str) + # Multiple LLM calls = first merge round + re-chunked merge rounds + assert mock_provision.await_count >= 2 + + @pytest.mark.asyncio + async def test_process_chunk_success(self): + """ + process_chunk produces a correctly-indexed ChunkResult. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from open_notebook.graphs.transformation import process_chunk + + fake_response = MagicMock() + fake_response.content = "chunk output" + fake_chain = AsyncMock() + fake_chain.ainvoke = AsyncMock(return_value=fake_response) + + state = { + "input_text": "source text", + "source": None, + "transformation": MagicMock(title="T", prompt="Summarize"), + "chunk_index": 2, + "chunk_text": "This is the third chunk content.", + "total_chunks": 5, + "output_budget": 1024, + } + + with ( + patch( + "open_notebook.graphs.transformation.DefaultPrompts", + return_value=MagicMock(transformation_instructions=None), + ), + patch( + "open_notebook.graphs.transformation.provision_langchain_model", + new=AsyncMock(return_value=fake_chain), + ), + ): + result = await process_chunk(cast(ChunkState, state), {"configurable": {}}) + + assert result == { + "chunk_results": [{"index": 2, "output": "chunk output"}] + } + + @pytest.mark.asyncio + async def test_process_chunk_error_classified(self): + """ + When process_chunk's LLM call fails, the error is classified and + stored in the ChunkResult (not re-raised), so sibling chunks can + still complete. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from open_notebook.graphs.transformation import process_chunk + + class FakeConnectionError(Exception): + pass + + fake_chain = AsyncMock() + fake_chain.ainvoke = AsyncMock( + side_effect=FakeConnectionError("connection reset") + ) + + state = { + "input_text": "src", + "source": None, + "transformation": MagicMock(title="T", prompt="Summarize"), + "chunk_index": 0, + "chunk_text": "text", + "total_chunks": 1, + "output_budget": 1024, + } + + with ( + patch( + "open_notebook.graphs.transformation.DefaultPrompts", + return_value=MagicMock(transformation_instructions=None), + ), + patch( + "open_notebook.graphs.transformation.provision_langchain_model", + new=AsyncMock(return_value=fake_chain), + ), + patch( + "open_notebook.graphs.transformation.classify_error", + return_value=(Exception, "Network error occurred"), + ), + ): + result = await process_chunk(cast(ChunkState, state), {"configurable": {}}) + + assert "chunk_results" in result + assert len(result["chunk_results"]) == 1 + err = result["chunk_results"][0] + assert err["index"] == 0 + assert err["output"] == "" + assert "error_message" in err + # ============================================================================ -# TEST SUITE 4: Source Graph - Title Preservation +# TEST SUITE 5: Source Graph - Title Preservation # ============================================================================ @@ -275,7 +668,7 @@ async def test_empty_title_replaced(self, mock_get): # ============================================================================ -# TEST SUITE 5: Source Graph - content_process (content-core 2.x) +# TEST SUITE 6: Source Graph - content_process (content-core 2.x) # ============================================================================ diff --git a/tests/test_token_utils.py b/tests/test_token_utils.py new file mode 100644 index 0000000000..af60d191f9 --- /dev/null +++ b/tests/test_token_utils.py @@ -0,0 +1,390 @@ +""" +Unit tests for the open_notebook.utils.token_utils module. + +Tests parse_context_limit_error with provider-specific error formats, +chunk_text_by_tokens edge cases (oversized words, overlap overhead), +and token counting utilities. +""" + + +import pytest + +from open_notebook.utils.token_utils import ( + SAFETY_BUFFER, + chunk_text_by_tokens, + parse_context_limit_error, + token_cost, + token_count, +) + +# ============================================================================ +# TEST SUITE 1: Token Counting +# ============================================================================ + + +class TestTokenCount: + """Test suite for token_count utility.""" + + def test_empty_string(self): + """Empty string should have 0 tokens.""" + assert token_count("") == 0 + + def test_simple_text(self): + """Simple English text should have reasonable token count.""" + count = token_count("Hello, world!") + assert count > 0 + assert count < 20 + + def test_longer_text(self): + """Longer text should have proportionally more tokens.""" + short = token_count("This is a short sentence.") + long = token_count("This is a much longer sentence with many more words in it for testing purposes.") + assert long > short + + def test_whitespace_only(self): + """Whitespace-only strings may have 0 or minimal tokens.""" + count = token_count(" \n\n \t ") + # Some tokenizers may count whitespace tokens + assert isinstance(count, int) + assert count >= 0 + + +# ============================================================================ +# TEST SUITE 2: Token Cost +# ============================================================================ + + +class TestTokenCost: + """Test suite for token_cost utility.""" + + def test_zero_tokens(self): + """Zero tokens should cost zero.""" + assert token_cost(0) == 0.0 + + def test_default_rate(self): + """Test default $0.150 per million tokens.""" + assert token_cost(1_000_000) == 0.150 + + def test_custom_rate(self): + """Test custom cost per million.""" + assert token_cost(1_000_000, 0.5) == 0.5 + + def test_partial_million(self): + """Test partial million calculation.""" + assert token_cost(500_000) == 0.075 # half of $0.150 + + +# ============================================================================ +# TEST SUITE 3: parse_context_limit_error +# ============================================================================ + + +class TestParseContextLimitError: + """Test suite for parse_context_limit_error with provider-specific formats.""" + + def test_anthropic_format(self): + """Anthropic: 'your request was X tokens ... model's maximum is Y'""" + error = ( + "your request was 5000 tokens, but the model's maximum " + "is 4097 tokens for this model." + ) + result = parse_context_limit_error(error) + assert result is not None + limit, sent = result + assert limit == 4097 # group(2) = limit + assert sent == 5000 # group(1) = sent + + def test_anthropic_variant(self): + """Anthropic variant: 'you have used X tokens ... limit is Y'""" + error = ( + "you have used 15000 tokens of context. The model's maximum " + "context length is 100000 tokens." + ) + result = parse_context_limit_error(error) + assert result is not None + limit, sent = result + assert limit == 100000 + assert sent == 15000 + + def test_openai_format(self): + """OpenAI: 'maximum context length is Y tokens ... you requested X'""" + error = ( + "This model's maximum context length is 8192 tokens. " + "However, you requested 10000 tokens (7500 in the messages, " + "2500 in the completion)." + ) + result = parse_context_limit_error(error) + assert result is not None + limit, sent = result + # OpenAI format: group(1)=limit, group(2)=sent + assert limit == 8192 + assert sent == 10000 + + def test_openai_long_context_format(self): + """OpenAI GPT-4-turbo: large context with large request.""" + error = ( + "This model's maximum context length is 128000 tokens. " + "However, you requested 150000 tokens." + ) + result = parse_context_limit_error(error) + assert result is not None + limit, sent = result + assert limit == 128000 + assert sent == 150000 + + def test_gemini_format(self): + """Google/Gemini: 'The number of tokens (X) exceeded the limit (Y)'""" + error = ( + "The number of tokens (5000) exceeded the limit of 4097. " + "Reduce the input or use a model with a larger context window." + ) + result = parse_context_limit_error(error) + assert result is not None + limit, sent = result + # Gemini format: the larger number (5000) is the sent tokens, + # the smaller number (4097) is the limit. The old max(g1,g2) + # heuristic would pick 5000 as limit — WRONG. + assert limit == 4097 # limit follows "limit of" + assert sent == 5000 # sent in parentheses + + def test_gemini_format_variant(self): + """Gemini variant: tokens exceed limit in reverse order.""" + error = ( + "The number of tokens (15000) exceeded the limit of 8192." + ) + result = parse_context_limit_error(error) + assert result is not None + limit, sent = result + assert limit == 8192 + assert sent == 15000 + + def test_generic_format(self): + """Generic: 'context_length_exceeded: X > Y'""" + error = "context_length_exceeded: 5000 > 4097" + result = parse_context_limit_error(error) + assert result is not None + limit, sent = result + assert limit == 4097 + assert sent == 5000 + + def test_generic_format_variant(self): + """Generic: 'context_length_exceeded: X exceeded Y'""" + error = "context_length_exceeded: 10000 exceeded 8192" + result = parse_context_limit_error(error) + assert result is not None + limit, sent = result + assert limit == 8192 + assert sent == 10000 + + def test_unrelated_error_returns_none(self): + """Non-context errors should return None.""" + error = "Authentication failed: invalid API key" + assert parse_context_limit_error(error) is None + + def test_rate_limit_error_returns_none(self): + """Rate limit errors should return None.""" + error = "Rate limit exceeded. Please wait and retry." + assert parse_context_limit_error(error) is None + + def test_empty_string_returns_none(self): + """Empty string should return None.""" + assert parse_context_limit_error("") is None + + def test_provider_model_not_found(self): + """Model-not-found errors should return None.""" + error = "model not found: gpt-4-fake-model" + assert parse_context_limit_error(error) is None + + +# ============================================================================ +# TEST SUITE 4: chunk_text_by_tokens +# ============================================================================ + + +class TestChunkTextByTokens: + """Test suite for chunk_text_by_tokens utility.""" + + def test_empty_text(self): + """Empty text returns empty list.""" + chunks = chunk_text_by_tokens("", max_chunk_tokens=1000) + # The function checks token_count("") <= max_chunk_tokens first, + # so empty string returns [""] — one empty chunk, not nothing. + assert len(chunks) == 1 + assert chunks[0] == "" + + def test_short_text_no_chunking(self): + """Text within budget returns unchanged.""" + text = "Short text." + chunks = chunk_text_by_tokens(text, max_chunk_tokens=1000) + assert len(chunks) == 1 + assert chunks[0] == text + + def test_text_at_token_limit(self): + """Text exactly at the budget boundary fits in one chunk.""" + text = "Hello world. " * 100 + max_tokens = token_count(text) + chunks = chunk_text_by_tokens(text, max_chunk_tokens=max_tokens) + assert len(chunks) == 1 + + def test_paragraph_splitting(self): + """Text with multiple paragraphs should split at paragraph boundaries.""" + # Build a text where each paragraph is within budget + # but the whole exceeds it. + para = "This is a paragraph with enough content to be meaningful. " * 50 + text = f"{para}\n\n{para}\n\n{para}" + max_tokens = int(token_count(para) * 1.5) # Fits ~1.5 paragraphs + chunks = chunk_text_by_tokens(text, max_chunk_tokens=max_tokens) + assert len(chunks) >= 2 + # Each chunk should be within the token budget (with some slack + # for the overlap overhead if added at the end). + for chunk in chunks: + assert token_count(chunk) <= max_tokens * 1.15, ( + f"Chunk has {token_count(chunk)} tokens, " + f"budget was {max_tokens}" + ) + + def test_oversized_paragraph_sentences(self): + """A single paragraph exceeding budget gets split by sentences.""" + # One paragraph that exceeds the budget + para = "Short sentence. " * 200 + max_tokens = int(token_count(para) * 0.3) + text = f"Intro.\n\n{para}\n\nOutro." + chunks = chunk_text_by_tokens(text, max_chunk_tokens=max_tokens) + assert len(chunks) >= 3 # intro + split para + outro + + def test_oversized_word_fallback(self): + """ + A single whitespace-free token exceeding the budget is split + at character boundaries (paragraph→sentence→word fallback). + """ + # A single "word" with no spaces that exceeds max_chunk_tokens + oversized = "a" * 10000 # ~2500 tokens for o200k + budget = 500 # tokens — the word should exceed this + chunks = chunk_text_by_tokens(oversized, max_chunk_tokens=budget) + assert len(chunks) > 1, ( + f"Oversized word should be split into multiple chunks, " + f"got {len(chunks)}" + ) + # Every chunk should be within budget or close + for chunk in chunks: + assert token_count(chunk) <= budget * 1.15, ( + f"Chunk has {token_count(chunk)} tokens, " + f"budget was {budget}" + ) + + def test_oversized_code_block(self): + """ + Minified code without spaces (e.g. base64, JSON) should + be split at character boundaries. + """ + minified = "abcdefghij" * 2000 # ~3000 chars, ~750 tokens + budget = 200 + chunks = chunk_text_by_tokens(minified, max_chunk_tokens=budget) + assert len(chunks) > 1 + for chunk in chunks: + assert token_count(chunk) <= budget * 1.15 + + def test_overlap_does_not_overshoot_budget(self): + """ + Overlap + continuation marker added between chunks should + not push chunks significantly over the token budget. + """ + para = "This paragraph has enough content to build up tokens. " * 30 + text = f"{para}\n\n{para}\n\n{para}\n\n{para}" + budget = int(token_count(para) * 1.8) + chunks = chunk_text_by_tokens( + text, max_chunk_tokens=budget, overlap_chars=200 + ) + assert len(chunks) >= 2 + for chunk in chunks: + assert token_count(chunk) <= budget * 1.15, ( + f"Chunk with overlap has {token_count(chunk)} tokens, " + f"budget was {budget}" + ) + + def test_no_overlap_when_overlap_chars_zero(self): + """Setting overlap_chars=0 should not add overlap.""" + text = "A chunk. " * 100 + "\n\n" + "Another chunk. " * 100 + budget = int(token_count("A chunk. " * 80)) + chunks_with = chunk_text_by_tokens( + text, max_chunk_tokens=budget, overlap_chars=200 + ) + chunks_without = chunk_text_by_tokens( + text, max_chunk_tokens=budget, overlap_chars=0 + ) + # Without overlap, chunks should be shorter + assert len(chunks_without) >= len(chunks_with) - 1 # overlap can create extra chunk + + def test_large_text_hierarchical(self): + """A long document should be split hierarchically.""" + paragraph = "This is a paragraph used for testing hierarchical splitting. " * 20 + text = "\n\n".join([paragraph] * 20) + budget = int(token_count(paragraph) * 2.5) + chunks = chunk_text_by_tokens(text, max_chunk_tokens=budget) + assert len(chunks) > 1 + # All chunks should contain text (not empty) + assert all(c.strip() for c in chunks) + + +# ============================================================================ +# TEST SUITE 5: Integration — parse + chunk pipeline +# ============================================================================ + + +class TestParseThenChunkIntegration: + """Test that parse_context_limit_error output feeds correctly into + chunk_text_by_tokens as it does in _handle_llm_error.""" + + def test_anthropic_parse_then_chunk(self): + """Simulate the _handle_llm_error pipeline with an Anthropic error.""" + error_msg = ( + "your request was 50000 tokens, but the model's maximum " + "is 100000 tokens for this model." + ) + result = parse_context_limit_error(error_msg) + assert result is not None + context_limit, _tokens_sent = result + + # Simulate the chunk size calculation from _handle_llm_error + system_prompt = "You are an AI assistant." + system_tokens = token_count(system_prompt) + output_budget = min(8192, int(context_limit * 0.10)) + max_chunk_tokens = context_limit - system_tokens - output_budget + max_chunk_tokens = int(max_chunk_tokens * SAFETY_BUFFER) + max_chunk_tokens = max(max_chunk_tokens, 512) + + # Create content that exceeds the chunk budget + # "Content paragraph. " * 50000 ≈ 150K tokens > 78K chunk budget + content = "Content paragraph. " * 50000 + chunks = chunk_text_by_tokens(content, max_chunk_tokens=max_chunk_tokens) + assert len(chunks) > 1 + for chunk in chunks: + assert token_count(chunk) <= max_chunk_tokens * 1.15 + + def test_gemini_parse_then_chunk(self): + """Simulate pipeline with a Gemini error — validates the fixed + context_limit ordering doesn't produce tiny/invalid chunks.""" + error_msg = ( + "The number of tokens (5000) exceeded the limit of 4097." + ) + result = parse_context_limit_error(error_msg) + assert result is not None + context_limit, _ = result + # With the old max-heuristic, context_limit would be 5000 (wrong). + # With the fix, it's 4097 (correct). + assert context_limit == 4097 + + def test_openai_parse_then_chunk(self): + """Simulate pipeline with an OpenAI error.""" + error_msg = ( + "This model's maximum context length is 8192 tokens. " + "However, you requested 10000 tokens." + ) + result = parse_context_limit_error(error_msg) + assert result is not None + context_limit, _ = result + assert context_limit == 8192 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])