diff --git a/open_notebook/utils/context_builder.py b/open_notebook/utils/context_builder.py index 7d3717ca63..69f83436b9 100644 --- a/open_notebook/utils/context_builder.py +++ b/open_notebook/utils/context_builder.py @@ -140,11 +140,16 @@ async def build_notebook_context( async def build_source_context( source_id: str, max_tokens: Optional[int] = None ) -> Dict[str, Any]: - """Assemble a single source's short context plus its insights. + """Assemble a single source's full content plus its insights. - Used by the source-chat graph. If `max_tokens` is given, insights are - dropped (last-fetched first) until the total fits — the source itself is - always kept. + Used by the source-chat graph. Uses the "long" context so the source's + full text is always included — a source with no insights/transformations + would otherwise leave the LLM with only the title. If `max_tokens` is + given, insights are dropped (last-fetched first) until the total fits. + The source itself is never dropped — a large full_text can alone exceed + the budget, and dropping it would leave an empty context; the caller + (source-chat's prompt formatter) already caps full_text at a fixed size + before it reaches the LLM. Returns a dict with "sources", "notes" (always empty), "insights", "total_tokens", "total_items" and per-type counts in "metadata". @@ -161,7 +166,13 @@ async def build_source_context( source = None if source: - source_context = await source.get_context(context_size="short") + # Pass insights=[] so get_context() doesn't fetch and embed them + # itself — this function fetches and represents insights + # separately below. Fetching them here too would hit the DB + # twice and double-count their tokens toward the budget. + source_context = await source.get_context( + context_size="long", insights=[] + ) sources.append(source_context) item_tokens.append(token_count(str(source_context))) @@ -177,16 +188,23 @@ async def build_source_context( else: logger.warning(f"Source {source_id} not found") - # Truncate to the token budget: drop insights from the end (the - # source, added first, is dropped only if it alone exceeds the budget). + # Truncate to the token budget: drop insights from the end. The + # source is never dropped, even if it alone exceeds the budget — + # an empty context is worse than a source whose full_text gets + # capped later by the prompt formatter (see docstring). If the + # source alone already exceeds the budget, trimming insights can't + # bring the total under it anyway, so keep them all instead of + # discarding useful context for nothing. total_tokens = sum(item_tokens) - if max_tokens: - while total_tokens > max_tokens and item_tokens: + source_tokens = item_tokens[0] if sources else 0 + if max_tokens is not None: + while ( + total_tokens > max_tokens + and insights + and source_tokens <= max_tokens + ): total_tokens -= item_tokens.pop() - if insights: - insights.pop() - else: - sources.pop() + insights.pop() total_items = len(sources) + len(insights) logger.info(f"Built context with {total_items} items, {total_tokens} tokens") diff --git a/tests/test_utils.py b/tests/test_utils.py index 5917fe2063..7357792b83 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -273,7 +273,7 @@ class TestBuildSourceContext: @pytest.mark.asyncio async def test_source_and_insights_shape(self): - """The response carries the source's short context and its insights.""" + """The response carries the source's full-text context and its insights.""" source = _mock_source([_insight("source_insight:1")]) with patch( @@ -283,7 +283,13 @@ async def test_source_and_insights_shape(self): result = await build_source_context("123") mock_get.assert_awaited_once_with("source:123") # bare id gets prefixed - source.get_context.assert_awaited_once_with(context_size="short") + # Long context is used so full_text is always included — a source + # with no insights would otherwise leave the LLM with only the + # title. insights=[] avoids get_context() double-fetching/counting + # insights that this function already fetches separately below. + source.get_context.assert_awaited_once_with( + context_size="long", insights=[] + ) assert result["sources"] == [ {"id": "source:123", "title": "T", "full_text": "body"} ] @@ -323,6 +329,37 @@ async def test_truncates_insights_to_token_budget(self): assert [i["id"] for i in result["insights"]] == ["source_insight:1"] assert result["total_tokens"] <= 600 + @pytest.mark.asyncio + async def test_keeps_source_when_it_alone_exceeds_budget(self): + """A source whose full_text alone exceeds max_tokens is still + returned (not dropped) — an empty context is worse than one over + budget, and the prompt formatter caps full_text length downstream. + Its insights are kept too: trimming them can't bring the total + under budget anyway (the source alone already exceeds it), so + dropping them would lose context for no benefit.""" + source = SimpleNamespace(id="source:123") + big_text = "word " * 20000 # far bigger than the token budget below + source.get_context = AsyncMock( + return_value={"id": "source:123", "title": "T", "full_text": big_text} + ) + source.get_insights = AsyncMock( + return_value=[_insight("source_insight:1"), _insight("source_insight:2")] + ) + + with patch( + "open_notebook.utils.context_builder.Source.get", + new=AsyncMock(return_value=source), + ): + result = await build_source_context("source:123", max_tokens=600) + + assert len(result["sources"]) == 1 + assert result["sources"][0]["full_text"] == big_text + assert [i["id"] for i in result["insights"]] == [ + "source_insight:1", + "source_insight:2", + ] + assert result["total_tokens"] > 600 + @pytest.mark.asyncio async def test_missing_source_yields_empty_context(self): """A missing source produces an empty context, not an error."""