Merge pull request 'Give the model full retrieved chunks instead of 400-character snippets' (#83) from perf/full-search-chunks into main

blueprints/conversation/agents.py

@@ -23,6 +23,7 @@ from blueprints.conversation.artifacts_util import (
     note_title,
     obsidian_source_ref,
     search_documents_artifact,
+    search_results_model_context,
     truncate,
     unique,
 )
@@ -162,7 +163,10 @@ async def simba_search(
         )
     except ValueError as exc:
         return str(exc), None
-    return serialized, search_documents_artifact(query, "paperless", docs)
+    artifact = search_documents_artifact(query, "paperless", docs)
+    if artifact is not None:
+        artifact["model_context"] = search_results_model_context(docs, "paperless")
+    return serialized, artifact
 
 
 MEALIE_NOT_CONFIGURED = "Mealie integration is not configured. Please set MEALIE_BASE_URL and MEALIE_API_TOKEN environment variables."
@@ -488,7 +492,10 @@ async def obsidian_search_notes(
         serialized, docs = await query_vector_store(
             query=query, source="obsidian", **kwargs
         )
-        return serialized, search_documents_artifact(query, "obsidian", docs)
+        artifact = search_documents_artifact(query, "obsidian", docs)
+        if artifact is not None:
+            artifact["model_context"] = search_results_model_context(docs, "obsidian")
+        return serialized, artifact
 
     except Exception as e:
         return f"Error searching Obsidian notes: {str(e)}", None

blueprints/conversation/artifacts_util.py

@@ -143,43 +143,64 @@ def vault_relative_path(filepath: str) -> str:
         return filepath
 
 
+def _result_identity(doc, source: str) -> dict:
+    """Identifying fields shared by the UI snippet and the model's full-text view."""
+    from blueprints.rag.retrieval import citation_url
+
+    meta = doc.metadata or {}
+    result = {}
+    if url := citation_url(doc):
+        result["url"] = url
+        result["source_id"] = str(doc.id)
+    if headings := [meta[k] for k in ("h1", "h2", "h3") if meta.get(k)]:
+        result["headings"] = headings
+    if source == "obsidian":
+        filepath = meta.get("filepath") or ""
+        path = vault_relative_path(filepath) if filepath else ""
+        result["title"] = note_title(path, meta) if path else "Untitled note"
+        if path:
+            result["path"] = path
+        created = meta.get("created_at")
+        if isinstance(created, str) and created:
+            result["date"] = created[:10]
+    else:
+        filename = meta.get("filename") or ""
+        result["title"] = Path(filename).stem or filename or "Document"
+        if filename:
+            result["filename"] = filename
+        if meta.get("document_type"):
+            result["document_type"] = str(meta["document_type"])
+        created = epoch_to_date(meta.get("created_date"))
+        if created:
+            result["date"] = created
+    return result
+
+
 def search_documents_artifact(query: str, source: str, docs) -> dict | None:
     """Build a `search.documents` artifact from vector-store Documents."""
     if not docs:
         return None
-    results = []
-    from blueprints.rag.retrieval import citation_url
-
-    for doc in docs:
-        meta = doc.metadata or {}
-        result = {"snippet": snippet(doc.page_content)}
-        if url := citation_url(doc):
-            result["url"] = url
-            result["source_id"] = str(doc.id)
-        if headings := [meta[k] for k in ("h1", "h2", "h3") if meta.get(k)]:
-            result["headings"] = headings
-        if source == "obsidian":
-            filepath = meta.get("filepath") or ""
-            path = vault_relative_path(filepath) if filepath else ""
-            result["title"] = note_title(path, meta) if path else "Untitled note"
-            if path:
-                result["path"] = path
-            created = meta.get("created_at")
-            if isinstance(created, str) and created:
-                result["date"] = created[:10]
-        else:
-            filename = meta.get("filename") or ""
-            result["title"] = Path(filename).stem or filename or "Document"
-            if filename:
-                result["filename"] = filename
-            if meta.get("document_type"):
-                result["document_type"] = str(meta["document_type"])
-            created = epoch_to_date(meta.get("created_date"))
-            if created:
-                result["date"] = created
-        results.append(result)
+    results = [
+        {"snippet": snippet(doc.page_content), **_result_identity(doc, source)}
+        for doc in docs
+    ]
     return make_artifact(
         "search.documents",
         f"Search: {query}",
         {"query": query, "source": source, "results": results},
     )
+
+
+def search_results_model_context(docs, source: str) -> dict:
+    """Full chunk text for the model, keyed like `search_documents_artifact`'s results.
+
+    Set as a tool's `artifact["model_context"]` so that under compaction (see
+    compact_tool_result in blueprints/conversation/context.py) the model still
+    sees the full chunk text it retrieved, not the UI's 400-char snippets.
+    """
+    return {
+        "results": [
+            {**_result_identity(doc, source), "content": doc.page_content}
+            for doc in docs
+        ]
+    }

blueprints/rag/logic.py

@@ -16,7 +16,7 @@ from langchain_text_splitters import (
 from sqlalchemy import create_engine, text
 
 from .fetchers import PaperlessNGXService
-from .retrieval import build_filter, citation_url, document_dates
+from .retrieval import build_filter, citation_url, document_date, document_dates
 from utils.obsidian_service import ObsidianService
 
 # Load environment variables
@@ -462,7 +462,7 @@ async def query_vector_store(
         ]
     serialized = "\n\n".join(
         (
-            f"Source: {doc.metadata}\n"
+            f"{_chunk_header(doc)}\n"
             + (f"Citation: {citation_url(doc)}\n" if citation_url(doc) else "")
             + f"Content: {doc.page_content}"
         )
@@ -471,6 +471,28 @@ async def query_vector_store(
     return serialized, retrieved_docs
 
 
+def _chunk_title(doc: Document) -> str:
+    from blueprints.conversation.artifacts_util import note_title, vault_relative_path
+
+    meta = doc.metadata or {}
+    if meta.get("source") == "obsidian":
+        path = vault_relative_path(meta.get("filepath") or "")
+        return note_title(path, meta) if path else "Untitled note"
+    filename = meta.get("filename") or ""
+    return Path(filename).stem or filename or "Document"
+
+
+def _chunk_header(doc: Document) -> str:
+    """Compact per-chunk header (title, date, source) replacing the raw metadata dict."""
+    meta = doc.metadata or {}
+    parts = [f"Title: {_chunk_title(doc)}"]
+    if date := document_date(meta):
+        parts.append(f"Date: {date}")
+    if source := meta.get("source"):
+        parts.append(f"Source: {source}")
+    return " | ".join(parts)
+
+
 def delete_all_documents():
     """Delete all documents from the vector store collection."""
     collection_id = _get_collection_id()

tests/unit/test_artifacts.py

@@ -85,6 +85,29 @@ class TestCreateArtifact:
         assert result is None
         assert await Artifact.all().count() == 0
 
+    async def test_drops_model_context_before_persisting(self, owner):
+        """model_context carries the full chunk text for the model's own view;
+        it must never reach the DB (or, from there, the UI)."""
+        user, conversation, message = owner
+        artifact = await create_artifact_from_output(
+            output=_output(
+                {
+                    "kind": "search.documents",
+                    "title": "Search: vet",
+                    "data": {"results": [{"snippet": "short"}]},
+                    "model_context": {"results": [{"content": "x" * 5000}]},
+                }
+            ),
+            tool="simba_search",
+            user=user,
+            conversation=conversation,
+            message=message,
+        )
+        assert artifact is not None
+        saved = await Artifact.get(id=artifact.id)
+        assert saved.data == {"results": [{"snippet": "short"}]}
+        assert "model_context" not in saved.data
+
     async def test_plain_string_output(self, owner):
         user, conversation, message = owner
         result = await create_artifact_from_output(

tests/unit/test_context_trim.py

@@ -464,6 +464,64 @@ def test_focused_note_heading_index_uses_absolute_content_lines():
     assert msg.artifact["data"]["markdown"] == "full original"
 
 
+def _search_artifact(chunk_text: str) -> dict:
+    return {
+        "kind": "search.documents",
+        "title": "Search: vet",
+        "data": {
+            "query": "vet",
+            "source": "paperless",
+            "results": [
+                {
+                    "snippet": chunk_text[:50],
+                    "title": "vet_invoice",
+                    "source_id": "chunk-1",
+                    "url": "/sources/chunk-1",
+                }
+            ],
+        },
+        "model_context": {
+            "results": [
+                {
+                    "content": chunk_text,
+                    "title": "vet_invoice",
+                    "source_id": "chunk-1",
+                    "url": "/sources/chunk-1",
+                }
+            ]
+        },
+    }
+
+
+def test_search_result_compaction_keeps_full_chunk_text_under_a_generous_limit():
+    import json
+
+    chunk_text = "Rabies vaccine administered. " + "detail " * 200
+    msg = ToolMessage(
+        content="display", artifact=_search_artifact(chunk_text), tool_call_id="a"
+    )
+    projected = context.compact_tool_result(msg, 100_000)
+    parsed = json.loads(projected.content)
+    (result,) = parsed["data"]["results"]
+    assert result["content"] == chunk_text
+    assert "snippet" not in result
+
+
+def test_search_result_compaction_shortens_content_but_keeps_source_ids_under_pressure():
+    import json
+
+    chunk_text = "Rabies vaccine administered. " + "detail " * 200
+    msg = ToolMessage(
+        content="display", artifact=_search_artifact(chunk_text), tool_call_id="a"
+    )
+    projected = context.compact_tool_result(msg, 700)
+    parsed = json.loads(projected.content)
+    (result,) = parsed["data"]["results"]
+    assert len(result["content"]) < len(chunk_text)
+    assert result["source_id"] == "chunk-1"
+    assert result["url"] == "/sources/chunk-1"
+
+
 def _conversation(turns, size=1000):
     messages = []
     for i in range(turns):
@@ -542,3 +600,18 @@ def test_trim_drops_to_low_watermark_and_prefix_stays_stable():
     # ...so the retained prefix is unchanged for several turns in between.
     assert all(b - a >= 3 for a, b in zip(cuts, cuts[1:]))
     assert len(cuts) <= 5
+
+
+def test_search_result_compaction_keeps_source_ids_when_no_rows_fit():
+    import json
+
+    chunk_text = "Rabies vaccine administered. " + "detail " * 200
+    msg = ToolMessage(
+        content="display", artifact=_search_artifact(chunk_text), tool_call_id="a"
+    )
+    projected = json.loads(context.compact_tool_result(msg, 300).content)
+    data = projected["data"]
+    assert data["results"] == []
+    assert data["results_compaction"]["omitted_source_refs"] == [
+        {"source_id": "chunk-1", "url": "/sources/chunk-1"}
+    ]

tests/unit/test_retrieval.py

@@ -94,6 +94,32 @@ def test_artifact_citations_identify_exact_retrieved_chunk(monkeypatch):
     assert result["headings"] == ["Budget", "Savings"]
 
 
+async def test_serialized_chunk_header_is_compact_not_raw_metadata(monkeypatch):
+    identifier = str(uuid4())
+    doc = Document(
+        id=identifier,
+        page_content="Rabies vaccine administered",
+        metadata={
+            "source": "paperless",
+            "filename": "vet_invoice.pdf",
+            "created_date": datetime(2026, 9, 1, tzinfo=timezone.utc).timestamp(),
+            "document_type": "Invoice",
+        },
+    )
+    store = SimpleNamespace(
+        amax_marginal_relevance_search=AsyncMock(return_value=[doc])
+    )
+    monkeypatch.setattr(logic, "_get_vector_store", lambda: store)
+    text, _ = await logic.query_vector_store("vet", source="paperless")
+    assert "{'source'" not in text
+    assert str(doc.metadata) not in text
+    assert "Title: vet_invoice" in text
+    assert "Date: 2026-09-01" in text
+    assert "Source: paperless" in text
+    assert f"Citation: /sources/{identifier}" in text
+    assert "Content: Rabies vaccine administered" in text
+
+
 def test_document_date_and_unknown_citations():
     assert (
         document_date(

tests/unit/test_tool_artifacts.py

@@ -521,6 +521,13 @@ async def test_simba_search(monkeypatch):
     assert result["snippet"].startswith("Rabies vaccine administered. ")
     assert len(result["snippet"]) <= SNIPPET_MAX_CHARS
 
+    (model_result,) = art["model_context"]["results"]
+    assert model_result["content"] == PAPERLESS_DOCS[0].page_content
+    assert len(model_result["content"]) > SNIPPET_MAX_CHARS
+    assert model_result["title"] == result["title"]
+    assert model_result["filename"] == result["filename"]
+    assert "snippet" not in model_result
+
 
 async def test_simba_search_empty(monkeypatch):
     monkeypatch.setattr(agents, "query_vector_store", AsyncMock(return_value=("", [])))
@@ -560,6 +567,14 @@ async def test_obsidian_search_notes(obsidian, monkeypatch):
             "date": "2026-01-02",
         }
     ]
+    assert art["model_context"]["results"] == [
+        {
+            "content": "Simba's vet is Dr Smith.",
+            "title": "simba",
+            "path": "pets/simba.md",
+            "date": "2026-01-02",
+        }
+    ]
 
 
 async def test_obsidian_search_notes_error(obsidian, monkeypatch):