From 0f7ed648cf069333c267c82e73d84add9b83885d Mon Sep 17 00:00:00 2001 From: Shine1i Date: Fri, 13 Feb 2026 16:45:00 +0100 Subject: [PATCH 1/4] feat: refactor chat runtime with modular APIs, state management, and runtime synchronization --- studio/frontend/src/features/chat/adapter.ts | 121 --------------- .../src/features/chat/api/chat-adapter.ts | 113 ++++++++++++++ .../src/features/chat/api/chat-api.ts | 139 +++++++++++++++++ .../frontend/src/features/chat/chat-page.tsx | 86 +++++------ .../src/features/chat/chat-settings-sheet.tsx | 60 +------- .../chat/hooks/use-chat-model-runtime.ts | 140 ++++++++++++++++++ studio/frontend/src/features/chat/index.ts | 2 + .../src/features/chat/runtime-provider.tsx | 20 +-- .../chat/stores/chat-runtime-store.ts | 40 +++++ .../frontend/src/features/chat/types/api.ts | 68 +++++++++ .../src/features/chat/types/runtime.ts | 27 ++++ .../chat/utils/parse-assistant-content.ts | 54 +++++++ 12 files changed, 630 insertions(+), 240 deletions(-) delete mode 100644 studio/frontend/src/features/chat/adapter.ts create mode 100644 studio/frontend/src/features/chat/api/chat-adapter.ts create mode 100644 studio/frontend/src/features/chat/api/chat-api.ts create mode 100644 studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts create mode 100644 studio/frontend/src/features/chat/stores/chat-runtime-store.ts create mode 100644 studio/frontend/src/features/chat/types/api.ts create mode 100644 studio/frontend/src/features/chat/types/runtime.ts create mode 100644 studio/frontend/src/features/chat/utils/parse-assistant-content.ts diff --git a/studio/frontend/src/features/chat/adapter.ts b/studio/frontend/src/features/chat/adapter.ts deleted file mode 100644 index 1c731cecaa..0000000000 --- a/studio/frontend/src/features/chat/adapter.ts +++ /dev/null @@ -1,121 +0,0 @@ -import type { ChatModelAdapter, ChatModelRunResult } from "@assistant-ui/react"; - -const API = import.meta.env.VITE_INFERENCE_URL || "/api/chat/generate"; -type ContentPart = NonNullable[number]; -type RunMessages = Parameters[0]["messages"]; -type RunMessage = RunMessages[number]; - -function collectTextParts(message: RunMessage): string[] { - const textParts = message.content - .filter((c) => c.type === "text") - .map((c) => c.text); - - if ("attachments" in message && (message.attachments?.length ?? 0) > 0) { - for (const att of message.attachments ?? []) { - for (const part of att.content ?? []) { - if (part.type === "text") { - textParts.push(part.text); - } - } - } - } - - return textParts; -} - -function messageToPayload(message: RunMessage): { - role: string; - content: string; -} { - return { - role: message.role, - content: collectTextParts(message).join("\n"), - }; -} - -function makeBody(messages: RunMessages): string { - const payloadMessages: Array<{ role: string; content: string }> = []; - for (const message of messages) { - payloadMessages.push(messageToPayload(message)); - } - return JSON.stringify({ messages: payloadMessages }); -} - -export function parseThinkTags(raw: string): ChatModelRunResult["content"] { - const parts: ContentPart[] = []; - const thinkStart = raw.indexOf(""); - if (thinkStart === -1) { - if (raw) { - parts.push({ type: "text", text: raw }); - } - return parts; - } - const before = raw.slice(0, thinkStart); - if (before.trim()) { - parts.push({ type: "text", text: before }); - } - - const thinkEnd = raw.indexOf(""); - if (thinkEnd === -1) { - const reasoning = raw.slice(thinkStart + 7); - if (reasoning) { - parts.push({ type: "reasoning", text: reasoning }); - } - return parts; - } - const reasoning = raw.slice(thinkStart + 7, thinkEnd); - if (reasoning) { - parts.push({ type: "reasoning", text: reasoning }); - } - - const after = raw.slice(thinkEnd + 8); - if (after) { - parts.push({ type: "text", text: after }); - } - return parts; -} - -export function createStreamAdapter(apiUrl: string = API): ChatModelAdapter { - return { - // biome-ignore lint/complexity/noExcessiveCognitiveComplexity: stream loop ok - async *run({ messages, abortSignal }) { - const res = await fetch(apiUrl, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: makeBody(messages), - signal: abortSignal, - }); - const reader = res.body?.getReader(); - if (!reader) { - throw new Error("Response body is empty"); - } - const decoder = new TextDecoder(); - let text = ""; - let reasoningStart: number | null = null; - let reasoningDuration = 0; - - while (true) { - const { done, value } = await reader.read(); - if (done) { - break; - } - text += decoder.decode(value, { stream: true }); - const parts = parseThinkTags(text) ?? []; - - if (parts.some((p) => p.type === "reasoning") && !reasoningStart) { - reasoningStart = Date.now(); - } - if (text.includes("") && reasoningStart && !reasoningDuration) { - reasoningDuration = Math.round((Date.now() - reasoningStart) / 1000); - } - - if (parts.length > 0) { - yield { - content: parts, - metadata: { custom: { reasoningDuration } }, - }; - } - } - }, - }; -} diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts new file mode 100644 index 0000000000..0c7bc827e3 --- /dev/null +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -0,0 +1,113 @@ +import type { ChatModelAdapter } from "@assistant-ui/react"; +import { streamChatCompletions } from "./chat-api"; +import { useChatRuntimeStore } from "../stores/chat-runtime-store"; +import { + hasClosedThinkTag, + parseAssistantContent, +} from "../utils/parse-assistant-content"; + +type RunMessages = Parameters[0]["messages"]; +type RunMessage = RunMessages[number]; + +function collectTextParts(message: RunMessage): string[] { + const textParts = message.content + .filter((part) => part.type === "text") + .map((part) => part.text); + + if ("attachments" in message && (message.attachments?.length ?? 0) > 0) { + for (const attachment of message.attachments ?? []) { + for (const part of attachment.content ?? []) { + if (part.type === "text") { + textParts.push(part.text); + } + } + } + } + + return textParts; +} + +function toOpenAIMessage(message: RunMessage): { + role: "system" | "user" | "assistant"; + content: string; +} | null { + if ( + message.role !== "system" && + message.role !== "user" && + message.role !== "assistant" + ) { + return null; + } + + return { + role: message.role, + content: collectTextParts(message).join("\n"), + }; +} + +export function createOpenAIStreamAdapter(): ChatModelAdapter { + return { + async *run({ messages, abortSignal }) { + const state = useChatRuntimeStore.getState(); + const { params } = state; + + if (!params.checkpoint) { + throw new Error("Load a model first."); + } + + const outboundMessages = messages + .map(toOpenAIMessage) + .filter((message): message is NonNullable => + Boolean(message), + ); + + if (params.systemPrompt.trim()) { + outboundMessages.unshift({ + role: "system", + content: params.systemPrompt.trim(), + }); + } + + const stream = streamChatCompletions( + { + model: params.checkpoint, + messages: outboundMessages, + stream: true, + temperature: params.temperature, + top_p: params.topP, + max_tokens: params.maxTokens, + top_k: params.topK, + repetition_penalty: params.repetitionPenalty, + }, + abortSignal, + ); + + let cumulativeText = ""; + let reasoningStartAt: number | null = null; + let reasoningDuration = 0; + + for await (const chunk of stream) { + const delta = chunk.choices?.[0]?.delta?.content; + if (!delta) { + continue; + } + cumulativeText += delta; + const parts = parseAssistantContent(cumulativeText); + + if (parts.some((part) => part.type === "reasoning") && !reasoningStartAt) { + reasoningStartAt = Date.now(); + } + if (hasClosedThinkTag(cumulativeText) && reasoningStartAt && !reasoningDuration) { + reasoningDuration = Math.round((Date.now() - reasoningStartAt) / 1000); + } + + if (parts.length > 0) { + yield { + content: parts, + metadata: { custom: { reasoningDuration } }, + }; + } + } + }, + }; +} diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts new file mode 100644 index 0000000000..cbd4e8cc6c --- /dev/null +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -0,0 +1,139 @@ +import { authFetch } from "@/features/auth"; +import type { + InferenceStatusResponse, + ListModelsResponse, + LoadModelRequest, + LoadModelResponse, + OpenAIChatCompletionsRequest, + OpenAIChatChunk, + UnloadModelRequest, +} from "../types/api"; + +function parseErrorText(status: number, body: unknown): string { + if ( + body && + typeof body === "object" && + "detail" in body && + typeof body.detail === "string" + ) { + return body.detail; + } + if ( + body && + typeof body === "object" && + "message" in body && + typeof body.message === "string" + ) { + return body.message; + } + return `Request failed (${status})`; +} + +async function parseJsonOrThrow(response: Response): Promise { + const body = await response.json().catch(() => null); + if (!response.ok) { + throw new Error(parseErrorText(response.status, body)); + } + return body as T; +} + +export async function listModels(): Promise { + const response = await authFetch("/api/models/list"); + return parseJsonOrThrow(response); +} + +export async function getInferenceStatus(): Promise { + const response = await authFetch("/api/inference/status"); + return parseJsonOrThrow(response); +} + +export async function loadModel( + payload: LoadModelRequest, +): Promise { + const response = await authFetch("/api/inference/load", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + }); + return parseJsonOrThrow(response); +} + +export async function unloadModel(payload: UnloadModelRequest): Promise { + const response = await authFetch("/api/inference/unload", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + }); + await parseJsonOrThrow(response); +} + +function parseSseEvent(rawEvent: string): string[] { + const dataLines: string[] = []; + for (const line of rawEvent.split(/\r?\n/)) { + if (line.startsWith("data:")) { + dataLines.push(line.slice(5).trimStart()); + } + } + return dataLines; +} + +export async function* streamChatCompletions( + payload: OpenAIChatCompletionsRequest, + signal: AbortSignal, +): AsyncGenerator { + const response = await authFetch("/api/inference/chat/completions", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + signal, + }); + + if (!response.ok) { + const body = await response.json().catch(() => null); + throw new Error(parseErrorText(response.status, body)); + } + + if (!response.body) { + throw new Error("Stream response missing body"); + } + + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + + while (true) { + const { done, value } = await reader.read(); + if (done) { + break; + } + + buffer += decoder.decode(value, { stream: true }); + + let separatorIndex = buffer.search(/\r?\n\r?\n/); + while (separatorIndex >= 0) { + const rawEvent = buffer.slice(0, separatorIndex); + const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2; + buffer = buffer.slice(separatorIndex + separatorLength); + + const dataLines = parseSseEvent(rawEvent); + if (dataLines.length === 0) { + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } + + const dataText = dataLines.join("\n"); + if (dataText === "[DONE]") { + return; + } + + const parsed = JSON.parse(dataText) as + | OpenAIChatChunk + | { error?: { message?: string } }; + if ("error" in parsed && parsed.error) { + throw new Error(parsed.error.message || "Stream error"); + } + yield parsed as OpenAIChatChunk; + separatorIndex = buffer.search(/\r?\n\r?\n/); + } + } +} diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index f29d5766ae..6d201f50de 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -28,16 +28,15 @@ import { memo, useCallback, useEffect, + useMemo, useRef, useState, } from "react"; -import { - ChatSettingsPanel, - type InferenceParams, - defaultInferenceParams, -} from "./chat-settings-sheet"; +import { ChatSettingsPanel } from "./chat-settings-sheet"; import { db } from "./db"; +import { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; import { ChatRuntimeProvider } from "./runtime-provider"; +import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import { type CompareHandle, CompareHandlesProvider, @@ -47,42 +46,6 @@ import { import { ThreadSidebar } from "./thread-sidebar"; import type { ChatView } from "./types"; -const LORA_MODELS: ModelOption[] = [ - { - id: "outputs/llama-3.1-8b-instruct-lora", - name: "meta-llama/Llama-3.1-8B-Instruct", - description: "LoRA v1", - }, - { - id: "outputs/qwen2.5-7b-lora", - name: "Qwen/Qwen2.5-7B-Instruct", - description: "LoRA v2", - }, - { - id: "outputs/mistral-7b-v0.3-lora", - name: "mistralai/Mistral-7B-Instruct-v0.3", - description: "LoRA v1", - }, -]; - -const GGUF_MODELS: ModelOption[] = [ - { - id: "models/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf", - name: "Meta-Llama-3.1-8B-Instruct", - description: "Q4_K_M", - }, - { - id: "models/Qwen2.5-7B-Instruct-Q5_K_M.gguf", - name: "Qwen2.5-7B-Instruct", - description: "Q5_K_M", - }, - { - id: "models/Mistral-7B-Instruct-v0.3-Q4_K_M.gguf", - name: "Mistral-7B-Instruct-v0.3", - description: "Q4_K_M", - }, -]; - const SingleContent = memo(function SingleContent({ threadId, }: { threadId?: string }): ReactElement { @@ -235,26 +198,40 @@ function TopBarActions({ export function ChatPage(): ReactElement { const [view, setView] = useState({ mode: "single" }); const [settingsOpen, setSettingsOpen] = useState(false); - const [inferenceParams, setInferenceParams] = useState( - defaultInferenceParams, - ); + const inferenceParams = useChatRuntimeStore((state) => state.params); + const setInferenceParams = useChatRuntimeStore((state) => state.setParams); + const modelsFromStore = useChatRuntimeStore((state) => state.models); + const modelsError = useChatRuntimeStore((state) => state.modelsError); + const { refresh, selectModel, ejectModel } = useChatModelRuntime(); const handleCheckpointChange = useCallback( - (v: string) => setInferenceParams((p) => ({ ...p, checkpoint: v })), - [], - ); - const handleEject = useCallback( - () => setInferenceParams((p) => ({ ...p, checkpoint: "" })), - [], + (value: string) => { + void selectModel(value); + }, + [selectModel], ); + const handleEject = useCallback(() => { + void ejectModel(); + }, [ejectModel]); const handleNewThread = useCallback(() => setView({ mode: "single" }), []); const handleNewCompare = useCallback( () => setView({ mode: "compare", pairId: crypto.randomUUID() }), [], ); - const models = - inferenceParams.inferenceEngine === "llama-cpp" ? GGUF_MODELS : LORA_MODELS; + const models = useMemo( + () => + modelsFromStore.map((model) => ({ + id: model.id, + name: model.name, + description: model.description, + })), + [modelsFromStore], + ); + + useEffect(() => { + void refresh(); + }, [refresh]); return ( + {modelsError && ( +
+ {modelsError} +
+ )}
- -
- - Backend - - -
- - state.params); + const models = useChatRuntimeStore((state) => state.models); + const setModels = useChatRuntimeStore((state) => state.setModels); + const setModelsError = useChatRuntimeStore((state) => state.setModelsError); + const setCheckpoint = useChatRuntimeStore((state) => state.setCheckpoint); + const clearCheckpoint = useChatRuntimeStore((state) => state.clearCheckpoint); + + const refresh = useCallback(async () => { + setModelsError(null); + try { + const [listRes, statusRes] = await Promise.all([ + listModels(), + getInferenceStatus(), + ]); + + const modelList = listRes.models.map(toChatModelSummary); + setModels(modelList); + + if (statusRes.active_model) { + setCheckpoint(statusRes.active_model); + } + } catch (error) { + const message = + error instanceof Error ? error.message : "Failed to load models"; + setModelsError(message); + } + }, [ + setCheckpoint, + setModels, + setModelsError, + ]); + + const selectModel = useCallback( + async (modelId: string) => { + if (!modelId || params.checkpoint === modelId) { + return; + } + const selected = models.find((model) => model.id === modelId); + if (!selected) { + setModelsError("Selected model was not found in model list."); + return; + } + + setModelsError(null); + try { + if (params.checkpoint) { + await unloadModel({ model_path: params.checkpoint }); + } + + await loadModel({ + model_path: selected.id, + hf_token: null, + max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH, + load_in_4bit: true, + is_lora: selected.isLora, + }); + + setCheckpoint(selected.id); + await refresh(); + } catch (error) { + const message = + error instanceof Error ? error.message : "Failed to load model"; + setModelsError(message); + } + }, + [ + models, + params.checkpoint, + refresh, + setCheckpoint, + setModelsError, + ], + ); + + const ejectModel = useCallback(async () => { + if (!params.checkpoint) { + return; + } + setModelsError(null); + try { + await unloadModel({ model_path: params.checkpoint }); + clearCheckpoint(); + await refresh(); + } catch (error) { + const message = + error instanceof Error ? error.message : "Failed to unload model"; + setModelsError(message); + } + }, [ + clearCheckpoint, + params.checkpoint, + refresh, + setModelsError, + ]); + + return { + refresh, + selectModel, + ejectModel, + }; +} diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index c1daee9e25..b7eaaf83ae 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -5,3 +5,5 @@ export { type InferenceParams, type Preset, } from "./chat-settings-sheet"; +export { useChatRuntimeStore } from "./stores/chat-runtime-store"; +export { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index afe948387e..dade143141 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -9,7 +9,6 @@ import { RuntimeAdapterProvider, SimpleImageAttachmentAdapter, SimpleTextAttachmentAdapter, - Suggestions, type ThreadHistoryAdapter, type ThreadMessage, type ThreadUserMessagePart, @@ -24,7 +23,7 @@ import { createAssistantStream } from "assistant-stream"; import mammoth from "mammoth"; import { type ReactElement, type ReactNode, useEffect, useMemo } from "react"; import { extractText, getDocumentProxy } from "unpdf"; -import { createStreamAdapter } from "./adapter"; +import { createOpenAIStreamAdapter } from "./api/chat-adapter"; import { db } from "./db"; import type { MessageRecord, ModelType } from "./types"; @@ -288,9 +287,11 @@ function ThreadHistoryProvider({ ); } -const chatAdapter = createStreamAdapter(); -const useRuntimeHook = (): ReturnType => - useLocalRuntime(chatAdapter); +const chatAdapter = createOpenAIStreamAdapter(); + +function useRuntimeHook(): ReturnType { + return useLocalRuntime(chatAdapter); +} function ThreadAutoSwitch({ threadId, @@ -327,14 +328,7 @@ export function ChatRuntimeProvider({ }, }); - const aui = useAui({ - suggestions: Suggestions([ - "Draw a simple flowchart of a login system using Mermaid", - "Solve the integral of x\u00B2\u00B7sin(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", - ]), - }); + const aui = useAui(); return ( diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts new file mode 100644 index 0000000000..94b4e3fdb3 --- /dev/null +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -0,0 +1,40 @@ +import { create } from "zustand"; +import { + DEFAULT_INFERENCE_PARAMS, + type ChatModelSummary, + type InferenceParams, +} from "../types/runtime"; + +type ChatRuntimeStore = { + params: InferenceParams; + models: ChatModelSummary[]; + modelsError: string | null; + setParams: (params: InferenceParams) => void; + setModels: (models: ChatModelSummary[]) => void; + setModelsError: (error: string | null) => void; + setCheckpoint: (modelId: string) => void; + clearCheckpoint: () => void; +}; + +export const useChatRuntimeStore = create((set) => ({ + params: DEFAULT_INFERENCE_PARAMS, + models: [], + modelsError: null, + setParams: (params) => set({ params }), + setModels: (models) => set({ models }), + setModelsError: (modelsError) => set({ modelsError }), + setCheckpoint: (modelId) => + set((state) => ({ + params: { + ...state.params, + checkpoint: modelId, + }, + })), + clearCheckpoint: () => + set((state) => ({ + params: { + ...state.params, + checkpoint: "", + }, + })), +})); diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts new file mode 100644 index 0000000000..6420020c06 --- /dev/null +++ b/studio/frontend/src/features/chat/types/api.ts @@ -0,0 +1,68 @@ +export interface BackendModelDetails { + id: string; + name?: string | null; + is_vision?: boolean; + is_lora?: boolean; +} + +export interface ListModelsResponse { + models: BackendModelDetails[]; + default_models: string[]; +} + +export interface LoadModelRequest { + model_path: string; + hf_token: string | null; + max_seq_length: number; + load_in_4bit: boolean; + is_lora: boolean; +} + +export interface LoadModelResponse { + status: string; + model: string; + display_name: string; + is_vision: boolean; + is_lora: boolean; +} + +export interface UnloadModelRequest { + model_path: string; +} + +export interface InferenceStatusResponse { + active_model: string | null; + is_vision: boolean; + loading: string[]; + loaded: string[]; +} + +export interface OpenAIChatMessage { + role: "system" | "user" | "assistant"; + content: string; +} + +export interface OpenAIChatCompletionsRequest { + model: string; + messages: OpenAIChatMessage[]; + stream: boolean; + temperature: number; + top_p: number; + max_tokens: number; + top_k: number; + repetition_penalty: number; +} + +export interface OpenAIChatDelta { + role?: string; + content?: string; +} + +export interface OpenAIChatChunkChoice { + delta?: OpenAIChatDelta; + finish_reason?: string | null; +} + +export interface OpenAIChatChunk { + choices?: OpenAIChatChunkChoice[]; +} diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts new file mode 100644 index 0000000000..67fa07eb14 --- /dev/null +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -0,0 +1,27 @@ +export interface InferenceParams { + temperature: number; + topP: number; + topK: number; + repetitionPenalty: number; + maxTokens: number; + systemPrompt: string; + checkpoint: string; +} + +export const DEFAULT_INFERENCE_PARAMS: InferenceParams = { + temperature: 0.7, + topP: 0.9, + topK: 50, + repetitionPenalty: 1.1, + maxTokens: 512, + systemPrompt: "", + checkpoint: "", +}; + +export interface ChatModelSummary { + id: string; + name: string; + description?: string; + isVision: boolean; + isLora: boolean; +} diff --git a/studio/frontend/src/features/chat/utils/parse-assistant-content.ts b/studio/frontend/src/features/chat/utils/parse-assistant-content.ts new file mode 100644 index 0000000000..fdf4ce051e --- /dev/null +++ b/studio/frontend/src/features/chat/utils/parse-assistant-content.ts @@ -0,0 +1,54 @@ +import type { ChatModelRunResult } from "@assistant-ui/react"; + +type ContentPart = NonNullable[number]; + +const THINK_OPEN_TAG = ""; +const THINK_CLOSE_TAG = ""; + +function appendTextPart(parts: ContentPart[], text: string): void { + if (text) { + parts.push({ type: "text", text }); + } +} + +function appendReasoningPart(parts: ContentPart[], text: string): void { + if (text) { + parts.push({ type: "reasoning", text }); + } +} + +export function parseAssistantContent( + raw: string, +): ContentPart[] { + const parts: ContentPart[] = []; + if (!raw) { + return parts; + } + + let cursor = 0; + while (cursor < raw.length) { + const openIndex = raw.indexOf(THINK_OPEN_TAG, cursor); + if (openIndex === -1) { + appendTextPart(parts, raw.slice(cursor)); + break; + } + + appendTextPart(parts, raw.slice(cursor, openIndex)); + + const reasoningStart = openIndex + THINK_OPEN_TAG.length; + const closeIndex = raw.indexOf(THINK_CLOSE_TAG, reasoningStart); + if (closeIndex === -1) { + appendReasoningPart(parts, raw.slice(reasoningStart)); + break; + } + + appendReasoningPart(parts, raw.slice(reasoningStart, closeIndex)); + cursor = closeIndex + THINK_CLOSE_TAG.length; + } + + return parts; +} + +export function hasClosedThinkTag(raw: string): boolean { + return raw.includes(THINK_CLOSE_TAG); +} From aaa8c1a816d6e02a519e33a2b268c7920ea973c9 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Fri, 13 Feb 2026 17:14:49 +0100 Subject: [PATCH 2/4] feat: integrate LoRA model management with UI and runtime synchronization --- studio/frontend/src/app/provider.tsx | 2 +- .../assistant-ui/model-selector.tsx | 436 ++++++++++++++---- .../src/features/chat/api/chat-api.ts | 9 +- .../frontend/src/features/chat/chat-page.tsx | 23 +- .../chat/hooks/use-chat-model-runtime.ts | 97 ++-- .../chat/stores/chat-runtime-store.ts | 5 + .../frontend/src/features/chat/types/api.ts | 11 + .../src/features/chat/types/runtime.ts | 7 + 8 files changed, 467 insertions(+), 123 deletions(-) diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index ce3af7ac54..5c2b93bcc9 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -10,7 +10,7 @@ export function AppProvider({ children }: AppProviderProps) { return ( {children} - + ); } diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 389dae424b..85bc1cacdb 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -1,14 +1,22 @@ "use client"; +import { Input } from "@/components/ui/input"; import { Popover, PopoverContent, PopoverTrigger, } from "@/components/ui/popover"; -import { cn } from "@/lib/utils"; -import { ArrowDown01Icon, Logout01Icon } from "@hugeicons/core-free-icons"; +import { Spinner } from "@/components/ui/spinner"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { useDebouncedValue, useHfModelSearch, useInfiniteScroll } from "@/hooks"; +import { cn, formatCompact } from "@/lib/utils"; +import { + ArrowDown01Icon, + Logout01Icon, + Search01Icon, +} from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { type ReactNode, useState } from "react"; +import { type ReactNode, useMemo, useState } from "react"; export interface ModelOption { id: string; @@ -17,11 +25,22 @@ export interface ModelOption { icon?: ReactNode; } +export interface LoraModelOption extends ModelOption { + baseModel?: string; + updatedAt?: number; +} + +export interface ModelSelectorChangeMeta { + source: "hub" | "lora"; + isLora: boolean; +} + interface ModelSelectorProps { models: ModelOption[]; + loraModels?: LoraModelOption[]; value?: string; defaultValue?: string; - onValueChange?: (value: string) => void; + onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void; onEject?: () => void; variant?: "outline" | "ghost" | "muted"; size?: "sm" | "default" | "lg"; @@ -29,7 +48,9 @@ interface ModelSelectorProps { contentClassName?: string; } -// --- Composable sub-components --- +function dedupe(values: string[]): string[] { + return [...new Set(values.filter(Boolean))]; +} function ModelSelectorTrigger({ currentModel, @@ -63,15 +84,11 @@ function ModelSelectorTrigger({ {isLoaded && ( )} - - {currentModel?.name ?? "Select a model\u2026"} + + {currentModel?.name ?? "Select model..."} {currentModel?.description && ( - - {currentModel.description} - + {currentModel.description} )} + {children} + + ); +} + +function ModelRow({ + label, + meta, + selected, + onClick, +}: { + label: string; + meta?: string; + selected?: boolean; + onClick: () => void; +}) { + return ( + + ); +} + +function HubModelPicker({ + models, + value, + onSelect, +}: { + models: ModelOption[]; + value?: string; + onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; +}) { + const [query, setQuery] = useState(""); + const debouncedQuery = useDebouncedValue(query); + const { results, isLoading, isLoadingMore, fetchMore } = useHfModelSearch( + debouncedQuery, + ); + + const recommendedIds = useMemo( + () => dedupe([...models.map((model) => model.id), value ?? ""]), + [models, value], + ); + + const showHfSection = debouncedQuery.trim().length > 0; + const recommendedSet = useMemo( + () => new Set(recommendedIds), + [recommendedIds], + ); + + const hfIds = useMemo(() => { + if (!showHfSection) { + return []; + } + return results + .map((result) => result.id) + .filter((id) => !recommendedSet.has(id)); + }, [recommendedSet, results, showHfSection]); + + const metricsById = useMemo( + () => + new Map( + results.map((result) => [ + result.id, + result.totalParams + ? formatCompact(result.totalParams) + : `↓${formatCompact(result.downloads)}`, + ]), + ), + [results], + ); + + const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length); + + return ( +
+
+ + setQuery(event.target.value)} + placeholder="Search Hugging Face models" + className="h-9 pl-8 pr-8" + /> + {isLoading && ( + + )} +
+ +
+
+ {!showHfSection ? ( + <> + Recommended + {recommendedIds.length === 0 ? ( +
+ No default models. +
+ ) : ( + recommendedIds.map((id) => ( + onSelect(id, { source: "hub", isLora: false })} + /> + )) + )} + + ) : null} + + {showHfSection ? ( + <> + Hugging Face + {hfIds.length === 0 && !isLoading ? ( +
+ No matching models. +
+ ) : ( + hfIds.map((id) => ( + onSelect(id, { source: "hub", isLora: false })} + /> + )) + )} +
+ {isLoadingMore ? ( +
+ +
+ ) : null} + + ) : null} +
+
+
+ ); +} + +function LoraModelPicker({ + loraModels, + value, + onSelect, +}: { + loraModels: LoraModelOption[]; + value?: string; + onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; +}) { + const [query, setQuery] = useState(""); + + const normalized = useMemo( + () => + loraModels + .map((model) => ({ + ...model, + baseModel: model.baseModel || model.description || "Unknown base model", + })) + .sort((a, b) => { + const aTime = a.updatedAt ?? -1; + const bTime = b.updatedAt ?? -1; + if (aTime !== bTime) { + return bTime - aTime; + } + const baseCmp = a.baseModel.localeCompare(b.baseModel); + if (baseCmp !== 0) { + return baseCmp; + } + return a.name.localeCompare(b.name); + }), + [loraModels], + ); + + const grouped = useMemo(() => { + const needle = query.trim().toLowerCase(); + const out = new Map(); + + for (const model of normalized) { + const searchText = `${model.name} ${model.baseModel} ${model.id}`.toLowerCase(); + if (needle && !searchText.includes(needle)) { + continue; + } + + const key = model.baseModel || "Unknown base model"; + const prev = out.get(key) ?? []; + prev.push(model); + out.set(key, prev); + } + + return [...out.entries()].sort((a, b) => { + const aLatest = Math.max(...a[1].map((model) => model.updatedAt ?? -1)); + const bLatest = Math.max(...b[1].map((model) => model.updatedAt ?? -1)); + if (aLatest !== bLatest) { + return bLatest - aLatest; + } + return a[0].localeCompare(b[0]); + }); + }, [normalized, query]); + + return ( +
+
+ + setQuery(event.target.value)} + placeholder="Search local adapters" + className="h-9 pl-8" + /> +
+ +
+
+ {grouped.length === 0 ? ( +
No adapters found.
+ ) : ( + grouped.map(([baseModel, adapters], index) => ( +
+ {index > 0 ?
: null} + {baseModel} + {adapters.map((adapter) => ( + onSelect(adapter.id, { source: "lora", isLora: true })} + /> + ))} +
+ )) + )} +
+
+
+ ); +} + function ModelSelectorContent({ models, + loraModels, value, onSelect, onEject, className, }: { models: ModelOption[]; + loraModels: LoraModelOption[]; value?: string; - onSelect: (id: string) => void; + onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; onEject?: () => void; className?: string; }) { + const hasSelection = Boolean(value); + return ( - {models.map((model) => ( - - ))} + + + Hub models + Fine-tuned + + + + + + + + + + + + {hasSelection && onEject ? ( +
+ +
+ ) : null} ); } -function ModelSelectorItem({ - model, - isActive, - onSelect, - onEject, -}: { - model: ModelOption; - isActive: boolean; - onSelect: (id: string) => void; - onEject?: () => void; -}) { - return ( - - )} - - ); -} - -// --- Main component --- - export function ModelSelector({ models, + loraModels = [], value, defaultValue, onValueChange, @@ -182,13 +433,31 @@ export function ModelSelector({ }: ModelSelectorProps) { const [open, setOpen] = useState(false); const [uncontrolled, setUncontrolled] = useState(defaultValue ?? ""); + const selected = value ?? uncontrolled; const isLoaded = selected !== ""; - const currentModel = models.find((m) => m.id === selected); - function handleSelect(id: string) { + const optionById = useMemo(() => { + const all = new Map(); + for (const model of models) { + all.set(model.id, model); + } + for (const lora of loraModels) { + all.set(lora.id, { + ...lora, + description: lora.baseModel || lora.description, + }); + } + return all; + }, [loraModels, models]); + + const currentModel = selected + ? optionById.get(selected) ?? { id: selected, name: selected } + : undefined; + + function handleSelect(id: string, meta: ModelSelectorChangeMeta) { if (onValueChange) { - onValueChange(id); + onValueChange(id, meta); } else { setUncontrolled(id); } @@ -211,6 +480,7 @@ export function ModelSelector({ /> { return parseJsonOrThrow(response); } +export async function listLoras(outputsDir = "./outputs"): Promise { + const query = new URLSearchParams({ outputs_dir: outputsDir }).toString(); + const response = await authFetch(`/api/models/loras?${query}`); + return parseJsonOrThrow(response); +} + export async function getInferenceStatus(): Promise { const response = await authFetch("/api/inference/status"); return parseJsonOrThrow(response); diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 6d201f50de..c787cf471c 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -1,4 +1,5 @@ import { + type LoraModelOption, type ModelOption, ModelSelector, } from "@/components/assistant-ui/model-selector"; @@ -201,12 +202,13 @@ export function ChatPage(): ReactElement { const inferenceParams = useChatRuntimeStore((state) => state.params); const setInferenceParams = useChatRuntimeStore((state) => state.setParams); const modelsFromStore = useChatRuntimeStore((state) => state.models); + const lorasFromStore = useChatRuntimeStore((state) => state.loras); const modelsError = useChatRuntimeStore((state) => state.modelsError); const { refresh, selectModel, ejectModel } = useChatModelRuntime(); const handleCheckpointChange = useCallback( - (value: string) => { - void selectModel(value); + (value: string, meta?: { isLora: boolean }) => { + void selectModel({ id: value, isLora: meta?.isLora }); }, [selectModel], ); @@ -229,6 +231,17 @@ export function ChatPage(): ReactElement { [modelsFromStore], ); + const loraModels = useMemo( + () => + lorasFromStore.map((lora) => ({ + id: lora.id, + name: lora.name, + baseModel: lora.baseModel, + updatedAt: lora.updatedAt, + })), + [lorasFromStore], + ); + useEffect(() => { void refresh(); }, [refresh]); @@ -263,6 +276,7 @@ export function ChatPage(): ReactElement { /> {view.mode === "single" ? ( - + ) : ( )} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index e0aaca9622..3eb9301302 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -1,15 +1,38 @@ import { useCallback } from "react"; +import { toast } from "sonner"; import { getInferenceStatus, + listLoras, listModels, loadModel, unloadModel, } from "../api/chat-api"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; -import type { ChatModelSummary } from "../types/runtime"; +import type { ChatLoraSummary, ChatModelSummary } from "../types/runtime"; const DEFAULT_MODEL_MAX_SEQ_LENGTH = 2048; +type SelectedModelInput = { + id: string; + isLora?: boolean; +}; + +const LORA_SUFFIX_RE = /_(\d{9,})$/; + +function parseTrailingEpoch(input: string): number | undefined { + const match = input.match(LORA_SUFFIX_RE); + if (!match) { + return undefined; + } + const parsed = Number.parseInt(match[1], 10); + return Number.isFinite(parsed) ? parsed : undefined; +} + +function stripTrailingEpoch(input: string): string { + const cleaned = input.replace(LORA_SUFFIX_RE, "").replace(/[_-]+$/, "").trim(); + return cleaned || input; +} + function describeModel(model: { is_lora?: boolean; is_vision?: boolean; @@ -36,10 +59,29 @@ function toChatModelSummary(model: { }; } +function toLoraSummary(lora: { + display_name: string; + adapter_path: string; + base_model?: string | null; +}): ChatLoraSummary { + const idTail = lora.adapter_path.split("/").filter(Boolean).at(-1) ?? ""; + const updatedAt = + parseTrailingEpoch(lora.display_name) ?? parseTrailingEpoch(idTail); + + return { + id: lora.adapter_path, + name: stripTrailingEpoch(lora.display_name), + baseModel: lora.base_model || "Unknown base model", + updatedAt, + }; +} + export function useChatModelRuntime() { const params = useChatRuntimeStore((state) => state.params); const models = useChatRuntimeStore((state) => state.models); + const loras = useChatRuntimeStore((state) => state.loras); const setModels = useChatRuntimeStore((state) => state.setModels); + const setLoras = useChatRuntimeStore((state) => state.setLoras); const setModelsError = useChatRuntimeStore((state) => state.setModelsError); const setCheckpoint = useChatRuntimeStore((state) => state.setCheckpoint); const clearCheckpoint = useChatRuntimeStore((state) => state.clearCheckpoint); @@ -47,13 +89,14 @@ export function useChatModelRuntime() { const refresh = useCallback(async () => { setModelsError(null); try { - const [listRes, statusRes] = await Promise.all([ + const [listRes, statusRes, lorasRes] = await Promise.all([ listModels(), getInferenceStatus(), + listLoras(), ]); - const modelList = listRes.models.map(toChatModelSummary); - setModels(modelList); + setModels(listRes.models.map(toChatModelSummary)); + setLoras(lorasRes.loras.map(toLoraSummary)); if (statusRes.active_model) { setCheckpoint(statusRes.active_model); @@ -63,22 +106,23 @@ export function useChatModelRuntime() { error instanceof Error ? error.message : "Failed to load models"; setModelsError(message); } - }, [ - setCheckpoint, - setModels, - setModelsError, - ]); + }, [setCheckpoint, setLoras, setModels, setModelsError]); const selectModel = useCallback( - async (modelId: string) => { + async (selection: string | SelectedModelInput) => { + const modelId = typeof selection === "string" ? selection : selection.id; if (!modelId || params.checkpoint === modelId) { return; } - const selected = models.find((model) => model.id === modelId); - if (!selected) { - setModelsError("Selected model was not found in model list."); - return; - } + + const explicitIsLora = + typeof selection === "string" ? undefined : selection.isLora; + const model = models.find((entry) => entry.id === modelId); + const lora = loras.find((entry) => entry.id === modelId); + const isLora = + explicitIsLora ?? model?.isLora ?? (lora ? true : false); + const displayName = model?.name || lora?.name || modelId; + const loadingToastId = toast.loading(`Loading ${displayName}...`); setModelsError(null); try { @@ -87,28 +131,24 @@ export function useChatModelRuntime() { } await loadModel({ - model_path: selected.id, + model_path: modelId, hf_token: null, max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH, load_in_4bit: true, - is_lora: selected.isLora, + is_lora: isLora, }); - setCheckpoint(selected.id); + setCheckpoint(modelId); await refresh(); + toast.success(`${displayName} loaded`, { id: loadingToastId }); } catch (error) { const message = error instanceof Error ? error.message : "Failed to load model"; setModelsError(message); + toast.error(message, { id: loadingToastId }); } }, - [ - models, - params.checkpoint, - refresh, - setCheckpoint, - setModelsError, - ], + [loras, models, params.checkpoint, refresh, setCheckpoint, setModelsError], ); const ejectModel = useCallback(async () => { @@ -125,12 +165,7 @@ export function useChatModelRuntime() { error instanceof Error ? error.message : "Failed to unload model"; setModelsError(message); } - }, [ - clearCheckpoint, - params.checkpoint, - refresh, - setModelsError, - ]); + }, [clearCheckpoint, params.checkpoint, refresh, setModelsError]); return { refresh, 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 94b4e3fdb3..76bf079ce7 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -1,6 +1,7 @@ import { create } from "zustand"; import { DEFAULT_INFERENCE_PARAMS, + type ChatLoraSummary, type ChatModelSummary, type InferenceParams, } from "../types/runtime"; @@ -8,9 +9,11 @@ import { type ChatRuntimeStore = { params: InferenceParams; models: ChatModelSummary[]; + loras: ChatLoraSummary[]; modelsError: string | null; setParams: (params: InferenceParams) => void; setModels: (models: ChatModelSummary[]) => void; + setLoras: (loras: ChatLoraSummary[]) => void; setModelsError: (error: string | null) => void; setCheckpoint: (modelId: string) => void; clearCheckpoint: () => void; @@ -19,9 +22,11 @@ type ChatRuntimeStore = { export const useChatRuntimeStore = create((set) => ({ params: DEFAULT_INFERENCE_PARAMS, models: [], + loras: [], modelsError: null, setParams: (params) => set({ params }), setModels: (models) => set({ models }), + setLoras: (loras) => set({ loras }), setModelsError: (modelsError) => set({ modelsError }), setCheckpoint: (modelId) => set((state) => ({ diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 6420020c06..9d21c854af 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -10,6 +10,17 @@ export interface ListModelsResponse { default_models: string[]; } +export interface BackendLoraInfo { + display_name: string; + adapter_path: string; + base_model?: string | null; +} + +export interface ListLorasResponse { + loras: BackendLoraInfo[]; + outputs_dir: string; +} + export interface LoadModelRequest { model_path: string; hf_token: string | null; diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts index 67fa07eb14..558f1a5464 100644 --- a/studio/frontend/src/features/chat/types/runtime.ts +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -25,3 +25,10 @@ export interface ChatModelSummary { isVision: boolean; isLora: boolean; } + +export interface ChatLoraSummary { + id: string; + name: string; + baseModel: string; + updatedAt?: number; +} From 76830db0cf16bd3d3c54700c69747398be097707 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Fri, 13 Feb 2026 17:28:01 +0100 Subject: [PATCH 3/4] feat: add warm-up indicator, new thread feature, and runtime improvements in chat UI --- .../src/components/assistant-ui/thread.tsx | 40 +++++++++- .../src/features/chat/api/chat-adapter.ts | 80 +++++++++++-------- .../frontend/src/features/chat/chat-page.tsx | 25 ++++-- .../src/features/chat/runtime-provider.tsx | 30 ++++++- .../chat/stores/chat-runtime-store.ts | 13 +++ studio/frontend/src/features/chat/types.ts | 2 +- 6 files changed, 149 insertions(+), 41 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 3bcaaea105..df77d3ea28 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -7,7 +7,9 @@ import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { AnimatedShinyText } from "@/components/ui/animated-shiny-text"; import { Button } from "@/components/ui/button"; +import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import { cn } from "@/lib/utils"; import { ActionBarMorePrimitive, @@ -20,6 +22,8 @@ import { SuggestionPrimitive, ThreadPrimitive, useAui, + useAuiEvent, + useAuiState, } from "@assistant-ui/react"; import { motion } from "framer-motion"; import { @@ -36,7 +40,7 @@ import { RefreshCwIcon, SquareIcon, } from "lucide-react"; -import type { FC } from "react"; +import { type FC, useRef } from "react"; export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ hideComposer, @@ -69,6 +73,7 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ + !thread.isEmpty}> {!hideComposer && } @@ -78,6 +83,28 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ ); }; +const WarmupIndicator: FC = () => { + const threadId = useAuiState(({ threads }) => threads.mainThreadId); + const isRunning = useAuiState(({ thread }) => thread.isRunning); + const isWarmingUp = useChatRuntimeStore((state) => + Boolean(state.warmingByThreadId[threadId ?? "__default"]), + ); + + if (!isRunning || !isWarmingUp) { + return null; + } + + return ( +
+
+ + Warming up model... + +
+
+ ); +}; + const ThreadScrollToBottom: FC = () => { return ( @@ -358,6 +385,15 @@ const UserActionBar: FC = () => { const EditComposer: FC = () => { const aui = useAui(); + const resendAfterCancelRef = useRef(false); + + useAuiEvent("thread.runEnd", () => { + if (!resendAfterCancelRef.current) { + return; + } + resendAfterCancelRef.current = false; + aui.composer().send(); + }); return ( @@ -384,7 +420,9 @@ const EditComposer: FC = () => { } if (aui.thread().getState().isRunning) { + resendAfterCancelRef.current = true; aui.thread().cancelRun(); + return; } aui.composer().send(); }} diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 0c7bc827e3..fa670ecda2 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -47,7 +47,7 @@ function toOpenAIMessage(message: RunMessage): { export function createOpenAIStreamAdapter(): ChatModelAdapter { return { - async *run({ messages, abortSignal }) { + async *run({ messages, abortSignal, unstable_threadId }) { const state = useChatRuntimeStore.getState(); const { params } = state; @@ -68,44 +68,58 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }); } - const stream = streamChatCompletions( - { - model: params.checkpoint, - messages: outboundMessages, - stream: true, - temperature: params.temperature, - top_p: params.topP, - max_tokens: params.maxTokens, - top_k: params.topK, - repetition_penalty: params.repetitionPenalty, - }, - abortSignal, - ); - + const threadKey = unstable_threadId || "__default"; + let waitingFirstChunk = true; + useChatRuntimeStore.getState().setThreadWarming(threadKey, true); let cumulativeText = ""; let reasoningStartAt: number | null = null; let reasoningDuration = 0; - for await (const chunk of stream) { - const delta = chunk.choices?.[0]?.delta?.content; - if (!delta) { - continue; - } - cumulativeText += delta; - const parts = parseAssistantContent(cumulativeText); + try { + const stream = streamChatCompletions( + { + model: params.checkpoint, + messages: outboundMessages, + stream: true, + temperature: params.temperature, + top_p: params.topP, + max_tokens: params.maxTokens, + top_k: params.topK, + repetition_penalty: params.repetitionPenalty, + }, + abortSignal, + ); - if (parts.some((part) => part.type === "reasoning") && !reasoningStartAt) { - reasoningStartAt = Date.now(); - } - if (hasClosedThinkTag(cumulativeText) && reasoningStartAt && !reasoningDuration) { - reasoningDuration = Math.round((Date.now() - reasoningStartAt) / 1000); - } + for await (const chunk of stream) { + const delta = chunk.choices?.[0]?.delta?.content; + if (!delta) { + continue; + } + if (waitingFirstChunk) { + waitingFirstChunk = false; + useChatRuntimeStore.getState().setThreadWarming(threadKey, false); + } - if (parts.length > 0) { - yield { - content: parts, - metadata: { custom: { reasoningDuration } }, - }; + cumulativeText += delta; + const parts = parseAssistantContent(cumulativeText); + + if (parts.some((part) => part.type === "reasoning") && !reasoningStartAt) { + reasoningStartAt = Date.now(); + } + if (hasClosedThinkTag(cumulativeText) && reasoningStartAt && !reasoningDuration) { + reasoningDuration = Math.round((Date.now() - reasoningStartAt) / 1000); + } + + if (parts.length > 0) { + yield { + content: parts, + metadata: { custom: { reasoningDuration } }, + }; + } + } + } finally { + if (waitingFirstChunk) { + useChatRuntimeStore.getState().setThreadWarming(threadKey, false); } } }, diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index c787cf471c..543ae9e865 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -49,9 +49,14 @@ import type { ChatView } from "./types"; const SingleContent = memo(function SingleContent({ threadId, -}: { threadId?: string }): ReactElement { + newThreadNonce, +}: { threadId?: string; newThreadNonce?: string }): ReactElement { return ( - +
@@ -197,7 +202,10 @@ function TopBarActions({ } export function ChatPage(): ReactElement { - const [view, setView] = useState({ mode: "single" }); + const [view, setView] = useState({ + mode: "single", + newThreadNonce: crypto.randomUUID(), + }); const [settingsOpen, setSettingsOpen] = useState(false); const inferenceParams = useChatRuntimeStore((state) => state.params); const setInferenceParams = useChatRuntimeStore((state) => state.setParams); @@ -215,7 +223,10 @@ export function ChatPage(): ReactElement { const handleEject = useCallback(() => { void ejectModel(); }, [ejectModel]); - const handleNewThread = useCallback(() => setView({ mode: "single" }), []); + const handleNewThread = useCallback( + () => setView({ mode: "single", newThreadNonce: crypto.randomUUID() }), + [], + ); const handleNewCompare = useCallback( () => setView({ mode: "compare", pairId: crypto.randomUUID() }), [], @@ -300,7 +311,11 @@ export function ChatPage(): ReactElement {
{view.mode === "single" ? ( - + ) : ( )} diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index dade143141..bf1b658c3d 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -7,6 +7,7 @@ import { type ExportedMessageRepositoryItem, type PendingAttachment, RuntimeAdapterProvider, + Suggestions, SimpleImageAttachmentAdapter, SimpleTextAttachmentAdapter, type ThreadHistoryAdapter, @@ -296,7 +297,14 @@ function useRuntimeHook(): ReturnType { function ThreadAutoSwitch({ threadId, }: { threadId: string }): ReactElement | null { - const aui = useAui(); + const aui = useAui({ + suggestions: Suggestions([ + "Draw a simple flowchart of a login system using Mermaid", + "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", + ]), + }); const isLoading = useAuiState(({ threads }) => threads.isLoading); const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId); @@ -309,16 +317,33 @@ function ThreadAutoSwitch({ return null; } +function ThreadNewChatSwitch({ + nonce, +}: { nonce: string }): ReactElement | null { + const aui = useAui(); + const isLoading = useAuiState(({ threads }) => threads.isLoading); + + useEffect(() => { + if (!isLoading) { + aui.threads().switchToNewThread(); + } + }, [aui, isLoading, nonce]); + + return null; +} + export function ChatRuntimeProvider({ children, modelType = "base", pairId, initialThreadId, + newThreadNonce, }: { children: ReactNode; modelType?: ModelType; pairId?: string; initialThreadId?: string; + newThreadNonce?: string; }): ReactElement { const runtime = useRemoteThreadListRuntime({ runtimeHook: useRuntimeHook, @@ -333,6 +358,9 @@ export function ChatRuntimeProvider({ return ( {initialThreadId && } + {!initialThreadId && newThreadNonce && ( + + )} {children} ); 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 76bf079ce7..74efc57134 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -10,10 +10,12 @@ type ChatRuntimeStore = { params: InferenceParams; models: ChatModelSummary[]; loras: ChatLoraSummary[]; + warmingByThreadId: Record; modelsError: string | null; setParams: (params: InferenceParams) => void; setModels: (models: ChatModelSummary[]) => void; setLoras: (loras: ChatLoraSummary[]) => void; + setThreadWarming: (threadId: string, warming: boolean) => void; setModelsError: (error: string | null) => void; setCheckpoint: (modelId: string) => void; clearCheckpoint: () => void; @@ -23,10 +25,21 @@ export const useChatRuntimeStore = create((set) => ({ params: DEFAULT_INFERENCE_PARAMS, models: [], loras: [], + warmingByThreadId: {}, modelsError: null, setParams: (params) => set({ params }), setModels: (models) => set({ models }), setLoras: (loras) => set({ loras }), + setThreadWarming: (threadId, warming) => + set((state) => { + const next = { ...state.warmingByThreadId }; + if (warming) { + next[threadId] = true; + } else { + delete next[threadId]; + } + return { warmingByThreadId: next }; + }), setModelsError: (modelsError) => set({ modelsError }), setCheckpoint: (modelId) => set((state) => ({ diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index 45d4f9e195..b0dccab307 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/studio/frontend/src/features/chat/types.ts @@ -1,7 +1,7 @@ export type ModelType = "base" | "lora"; export type ChatView = - | { mode: "single"; threadId?: string } + | { mode: "single"; threadId?: string; newThreadNonce?: string } | { mode: "compare"; pairId: string }; export interface ThreadRecord { From 95c56ed465895d6e6d6ed4b5cea87e25598afdc2 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Fri, 13 Feb 2026 17:42:45 +0100 Subject: [PATCH 4/4] feat: refactor suggestion handling and centralize defaults for thread UI --- .../src/components/assistant-ui/thread.tsx | 19 +++++++++++++++--- .../src/features/chat/runtime-provider.tsx | 20 ++++++++++--------- 2 files changed, 27 insertions(+), 12 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index df77d3ea28..9c94e4eb22 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -120,13 +120,26 @@ const ThreadScrollToBottom: FC = () => { }; const SuggestionItem: FC = () => { + const aui = useAui(); + const prompt = useAuiState(({ suggestion }) => suggestion.prompt); + const isDisabled = useAuiState(({ thread }) => thread.isDisabled); + const isRunning = useAuiState(({ thread }) => thread.isRunning); + return ( - { + if (!isDisabled && !isRunning) { + 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" > - + ); }; diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index bf1b658c3d..7dcc6644e8 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -28,6 +28,13 @@ import { createOpenAIStreamAdapter } from "./api/chat-adapter"; import { db } from "./db"; import type { MessageRecord, ModelType } from "./types"; +const DEFAULT_SUGGESTIONS = [ + "Draw a simple flowchart of a login system using Mermaid", + "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", +]; + class PDFAttachmentAdapter implements AttachmentAdapter { accept = "application/pdf"; @@ -297,14 +304,7 @@ function useRuntimeHook(): ReturnType { function ThreadAutoSwitch({ threadId, }: { threadId: string }): ReactElement | null { - const aui = useAui({ - suggestions: Suggestions([ - "Draw a simple flowchart of a login system using Mermaid", - "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", - ]), - }); + const aui = useAui(); const isLoading = useAuiState(({ threads }) => threads.isLoading); const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId); @@ -353,7 +353,9 @@ export function ChatRuntimeProvider({ }, }); - const aui = useAui(); + const aui = useAui({ + suggestions: Suggestions(DEFAULT_SUGGESTIONS), + }); return (