feat(chat): server-side timings, context display & source hover cards (#4467)
* feat(chat): add server-side timings and context display for GGUF Extract timings/usage metadata from llama-server SSE stream and forward through the full stack. Replace client-side estimates with accurate server-reported metrics (prompt eval, tok/s, token counts, cache hits). Add context window usage bar to chat top nav. * feat(chat): source badges with hover cards and 2-row collapse - Add hover cards to source badges showing favicon, title, URL and snippet description on hover - Limit source badges to 2 rows with +X more expand/collapse - Parse snippet from web search results for hover card descriptions - Replace individual Source rendering with grouped SourcesGroup component * fix(chat): add null guards for server timings edge cases * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(chat): reset contextUsage on thread switch, remove unused context-display * fix(chat): stop double-counting completion tokens in tool-calling path * fix(chat): skip metadata events in llm_assist consumers * fix(chat): hide context usage bar in compare mode * fix(chat): harden timings pipeline and context usage persistence Accumulate prompt_ms, predicted_ms, and predicted_n from intermediate tool-detection passes so the final metadata reflects total server work. Persist contextUsage in message metadata (Dexie) and restore on thread load. Add type guard in gguf_stream_chunks for unexpected dict events. Clear contextUsage when entering compare mode. * feat(chat): make GGUF stream metadata OpenAI-compatible * fix(chat): address PR review feedback * feat(chat): address PR review feedback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Wasim Yousef Said
GitHub
parent
dd283b0605
commit
50cccfd55e
17 changed files with 840 additions and 91 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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 ───────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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=="],
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
*
|
||||
* <ActionBarPrimitive.Root >
|
||||
* <ActionBarPrimitive.Copy />
|
||||
* <ActionBarPrimitive.Reload />
|
||||
* <MessageTiming /> // <-- add this
|
||||
* </ActionBarPrimitive.Root>
|
||||
* ```
|
||||
*
|
||||
* @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<string, unknown> | undefined
|
||||
)?.custom as { serverTimings?: Record<string, number> } | 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 (
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild>
|
||||
|
|
@ -53,7 +55,7 @@ export const MessageTiming: FC<{
|
|||
className,
|
||||
)}
|
||||
>
|
||||
{formatTimingMs(timing.totalStreamTime)}
|
||||
{badgeText}
|
||||
</button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent
|
||||
|
|
@ -62,25 +64,97 @@ export const MessageTiming: FC<{
|
|||
data-slot="message-timing-popover"
|
||||
className="[&_span>svg]:hidden! rounded-lg border bg-popover px-3 py-2 text-popover-foreground shadow-md"
|
||||
>
|
||||
<div className="grid min-w-35 gap-1.5 text-xs">
|
||||
{timing.firstTokenTime !== undefined && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">First token</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTimingMs(timing.firstTokenTime)}
|
||||
</span>
|
||||
</div>
|
||||
<div className="grid min-w-40 gap-1.5 text-xs">
|
||||
{st ? (
|
||||
<>
|
||||
{/* Server-side metrics (GGUF) */}
|
||||
{st?.prompt_ms != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Prompt eval</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTimingMs(st.prompt_ms)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{(st?.prompt_n ?? 0) > 1 && st?.prompt_per_second != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Prompt speed</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{st.prompt_per_second.toFixed(1)} tok/s
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{st?.predicted_ms != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Generation</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTimingMs(st.predicted_ms)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{st?.predicted_per_second != null && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Speed</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{st.predicted_per_second.toFixed(1)} tok/s
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{timing.tokenCount !== undefined && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Tokens</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatNumber(timing.tokenCount)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{(st?.cache_n ?? 0) > 0 && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Cache hits</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatNumber(st!.cache_n)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
<div className="my-0.5 border-t border-border/40" />
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Total</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTimingMs(timing.totalStreamTime)}
|
||||
</span>
|
||||
</div>
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Chunks</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{timing.totalChunks}
|
||||
</span>
|
||||
</div>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
{/* Client-side metrics (safetensors fallback) */}
|
||||
{timing.firstTokenTime !== undefined && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">First token</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTimingMs(timing.firstTokenTime)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Total</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTimingMs(timing.totalStreamTime)}
|
||||
</span>
|
||||
</div>
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Chunks</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{timing.totalChunks}
|
||||
</span>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Total</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTimingMs(timing.totalStreamTime)}
|
||||
</span>
|
||||
</div>
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Chunks</span>
|
||||
<span className="font-mono tabular-nums">{timing.totalChunks}</span>
|
||||
</div>
|
||||
</div>
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<span
|
||||
data-slot="source-icon-fallback"
|
||||
className={cn(
|
||||
"flex size-3 shrink-0 items-center justify-center rounded-sm bg-muted font-medium text-[10px]",
|
||||
`flex ${sizeClass} shrink-0 items-center justify-center rounded-sm bg-muted font-medium text-[10px]`,
|
||||
className,
|
||||
)}
|
||||
{...props}
|
||||
|
|
@ -46,7 +65,7 @@ function SourceIcon({
|
|||
data-slot="source-icon"
|
||||
src={`https://www.google.com/s2/favicons?domain=${domain}&sz=32`}
|
||||
alt=""
|
||||
className={cn("size-3 shrink-0 rounded-sm", className)}
|
||||
className={cn(`${sizeClass} shrink-0 rounded-sm`, className)}
|
||||
onError={() => 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 (
|
||||
<span className="mr-1 mt-1 inline-block first:mt-2">
|
||||
<Source href={url}>
|
||||
<SourceIcon url={url} />
|
||||
<SourceTitle>{displayTitle}</SourceTitle>
|
||||
</Source>
|
||||
</span>
|
||||
<HoverCard openDelay={300} closeDelay={100}>
|
||||
<HoverCardTrigger asChild>
|
||||
<span className="inline-block">
|
||||
<Source href={source.url}>
|
||||
<SourceIcon url={source.url} />
|
||||
<SourceTitle>{displayTitle}</SourceTitle>
|
||||
</Source>
|
||||
</span>
|
||||
</HoverCardTrigger>
|
||||
<HoverCardContent side="top" align="start" className="w-72 p-3">
|
||||
<div className="flex gap-2.5">
|
||||
<SourceIcon url={source.url} size={4} className="mt-0.5 shrink-0" />
|
||||
<div className="min-w-0 space-y-1">
|
||||
<p className="text-sm font-semibold leading-tight truncate">
|
||||
{source.title || domain}
|
||||
</p>
|
||||
<p className="text-xs text-muted-foreground truncate">{domain}</p>
|
||||
{source.description && (
|
||||
<p className="text-xs text-muted-foreground leading-relaxed line-clamp-3">
|
||||
{source.description}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</HoverCardContent>
|
||||
</HoverCard>
|
||||
);
|
||||
};
|
||||
|
||||
const Sources = memo(SourcesImpl) as unknown as SourceMessagePartComponent & {
|
||||
// ── Grouped sources with 2-row collapse ─────────────────────
|
||||
|
||||
const SourcesGroup: FC = () => {
|
||||
const message = useMessage();
|
||||
const containerRef = useRef<HTMLDivElement>(null);
|
||||
const [visibleCount, setVisibleCount] = useState<number | null>(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 (
|
||||
<div className="relative mt-2">
|
||||
{/* Hidden measurement container — renders all badges to measure row positions */}
|
||||
<div
|
||||
ref={containerRef}
|
||||
aria-hidden
|
||||
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>
|
||||
))}
|
||||
</div>
|
||||
|
||||
{/* Visible container */}
|
||||
<div className="flex flex-wrap gap-1">
|
||||
{displayedSources.map((source) => (
|
||||
<SourceBadge key={source.url} source={source} />
|
||||
))}
|
||||
{shouldCollapse && !expanded && (
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setExpanded(true)}
|
||||
className={cn(
|
||||
badgeVariants({ variant: "outline", size: "default" }),
|
||||
"cursor-pointer text-muted-foreground hover:text-foreground",
|
||||
)}
|
||||
>
|
||||
+{hiddenCount} more
|
||||
</button>
|
||||
)}
|
||||
{shouldCollapse && expanded && (
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setExpanded(false)}
|
||||
className={cn(
|
||||
badgeVariants({ variant: "outline", size: "default" }),
|
||||
"cursor-pointer text-muted-foreground hover:text-foreground",
|
||||
)}
|
||||
>
|
||||
Show less
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
// ── Individual source (renders null — SourcesGroup handles all) ──
|
||||
|
||||
const SourcesNoop: FC<Record<string, unknown>> = () => null;
|
||||
|
||||
// ── Exports ──────────────────────────────────────────────────
|
||||
|
||||
const Sources = memo(SourcesNoop) as unknown as FC<Record<string, unknown>> & {
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -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<string, Array<"thinking" | "search" | "code">> = {
|
||||
"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 (
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
if (!isDisabled && !isRunning) {
|
||||
const store = useChatRuntimeStore.getState();
|
||||
if (store.supportsReasoning) {
|
||||
store.setReasoningEnabled(tools.includes("thinking"));
|
||||
}
|
||||
if (store.supportsTools) {
|
||||
store.setToolsEnabled(tools.includes("search"));
|
||||
store.setCodeToolsEnabled(tools.includes("code"));
|
||||
}
|
||||
aui.thread().append(prompt);
|
||||
aui.composer().setText("");
|
||||
return;
|
||||
}
|
||||
aui.composer().setText(prompt);
|
||||
}}
|
||||
className="fade-in slide-in-from-bottom-1 animate-in cursor-pointer corner-squircle rounded-xl border bg-background px-4 py-2.5 text-left text-sm text-foreground shadow-sm transition-colors duration-150 hover:bg-accent"
|
||||
className="fade-in slide-in-from-bottom-1 animate-in relative cursor-pointer corner-squircle rounded-xl border bg-background px-4 py-2.5 pr-12 text-left text-sm text-foreground shadow-sm transition-colors duration-150 hover:bg-accent"
|
||||
>
|
||||
<SuggestionPrimitive.Title />
|
||||
{tools.length > 0 && (
|
||||
<div className="absolute bottom-2.5 right-3 flex items-center gap-1">
|
||||
{tools.map((tool) => {
|
||||
const { icon: Icon, label } = toolIconMap[tool];
|
||||
return (
|
||||
<Icon
|
||||
key={tool}
|
||||
className="size-3 text-muted-foreground/60"
|
||||
aria-label={label}
|
||||
/>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
)}
|
||||
</button>
|
||||
);
|
||||
};
|
||||
|
|
@ -517,6 +559,7 @@ const AssistantMessage: FC = () => {
|
|||
},
|
||||
}}
|
||||
/>
|
||||
<SourcesGroup />
|
||||
<MessageError />
|
||||
</div>
|
||||
|
||||
|
|
|
|||
|
|
@ -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 ? (
|
||||
<LoaderIcon
|
||||
data-slot="tool-group-trigger-loader"
|
||||
className="aui-tool-group-trigger-loader size-4 shrink-0 animate-spin"
|
||||
/>
|
||||
) : (
|
||||
<HugeiconsIcon
|
||||
icon={Wrench01Icon}
|
||||
data-slot="tool-group-trigger-wrench"
|
||||
className="size-4 shrink-0 text-muted-foreground"
|
||||
strokeWidth={2}
|
||||
/>
|
||||
)}
|
||||
<span
|
||||
data-slot="tool-group-trigger-label"
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { HoverCard as HoverCardPrimitive } from "radix-ui";
|
||||
import type * as React from "react";
|
||||
|
||||
|
|
@ -33,7 +33,7 @@ function HoverCardContent({
|
|||
align={align}
|
||||
sideOffset={sideOffset}
|
||||
className={cn(
|
||||
"data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 ring-foreground/5 bg-popover text-popover-foreground w-72 rounded-2xl p-4 text-sm shadow-2xl ring-1 duration-100 z-50 origin-(--radix-hover-card-content-transform-origin) outline-hidden",
|
||||
"data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 ring-foreground/5 bg-popover text-popover-foreground w-72 corner-squircle rounded-2xl p-4 text-sm shadow-2xl ring-1 duration-100 z-50 origin-(--radix-hover-card-content-transform-origin) outline-hidden",
|
||||
className,
|
||||
)}
|
||||
{...props}
|
||||
|
|
|
|||
|
|
@ -20,6 +20,26 @@ import {
|
|||
parseAssistantContent,
|
||||
} from "../utils/parse-assistant-content";
|
||||
|
||||
/** Server-side usage data from llama-server (via stream_options.include_usage). */
|
||||
interface ServerUsage {
|
||||
prompt_tokens: number;
|
||||
completion_tokens: number;
|
||||
total_tokens: number;
|
||||
}
|
||||
|
||||
/** Server-side timing data from llama-server's timings object. */
|
||||
interface ServerTimings {
|
||||
prompt_n: number;
|
||||
cache_n: number;
|
||||
prompt_ms: number;
|
||||
prompt_per_token_ms: number;
|
||||
prompt_per_second: number;
|
||||
predicted_n: number;
|
||||
predicted_ms: number;
|
||||
predicted_per_token_ms: number;
|
||||
predicted_per_second: number;
|
||||
}
|
||||
|
||||
type RunMessages = Parameters<ChatModelAdapter["run"]>[0]["messages"];
|
||||
type RunMessage = RunMessages[number];
|
||||
|
||||
|
|
@ -27,21 +47,24 @@ type RunMessage = RunMessages[number];
|
|||
export const sentAudioNames = new Map<string, string>();
|
||||
|
||||
/** 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<string, unknown>).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) {
|
||||
|
|
|
|||
|
|
@ -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<string, unknown> | 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 {
|
|||
</div>
|
||||
)}
|
||||
<div className="flex-1" />
|
||||
{view.mode === "single" && ggufContextLength && contextUsage ? (
|
||||
<ContextUsageBar
|
||||
used={contextUsage.totalTokens}
|
||||
total={ggufContextLength}
|
||||
cached={contextUsage.cachedTokens}
|
||||
promptTokens={contextUsage.promptTokens}
|
||||
completionTokens={contextUsage.completionTokens}
|
||||
/>
|
||||
) : null}
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setSettingsOpen((o) => !o)}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,111 @@
|
|||
"use client";
|
||||
|
||||
import {
|
||||
Tooltip,
|
||||
TooltipContent,
|
||||
TooltipTrigger,
|
||||
} from "@/components/ui/tooltip";
|
||||
import { cn } from "@/lib/utils";
|
||||
import type { FC } from "react";
|
||||
|
||||
const formatTokenCount = (n: number): string => {
|
||||
if (n >= 1000) return `${(n / 1000).toFixed(1)}k`;
|
||||
return String(n);
|
||||
};
|
||||
|
||||
const formatTokenCountFull = (n: number): string => {
|
||||
return n.toLocaleString();
|
||||
};
|
||||
|
||||
function getSeverityColor(percent: number): {
|
||||
bar: string;
|
||||
text: string;
|
||||
} {
|
||||
if (percent > 85) return { bar: "bg-red-500", text: "text-red-500" };
|
||||
if (percent > 65) return { bar: "bg-amber-500", text: "text-amber-500" };
|
||||
return { bar: "bg-emerald-500", text: "text-emerald-500" };
|
||||
}
|
||||
|
||||
export const ContextUsageBar: FC<{
|
||||
used: number;
|
||||
total: number;
|
||||
cached?: number;
|
||||
promptTokens?: number;
|
||||
completionTokens?: number;
|
||||
className?: string;
|
||||
}> = ({ used, total, cached, promptTokens, completionTokens, className }) => {
|
||||
if (total <= 0) return null;
|
||||
|
||||
const percent = Math.min((used / total) * 100, 100);
|
||||
const severity = getSeverityColor(percent);
|
||||
|
||||
return (
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild>
|
||||
<button
|
||||
type="button"
|
||||
aria-label={`Context usage: ${formatTokenCount(used)} of ${formatTokenCount(total)} tokens`}
|
||||
className={cn(
|
||||
"flex items-center gap-2 rounded-md px-2 py-1 text-xs font-mono tabular-nums text-muted-foreground transition-colors hover:bg-accent hover:text-accent-foreground",
|
||||
className,
|
||||
)}
|
||||
>
|
||||
<span>
|
||||
{formatTokenCount(used)} / {formatTokenCount(total)}
|
||||
</span>
|
||||
<div className="h-1.5 w-16 rounded-full bg-muted overflow-hidden">
|
||||
<div
|
||||
className={cn("h-full rounded-full transition-all", severity.bar)}
|
||||
style={{ width: `${percent}%` }}
|
||||
/>
|
||||
</div>
|
||||
</button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent
|
||||
side="bottom"
|
||||
sideOffset={8}
|
||||
className="[&_span>svg]:hidden! rounded-lg border bg-popover px-3 py-2 text-popover-foreground shadow-md"
|
||||
>
|
||||
<div className="grid min-w-44 gap-1.5 text-xs">
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Context usage</span>
|
||||
<span className={cn("font-mono tabular-nums font-medium", severity.text)}>
|
||||
{percent.toFixed(1)}%
|
||||
</span>
|
||||
</div>
|
||||
{promptTokens !== undefined && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Prompt tokens</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTokenCountFull(promptTokens)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{completionTokens !== undefined && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Completion</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTokenCountFull(completionTokens)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{cached !== undefined && cached > 0 && (
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Cache hits</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTokenCountFull(cached)}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
<div className="my-0.5 border-t border-border/40" />
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="text-muted-foreground">Total</span>
|
||||
<span className="font-mono tabular-nums">
|
||||
{formatTokenCountFull(used)} / {formatTokenCountFull(total)}
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
);
|
||||
};
|
||||
|
|
@ -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<string, unknown>) ?? {};
|
||||
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<ThreadMessage, { role: "assistant" }>["content"],
|
||||
status: { type: "complete" as const, reason: "unknown" as const },
|
||||
metadata: {
|
||||
custom: (m.metadata as Record<string, unknown>) ?? {},
|
||||
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<string, unknown>)?.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));
|
||||
},
|
||||
|
||||
|
|
|
|||
|
|
@ -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<ChatRuntimeStore>((set) => ({
|
||||
|
|
@ -134,6 +141,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((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<ChatRuntimeStore>((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<ChatRuntimeStore>((set) => ({
|
|||
},
|
||||
activeGgufVariant: null,
|
||||
ggufContextLength: null,
|
||||
contextUsage: null,
|
||||
supportsReasoning: false,
|
||||
reasoningEnabled: true,
|
||||
supportsTools: false,
|
||||
|
|
@ -208,4 +217,5 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
set({ pendingAudioBase64: base64, pendingAudioName: name }),
|
||||
clearPendingAudio: () =>
|
||||
set({ pendingAudioBase64: null, pendingAudioName: null }),
|
||||
setContextUsage: (contextUsage) => set({ contextUsage }),
|
||||
}));
|
||||
|
|
|
|||
|
|
@ -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({
|
|||
</SidebarGroupContent>
|
||||
</SidebarGroup>
|
||||
</SidebarContent>
|
||||
<SidebarFooter className="space-y-1 px-4 pb-3">
|
||||
<a
|
||||
href="https://unsloth.ai/docs/new/studio/chat"
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="flex items-center gap-2 corner-squircle rounded-md px-2 py-1.5 text-xs font-medium text-primary bg-primary/10 transition-colors hover:bg-primary/20"
|
||||
>
|
||||
<HugeiconsIcon icon={BookOpen02Icon} className="size-4 shrink-0" strokeWidth={2} />
|
||||
<span>Learn more in docs</span>
|
||||
</a>
|
||||
<a
|
||||
href="https://unsloth.ai/blog"
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="flex items-center gap-2 rounded-md px-2 py-1.5 text-xs text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
|
||||
>
|
||||
<HugeiconsIcon icon={NewReleasesIcon} className="size-4 shrink-0" strokeWidth={2} />
|
||||
<span>What's new</span>
|
||||
</a>
|
||||
</SidebarFooter>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -175,4 +175,10 @@ export interface OpenAIChatChunkChoice {
|
|||
|
||||
export interface OpenAIChatChunk {
|
||||
choices?: OpenAIChatChunkChoice[];
|
||||
usage?: {
|
||||
prompt_tokens: number;
|
||||
completion_tokens: number;
|
||||
total_tokens: number;
|
||||
};
|
||||
timings?: Record<string, number>;
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue