feat: refactor chat runtime with modular APIs, state management, and runtime synchronization
This commit is contained in:
Shine1i
parent
a7c6432ffd
commit
23d2cfd09d
12 changed files with 630 additions and 240 deletions
|
|
@ -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<ChatModelRunResult["content"]>[number];
|
||||
type RunMessages = Parameters<ChatModelAdapter["run"]>[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("<think>");
|
||||
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("</think>");
|
||||
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("</think>") && reasoningStart && !reasoningDuration) {
|
||||
reasoningDuration = Math.round((Date.now() - reasoningStart) / 1000);
|
||||
}
|
||||
|
||||
if (parts.length > 0) {
|
||||
yield {
|
||||
content: parts,
|
||||
metadata: { custom: { reasoningDuration } },
|
||||
};
|
||||
}
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
113
studio/frontend/src/features/chat/api/chat-adapter.ts
Normal file
113
studio/frontend/src/features/chat/api/chat-adapter.ts
Normal file
|
|
@ -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<ChatModelAdapter["run"]>[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<typeof message> =>
|
||||
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 } },
|
||||
};
|
||||
}
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
139
studio/frontend/src/features/chat/api/chat-api.ts
Normal file
139
studio/frontend/src/features/chat/api/chat-api.ts
Normal file
|
|
@ -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<T>(response: Response): Promise<T> {
|
||||
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<ListModelsResponse> {
|
||||
const response = await authFetch("/api/models/list");
|
||||
return parseJsonOrThrow<ListModelsResponse>(response);
|
||||
}
|
||||
|
||||
export async function getInferenceStatus(): Promise<InferenceStatusResponse> {
|
||||
const response = await authFetch("/api/inference/status");
|
||||
return parseJsonOrThrow<InferenceStatusResponse>(response);
|
||||
}
|
||||
|
||||
export async function loadModel(
|
||||
payload: LoadModelRequest,
|
||||
): Promise<LoadModelResponse> {
|
||||
const response = await authFetch("/api/inference/load", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(payload),
|
||||
});
|
||||
return parseJsonOrThrow<LoadModelResponse>(response);
|
||||
}
|
||||
|
||||
export async function unloadModel(payload: UnloadModelRequest): Promise<void> {
|
||||
const response = await authFetch("/api/inference/unload", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(payload),
|
||||
});
|
||||
await parseJsonOrThrow<unknown>(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<OpenAIChatChunk> {
|
||||
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/);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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<ChatView>({ mode: "single" });
|
||||
const [settingsOpen, setSettingsOpen] = useState(false);
|
||||
const [inferenceParams, setInferenceParams] = useState<InferenceParams>(
|
||||
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<ModelOption[]>(
|
||||
() =>
|
||||
modelsFromStore.map((model) => ({
|
||||
id: model.id,
|
||||
name: model.name,
|
||||
description: model.description,
|
||||
})),
|
||||
[modelsFromStore],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
void refresh();
|
||||
}, [refresh]);
|
||||
|
||||
return (
|
||||
<SidebarProvider
|
||||
|
|
@ -292,6 +269,11 @@ export function ChatPage(): ReactElement {
|
|||
variant="ghost"
|
||||
/>
|
||||
</div>
|
||||
{modelsError && (
|
||||
<div className="ml-2 text-xs text-destructive truncate max-w-[28rem]">
|
||||
{modelsError}
|
||||
</div>
|
||||
)}
|
||||
<div className="flex-1" />
|
||||
<button
|
||||
type="button"
|
||||
|
|
|
|||
|
|
@ -10,7 +10,6 @@ import { Textarea } from "@/components/ui/textarea";
|
|||
import {
|
||||
ArrowDown01Icon,
|
||||
Delete02Icon,
|
||||
EngineIcon,
|
||||
FloppyDiskIcon,
|
||||
PencilEdit01Icon,
|
||||
Settings02Icon,
|
||||
|
|
@ -20,28 +19,13 @@ import { HugeiconsIcon } from "@hugeicons/react";
|
|||
import { AnimatePresence, motion } from "motion/react";
|
||||
import type { ReactNode } from "react";
|
||||
import { useState } from "react";
|
||||
import {
|
||||
DEFAULT_INFERENCE_PARAMS,
|
||||
type InferenceParams,
|
||||
} from "./types/runtime";
|
||||
|
||||
export interface InferenceParams {
|
||||
temperature: number;
|
||||
topP: number;
|
||||
topK: number;
|
||||
repetitionPenalty: number;
|
||||
maxTokens: number;
|
||||
systemPrompt: string;
|
||||
inferenceEngine: string;
|
||||
checkpoint: string;
|
||||
}
|
||||
|
||||
export const defaultInferenceParams: InferenceParams = {
|
||||
temperature: 0.7,
|
||||
topP: 0.9,
|
||||
topK: 50,
|
||||
repetitionPenalty: 1.1,
|
||||
maxTokens: 512,
|
||||
systemPrompt: "",
|
||||
inferenceEngine: "unsloth",
|
||||
checkpoint: "outputs/llama-3.1-8b-instruct-lora",
|
||||
};
|
||||
export const defaultInferenceParams = DEFAULT_INFERENCE_PARAMS;
|
||||
export type { InferenceParams } from "./types/runtime";
|
||||
|
||||
export interface Preset {
|
||||
name: string;
|
||||
|
|
@ -72,11 +56,6 @@ const BUILTIN_PRESETS: Preset[] = [
|
|||
},
|
||||
];
|
||||
|
||||
const ENGINE_OPTIONS = [
|
||||
{ value: "unsloth", label: "Unsloth" },
|
||||
{ value: "llama-cpp", label: "llama.cpp (GGUF)" },
|
||||
];
|
||||
|
||||
function ParamSlider({
|
||||
label,
|
||||
value,
|
||||
|
|
@ -285,33 +264,6 @@ export function ChatSettingsPanel({
|
|||
/>
|
||||
</div>
|
||||
|
||||
<CollapsibleSection
|
||||
icon={EngineIcon}
|
||||
label="Inference Engine"
|
||||
defaultOpen={true}
|
||||
>
|
||||
<div>
|
||||
<span className="mb-1 block text-[11px] text-muted-foreground">
|
||||
Backend
|
||||
</span>
|
||||
<Select
|
||||
value={params.inferenceEngine}
|
||||
onValueChange={set("inferenceEngine")}
|
||||
>
|
||||
<SelectTrigger className="h-8 w-full text-xs corner-squircle">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{ENGINE_OPTIONS.map((o) => (
|
||||
<SelectItem key={o.value} value={o.value}>
|
||||
{o.label}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
</CollapsibleSection>
|
||||
|
||||
<CollapsibleSection
|
||||
icon={SlidersHorizontalIcon}
|
||||
label="Sampling"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,140 @@
|
|||
import { useCallback } from "react";
|
||||
import {
|
||||
getInferenceStatus,
|
||||
listModels,
|
||||
loadModel,
|
||||
unloadModel,
|
||||
} from "../api/chat-api";
|
||||
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
|
||||
import type { ChatModelSummary } from "../types/runtime";
|
||||
|
||||
const DEFAULT_MODEL_MAX_SEQ_LENGTH = 2048;
|
||||
|
||||
function describeModel(model: {
|
||||
is_lora?: boolean;
|
||||
is_vision?: boolean;
|
||||
}): string | undefined {
|
||||
const tags: string[] = [];
|
||||
if (model.is_lora) tags.push("LoRA");
|
||||
if (model.is_vision) tags.push("Vision");
|
||||
if (!model.is_lora && !model.is_vision) tags.push("Base");
|
||||
return tags.join(" · ");
|
||||
}
|
||||
|
||||
function toChatModelSummary(model: {
|
||||
id: string;
|
||||
name?: string | null;
|
||||
is_lora?: boolean;
|
||||
is_vision?: boolean;
|
||||
}): ChatModelSummary {
|
||||
return {
|
||||
id: model.id,
|
||||
name: model.name || model.id,
|
||||
description: describeModel(model),
|
||||
isLora: Boolean(model.is_lora),
|
||||
isVision: Boolean(model.is_vision),
|
||||
};
|
||||
}
|
||||
|
||||
export function useChatModelRuntime() {
|
||||
const params = useChatRuntimeStore((state) => 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,
|
||||
};
|
||||
}
|
||||
|
|
@ -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";
|
||||
|
|
|
|||
|
|
@ -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<typeof useLocalRuntime> =>
|
||||
useLocalRuntime(chatAdapter);
|
||||
const chatAdapter = createOpenAIStreamAdapter();
|
||||
|
||||
function useRuntimeHook(): ReturnType<typeof useLocalRuntime> {
|
||||
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 (
|
||||
<AssistantRuntimeProvider runtime={runtime} aui={aui}>
|
||||
|
|
|
|||
|
|
@ -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<ChatRuntimeStore>((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: "",
|
||||
},
|
||||
})),
|
||||
}));
|
||||
68
studio/frontend/src/features/chat/types/api.ts
Normal file
68
studio/frontend/src/features/chat/types/api.ts
Normal file
|
|
@ -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[];
|
||||
}
|
||||
27
studio/frontend/src/features/chat/types/runtime.ts
Normal file
27
studio/frontend/src/features/chat/types/runtime.ts
Normal file
|
|
@ -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;
|
||||
}
|
||||
|
|
@ -0,0 +1,54 @@
|
|||
import type { ChatModelRunResult } from "@assistant-ui/react";
|
||||
|
||||
type ContentPart = NonNullable<ChatModelRunResult["content"]>[number];
|
||||
|
||||
const THINK_OPEN_TAG = "<think>";
|
||||
const THINK_CLOSE_TAG = "</think>";
|
||||
|
||||
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);
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue