feat: track and display reasoning duration, enhance runtime with adapter for copying during inference and edit and UI integration

This commit is contained in:
shine1i 2026-02-01 09:08:23 +01:00
commit aeb5382f0a
5 changed files with 67 additions and 10 deletions

View file

@ -165,7 +165,7 @@ function ReasoningTrigger({
Thinking...
</AnimatedShinyText>
) : (
<span>Thought for {duration ?? 0}s</span>
<span>Thought for {duration ?? 0} seconds</span>
)}
</span>
<ChevronDownIcon
@ -276,6 +276,12 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
return lastIndex >= startIndex && lastIndex <= endIndex;
});
const persistedDuration = useAuiState(({ message }) => {
const d = (message.metadata?.custom as Record<string, unknown>)
?.reasoningDuration;
return typeof d === "number" ? d : 0;
});
const [manualOpen, setManualOpen] = useState(false);
const [duration, setDuration] = useState<number>(0);
const startTimeRef = useRef<number | null>(null);
@ -319,7 +325,7 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
>
<ReasoningTrigger
active={isReasoningStreaming}
duration={duration}
duration={duration || persistedDuration}
/>
<ReasoningContent
aria-busy={isReasoningStreaming}

View file

@ -19,6 +19,9 @@ import {
MessagePrimitive,
SuggestionPrimitive,
ThreadPrimitive,
useComposerRuntime,
useMessageRuntime,
useThreadRuntime,
} from "@assistant-ui/react";
import { motion } from "framer-motion";
import {
@ -333,12 +336,21 @@ const UserMessage: FC = () => {
const UserActionBar: FC = () => {
return (
<ActionBarPrimitive.Root
hideWhenRunning={true}
autohide="not-last"
className="aui-user-action-bar-root flex flex-col items-end"
className="aui-user-action-bar-root flex items-center"
>
<ActionBarPrimitive.Copy asChild={true}>
<TooltipIconButton tooltip="Copy">
<AuiIf condition={({ message }) => message.isCopied}>
<CheckIcon />
</AuiIf>
<AuiIf condition={({ message }) => !message.isCopied}>
<CopyIcon />
</AuiIf>
</TooltipIconButton>
</ActionBarPrimitive.Copy>
<ActionBarPrimitive.Edit asChild={true}>
<TooltipIconButton tooltip="Edit" className="aui-user-action-edit p-4">
<TooltipIconButton tooltip="Edit" className="aui-user-action-edit">
<PencilIcon />
</TooltipIconButton>
</ActionBarPrimitive.Edit>
@ -347,6 +359,10 @@ const UserActionBar: FC = () => {
};
const EditComposer: FC = () => {
const threadRuntime = useThreadRuntime();
const composerRuntime = useComposerRuntime();
const messageRuntime = useMessageRuntime();
return (
<MessagePrimitive.Root className="aui-edit-composer-wrapper mx-auto flex w-full max-w-(--thread-max-width) flex-col px-2 py-3">
<ComposerPrimitive.Root className="aui-edit-composer-root ml-auto flex w-full max-w-[85%] flex-col rounded-2xl bg-muted">
@ -360,9 +376,25 @@ const EditComposer: FC = () => {
Cancel
</Button>
</ComposerPrimitive.Cancel>
<ComposerPrimitive.Send asChild={true}>
<Button size="sm">Update</Button>
</ComposerPrimitive.Send>
<Button
size="sm"
onClick={() => {
const newText = composerRuntime.getState().text;
const originalText = messageRuntime.unstable_getCopyText();
if (newText === originalText) {
composerRuntime.cancel();
return;
}
if (threadRuntime.getState().isRunning) {
threadRuntime.cancelRun();
}
composerRuntime.send();
}}
>
Update
</Button>
</div>
</ComposerPrimitive.Root>
</MessagePrimitive.Root>

View file

@ -65,6 +65,9 @@ export function createStreamAdapter(apiUrl: string = API): ChatModelAdapter {
}
const decoder = new TextDecoder();
let text = "";
let reasoningStart: number | null = null;
let reasoningDuration = 0;
while (true) {
const { done, value } = await reader.read();
if (done) {
@ -72,8 +75,21 @@ export function createStreamAdapter(apiUrl: string = API): ChatModelAdapter {
}
text += decoder.decode(value, { stream: true });
const parts = parseThinkTags(text) ?? [];
if (parts.some((p) => p.type === "reasoning") && !reasoningStart) {
reasoningStart = Date.now();
}
if (text.includes("</think>") && reasoningStart && !reasoningDuration) {
reasoningDuration = Math.round(
(Date.now() - reasoningStart) / 1000,
);
}
if (parts.length > 0) {
yield { content: parts };
yield {
content: parts,
metadata: { custom: { reasoningDuration } },
};
}
}
},

View file

@ -52,7 +52,7 @@ function toThreadMessage(m: MessageRecord): ThreadMessage {
role: "assistant" as const,
status: { type: "complete" as const, reason: "unknown" as const },
metadata: {
custom: {},
custom: (m.metadata as Record<string, unknown>) ?? {},
steps: [],
unstable_annotations: [],
unstable_data: [],
@ -183,11 +183,13 @@ function ThreadHistoryProvider({
const content = Array.isArray(message.content)
? JSON.parse(JSON.stringify(message.content))
: [];
const custom = message.metadata?.custom;
await db.messages.put({
id: message.id,
threadId: remoteId,
role: message.role,
content,
...(custom && Object.keys(custom).length > 0 && { metadata: custom }),
createdAt: message.createdAt?.getTime() ?? Date.now(),
});
},

View file

@ -18,6 +18,7 @@ export interface MessageRecord {
threadId: string;
role: import("@assistant-ui/react").ThreadMessage["role"];
content: import("@assistant-ui/react").ThreadMessage["content"];
metadata?: Record<string, unknown>;
createdAt: number;
}