ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: render only LLM-cited RAG chunks as Source badges; globally unique chunk IDs

Browse source
This commit is contained in:
Roland Tannous 2026-05-26 17:52:58 +04:00
commit 7c1f8efe99
5 changed files with 244 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

26
studio/backend/core/rag/tool.py
View file

@ -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 "
'<chunk id="N" source="..." page="..." score="...">...</chunk> '
"tags; cite them in your reply as [1], [2], etc."
"tags. CITE each chunk you use with its LITERAL id attribute, "
'e.g. `<chunk id="7">` 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 <chunk> blocks with metadata."""
def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str:
"""Render hits as fenced <chunk> 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

132
studio/frontend/src/components/assistant-ui/sources.tsx
View file

@ -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 (
<HoverCard openDelay={0} closeDelay={0}>
<HoverCardTrigger asChild>
<span className="inline-block">
<Badge
variant="outline"
className="rounded-full cursor-default inline-flex items-center gap-1.5 outline-none"
>
<span className="font-mono text-[10px] font-semibold text-muted-foreground">
[{source.chunkId}]
</span>
<FileTextIcon className="size-3 shrink-0 text-muted-foreground" />
<SourceTitle>{source.filename}</SourceTitle>
</Badge>
</span>
</HoverCardTrigger>
<HoverCardContent
side="top"
align="start"
className="!bg-black !text-white !w-72 !p-3 !rounded-2xl !shadow-md !ring-0 !duration-0"
style={{ animation: "none" }}
>
<div className="flex gap-2.5">
<FileTextIcon className="size-4 mt-0.5 shrink-0 text-white/70" />
<div className="min-w-0 space-y-1">
<p className="text-sm font-semibold leading-tight truncate">
{source.filename}
</p>
{metaParts.length > 0 ? (
<p className="text-xs text-white/60 truncate">
{metaParts.join(" · ")}
</p>
) : null}
<p className="text-xs text-white/70 leading-relaxed line-clamp-3 whitespace-pre-wrap">
{source.text}
</p>
</div>
</div>
</HoverCardContent>
</HoverCard>
);
};
// ── Grouped sources with 2-row collapse ─────────────────────
const SourcesGroup: FC = () => {
@ -179,7 +246,7 @@ const SourcesGroup: FC = () => {
const [visibleCount, setVisibleCount] = useState<number | null>(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) => (
<span key={source.url} className="inline-block">
<Source href={source.url}>
<SourceIcon url={source.url} />
<SourceTitle>{source.title || extractDomain(source.url)}</SourceTitle>
</Source>
<span key={sourceKey(source)} className="inline-block">
{source.kind === "url" ? (
<Source href={source.url}>
<SourceIcon url={source.url} />
<SourceTitle>
{source.title || extractDomain(source.url)}
</SourceTitle>
</Source>
) : (
<Badge
variant="outline"
className="rounded-full inline-flex items-center gap-1.5"
>
<span className="font-mono text-[10px] font-semibold text-muted-foreground">
[{source.chunkId}]
</span>
<FileTextIcon className="size-3 shrink-0 text-muted-foreground" />
<SourceTitle>{source.filename}</SourceTitle>
</Badge>
)}
</span>
))}
</div>
{/* Visible container */}
<div className="flex flex-wrap gap-1">
{displayedSources.map((source) => (
<SourceBadge key={source.url} source={source} />
))}
{displayedSources.map((source) =>
source.kind === "url" ? (
<SourceBadge key={sourceKey(source)} source={source} />
) : (
<DocumentSourceBadge key={sourceKey(source)} source={source} />
),
)}
{shouldCollapse && !expanded && (
<button
type="button"
@ -326,5 +435,6 @@ export {
Source,
SourceIcon,
SourceTitle,
DocumentSourceBadge,
badgeVariants as sourceVariants,
};

4
studio/frontend/src/components/assistant-ui/tool-ui-search-knowledge-base.tsx
View file

@ -16,7 +16,7 @@ import {
ToolFallbackTrigger,
} from "./tool-fallback";
interface ParsedChunk {
export interface ParsedChunk {
id: string;
source: string;
page?: string;
@ -39,7 +39,7 @@ function decodeXml(value: string): string {
.replace(/&amp;/g, "&");
}
function parseChunks(raw: string): ParsedChunk[] {
export function parseChunks(raw: string): ParsedChunk[] {
if (!raw) return [];
const out: ParsedChunk[] = [];
let match: RegExpExecArray | null = CHUNK_RE.exec(raw);

90
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -61,6 +61,10 @@ import {
type SearchRequest,
search as ragSearch,
} from "@/features/rag/api/rag-api";
import {
type ParsedChunk,
parseChunks,
} from "@/components/assistant-ui/tool-ui-search-knowledge-base";
import type { RagMode, RagSource } from "./chat-settings-api";
import {
createOpenAIContainer,
@ -241,6 +245,64 @@ function parseSourcesFromResult(raw: string): {
return sources;
}
interface DocumentSourcePart {
type: "source";
sourceType: "document";
id: string;
chunkId: string;
filename: string;
page?: string;
score?: string;
text: string;
}
/** Pull every `[N]` token the model wrote in its reply.
* Naive: regex over the whole text. False positives (e.g. `[1]` inside a
* code fence or list marker) are tolerated — the worst case is a stray
* badge for an id that exists in the retrieval set. */
const CITATION_RE = /\[(\d+)\]/g;
function extractCitedIds(text: string): Set<string> {
const ids = new Set<string>();
let match: RegExpExecArray | null = CITATION_RE.exec(text);
while (match !== null) {
ids.add(match[1]);
match = CITATION_RE.exec(text);
}
CITATION_RE.lastIndex = 0;
return ids;
}
/** Build doc-shaped source parts for chunks the model actually cited.
* `allChunks` is the flat union of every search_knowledge_base tool
* result in this turn (deduped by id). Returns one part per unique
* cited id that maps to a real chunk; hallucinated `[99]` refs without
* a matching chunk are silently dropped. */
function buildDocumentSourceParts(
allChunks: ParsedChunk[],
citedIds: Set<string>,
): DocumentSourcePart[] {
const byId = new Map<string, ParsedChunk>();
for (const chunk of allChunks) {
if (!byId.has(chunk.id)) byId.set(chunk.id, chunk);
}
const out: DocumentSourcePart[] = [];
for (const id of citedIds) {
const chunk = byId.get(id);
if (!chunk) continue;
out.push({
type: "source",
sourceType: "document",
id: `rag-${id}`,
chunkId: id,
filename: chunk.source,
...(chunk.page ? { page: chunk.page } : {}),
...(chunk.score ? { score: chunk.score } : {}),
text: chunk.text,
});
}
return out;
}
function estimateTokenCount(text: string): number | undefined {
const trimmed = text.trim();
if (!trimmed) {
@ -1013,8 +1075,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
"tool calls). Phrase each query as a focused question, not a " +
"keyword bag.\n" +
"3. After the tool calls return, ground your reply in the " +
"returned <chunk> blocks and cite them as [1], [2], etc. Do not " +
"make additional tool calls beyond the planned 3.",
"returned <chunk> blocks. CITE each chunk you use with its " +
'LITERAL id attribute — e.g. `<chunk id="7">` is cited as ' +
"`[7]`. IDs are unique across the whole turn; never renumber, " +
"never reuse a different number. Do not make additional tool " +
"calls beyond the planned 3.",
);
}
if (systemPromptParts.length > 0) {
@ -2021,7 +2086,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
// tool calls. Both emit the same `Title:` / `URL:` / `Snippet:`
// block shape from the Anthropic backend, so the parser does
// not need to branch on tool name.
const sourceParts = toolCallParts.flatMap((tc) => {
const urlSourceParts = toolCallParts.flatMap((tc) => {
if (
(tc.toolName !== "web_search" && tc.toolName !== "web_fetch") ||
!tc.result
@ -2033,6 +2098,25 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
);
});
// RAG: flatten chunks across every search_knowledge_base call this
// turn, then emit doc-source parts only for ids the model actually
// cited as [N] in its final reply.
const ragChunks = toolCallParts.flatMap((tc) => {
if (tc.toolName !== "search_knowledge_base" || !tc.result) {
return [];
}
return parseChunks(typeof tc.result === "string" ? tc.result : "");
});
const documentSourceParts =
ragChunks.length > 0
? buildDocumentSourceParts(
ragChunks,
extractCitedIds(cumulativeText),
)
: [];
const sourceParts = [...urlSourceParts, ...documentSourceParts];
const meta = serverMetadata;
const finalTokenCount =
meta?.usage?.completion_tokens ?? estimateTokenCount(cumulativeText);

13
tests/python/test_rag_tool_handler.py
View file

@ -140,6 +140,19 @@ def test_format_hits_handles_unknown_source():
assert "\norphan\n</chunk>" in result
def test_format_hits_offsets_ids_by_start_id():
from core.rag.tool import _format_hits_for_llm
hits = [
{"filename": "a.pdf", "text": "first"},
{"filename": "b.pdf", "text": "second"},
]
result = _format_hits_for_llm(hits, start_id = 5)
assert '<chunk id="6"' in result
assert '<chunk id="7"' in result
assert '<chunk id="1"' not in result
def test_format_hits_escapes_xml_in_source():
from core.rag.tool import _format_hits_for_llm