Hugging Face's logo

Spaces:
Etiya
/
d2l-ui
Sleeping

App Files Community
d2l-ui / lib /chatStore.ts
Berkkirik's picture
Berkkirik
feat: streaming + multi-turn + view source + doc-scoped + UX fixes
df790cc
Raw History Blame Contribute Delete
9.79 kB
"use client";
import { create } from "zustand";
import { persist, createJSONStorage } from "zustand/middleware";
import type { AskSmartRequest, AskSmartResponse } from "./types";
export type ChatTurn = {
id: string;
question: string;
response?: AskSmartResponse;
error?: { status: number; message: string };
pending?: boolean;
timestamp: number;
};
export type Conversation = {
id: string;
title: string;
turns: ChatTurn[];
createdAt: number;
updatedAt: number;
/** When set, follow-up questions go to /ask scoped to this doc instead of /ask_smart. */
pinnedDocId?: string | null;
pinnedDocName?: string | null;
};
export type Settings = Required<
Omit<AskSmartRequest, "question" | "history" | "stream">
>;
export const DEFAULT_SETTINGS: Settings = {
top_k: 1,
max_new_tokens: 200,
similarity_threshold: 0.45,
rerank_threshold: 0.6,
anchor_threshold: 0.2,
use_grounding: true,
repetition_penalty: 1.15,
no_repeat_ngram_size: 4,
scaler: 1.0,
bias_scaler: 1.0,
};
type ChatStore = {
conversations: Record<string, Conversation>;
activeId: string | null;
settings: Settings;
sidebarOpen: boolean;
// sync state
hydrated: boolean;
syncStatus: "idle" | "saving" | "error";
lastSavedAt: number;
// selectors
active: () => Conversation | null;
list: () => Conversation[];
// mutations
newConversation: () => string;
newConversationPinned: (doc_id: string, doc_name: string) => string;
openConversation: (id: string) => void;
deleteConversation: (id: string) => void;
renameConversation: (id: string, title: string) => void;
pinDoc: (id: string, doc_id: string | null, doc_name?: string | null) => void;
appendTurn: (turn: ChatTurn) => void;
patchTurn: (id: string, patch: Partial<ChatTurn>) => void;
clearActive: () => void;
setSettings: (next: Settings) => void;
resetSettings: () => void;
toggleSidebar: () => void;
setSidebarOpen: (open: boolean) => void;
// sync helpers (used by useBackendSync hook)
_setHydrated: (v: boolean) => void;
_setSyncStatus: (v: "idle" | "saving" | "error") => void;
_hydrateFromRemote: (payload: {
conversations: Record<string, Conversation>;
activeId: string | null;
}) => void;
};
const newId = () =>
(typeof crypto !== "undefined" && crypto.randomUUID
? crypto.randomUUID()
: Math.random().toString(36).slice(2, 12));
/**
* Build a multi-turn history payload for the backend from a conversation.
* Excludes pending and errored turns; caps to last `max` Q→A pairs.
*/
export function buildHistory(
conv: Conversation | null | undefined,
max = 5
): { question: string; answer: string }[] {
if (!conv) return [];
const pairs: { question: string; answer: string }[] = [];
for (const t of conv.turns) {
if (t.pending || t.error || !t.response) continue;
pairs.push({ question: t.question, answer: t.response.answer });
}
return pairs.slice(-max);
}
const seedConversation = (): Conversation => ({
id: newId(),
title: "New conversation",
turns: [],
createdAt: Date.now(),
updatedAt: Date.now(),
});
export const useChatStore = create<ChatStore>()(
persist(
(set, get) => ({
conversations: {},
activeId: null,
settings: DEFAULT_SETTINGS,
sidebarOpen: true,
hydrated: false,
syncStatus: "idle",
lastSavedAt: 0,
active: () => {
const { activeId, conversations } = get();
return activeId ? conversations[activeId] ?? null : null;
},
list: () =>
Object.values(get().conversations).sort(
(a, b) => b.updatedAt - a.updatedAt
),
newConversation: () => {
const c = seedConversation();
set((s) => ({
conversations: { ...s.conversations, [c.id]: c },
activeId: c.id,
}));
return c.id;
},
newConversationPinned: (doc_id, doc_name) => {
const c: Conversation = {
...seedConversation(),
title: doc_name ? `Scoped: ${doc_name}` : "Scoped conversation",
pinnedDocId: doc_id,
pinnedDocName: doc_name ?? null,
};
set((s) => ({
conversations: { ...s.conversations, [c.id]: c },
activeId: c.id,
}));
return c.id;
},
pinDoc: (id, doc_id, doc_name) =>
set((s) => {
const conv = s.conversations[id];
if (!conv) return s;
return {
conversations: {
...s.conversations,
[id]: {
...conv,
pinnedDocId: doc_id,
pinnedDocName: doc_name ?? conv.pinnedDocName ?? null,
updatedAt: Date.now(),
},
},
};
}),
openConversation: (id) => set({ activeId: id }),
deleteConversation: (id) =>
set((s) => {
const next = { ...s.conversations };
delete next[id];
const remainingIds = Object.keys(next);
return {
conversations: next,
activeId:
s.activeId === id
? remainingIds[0] ?? null
: s.activeId,
};
}),
renameConversation: (id, title) =>
set((s) => ({
conversations: {
...s.conversations,
[id]: s.conversations[id]
? { ...s.conversations[id], title, updatedAt: Date.now() }
: s.conversations[id],
},
})),
appendTurn: (turn) =>
set((s) => {
let { activeId } = s;
let conversations = s.conversations;
if (!activeId || !conversations[activeId]) {
const c = seedConversation();
activeId = c.id;
conversations = { ...conversations, [c.id]: c };
}
const conv = conversations[activeId];
const isFirstTurn = conv.turns.length === 0;
const updated: Conversation = {
...conv,
title: isFirstTurn
? turn.question.slice(0, 60).trim() || conv.title
: conv.title,
turns: [...conv.turns, turn],
updatedAt: Date.now(),
};
return {
activeId,
conversations: { ...conversations, [activeId]: updated },
};
}),
patchTurn: (id, patch) =>
set((s) => {
if (!s.activeId) return s;
const conv = s.conversations[s.activeId];
if (!conv) return s;
const turns = conv.turns.map((t) =>
t.id === id ? { ...t, ...patch } : t
);
return {
conversations: {
...s.conversations,
[s.activeId]: { ...conv, turns, updatedAt: Date.now() },
},
};
}),
clearActive: () =>
set((s) => {
if (!s.activeId) return s;
const conv = s.conversations[s.activeId];
if (!conv) return s;
return {
conversations: {
...s.conversations,
[s.activeId]: { ...conv, turns: [], updatedAt: Date.now() },
},
};
}),
setSettings: (next) => set({ settings: next }),
resetSettings: () => set({ settings: DEFAULT_SETTINGS }),
toggleSidebar: () => set((s) => ({ sidebarOpen: !s.sidebarOpen })),
setSidebarOpen: (open) => set({ sidebarOpen: open }),
// ── sync helpers ─────────────────────────────────────────────
_setHydrated: (v) => set({ hydrated: v }),
_setSyncStatus: (v) => set({ syncStatus: v }),
_hydrateFromRemote: ({ conversations, activeId }) =>
set((s) => {
// Merge remote into local: remote wins per-id (server is the
// source of truth on hydrate). If remote is empty but local has
// pending in-flight turns, keep local — avoids losing a turn
// that's mid-flight when hydration races with the request.
const remoteIds = Object.keys(conversations);
if (remoteIds.length === 0) {
// Server is empty — keep local; the next save will populate it.
return { hydrated: true };
}
// Preserve any locally-pending turns that aren't yet on server.
const merged: Record<string, Conversation> = { ...conversations };
for (const [id, local] of Object.entries(s.conversations)) {
if (!merged[id]) {
merged[id] = local;
continue;
}
const localPending = local.turns.filter((t) => t.pending);
if (localPending.length === 0) continue;
const serverTurnIds = new Set(merged[id].turns.map((t) => t.id));
const extras = localPending.filter(
(t) => !serverTurnIds.has(t.id)
);
if (extras.length > 0) {
merged[id] = {
...merged[id],
turns: [...merged[id].turns, ...extras],
};
}
}
const nextActive =
activeId && merged[activeId]
? activeId
: s.activeId && merged[s.activeId]
? s.activeId
: Object.keys(merged)[0] ?? null;
return {
conversations: merged,
activeId: nextActive,
hydrated: true,
};
}),
}),
{
name: "etiya-d2l-chat",
storage: createJSONStorage(() => localStorage),
partialize: (state) => ({
conversations: state.conversations,
activeId: state.activeId,
settings: state.settings,
sidebarOpen: state.sidebarOpen,
}),
}
)
);