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():