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({ /> = ({ 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 ( @@ -93,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" > - + ); }; @@ -358,6 +398,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 +433,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/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..fa670ecda2 --- /dev/null +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -0,0 +1,127 @@ +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, unstable_threadId }) { + 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 threadKey = unstable_threadId || "__default"; + let waitingFirstChunk = true; + useChatRuntimeStore.getState().setThreadWarming(threadKey, true); + let cumulativeText = ""; + let reasoningStartAt: number | null = null; + let reasoningDuration = 0; + + 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, + ); + + 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); + } + + 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/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts new file mode 100644 index 0000000000..72baf9a6f6 --- /dev/null +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -0,0 +1,146 @@ +import { authFetch } from "@/features/auth"; +import type { + InferenceStatusResponse, + ListLorasResponse, + ListModelsResponse, + LoadModelRequest, + LoadModelResponse, + OpenAIChatChunk, + OpenAIChatCompletionsRequest, + 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 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); +} + +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..543ae9e865 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"; @@ -28,16 +29,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,47 +47,16 @@ 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 { + newThreadNonce, +}: { threadId?: string; newThreadNonce?: string }): ReactElement { return ( - +
@@ -233,28 +202,60 @@ 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, setInferenceParams] = useState( - defaultInferenceParams, - ); + 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( - (v: string) => setInferenceParams((p) => ({ ...p, checkpoint: v })), + (value: string, meta?: { isLora: boolean }) => { + void selectModel({ id: value, isLora: meta?.isLora }); + }, + [selectModel], + ); + const handleEject = useCallback(() => { + void ejectModel(); + }, [ejectModel]); + const handleNewThread = useCallback( + () => setView({ mode: "single", newThreadNonce: crypto.randomUUID() }), [], ); - const handleEject = useCallback( - () => setInferenceParams((p) => ({ ...p, checkpoint: "" })), - [], - ); - 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], + ); + + const loraModels = useMemo( + () => + lorasFromStore.map((lora) => ({ + id: lora.id, + name: lora.name, + baseModel: lora.baseModel, + updatedAt: lora.updatedAt, + })), + [lorasFromStore], + ); + + useEffect(() => { + void refresh(); + }, [refresh]); return (
+ {modelsError && ( +
+ {modelsError} +
+ )}
- -
- - Backend - - -
- - 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); + + const refresh = useCallback(async () => { + setModelsError(null); + try { + const [listRes, statusRes, lorasRes] = await Promise.all([ + listModels(), + getInferenceStatus(), + listLoras(), + ]); + + setModels(listRes.models.map(toChatModelSummary)); + setLoras(lorasRes.loras.map(toLoraSummary)); + + if (statusRes.active_model) { + setCheckpoint(statusRes.active_model); + } + } catch (error) { + const message = + error instanceof Error ? error.message : "Failed to load models"; + setModelsError(message); + } + }, [setCheckpoint, setLoras, setModels, setModelsError]); + + const selectModel = useCallback( + async (selection: string | SelectedModelInput) => { + const modelId = typeof selection === "string" ? selection : selection.id; + if (!modelId || params.checkpoint === modelId) { + 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 { + if (params.checkpoint) { + await unloadModel({ model_path: params.checkpoint }); + } + + await loadModel({ + model_path: modelId, + hf_token: null, + max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH, + load_in_4bit: true, + is_lora: isLora, + }); + + 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 }); + } + }, + [loras, 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..7dcc6644e8 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -7,9 +7,9 @@ import { type ExportedMessageRepositoryItem, type PendingAttachment, RuntimeAdapterProvider, + Suggestions, SimpleImageAttachmentAdapter, SimpleTextAttachmentAdapter, - Suggestions, type ThreadHistoryAdapter, type ThreadMessage, type ThreadUserMessagePart, @@ -24,10 +24,17 @@ 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"; +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"; @@ -288,9 +295,11 @@ function ThreadHistoryProvider({ ); } -const chatAdapter = createStreamAdapter(); -const useRuntimeHook = (): ReturnType => - useLocalRuntime(chatAdapter); +const chatAdapter = createOpenAIStreamAdapter(); + +function useRuntimeHook(): ReturnType { + return useLocalRuntime(chatAdapter); +} function ThreadAutoSwitch({ threadId, @@ -308,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, @@ -328,17 +354,15 @@ 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", - ]), + suggestions: Suggestions(DEFAULT_SUGGESTIONS), }); 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 new file mode 100644 index 0000000000..74efc57134 --- /dev/null +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -0,0 +1,58 @@ +import { create } from "zustand"; +import { + DEFAULT_INFERENCE_PARAMS, + type ChatLoraSummary, + type ChatModelSummary, + type InferenceParams, +} from "../types/runtime"; + +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; +}; + +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) => ({ + params: { + ...state.params, + checkpoint: modelId, + }, + })), + clearCheckpoint: () => + set((state) => ({ + params: { + ...state.params, + checkpoint: "", + }, + })), +})); 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 { 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..9d21c854af --- /dev/null +++ b/studio/frontend/src/features/chat/types/api.ts @@ -0,0 +1,79 @@ +export interface BackendModelDetails { + id: string; + name?: string | null; + is_vision?: boolean; + is_lora?: boolean; +} + +export interface ListModelsResponse { + models: BackendModelDetails[]; + 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; + 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..558f1a5464 --- /dev/null +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -0,0 +1,34 @@ +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; +} + +export interface ChatLoraSummary { + id: string; + name: string; + baseModel: string; + updatedAt?: number; +} 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); +}