Studio: format search_knowledge_base hits as fenced <chunk> blocks with score/page/tokens

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 "
+            '<chunk id="N" source="..." page="..." score="...">...</chunk> '
+            "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("&", "&amp;")
+        .replace('"', "&quot;")
+        .replace("<", "&lt;")
+        .replace(">", "&gt;")
+    )
+
+
+def _format_hits_for_llm(hits: list[dict]) -> str:
+    """Render hits as fenced <chunk> 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"<chunk {' '.join(attrs)}>\n{text}\n</chunk>")
+    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)

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 '<chunk id="1" source="alpha.pdf" page="3" score="0.780"' in result
+    assert 'chunk_index="12"' in result
+    assert 'tokens="42"' in result
+    assert "first body\n</chunk>" in result
+    assert '<chunk id="2" source="beta.md" score="0.610">' in result
+    # Blocks separated by a blank line so the model can scan the list.
+    assert "</chunk>\n\n<chunk" in result
 
 
 def test_format_hits_handles_unknown_source():
@@ -121,7 +136,16 @@ def test_format_hits_handles_unknown_source():
 
     hits = [{"filename": None, "page_number": None, "text": "orphan"}]
     result = _format_hits_for_llm(hits)
-    assert "[1] unknown source: orphan" in result
+    assert '<chunk id="1" source="unknown">' in result
+    assert "\norphan\n</chunk>" 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&quot;name&lt;.pdf"' in result
 
 
 def test_tool_spec_shape_is_openai_compatible():