Give the model full retrieved chunks instead of 400-character snippets #83

Merged
ryan merged 1 commits from perf/full-search-chunks into main 2026-10-13 03:41:12 -04:00
7 changed files with 222 additions and 35 deletions
+9 -2
View File
@@ -23,6 +23,7 @@ from blueprints.conversation.artifacts_util import (
note_title, note_title,
obsidian_source_ref, obsidian_source_ref,
search_documents_artifact, search_documents_artifact,
search_results_model_context,
truncate, truncate,
unique, unique,
) )
@@ -162,7 +163,10 @@ async def simba_search(
) )
except ValueError as exc: except ValueError as exc:
return str(exc), None 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." 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( serialized, docs = await query_vector_store(
query=query, source="obsidian", **kwargs 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: except Exception as e:
return f"Error searching Obsidian notes: {str(e)}", None return f"Error searching Obsidian notes: {str(e)}", None
+52 -31
View File
@@ -143,43 +143,64 @@ def vault_relative_path(filepath: str) -> str:
return filepath 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: def search_documents_artifact(query: str, source: str, docs) -> dict | None:
"""Build a `search.documents` artifact from vector-store Documents.""" """Build a `search.documents` artifact from vector-store Documents."""
if not docs: if not docs:
return None return None
results = [] results = [
from blueprints.rag.retrieval import citation_url {"snippet": snippet(doc.page_content), **_result_identity(doc, source)}
for doc in docs
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)
return make_artifact( return make_artifact(
"search.documents", "search.documents",
f"Search: {query}", f"Search: {query}",
{"query": query, "source": source, "results": results}, {"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
]
}
+24 -2
View File
@@ -16,7 +16,7 @@ from langchain_text_splitters import (
from sqlalchemy import create_engine, text from sqlalchemy import create_engine, text
from .fetchers import PaperlessNGXService 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 from utils.obsidian_service import ObsidianService
# Load environment variables # Load environment variables
@@ -462,7 +462,7 @@ async def query_vector_store(
] ]
serialized = "\n\n".join( 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"Citation: {citation_url(doc)}\n" if citation_url(doc) else "")
+ f"Content: {doc.page_content}" + f"Content: {doc.page_content}"
) )
@@ -471,6 +471,28 @@ async def query_vector_store(
return serialized, retrieved_docs 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(): def delete_all_documents():
"""Delete all documents from the vector store collection.""" """Delete all documents from the vector store collection."""
collection_id = _get_collection_id() collection_id = _get_collection_id()
+23
View File
@@ -85,6 +85,29 @@ class TestCreateArtifact:
assert result is None assert result is None
assert await Artifact.all().count() == 0 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): async def test_plain_string_output(self, owner):
user, conversation, message = owner user, conversation, message = owner
result = await create_artifact_from_output( result = await create_artifact_from_output(
+73
View File
@@ -464,6 +464,64 @@ def test_focused_note_heading_index_uses_absolute_content_lines():
assert msg.artifact["data"]["markdown"] == "full original" 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): def _conversation(turns, size=1000):
messages = [] messages = []
for i in range(turns): 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. # ...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 all(b - a >= 3 for a, b in zip(cuts, cuts[1:]))
assert len(cuts) <= 5 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"}
]
+26
View File
@@ -94,6 +94,32 @@ def test_artifact_citations_identify_exact_retrieved_chunk(monkeypatch):
assert result["headings"] == ["Budget", "Savings"] 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(): def test_document_date_and_unknown_citations():
assert ( assert (
document_date( document_date(
+15
View File
@@ -521,6 +521,13 @@ async def test_simba_search(monkeypatch):
assert result["snippet"].startswith("Rabies vaccine administered. ") assert result["snippet"].startswith("Rabies vaccine administered. ")
assert len(result["snippet"]) <= SNIPPET_MAX_CHARS 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): async def test_simba_search_empty(monkeypatch):
monkeypatch.setattr(agents, "query_vector_store", AsyncMock(return_value=("", []))) 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", "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): async def test_obsidian_search_notes_error(obsidian, monkeypatch):