apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-07-22 23:31:12 +02:00
Code Issues Projects Releases Packages Wiki Activity

feat: integrate token usage context and enhance message info display in chat UI

Browse source
This commit is contained in:
Anish Sarkar 2026-04-14 14:47:59 +05:30
parent 5af6005163
commit 5510c1de03
3 changed files with 145 additions and 67 deletions
Download patch file Download diff file Expand all files Collapse all files

73
surfsense_web/components/assistant-ui/assistant-message.tsx
View file

@ -45,12 +45,14 @@ import {
DropdownMenuContent,
DropdownMenuItem,
DropdownMenuLabel,
DropdownMenuSeparator,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Button } from "@/components/ui/button";
import { useComments } from "@/hooks/use-comments";
import { useMediaQuery } from "@/hooks/use-media-query";
import { useElectronAPI } from "@/hooks/use-platform";
import { useTokenUsage } from "@/components/assistant-ui/token-usage-context";
import { cn } from "@/lib/utils";
// Captured once at module load — survives client-side navigations that strip the query param.
@ -375,22 +377,24 @@ export const MessageError: FC = () => {
);
};
const TokenUsageDropdown: FC = () => {
const usage = useAuiState(({ message }) => {
const custom = message?.metadata?.custom as Record<string, unknown> | undefined;
return custom?.usage as Record<string, unknown> | undefined;
function formatMessageDate(date: Date): string {
return date.toLocaleDateString(undefined, {
month: "short",
day: "numeric",
hour: "numeric",
minute: "2-digit",
hour12: true,
});
}
if (!usage) return null;
const totalTokens = (usage.total_tokens as number) ?? 0;
if (totalTokens === 0) return null;
const modelBreakdown = (usage.usage ?? usage.model_breakdown) as
| Record<string, { prompt_tokens: number; completion_tokens: number; total_tokens: number }>
| undefined;
const MessageInfoDropdown: FC = () => {
const messageId = useAuiState(({ message }) => message?.id);
const createdAt = useAuiState(({ message }) => message?.createdAt);
const usage = useTokenUsage(messageId);
const modelBreakdown = usage ? (usage.usage ?? usage.model_breakdown) : undefined;
const models = modelBreakdown ? Object.entries(modelBreakdown) : [];
const hasUsage = usage && usage.total_tokens > 0;
return (
<DropdownMenu>
@ -401,24 +405,31 @@ const TokenUsageDropdown: FC = () => {
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent align="start" className="min-w-[180px]">
<DropdownMenuLabel className="text-xs text-muted-foreground font-normal">
Token Usage
</DropdownMenuLabel>
{models.length > 0 ? (
models.map(([model, counts]) => (
<DropdownMenuItem key={model} className="flex-col items-start gap-0.5 cursor-default" onSelect={(e) => e.preventDefault()}>
<span className="text-xs font-medium">{model}</span>
<span className="text-xs text-muted-foreground">
{counts.total_tokens.toLocaleString()} tokens
</span>
</DropdownMenuItem>
))
) : (
<DropdownMenuItem className="flex-col items-start gap-0.5 cursor-default" onSelect={(e) => e.preventDefault()}>
<span className="text-xs text-muted-foreground">
{totalTokens.toLocaleString()} tokens
</span>
</DropdownMenuItem>
{createdAt && (
<DropdownMenuLabel className="text-xs text-muted-foreground font-normal select-none">
{formatMessageDate(createdAt)}
</DropdownMenuLabel>
)}
{hasUsage && (
<>
<DropdownMenuSeparator />
{models.length > 0 ? (
models.map(([model, counts]) => (
<DropdownMenuItem key={model} className="flex-col items-start gap-0.5 cursor-default" onSelect={(e) => e.preventDefault()}>
<span className="text-xs font-medium">{model}</span>
<span className="text-xs text-muted-foreground">
{counts.total_tokens.toLocaleString()} tokens
</span>
</DropdownMenuItem>
))
) : (
<DropdownMenuItem className="flex-col items-start gap-0.5 cursor-default" onSelect={(e) => e.preventDefault()}>
<span className="text-xs text-muted-foreground">
{usage.total_tokens.toLocaleString()} tokens
</span>
</DropdownMenuItem>
)}
</>
)}
</DropdownMenuContent>
</DropdownMenu>
@ -683,7 +694,7 @@ const AssistantActionBar: FC = () => {
<ClipboardPaste />
</TooltipIconButton>
)}
<TokenUsageDropdown />
<MessageInfoDropdown />
</ActionBarPrimitive.Root>
);
};

79
surfsense_web/components/assistant-ui/token-usage-context.tsx Normal file
View file

@ -0,0 +1,79 @@
"use client";
import { createContext, useContext, useCallback, useSyncExternalStore, type FC, type ReactNode } from "react";
export interface TokenUsageData {
prompt_tokens: number;
completion_tokens: number;
total_tokens: number;
usage?: Record<string, { prompt_tokens: number; completion_tokens: number; total_tokens: number }>;
model_breakdown?: Record<string, { prompt_tokens: number; completion_tokens: number; total_tokens: number }>;
}
type Listener = () => void;
class TokenUsageStore {
private data = new Map<string, TokenUsageData>();
private listeners = new Set<Listener>();
get(messageId: string): TokenUsageData | undefined {
return this.data.get(messageId);
}
set(messageId: string, usage: TokenUsageData): void {
this.data.set(messageId, usage);
this.notify();
}
rename(oldId: string, newId: string): void {
const usage = this.data.get(oldId);
if (usage) {
this.data.delete(oldId);
this.data.set(newId, usage);
this.notify();
}
}
clear(): void {
this.data.clear();
this.notify();
}
subscribe = (listener: Listener): (() => void) => {
this.listeners.add(listener);
return () => this.listeners.delete(listener);
};
private notify(): void {
for (const l of this.listeners) l();
}
}
const TokenUsageContext = createContext<TokenUsageStore | null>(null);
export const TokenUsageProvider: FC<{ store: TokenUsageStore; children: ReactNode }> = ({ store, children }) => (
<TokenUsageContext.Provider value={store}>{children}</TokenUsageContext.Provider>
);
export function useTokenUsageStore(): TokenUsageStore {
const store = useContext(TokenUsageContext);
if (!store) throw new Error("useTokenUsageStore must be used within TokenUsageProvider");
return store;
}
export function useTokenUsage(messageId: string | undefined): TokenUsageData | undefined {
const store = useContext(TokenUsageContext);
const getSnapshot = useCallback(
() => (store && messageId ? store.get(messageId) : undefined),
[store, messageId],
);
const subscribe = useCallback(
(onStoreChange: () => void) => (store ? store.subscribe(onStoreChange) : () => {}),
[store],
);
return useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
}
export function createTokenUsageStore(): TokenUsageStore {
return new TokenUsageStore();
}