diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8054bd2a19..7b1db8fd04 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1459,7 +1459,7 @@ class LlamaCppBackend: stop: Optional[list[str]] = None, cancel_event: Optional[threading.Event] = None, enable_thinking: Optional[bool] = None, - ) -> Generator[str, None, None]: + ) -> Generator[str | dict, None, None]: """ Send a chat completion request to llama-server and stream tokens back. @@ -1490,10 +1490,14 @@ class LlamaCppBackend: payload["max_tokens"] = max_tokens if stop: payload["stop"] = stop + payload["stream_options"] = {"include_usage": True} url = f"{self.base_url}/v1/chat/completions" cumulative = "" in_thinking = False + _stream_done = False + _metadata_usage = None + _metadata_timings = None try: # _stream_with_retry uses a 120 s read timeout so prefill @@ -1536,12 +1540,20 @@ class LlamaCppBackend: # as the main response, not as a thinking block. cumulative = reasoning_text yield cumulative - return + _stream_done = True + break # exit inner while if not line.startswith("data: "): continue try: data = json.loads(line[6:]) + # Capture server timings/usage from final chunks + _chunk_timings = data.get("timings") + if _chunk_timings: + _metadata_timings = _chunk_timings + _chunk_usage = data.get("usage") + if _chunk_usage: + _metadata_usage = _chunk_usage choices = data.get("choices", []) if choices: delta = choices[0].get("delta", {}) @@ -1570,6 +1582,14 @@ class LlamaCppBackend: logger.debug( f"Skipping malformed SSE line: {line[:100]}" ) + if _stream_done: + break # exit outer for + if _metadata_usage or _metadata_timings: + yield { + "type": "metadata", + "usage": _metadata_usage, + "timings": _metadata_timings, + } except httpx.ConnectError: raise RuntimeError("Lost connection to llama-server") @@ -1614,6 +1634,9 @@ class LlamaCppBackend: conversation = list(messages) url = f"{self.base_url}/v1/chat/completions" + _accumulated_completion_tokens = 0 + _accumulated_predicted_ms = 0.0 + _accumulated_predicted_n = 0 for iteration in range(max_tool_iterations): if cancel_event is not None and cancel_event.is_set(): @@ -1710,6 +1733,13 @@ class LlamaCppBackend: ) if finish_reason == "tool_calls" or (tool_calls and len(tool_calls) > 0): + # Only accumulate metrics for responses that are actually used + _accumulated_completion_tokens += data.get("usage", {}).get( + "completion_tokens", 0 + ) + _iter_timings = data.get("timings", {}) + _accumulated_predicted_ms += _iter_timings.get("predicted_ms", 0) + _accumulated_predicted_n += _iter_timings.get("predicted_n", 0) # Append the assistant message with tool_calls to conversation assistant_msg = {"role": "assistant", "content": content_text} if tool_calls: @@ -1805,6 +1835,14 @@ class LlamaCppBackend: if iteration == 0 and content_text: yield {"type": "status", "text": ""} yield {"type": "content", "text": content_text} + _direct_usage = data.get("usage") + _direct_timings = data.get("timings") + if _direct_usage or _direct_timings: + yield { + "type": "metadata", + "usage": _direct_usage, + "timings": _direct_timings, + } return # Tools were called in previous iterations; do a final @@ -1834,6 +1872,7 @@ class LlamaCppBackend: stream_payload["max_tokens"] = max_tokens if stop: stream_payload["stop"] = stop + stream_payload["stream_options"] = {"include_usage": True} import re as _re_final @@ -1862,6 +1901,9 @@ class LlamaCppBackend: in_thinking = False has_content_tokens = False reasoning_text = "" + _metadata_usage = None + _metadata_timings = None + _stream_done = False try: stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) @@ -1899,12 +1941,20 @@ class LlamaCppBackend: else: cumulative = reasoning_text yield {"type": "content", "text": cumulative} - return + _stream_done = True + break # exit inner while if not line.startswith("data: "): continue try: chunk_data = json.loads(line[6:]) + # Capture server timings/usage from final chunks + _chunk_timings = chunk_data.get("timings") + if _chunk_timings: + _metadata_timings = _chunk_timings + _chunk_usage = chunk_data.get("usage") + if _chunk_usage: + _metadata_usage = _chunk_usage choices = chunk_data.get("choices", []) if choices: delta = choices[0].get("delta", {}) @@ -1934,6 +1984,42 @@ class LlamaCppBackend: logger.debug( f"Skipping malformed SSE line: {line[:100]}" ) + if _stream_done: + break # exit outer for + _final_usage = _metadata_usage or {} + _final_completion = _final_usage.get("completion_tokens", 0) + _final_prompt = _final_usage.get("prompt_tokens", 0) + _total_completion = ( + _final_completion + _accumulated_completion_tokens + ) + if _metadata_usage or _metadata_timings: + _merged_timings = ( + dict(_metadata_timings) if _metadata_timings else {} + ) + if _accumulated_predicted_ms or _accumulated_predicted_n: + _merged_timings["predicted_ms"] = ( + _merged_timings.get("predicted_ms", 0) + + _accumulated_predicted_ms + ) + _total_predicted_n = ( + _merged_timings.get("predicted_n", 0) + + _accumulated_predicted_n + ) + _merged_timings["predicted_n"] = _total_predicted_n + _total_predicted_ms = _merged_timings["predicted_ms"] + if _total_predicted_ms > 0: + _merged_timings["predicted_per_second"] = ( + _total_predicted_n / (_total_predicted_ms / 1000.0) + ) + yield { + "type": "metadata", + "usage": { + "prompt_tokens": _final_prompt, + "completion_tokens": _total_completion, + "total_tokens": _final_prompt + _total_completion, + }, + "timings": _merged_timings, + } except httpx.ConnectError: raise RuntimeError("Lost connection to llama-server") diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 41a942d217..b0498319ca 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -364,6 +364,8 @@ class ChatCompletionChunk(BaseModel): created: int = Field(default_factory = lambda: int(time.time())) model: str = "default" choices: list[ChunkChoice] + usage: Optional[CompletionUsage] = None + timings: Optional[dict] = None # ── Non-streaming response ─────────────────────────────────────── diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4c98dc6d24..aa8c34a3c5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -55,6 +55,7 @@ from models.inference import ( ChoiceDelta, CompletionChoice, CompletionMessage, + CompletionUsage, ValidateModelRequest, ValidateModelResponse, ) @@ -1083,6 +1084,8 @@ async def openai_chat_completions( # the event loop stays free for disconnect detection. gen = gguf_generate_with_tools() prev_text = "" + _stream_usage = None + _stream_timings = None while True: if await request.is_disconnected(): cancel_event.set() @@ -1107,6 +1110,11 @@ async def openai_chat_completions( yield f"data: {json.dumps(event)}\n\n" continue + if event["type"] == "metadata": + _stream_usage = event.get("usage") + _stream_timings = event.get("timings") + continue + # "content" type -- cumulative text cumulative = event.get("text", "") new_text = cumulative[len(prev_text) :] @@ -1138,6 +1146,24 @@ async def openai_chat_completions( ], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" + # Usage chunk (OpenAI-standard: choices=[], usage populated) + if _stream_usage or _stream_timings: + usage_obj = CompletionUsage( + prompt_tokens = (_stream_usage or {}).get("prompt_tokens", 0), + completion_tokens = (_stream_usage or {}).get( + "completion_tokens", 0 + ), + total_tokens = (_stream_usage or {}).get("total_tokens", 0), + ) + usage_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [], + usage = usage_obj, + timings = _stream_timings, + ) + yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -1207,6 +1233,8 @@ async def openai_chat_completions( # the event loop stays free for disconnect detection. gen = gguf_generate() prev_text = "" + _stream_usage = None + _stream_timings = None while True: if await request.is_disconnected(): cancel_event.set() @@ -1214,6 +1242,21 @@ async def openai_chat_completions( cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) if cumulative is _gguf_sentinel: break + # Capture server metadata for final usage chunk + if isinstance(cumulative, dict): + if cumulative.get("type") == "metadata": + _stream_usage = cumulative.get("usage") + _stream_timings = cumulative.get("timings") + else: + logger.warning( + "gguf_stream_chunks: unexpected dict event: %s", + { + k: v + for k, v in cumulative.items() + if k != "timings" + }, + ) + continue new_text = cumulative[len(prev_text) :] prev_text = cumulative if not new_text: @@ -1244,6 +1287,24 @@ async def openai_chat_completions( ], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" + # Usage chunk (OpenAI-standard: choices=[], usage populated) + if _stream_usage or _stream_timings: + usage_obj = CompletionUsage( + prompt_tokens = (_stream_usage or {}).get("prompt_tokens", 0), + completion_tokens = (_stream_usage or {}).get( + "completion_tokens", 0 + ), + total_tokens = (_stream_usage or {}).get("total_tokens", 0), + ) + usage_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [], + usage = usage_obj, + timings = _stream_timings, + ) + yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -1272,6 +1333,8 @@ async def openai_chat_completions( try: full_text = "" for token in gguf_generate(): + if isinstance(token, dict): + continue # skip metadata dict in non-streaming path full_text = token response = ChatCompletion( diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 51541774d9..fdc4f374ab 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -147,7 +147,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]: "Helper model request: enable_thinking=False (per-request override)" ) cumulative = "" - for text in backend.generate_chat_completion( + for chunk in backend.generate_chat_completion( messages = messages, temperature = 0.1, top_p = 0.9, @@ -156,7 +156,9 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]: repetition_penalty = 1.0, enable_thinking = False, # Always disable thinking for AI Assist ): - cumulative = text # cumulative — last value is full text + if isinstance(chunk, dict): + continue # skip metadata events + cumulative = chunk # cumulative — last value is full text result = cumulative.strip() result = _strip_think_tags(result) @@ -422,7 +424,7 @@ def _generate_with_backend(backend, messages: list[dict], max_tokens: int = 512) """Run one chat completion on an already-loaded backend. Returns raw text.""" logger.info("Advisor request: enable_thinking=False (per-request override)") cumulative = "" - for text in backend.generate_chat_completion( + for chunk in backend.generate_chat_completion( messages = messages, temperature = 0.1, top_p = 0.9, @@ -431,7 +433,9 @@ def _generate_with_backend(backend, messages: list[dict], max_tokens: int = 512) repetition_penalty = 1.0, enable_thinking = False, # Always disable thinking for AI Assist ): - cumulative = text + if isinstance(chunk, dict): + continue # skip metadata events + cumulative = chunk result = cumulative.strip() result = _strip_think_tags(result) return result diff --git a/studio/frontend/bun.lock b/studio/frontend/bun.lock index 7e3b0ac51e..5504aea3d3 100644 --- a/studio/frontend/bun.lock +++ b/studio/frontend/bun.lock @@ -430,6 +430,8 @@ "@open-draft/until": ["@open-draft/until@2.1.0", "", {}, "sha512-U69T3ItWHvLwGg5eJ0n3I62nWuE6ilHlmz7zM0npLBRvPRd7e6NYmg54vvRtP5mZG7kZqZCFVdsTWo7BPtBujg=="], + "@opentelemetry/api": ["@opentelemetry/api@1.9.0", "", {}, "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg=="], + "@radix-ui/number": ["@radix-ui/number@1.1.1", "", {}, "sha512-MkKCwxlXTgz6CFoJx3pCwn07GKp36+aZyu/u2Ln2VrA5DcdyCZkASEDBTd8x5whTQQL5CiYf4prXKLcgQdv29g=="], "@radix-ui/primitive": ["@radix-ui/primitive@1.1.3", "", {}, "sha512-JTF99U/6XIjCBo0wqkU5sK10glYe27MRRsfwoiq5zzOEZLHU3A3KCMa5X/azekYRCJ0HlwI0crAXS/5dEHTzDg=="], diff --git a/studio/frontend/src/components/assistant-ui/message-timing.tsx b/studio/frontend/src/components/assistant-ui/message-timing.tsx index 567e8468d0..df233812b4 100644 --- a/studio/frontend/src/components/assistant-ui/message-timing.tsx +++ b/studio/frontend/src/components/assistant-ui/message-timing.tsx @@ -1,6 +1,6 @@ "use client"; -import { useMessageTiming } from "@assistant-ui/react"; +import { useMessageTiming, useMessage } from "@assistant-ui/react"; import { Tooltip, TooltipContent, @@ -15,32 +15,34 @@ const formatTimingMs = (ms: number | undefined): string => { return `${(ms / 1000).toFixed(2)}s`; }; +const formatNumber = (n: number): string => { + return n.toLocaleString(); +}; + /** - * Shows streaming stats (TTFT, total time, chunks) as a badge with a - * hover/focus tooltip. Renders nothing until the stream completes. - * - * Place it inside `ActionBarPrimitive.Root` in your `thread.tsx` so it - * inherits the action bar's autohide behaviour: - * - * ```tsx - * import { MessageTiming } from "@/components/assistant-ui/message-timing"; - * - * - * - * - * // <-- add this - * - * ``` - * - * @param side - Side of the tooltip relative to the badge trigger. Defaults to `"right"`. + * Shows streaming stats as a badge with hover tooltip. + * When server timings are available (GGUF), shows prompt eval, generation, + * speed, tokens, and cache hits. Falls back to client-side metrics otherwise. */ export const MessageTiming: FC<{ className?: string; side?: "top" | "right" | "bottom" | "left"; }> = ({ className, side = "right" }) => { const timing = useMessageTiming(); + const message = useMessage(); + if (timing?.totalStreamTime === undefined) return null; + const serverTimings = ( + message.metadata as Record | undefined + )?.custom as { serverTimings?: Record } | undefined; + const st = serverTimings?.serverTimings; + + // Badge text: show tok/s if available, otherwise total time + const badgeText = st?.predicted_per_second != null + ? `${st.predicted_per_second.toFixed(1)} tok/s` + : formatTimingMs(timing.totalStreamTime); + return ( @@ -53,7 +55,7 @@ export const MessageTiming: FC<{ className, )} > - {formatTimingMs(timing.totalStreamTime)} + {badgeText} -
- {timing.firstTokenTime !== undefined && ( -
- First token - - {formatTimingMs(timing.firstTokenTime)} - -
+
+ {st ? ( + <> + {/* Server-side metrics (GGUF) */} + {st?.prompt_ms != null && ( +
+ Prompt eval + + {formatTimingMs(st.prompt_ms)} + +
+ )} + {(st?.prompt_n ?? 0) > 1 && st?.prompt_per_second != null && ( +
+ Prompt speed + + {st.prompt_per_second.toFixed(1)} tok/s + +
+ )} + {st?.predicted_ms != null && ( +
+ Generation + + {formatTimingMs(st.predicted_ms)} + +
+ )} + {st?.predicted_per_second != null && ( +
+ Speed + + {st.predicted_per_second.toFixed(1)} tok/s + +
+ )} + {timing.tokenCount !== undefined && ( +
+ Tokens + + {formatNumber(timing.tokenCount)} + +
+ )} + {(st?.cache_n ?? 0) > 0 && ( +
+ Cache hits + + {formatNumber(st!.cache_n)} + +
+ )} +
+
+ Total + + {formatTimingMs(timing.totalStreamTime)} + +
+
+ Chunks + + {timing.totalChunks} + +
+ + ) : ( + <> + {/* Client-side metrics (safetensors fallback) */} + {timing.firstTokenTime !== undefined && ( +
+ First token + + {formatTimingMs(timing.firstTokenTime)} + +
+ )} +
+ Total + + {formatTimingMs(timing.totalStreamTime)} + +
+
+ Chunks + + {timing.totalChunks} + +
+ )} -
- Total - - {formatTimingMs(timing.totalStreamTime)} - -
-
- Chunks - {timing.totalChunks} -
diff --git a/studio/frontend/src/components/assistant-ui/sources.tsx b/studio/frontend/src/components/assistant-ui/sources.tsx index ada503fd54..18ee03cad0 100644 --- a/studio/frontend/src/components/assistant-ui/sources.tsx +++ b/studio/frontend/src/components/assistant-ui/sources.tsx @@ -1,9 +1,24 @@ "use client"; -import { memo, useState, type ComponentProps } from "react"; -import type { SourceMessagePartComponent } from "@assistant-ui/react"; +import { + memo, + useState, + useRef, + useEffect, + useCallback, + type ComponentProps, + type FC, +} from "react"; +import { useMessage } from "@assistant-ui/react"; import { cn } from "@/lib/utils"; import { Badge, badgeVariants, type BadgeProps } from "./badge"; +import { + HoverCard, + HoverCardTrigger, + HoverCardContent, +} from "@/components/ui/hover-card"; + +// ── Helpers ────────────────────────────────────────────────── const extractDomain = (url: string): string => { try { @@ -18,20 +33,24 @@ const getDomainInitial = (url: string): string => { return domain.charAt(0).toUpperCase(); }; +// ── Sub-components ─────────────────────────────────────────── + function SourceIcon({ url, className, + size = 3, ...props -}: ComponentProps<"span"> & { url: string }) { +}: ComponentProps<"span"> & { url: string; size?: number }) { const [hasError, setHasError] = useState(false); const domain = extractDomain(url); + const sizeClass = `size-${size}`; if (hasError) { return ( setHasError(true)} {...(props as ComponentProps<"img">)} /> @@ -97,39 +116,200 @@ function Source({ ); } -const SourcesImpl: SourceMessagePartComponent = ({ - url, - title, - sourceType, -}) => { - if (sourceType !== "url" || !url) return null; +// ── Source badge with hover card ───────────────────────────── - const domain = extractDomain(url); - const displayTitle = title || domain; +interface SourceData { + url: string; + title: string; + description?: string; +} + +const SourceBadge: FC<{ source: SourceData }> = ({ source }) => { + const domain = extractDomain(source.url); + const displayTitle = source.title || domain; return ( - - - - {displayTitle} - - + + + + + + {displayTitle} + + + + +
+ +
+

+ {source.title || domain} +

+

{domain}

+ {source.description && ( +

+ {source.description} +

+ )} +
+
+ + ); }; -const Sources = memo(SourcesImpl) as unknown as SourceMessagePartComponent & { +// ── Grouped sources with 2-row collapse ───────────────────── + +const SourcesGroup: FC = () => { + const message = useMessage(); + const containerRef = useRef(null); + const [visibleCount, setVisibleCount] = useState(null); + const [expanded, setExpanded] = useState(false); + + // Extract source parts from the message + const sources: SourceData[] = []; + if (message.content) { + for (const part of message.content) { + if ( + part.type === "source" && + "sourceType" in part && + part.sourceType === "url" && + "url" in part && + part.url + ) { + sources.push({ + url: part.url as string, + title: (part as { title?: string }).title || "", + description: (part as { metadata?: { description?: string } }) + .metadata?.description, + }); + } + } + } + + // Measure how many badges fit in 2 rows + const measure = useCallback(() => { + const container = containerRef.current; + if (!container || sources.length === 0) return; + + const children = Array.from(container.children) as HTMLElement[]; + if (children.length === 0) return; + + // Find the top of the first child as baseline + const firstTop = children[0].offsetTop; + let rowCount = 1; + let prevTop = firstTop; + let cutoff = children.length; + + for (let i = 1; i < children.length; i++) { + const childTop = children[i].offsetTop; + if (childTop > prevTop) { + rowCount++; + prevTop = childTop; + if (rowCount > 2) { + cutoff = i; + break; + } + } + } + + setVisibleCount(rowCount > 2 ? cutoff : null); + }, [sources.length]); + + useEffect(() => { + measure(); + }, [measure]); + + // Re-measure on resize + useEffect(() => { + const container = containerRef.current; + if (!container) return; + const observer = new ResizeObserver(measure); + observer.observe(container); + return () => observer.disconnect(); + }, [measure]); + + if (sources.length === 0) return null; + + const shouldCollapse = visibleCount !== null && visibleCount < sources.length; + const displayedSources = + expanded || !shouldCollapse ? sources : sources.slice(0, visibleCount); + const hiddenCount = sources.length - (visibleCount ?? sources.length); + + return ( +
+ {/* Hidden measurement container — renders all badges to measure row positions */} +
+ {sources.map((source) => ( + + + + {source.title || extractDomain(source.url)} + + + ))} +
+ + {/* Visible container */} +
+ {displayedSources.map((source) => ( + + ))} + {shouldCollapse && !expanded && ( + + )} + {shouldCollapse && expanded && ( + + )} +
+
+ ); +}; + +// ── Individual source (renders null — SourcesGroup handles all) ── + +const SourcesNoop: FC> = () => null; + +// ── Exports ────────────────────────────────────────────────── + +const Sources = memo(SourcesNoop) as unknown as FC> & { Root: typeof Source; Icon: typeof SourceIcon; Title: typeof SourceTitle; + Group: typeof SourcesGroup; }; Sources.displayName = "Sources"; Sources.Root = Source; Sources.Icon = SourceIcon; Sources.Title = SourceTitle; +Sources.Group = SourcesGroup; export { Sources, + SourcesGroup, Source, SourceIcon, SourceTitle, diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 567c25b53f..4d472ba208 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -9,7 +9,7 @@ import { import { MessageTiming } from "@/components/assistant-ui/message-timing"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; -import { Sources } from "@/components/assistant-ui/sources"; +import { Sources, SourcesGroup } from "@/components/assistant-ui/sources"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { ToolGroup } from "@/components/assistant-ui/tool-group"; import { WebSearchToolUI } from "@/components/assistant-ui/tool-ui-web-search"; @@ -114,26 +114,68 @@ const ThreadScrollToBottom: FC = () => { ); }; +const SUGGESTION_TOOLS: Record> = { + "How do you fine-tune an audio model with Unsloth?": ["thinking", "search"], + "Show me a live weather dashboard, no API key needed": ["thinking", "code", "search"], + "Solve the integral of x·sin(x), and verify it step by step": ["thinking", "code"], + "Draw an SVG of a cute sloth": ["thinking", "code", "search"], +}; + +const toolIconMap = { + thinking: { icon: LightbulbIcon, label: "Thinking" }, + search: { icon: GlobeIcon, label: "Web search" }, + code: { icon: TerminalIcon, label: "Code" }, +} as const; + const SuggestionItem: FC = () => { const aui = useAui(); const prompt = useAuiState(({ suggestion }) => suggestion.prompt); const isDisabled = useAuiState(({ thread }) => thread.isDisabled); const isRunning = useAuiState(({ thread }) => thread.isRunning); + const supportsTools = useChatRuntimeStore((s) => s.supportsTools); + const supportsReasoning = useChatRuntimeStore((s) => s.supportsReasoning); + const allTools = SUGGESTION_TOOLS[prompt] ?? []; + const tools = allTools.filter((tool) => { + if (tool === "thinking") return supportsReasoning; + return supportsTools; + }); return ( ); }; @@ -517,6 +559,7 @@ const AssistantMessage: FC = () => { }, }} /> +
diff --git a/studio/frontend/src/components/assistant-ui/tool-group.tsx b/studio/frontend/src/components/assistant-ui/tool-group.tsx index d1ad01335e..f29adb510f 100644 --- a/studio/frontend/src/components/assistant-ui/tool-group.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-group.tsx @@ -9,6 +9,8 @@ import { type PropsWithChildren, } from "react"; import { ChevronDownIcon, LoaderIcon } from "lucide-react"; +import { Wrench01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import { cva, type VariantProps } from "class-variance-authority"; import { useScrollLock } from "@assistant-ui/react"; import { @@ -116,11 +118,18 @@ function ToolGroupTrigger({ )} {...props} > - {active && ( + {active ? ( + ) : ( + )} [0]["messages"]; type RunMessage = RunMessages[number]; @@ -27,21 +47,24 @@ type RunMessage = RunMessages[number]; export const sentAudioNames = new Map(); /** Parse "Title: ...\nURL: ...\nSnippet: ..." blocks into source content parts. */ -function parseSourcesFromResult(raw: string): { type: "source"; sourceType: "url"; id: string; url: string; title: string }[] { +function parseSourcesFromResult(raw: string): { type: "source"; sourceType: "url"; id: string; url: string; title: string; metadata?: { description: string } }[] { if (!raw) return []; const blocks = raw.split(/\n---\n/).filter(Boolean); - const sources: { type: "source"; sourceType: "url"; id: string; url: string; title: string }[] = []; + const sources: { type: "source"; sourceType: "url"; id: string; url: string; title: string; metadata?: { description: string } }[] = []; for (const block of blocks) { const titleMatch = block.match(/Title:\s*(.+)/); const urlMatch = block.match(/URL:\s*(.+)/); + const snippetMatch = block.match(/Snippet:\s*(.+)/); if (titleMatch && urlMatch) { const url = urlMatch[1].trim(); + const snippet = snippetMatch?.[1]?.trim(); sources.push({ type: "source" as const, sourceType: "url" as const, id: url, url, title: titleMatch[1].trim(), + ...(snippet ? { metadata: { description: snippet } } : {}), }); } } @@ -63,6 +86,7 @@ function buildTiming( totalStreamTime?: number, tokenCount?: number, toolCallCount = 0, + tokensPerSecondOverride?: number, ): MessageTiming { return { streamStartTime, @@ -70,11 +94,12 @@ function buildTiming( totalStreamTime, tokenCount, tokensPerSecond: - typeof totalStreamTime === "number" && + tokensPerSecondOverride ?? + (typeof totalStreamTime === "number" && totalStreamTime > 0 && typeof tokenCount === "number" ? tokenCount / (totalStreamTime / 1000) - : undefined, + : undefined), totalChunks, toolCallCount, }; @@ -528,6 +553,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // Tool call content parts — accumulated and yielded cumulatively. // result is set directly on the tool-call part when tool_end arrives. const toolCallParts: ToolCallMessagePart[] = []; + let serverMetadata: { usage?: ServerUsage; timings?: ServerTimings } | null = null; try { const { supportsReasoning, reasoningEnabled } = runtime; @@ -610,6 +636,15 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { continue; } + // OpenAI-standard usage chunk: choices=[], usage populated + if (chunk.choices?.length === 0 && chunk.usage) { + serverMetadata = { + usage: chunk.usage, + timings: (chunk as Record).timings as ServerTimings | undefined, + }; + continue; + } + totalChunks += 1; const delta = chunk.choices?.[0]?.delta?.content; if (!delta) { @@ -654,6 +689,37 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { return parseSourcesFromResult(typeof tc.result === "string" ? tc.result : ""); }); + const meta = serverMetadata; + const finalTokenCount = meta?.usage?.completion_tokens + ?? estimateTokenCount(cumulativeText); + const finalTokPerSec = meta?.timings?.predicted_per_second; + const serverPromptEvalTime = meta?.timings?.prompt_ms; + + // Update context usage in store if we got valid server data + if ( + meta?.usage && + typeof meta.usage.prompt_tokens === "number" && + typeof meta.usage.completion_tokens === "number" && + typeof meta.usage.total_tokens === "number" + ) { + useChatRuntimeStore.getState().setContextUsage({ + promptTokens: meta.usage.prompt_tokens, + completionTokens: meta.usage.completion_tokens, + totalTokens: meta.usage.total_tokens, + cachedTokens: meta.timings?.cache_n ?? 0, + }); + } + + const finalTiming = buildTiming( + streamStartTime, + totalChunks, + serverPromptEvalTime ?? firstTokenTime, + Date.now() - streamStartTime, + finalTokenCount, + toolCallParts.length, + finalTokPerSec, + ); + yield { content: [ ...toolCallParts, @@ -661,15 +727,19 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ...sourceParts, ], metadata: { - timing: buildTiming( - streamStartTime, - totalChunks, - firstTokenTime, - Date.now() - streamStartTime, - estimateTokenCount(cumulativeText), - toolCallParts.length, - ), - custom: { reasoningDuration }, + timing: finalTiming, + custom: { + reasoningDuration, + serverTimings: meta?.timings ?? undefined, + contextUsage: meta?.usage ? { + promptTokens: meta.usage.prompt_tokens, + completionTokens: meta.usage.completion_tokens, + totalTokens: meta.usage.total_tokens, + cachedTokens: meta.timings?.cache_n ?? 0, + modelId: params.checkpoint, + } : undefined, + timing: finalTiming, + }, }, }; } catch (err) { diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 33c9b5ba39..c04cfbc89c 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -38,6 +38,7 @@ import { import { toast } from "sonner"; import { GuidedTour, useGuidedTourController } from "@/features/tour"; import { ChatSettingsPanel } from "./chat-settings-sheet"; +import { ContextUsageBar } from "./components/context-usage-bar"; import { ModelLoadInlineStatus } from "./components/model-load-status"; import { db } from "./db"; import { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; @@ -423,6 +424,8 @@ export function ChatPage(): ReactElement { const inferenceParams = useChatRuntimeStore((state) => state.params); const setInferenceParams = useChatRuntimeStore((state) => state.setParams); const activeGgufVariant = useChatRuntimeStore((state) => state.activeGgufVariant); + const ggufContextLength = useChatRuntimeStore((state) => state.ggufContextLength); + const contextUsage = useChatRuntimeStore((state) => state.contextUsage); const autoTitle = useChatRuntimeStore((state) => state.autoTitle); const setAutoTitle = useChatRuntimeStore((state) => state.setAutoTitle); const modelsFromStore = useChatRuntimeStore((state) => state.models); @@ -503,7 +506,10 @@ export function ChatPage(): ReactElement { [], ); const handleNewCompare = useCallback( - () => setView({ mode: "compare", pairId: crypto.randomUUID() }), + () => { + setView({ mode: "compare", pairId: crypto.randomUUID() }); + useChatRuntimeStore.getState().setContextUsage(null); + }, [], ); @@ -531,12 +537,28 @@ export function ChatPage(): ReactElement { const enterCompare = useCallback(() => { setViewBeforeCompare((prev) => prev ?? view); setView({ mode: "compare", pairId: crypto.randomUUID() }); + useChatRuntimeStore.getState().setContextUsage(null); }, [view]); const exitCompare = useCallback(() => { if (!viewBeforeCompare) return; setView(viewBeforeCompare); setViewBeforeCompare(null); + // Restore context usage from the active thread's last assistant message + const store = useChatRuntimeStore.getState(); + const threadId = store.activeThreadId; + if (threadId) { + void db.messages + .where("threadId") + .equals(threadId) + .reverse() + .first() + .then((msg) => { + const saved = msg?.metadata as Record | undefined; + const usage = saved?.contextUsage as typeof store.contextUsage | undefined; + if (usage) store.setContextUsage(usage); + }); + } }, [viewBeforeCompare]); const handleThreadSelect = useCallback( @@ -599,6 +621,7 @@ export function ChatPage(): ReactElement { await selectModelRef.current({ id: targetLora.id, isLora: true }); if (canceled) return; setView({ mode: "compare", pairId: crypto.randomUUID() }); + useChatRuntimeStore.getState().setContextUsage(null); clearHandoff(); console.info("[chat-handoff] loaded lora + opened compare"); return; @@ -744,6 +767,15 @@ export function ChatPage(): ReactElement {
)}
+ {view.mode === "single" && ggufContextLength && contextUsage ? ( + + ) : null} + + +
+
+ Context usage + + {percent.toFixed(1)}% + +
+ {promptTokens !== undefined && ( +
+ Prompt tokens + + {formatTokenCountFull(promptTokens)} + +
+ )} + {completionTokens !== undefined && ( +
+ Completion + + {formatTokenCountFull(completionTokens)} + +
+ )} + {cached !== undefined && cached > 0 && ( +
+ Cache hits + + {formatTokenCountFull(cached)} + +
+ )} +
+
+ Total + + {formatTokenCountFull(used)} / {formatTokenCountFull(total)} + +
+
+ + + ); +}; diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index b4b7b6dc7b..b12071639c 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -31,10 +31,26 @@ import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import type { MessageRecord, ModelType } from "./types"; const DEFAULT_SUGGESTIONS = [ - "Draw an SVG of a cute sloth", - "Solve the integral of x²·sin(x) step by step", - "Write a Python function that finds the longest palindrome in a string", - "Format a comparison of 3 databases as a markdown table with pros and cons", + { + title: "How do you fine-tune an audio model with Unsloth?", + label: "Audio fine-tuning", + prompt: "How do you fine-tune an audio model with Unsloth?", + }, + { + title: "Show me a live weather dashboard, no API key needed", + label: "Weather dashboard", + prompt: "Show me a live weather dashboard, no API key needed", + }, + { + title: "Solve the integral of x·sin(x), and verify it", + label: "Integral", + prompt: "Solve the integral of x·sin(x), and verify it step by step", + }, + { + title: "Draw an SVG of a cute sloth", + label: "SVG sloth", + prompt: "Draw an SVG of a cute sloth", + }, ]; type TitleResponse = { @@ -345,6 +361,8 @@ function toThreadMessage(m: MessageRecord): ThreadMessage { metadata: { custom: {} }, }; } + const custom = (m.metadata as Record) ?? {}; + const savedTiming = custom.timing as import("@assistant-ui/react").MessageTiming | undefined; return { id: m.id, createdAt: new Date(m.createdAt), @@ -352,7 +370,8 @@ function toThreadMessage(m: MessageRecord): ThreadMessage { content: content as Extract["content"], status: { type: "complete" as const, reason: "unknown" as const }, metadata: { - custom: (m.metadata as Record) ?? {}, + custom, + ...(savedTiming ? { timing: savedTiming } : {}), steps: [], unstable_annotations: [], unstable_data: [], @@ -535,6 +554,21 @@ function ThreadHistoryProvider({ return a.id < b.id ? -1 : a.id > b.id ? 1 : 0; }); + // Restore context usage from last assistant message if model matches + const lastAssistant = [...msgs].reverse().find((m) => m.role === "assistant"); + const savedUsage = (lastAssistant?.metadata as Record)?.contextUsage as + | { promptTokens: number; completionTokens: number; totalTokens: number; cachedTokens: number; modelId?: string } + | undefined; + const store = useChatRuntimeStore.getState(); + if ( + savedUsage && + store.ggufContextLength && + savedUsage.totalTokens <= store.ggufContextLength && + (!savedUsage.modelId || savedUsage.modelId === store.params.checkpoint) + ) { + store.setContextUsage(savedUsage); + } + return ExportedMessageRepository.fromArray(msgs.map(toThreadMessage)); }, diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index b7f222e3bf..fea5442187 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -84,6 +84,12 @@ type ChatRuntimeStore = { activeThreadId: string | null; pendingAudioBase64: string | null; pendingAudioName: string | null; + contextUsage: { + promptTokens: number; + completionTokens: number; + totalTokens: number; + cachedTokens: number; + } | null; modelLoading: boolean; setModelLoading: (loading: boolean) => void; setParams: (params: InferenceParams) => void; @@ -107,6 +113,7 @@ type ChatRuntimeStore = { setChatTemplateOverride: (template: string | null) => void; setPendingAudio: (base64: string, name: string) => void; clearPendingAudio: () => void; + setContextUsage: (usage: ChatRuntimeStore["contextUsage"]) => void; }; export const useChatRuntimeStore = create((set) => ({ @@ -134,6 +141,7 @@ export const useChatRuntimeStore = create((set) => ({ activeThreadId: null, pendingAudioBase64: null, pendingAudioName: null, + contextUsage: null, modelLoading: false, setModelLoading: (loading) => set({ modelLoading: loading }), setParams: (params) => set({ params }), @@ -163,7 +171,7 @@ export const useChatRuntimeStore = create((set) => ({ }, activeGgufVariant: ggufVariant ?? null, })), - setActiveThreadId: (activeThreadId) => set({ activeThreadId }), + setActiveThreadId: (activeThreadId) => set({ activeThreadId, contextUsage: null }), clearCheckpoint: () => set((state) => ({ params: { @@ -172,6 +180,7 @@ export const useChatRuntimeStore = create((set) => ({ }, activeGgufVariant: null, ggufContextLength: null, + contextUsage: null, supportsReasoning: false, reasoningEnabled: true, supportsTools: false, @@ -208,4 +217,5 @@ export const useChatRuntimeStore = create((set) => ({ set({ pendingAudioBase64: base64, pendingAudioName: name }), clearPendingAudio: () => set({ pendingAudioBase64: null, pendingAudioName: null }), + setContextUsage: (contextUsage) => set({ contextUsage }), })); diff --git a/studio/frontend/src/features/chat/thread-sidebar.tsx b/studio/frontend/src/features/chat/thread-sidebar.tsx index 751f04c171..53c5521dc7 100644 --- a/studio/frontend/src/features/chat/thread-sidebar.tsx +++ b/studio/frontend/src/features/chat/thread-sidebar.tsx @@ -3,6 +3,7 @@ import { SidebarContent, + SidebarFooter, SidebarGroup, SidebarGroupContent, SidebarGroupLabel, @@ -13,8 +14,10 @@ import { SidebarMenuItem, } from "@/components/ui/sidebar"; import { + BookOpen02Icon, ColumnInsertIcon, Delete02Icon, + NewReleasesIcon, PencilEdit02Icon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -158,6 +161,26 @@ export function ThreadSidebar({ + + + + Learn more in docs + + + + What's new + + ); } diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 013e658105..ff5ebe50ca 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -175,4 +175,10 @@ export interface OpenAIChatChunkChoice { export interface OpenAIChatChunk { choices?: OpenAIChatChunkChoice[]; + usage?: { + prompt_tokens: number; + completion_tokens: number; + total_tokens: number; + }; + timings?: Record; }