Give the model full retrieved chunks instead of 400-character snippets
Compaction replaces a tool's text with its artifact, and the search artifact only carries UI snippets, so the full chunk the retriever found never reached the model. Search tools now attach the full text as model_context, which compaction prefers and shortens under budget pressure while keeping source ids; the persisted artifact and UI are unchanged. The per-chunk header drops the raw metadata dict (about a fifth of search output tokens) for title, date and source. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
]
|
||||
}
|
||||
|
||||
+24
-2
@@ -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()
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"}
|
||||
]
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user