diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index 913bddc5b3..836750df78 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -9,12 +9,19 @@ so the model never sees KB UUIDs. from __future__ import annotations +from contextvars import ContextVar from typing import Any, Literal from loggers import get_logger logger = get_logger(__name__) +# Per-request chunk-id counter. Task-local (FastAPI runs each request in +# its own asyncio task → its own Context). Lets the model cite chunks +# unambiguously when multiple search_knowledge_base calls run in the +# same chat turn: call 1 returns ids 1..N, call 2 returns N+1..N+M, etc. +_chunk_id_counter: ContextVar[int] = ContextVar("rag_chunk_id_counter", default = 0) + SEARCH_KNOWLEDGE_BASE_TOOL = { "type": "function", @@ -27,7 +34,9 @@ SEARCH_KNOWLEDGE_BASE_TOOL = { "knowledge until you have called this tool with a focused query " "derived from the user's latest message. Returns chunks wrapped in " '... ' - "tags; cite them in your reply as [1], [2], etc." + "tags. CITE each chunk you use with its LITERAL id attribute, " + 'e.g. `` is cited as `[7]`. IDs are unique across ' + "all calls in this turn — never renumber, never reuse." ), "parameters": { "type": "object", @@ -65,8 +74,12 @@ def _xml_attr(value: Any) -> str: ) -def _format_hits_for_llm(hits: list[dict]) -> str: - """Render hits as fenced blocks with metadata.""" +def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str: + """Render hits as fenced blocks with metadata. + + ``start_id`` offsets the citation id so multiple calls in the same + request produce globally unique ids (call 1: 1..N, call 2: N+1..N+M). + """ if not hits: return ( "No matching chunks were found in the attached documents. " @@ -74,7 +87,7 @@ def _format_hits_for_llm(hits: list[dict]) -> str: "have been ingested yet." ) blocks: list[str] = [] - for index, hit in enumerate(hits, start = 1): + for index, hit in enumerate(hits, start = start_id + 1): attrs = [ f'id="{index}"', f'source="{_xml_attr(hit.get("filename") or "unknown")}"', @@ -244,4 +257,7 @@ def search_knowledge_base( "chunk_index": hit.chunk_index, } ) - return _format_hits_for_llm(formatted) + start_id = _chunk_id_counter.get() + rendered = _format_hits_for_llm(formatted, start_id = start_id) + _chunk_id_counter.set(start_id + len(formatted)) + return rendered diff --git a/studio/frontend/src/components/assistant-ui/sources.tsx b/studio/frontend/src/components/assistant-ui/sources.tsx index 3a55c3fa78..3593474c2c 100644 --- a/studio/frontend/src/components/assistant-ui/sources.tsx +++ b/studio/frontend/src/components/assistant-ui/sources.tsx @@ -10,6 +10,7 @@ import { type ComponentProps, type FC, } from "react"; +import { FileTextIcon } from "lucide-react"; import { useMessage } from "@assistant-ui/react"; import { cn } from "@/lib/utils"; import { Badge, badgeVariants, type BadgeProps } from "./badge"; @@ -126,13 +127,31 @@ function Source({ // ── Source badge with hover card ───────────────────────────── -interface SourceData { +interface UrlSourceData { + kind: "url"; url: string; title: string; description?: string; } -const SourceBadge: FC<{ source: SourceData }> = ({ source }) => { +interface DocSourceData { + kind: "document"; + chunkId: string; + filename: string; + page?: string; + score?: string; + text: string; +} + +type SourceData = UrlSourceData | DocSourceData; + +function sourceKey(source: SourceData): string { + return source.kind === "url" + ? `url:${source.url}` + : `doc:${source.chunkId}`; +} + +const SourceBadge: FC<{ source: UrlSourceData }> = ({ source }) => { const domain = extractDomain(source.url); const displayTitle = source.title || domain; @@ -171,6 +190,54 @@ const SourceBadge: FC<{ source: SourceData }> = ({ source }) => { ); }; +const DocumentSourceBadge: FC<{ source: DocSourceData }> = ({ source }) => { + const metaParts: string[] = []; + if (source.page) metaParts.push(`page ${source.page}`); + if (source.score) metaParts.push(`score ${source.score}`); + + return ( + + + + + + [{source.chunkId}] + + + {source.filename} + + + + +
+ +
+

+ {source.filename} +

+ {metaParts.length > 0 ? ( +

+ {metaParts.join(" · ")} +

+ ) : null} +

+ {source.text} +

+
+
+ + + ); +}; + // ── Grouped sources with 2-row collapse ───────────────────── const SourcesGroup: FC = () => { @@ -179,7 +246,7 @@ const SourcesGroup: FC = () => { const [visibleCount, setVisibleCount] = useState(null); const [expanded, setExpanded] = useState(false); - // Extract source parts from the message + // Extract source parts (both URL and document) from the message const sources: SourceData[] = []; if (message.content) { for (const part of message.content) { @@ -191,11 +258,34 @@ const SourcesGroup: FC = () => { part.url ) { sources.push({ + kind: "url", url: part.url as string, title: (part as { title?: string }).title || "", description: (part as { metadata?: { description?: string } }) .metadata?.description, }); + } else if ( + part.type === "source" && + "sourceType" in part && + (part as { sourceType?: string }).sourceType === "document" + ) { + const docPart = part as { + chunkId?: string; + filename?: string; + page?: string; + score?: string; + text?: string; + }; + if (docPart.chunkId && docPart.filename) { + sources.push({ + kind: "document", + chunkId: docPart.chunkId, + filename: docPart.filename, + page: docPart.page, + score: docPart.score, + text: docPart.text ?? "", + }); + } } } } @@ -258,20 +348,39 @@ const SourcesGroup: FC = () => { className="flex w-full flex-wrap gap-1 invisible absolute pointer-events-none" > {sources.map((source) => ( - - - - {source.title || extractDomain(source.url)} - + + {source.kind === "url" ? ( + + + + {source.title || extractDomain(source.url)} + + + ) : ( + + + [{source.chunkId}] + + + {source.filename} + + )} ))} {/* Visible container */}
- {displayedSources.map((source) => ( - - ))} + {displayedSources.map((source) => + source.kind === "url" ? ( + + ) : ( + + ), + )} {shouldCollapse && !expanded && (