From 23d2cfd09d687bf652a4f16bf94422767aaf44eb Mon Sep 17 00:00:00 2001 From: Shine1i Date: Fri, 13 Feb 2026 16:45:00 +0100 Subject: [PATCH] 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); +}