feat: refactor chat runtime with modular APIs, state management, and runtime synchronization

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<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 } },
+          };
+        }
+      }
+    },
+  };
+}

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<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/);
+    }
+  }
+}