Give the model full retrieved chunks instead of 400-character snippets #83
@@ -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
|
||||||
|
|||||||
@@ -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
@@ -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()
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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"}
|
||||||
|
]
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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):
|
||||||
|
|||||||
Reference in New Issue
Block a user