diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py
index 9db3fa7fd2..6a5cb18486 100644
--- a/studio/backend/core/rag/tool.py
+++ b/studio/backend/core/rag/tool.py
@@ -21,11 +21,11 @@ SEARCH_KNOWLEDGE_BASE_TOOL = {
"function": {
"name": "search_knowledge_base",
"description": (
- "Search the user's attached documents for information relevant to "
- "the user's question. Call this when the user references content "
- "from their docs, asks fact-heavy questions, or needs grounded "
- "citations. Returns numbered chunks with source filenames; cite "
- "them in your reply as [1], [2], etc."
+ "Search the user's attached documents. Call this when the user "
+ "references content from their docs, asks fact-heavy questions, "
+ "or needs grounded citations. Returns chunks wrapped in "
+ '... '
+ "tags; cite them in your reply as [1], [2], etc."
),
"parameters": {
"type": "object",
@@ -53,22 +53,51 @@ SEARCH_KNOWLEDGE_BASE_TOOL = {
}
-def _format_hits_for_llm(hits: list[Any]) -> str:
- """Render hits as numbered Markdown citations; empty results return a message, not ''."""
+def _xml_attr(value: Any) -> str:
+ return (
+ str(value)
+ .replace("&", "&")
+ .replace('"', """)
+ .replace("<", "<")
+ .replace(">", ">")
+ )
+
+
+def _format_hits_for_llm(hits: list[dict]) -> str:
+ """Render hits as fenced blocks with metadata."""
if not hits:
return (
"No matching chunks were found in the attached documents. "
"Either nothing in this scope is relevant, or no documents "
"have been ingested yet."
)
- lines: list[str] = []
+ blocks: list[str] = []
for index, hit in enumerate(hits, start = 1):
- name = hit.get("filename") or "unknown source"
+ attrs = [
+ f'id="{index}"',
+ f'source="{_xml_attr(hit.get("filename") or "unknown")}"',
+ ]
page = hit.get("page_number")
- suffix = f" (page {page})" if page is not None else ""
+ if page is not None:
+ attrs.append(f'page="{page}"')
+ score = hit.get("score")
+ if score is not None:
+ attrs.append(f'score="{float(score):.3f}"')
+ dense = hit.get("dense_score")
+ if dense is not None and dense != score:
+ attrs.append(f'dense_score="{float(dense):.3f}"')
+ chunk_index = hit.get("chunk_index")
+ if chunk_index is not None:
+ attrs.append(f'chunk_index="{chunk_index}"')
+ tokens = hit.get("token_count")
+ if tokens:
+ attrs.append(f'tokens="{tokens}"')
+ kind = hit.get("kind")
+ if kind and kind != "text":
+ attrs.append(f'kind="{_xml_attr(kind)}"')
text = (hit.get("text") or "").strip()
- lines.append(f"[{index}] {name}{suffix}: {text}")
- return "\n\n".join(lines)
+ blocks.append(f"\n{text}\n")
+ return "\n\n".join(blocks)
def search_knowledge_base(
@@ -166,7 +195,7 @@ def search_knowledge_base(
rows = conn.execute(
f"""
SELECT c.id AS chunk_id, c.text, c.page_number,
- c.kind, d.filename
+ c.token_count, c.kind, d.filename
FROM rag_chunks c
JOIN rag_documents d ON d.id = c.document_id
WHERE c.id IN ({placeholders})
@@ -198,9 +227,19 @@ def search_knowledge_base(
hits = hits[:k]
# Skip image-kind hits; the paired caption surfaces separately.
- formatted = [
- lookup[hit.chunk_id]
- for hit in hits
- if hit.chunk_id in lookup and lookup[hit.chunk_id].get("kind") != "image"
- ]
+ # Merge Hit-side metadata (score, dense_score, chunk_index) into the
+ # sqlite-side row so the formatter sees one flat dict per chunk.
+ formatted: list[dict] = []
+ for hit in hits:
+ row = lookup.get(hit.chunk_id)
+ if row is None or row.get("kind") == "image":
+ continue
+ formatted.append(
+ {
+ **row,
+ "score": hit.score,
+ "dense_score": hit.dense_score,
+ "chunk_index": hit.chunk_index,
+ }
+ )
return _format_hits_for_llm(formatted)
diff --git a/tests/python/test_rag_tool_handler.py b/tests/python/test_rag_tool_handler.py
index c28b49cbbe..e7d45b5423 100644
--- a/tests/python/test_rag_tool_handler.py
+++ b/tests/python/test_rag_tool_handler.py
@@ -102,18 +102,33 @@ def test_empty_results_message_is_user_facing():
assert "No matching chunks" in result
-def test_format_hits_produces_numbered_citations():
+def test_format_hits_produces_fenced_chunks():
from core.rag.tool import _format_hits_for_llm
hits = [
- {"filename": "alpha.pdf", "page_number": 3, "text": "first body"},
- {"filename": "beta.md", "page_number": None, "text": "second body"},
+ {
+ "filename": "alpha.pdf",
+ "page_number": 3,
+ "text": "first body",
+ "score": 0.78,
+ "chunk_index": 12,
+ "token_count": 42,
+ },
+ {
+ "filename": "beta.md",
+ "page_number": None,
+ "text": "second body",
+ "score": 0.61,
+ },
]
result = _format_hits_for_llm(hits)
- assert "[1] alpha.pdf (page 3): first body" in result
- assert "[2] beta.md: second body" in result
- # Each hit on its own paragraph so the LLM can cite cleanly.
- assert "\n\n" in result
+ assert '" in result
+ assert '' in result
+ # Blocks separated by a blank line so the model can scan the list.
+ assert "\n\n' in result
+ assert "\norphan\n" in result
+
+
+def test_format_hits_escapes_xml_in_source():
+ from core.rag.tool import _format_hits_for_llm
+
+ hits = [{"filename": 'weird"name<.pdf', "text": "body"}]
+ result = _format_hits_for_llm(hits)
+ assert 'source="weird"name<.pdf"' in result
def test_tool_spec_shape_is_openai_compatible():