SPB Git
11.6 KB · 129 lines typescript
1"use client";2import { useCallback, useEffect, useSyncExternalStore } from "react";3import { api, readSSE } from "./fetcher";4import type { MessageDTO, ConversationDTO, ChatSendOptions, Part, AttachmentDTO, ArtifactDTO } from "./types";5import { toast } from "@/components/ui";6import { mutate } from "swr";78export interface ChatState {9  conversationId: string | null;10  conversation: ConversationDTO | null;11  messages: MessageDTO[];12  allMessages: MessageDTO[];13  siblings: { parent_id: string | null; role: string; n: number }[];14  attachments: AttachmentDTO[];15  artifacts: ArtifactDTO[];16  loading: boolean;17  streaming: boolean;18  error: string | null;19  warning: string | null;20  abort: AbortController | null;21}2223const empty = (id: string | null): ChatState => ({ conversationId: id, conversation: null, messages: [], allMessages: [], siblings: [], attachments: [], artifacts: [], loading: !!id, streaming: false, error: null, warning: null, abort: null });2425// ---------------------------------------------------------------- tiny external store keyed by conversation id ("new" for drafts)26const store = new Map<string, ChatState>();27const alias = new Map<string, string>(); // "new" → id of the conversation created from the draft (so the / page keeps showing it)28const listeners = new Set<() => void>();29const key = (id: string | null) => { const k = id ?? "new"; return alias.get(k) ?? k; };30function get(id: string | null): ChatState { const k = key(id); if (!store.has(k)) store.set(k, empty(k === "new" ? null : k)); return store.get(k)!; }31function set(id: string | null, patch: Partial<ChatState> | ((s: ChatState) => Partial<ChatState>)) { const k = key(id); const cur = get(id); const p = typeof patch === "function" ? patch(cur) : patch; store.set(k, { ...cur, ...p }); listeners.forEach(l => l()); }32function subscribe(l: () => void) { listeners.add(l); return () => { listeners.delete(l); }; }3334export function resetDraft() { alias.delete("new"); store.set("new", empty(null)); listeners.forEach(l => l()); }3536async function load(id: string) {37  set(id, { loading: true, error: null });38  try {39    const r = await api<{ conversation: ConversationDTO; messages: MessageDTO[]; siblings: ChatState["siblings"]; attachments: AttachmentDTO[]; artifacts: ArtifactDTO[] }>(`/api/conversations/${id}?all=1`);40    set(id, { conversation: r.conversation, messages: r.messages.filter(m => m.active), allMessages: r.messages, siblings: r.siblings, attachments: r.attachments, artifacts: r.artifacts, loading: false });41  } catch (e) { set(id, { loading: false, error: (e as Error).message }); }42}4344/** Streams one assistant turn; `endpoint` is one of the chat SSE routes. */45async function stream(fromId: string | null, endpoint: string, payload: Record<string, unknown>, opts: { optimisticUser?: MessageDTO; replaceAssistantId?: string } = {}) {46  const abort = new AbortController();47  let id = fromId;48  set(id, s => ({ streaming: true, error: null, warning: null, abort, messages: opts.optimisticUser ? [...s.messages, opts.optimisticUser] : s.messages }));49  let assistant: MessageDTO | null = null;50  let textBuf = ""; let reasoningBuf = ""; let raf = 0;51  const flush = () => { raf = 0; if (!assistant) return; const a = assistant; set(id, s => ({ messages: s.messages.map(m => m.id === a.id ? { ...a, parts: [...a.parts] } : m) })); };52  const schedule = () => { if (!raf) raf = requestAnimationFrame(flush); };53  try {54    const res = await fetch(endpoint, { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify(payload), signal: abort.signal, credentials: "same-origin" });55    for await (const { event, data } of readSSE(res)) {56      const d = data as Record<string, unknown>;57      switch (event) {58        case "conversation": {59          const newId = d.id as string;60          if (id !== newId) { // draft → real conversation: move state61            const s = get(id); store.set(newId, { ...s, conversationId: newId }); if (id === null) alias.set("new", newId);62            id = newId; listeners.forEach(l => l());63            if (d.created) { window.history.replaceState(null, "", `/c/${newId}`); void mutate("/api/conversations"); }64          }65          if (opts.optimisticUser && d.userMessageId) { const uid = d.userMessageId as string; const tmp = opts.optimisticUser.id; set(id, s => ({ messages: s.messages.map(m => m.id === tmp ? { ...m, id: uid, conversation_id: id! } : m) })); }66          if (d.regenerated || d.edited) { const old = (d.regenerated ?? d.edited) as string; set(id, s => { const idx = s.messages.findIndex(m => m.id === old); return { messages: idx >= 0 ? s.messages.slice(0, d.edited ? idx : idx) : s.messages }; }); if (d.edited && opts.optimisticUser) set(id, s => ({ messages: [...s.messages, { ...opts.optimisticUser!, id: d.userMessageId as string }] })); }67          break;68        }69        case "start": {70          assistant = { id: d.messageId as string, conversation_id: id!, parent_id: (payload.parentMessageId as string) ?? null, role: "assistant", parts: [], text: "", model_id: d.model as string, provider: d.provider as string, settings: null, routing: d.routing as MessageDTO["routing"], usage: null, cost_usd: null, latency_ms: null, ttft_ms: null, status: "streaming", error: null, active: true, created_at: new Date().toISOString() };71          textBuf = ""; reasoningBuf = "";72          set(id, s => ({ messages: [...s.messages.filter(m => m.id !== opts.replaceAssistantId), assistant!] }));73          break;74        }75        case "text": { if (!assistant) break; textBuf += d.delta as string; let tp = assistant.parts.find(p => p.type === "text") as Extract<Part, { type: "text" }> | undefined; if (!tp) { tp = { type: "text", text: "" }; assistant.parts.push(tp); } tp.text = textBuf; schedule(); break; }76        case "reasoning": { if (!assistant) break; reasoningBuf += d.delta as string; let rp = assistant.parts.find(p => p.type === "reasoning") as Extract<Part, { type: "reasoning" }> | undefined; if (!rp) { rp = { type: "reasoning", text: "", summarized: true }; assistant.parts.unshift(rp); } rp.text = reasoningBuf; schedule(); break; }77        case "tool_call": { if (!assistant) break; assistant.parts.push({ type: "tool_call", id: d.id as string, name: d.name as string, input: d.input, status: (d.status as "running") ?? "running" }); flush(); break; }78        case "tool_result": { if (!assistant) break; const tc = assistant.parts.find(p => p.type === "tool_call" && p.id === d.id) as Extract<Part, { type: "tool_call" }> | undefined; if (tc) { tc.status = d.status as "ok"; tc.output = d.output; tc.error = d.error as string | undefined; tc.durationMs = d.durationMs as number; } flush(); break; }79        case "part": { if (!assistant) break; assistant.parts.push(d as unknown as Part); flush(); break; }80        case "artifacts": { if (!assistant) break; const arts = data as { id: string; title: string; type: string; language: string }[]; assistant.parts.push(...arts.map(a => ({ type: "artifact" as const, artifactId: a.id, title: a.title, artifactType: a.type, language: a.language }))); set(id, s => ({ artifacts: [...s.artifacts, ...arts.map(a => ({ id: a.id, root_id: a.id, version: 1, type: a.type, language: a.language, title: a.title, saved: false, message_id: assistant!.id, created_at: new Date().toISOString() }))] })); flush(); break; }81        case "usage": { if (!assistant) break; assistant.usage = d.usage as MessageDTO["usage"]; assistant.cost_usd = d.costUsd as number | null; assistant.latency_ms = d.latencyMs as number; assistant.ttft_ms = d.ttftMs as number | null; assistant.status = "complete"; flush(); break; }82        case "meta": { if (assistant && d.servedBy) { assistant.routing = { ...(assistant.routing ?? {}), servedBy: d.servedBy as string }; flush(); } break; }83        case "title": { set(id, s => ({ conversation: s.conversation ? { ...s.conversation, title: d.title as string } : s.conversation })); if (!get(id).conversation && id) void load(id); void mutate("/api/conversations"); break; }84        case "warning": { set(id, { warning: d.message as string }); break; }85        case "error": { if (assistant) { assistant.status = "error"; assistant.error = d.message as string; flush(); } set(id, { error: d.message as string }); break; }86        case "done": { if (assistant && assistant.status === "streaming") { assistant.status = (d.status as string) ?? "complete"; flush(); } break; }87      }88    }89  } catch (e) {90    if ((e as Error).name !== "AbortError") { set(id, { error: (e as Error).message }); toast((e as Error).message, "error"); }91    else if (assistant) { assistant.status = "cancelled"; flush(); }92  } finally {93    if (raf) cancelAnimationFrame(raf);94    set(id, { streaming: false, abort: null });95    if (id) void load(id).then(() => mutate("/api/conversations"));96  }97}9899export function useChat(conversationId: string | null) {100  const state = useSyncExternalStore(subscribe, () => get(conversationId), () => get(conversationId));101  useEffect(() => { if (conversationId && !get(conversationId).conversation && !get(conversationId).streaming) void load(conversationId); }, [conversationId]);102103  const send = useCallback(async (text: string, o: ChatSendOptions = {}, atts: AttachmentDTO[] = []) => {104    const parts: Part[] = [...atts.map(a => a.mime.startsWith("image/") ? ({ type: "image" as const, attachmentId: a.id, mime: a.mime, name: a.filename }) : ({ type: "file" as const, attachmentId: a.id, mime: a.mime, name: a.filename, size: a.size ?? a.size_bytes, pages: a.pages ?? a.page_count ?? undefined })), ...(text.trim() ? [{ type: "text" as const, text }] : [])];105    const optimistic: MessageDTO = { id: `tmp_${Date.now()}`, conversation_id: conversationId ?? "", parent_id: null, role: "user", parts, text, model_id: null, provider: null, settings: null, routing: null, usage: null, cost_usd: null, latency_ms: null, ttft_ms: null, status: "complete", error: null, active: true, created_at: new Date().toISOString() };106    await stream(conversationId, "/api/chat", { conversationId: conversationId ?? undefined, text, attachmentIds: atts.map(a => a.id), ...o }, { optimisticUser: optimistic });107  }, [conversationId]);108109  const regenerate = useCallback(async (messageId: string, o: Pick<ChatSendOptions, "modelId" | "routingMode" | "settings" | "toolsEnabled"> = {}) => {110    await stream(conversationId, "/api/chat/regenerate", { messageId, ...o }, { replaceAssistantId: messageId });111  }, [conversationId]);112113  const edit = useCallback(async (messageId: string, text: string, o: Pick<ChatSendOptions, "modelId" | "routingMode" | "toolsEnabled"> = {}) => {114    const old = get(conversationId).messages.find(m => m.id === messageId);115    const optimistic: MessageDTO = { ...(old as MessageDTO), id: `tmp_${Date.now()}`, parts: [...(old?.parts.filter(p => p.type !== "text") ?? []), { type: "text", text }], text };116    await stream(conversationId, "/api/chat/edit", { messageId, text, ...o }, { optimisticUser: optimistic });117  }, [conversationId]);118119  const approve = useCallback(async (messageId: string, approveIds: string[], denyIds: string[] = []) => {120    await stream(conversationId, "/api/chat/approve", { messageId, approve: approveIds, deny: denyIds });121  }, [conversationId]);122123  const stop = useCallback(() => { get(conversationId).abort?.abort(); }, [conversationId]);124  const reload = useCallback(() => { if (conversationId) void load(conversationId); }, [conversationId]);125  const switchBranch = useCallback(async (targetId: string) => { if (!conversationId) return; await api(`/api/conversations/${conversationId}/branch`, { method: "POST", json: { activate: targetId } }); await load(conversationId); }, [conversationId]);126127  return { ...state, send, regenerate, edit, approve, stop, reload, switchBranch };128}129
Files · main
Menu

Theme