Fix single-source chat missing full document content

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")

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