diff --git a/blueprints/conversation/agents.py b/blueprints/conversation/agents.py index e7927ef..b0ac9c8 100644 --- a/blueprints/conversation/agents.py +++ b/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 diff --git a/blueprints/conversation/artifacts_util.py b/blueprints/conversation/artifacts_util.py index 48f4e58..d878161 100644 --- a/blueprints/conversation/artifacts_util.py +++ b/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 + ] + } diff --git a/blueprints/rag/logic.py b/blueprints/rag/logic.py index 227b7f8..94ee7ae 100644 --- a/blueprints/rag/logic.py +++ b/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() diff --git a/tests/unit/test_artifacts.py b/tests/unit/test_artifacts.py index 9ac1b78..5a740f6 100644 --- a/tests/unit/test_artifacts.py +++ b/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( diff --git a/tests/unit/test_context_trim.py b/tests/unit/test_context_trim.py index 762c2ae..5cfdf23 100644 --- a/tests/unit/test_context_trim.py +++ b/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"} + ] diff --git a/tests/unit/test_retrieval.py b/tests/unit/test_retrieval.py index d8f1617..762ca81 100644 --- a/tests/unit/test_retrieval.py +++ b/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( diff --git a/tests/unit/test_tool_artifacts.py b/tests/unit/test_tool_artifacts.py index 3ed6b0e..2bf5b44 100644 --- a/tests/unit/test_tool_artifacts.py +++ b/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):