130 lines
3.4 KiB
TypeScript
130 lines
3.4 KiB
TypeScript
"use client";
|
|
|
|
import { useCallback, useEffect, useMemo, useRef } from "react";
|
|
import { useStorage } from "./use-storage.ts";
|
|
import { UIMessage } from "ai";
|
|
|
|
export type ChatHistory<T extends UIMessage> = Record<string, T[]>;
|
|
|
|
const CHAT_HISTORY_KEY = "storytellurx_chat_history";
|
|
|
|
function downloadJSON(data: any, filename: string) {
|
|
const dataStr = "data:text/json;charset=utf-8," +
|
|
encodeURIComponent(JSON.stringify(data, null, 2));
|
|
const downloadAnchorNode = document.createElement("a");
|
|
downloadAnchorNode.setAttribute("href", dataStr);
|
|
downloadAnchorNode.setAttribute("download", filename);
|
|
document.body.appendChild(downloadAnchorNode);
|
|
downloadAnchorNode.click();
|
|
downloadAnchorNode.remove();
|
|
}
|
|
|
|
const acceptableMessagesFilter = (message: UIMessage<unknown, any>) =>
|
|
message.parts.length > 0 &&
|
|
(message.parts.every((v) => v.type !== "text") ||
|
|
message.parts.some((v) => v.type === "text" && v.text.trim().length > 0));
|
|
|
|
const getTitleFromMessages = (messages: UIMessage<unknown, any>[]) =>
|
|
messages.at(0)?.parts.filter((v) => v.type === "text").at(0)
|
|
?.text;
|
|
|
|
export function useChatHistory<
|
|
T extends UIMessage<unknown, any>,
|
|
>(options?: { onHistoryInit?: (history: ChatHistory<T>) => void }) {
|
|
const [rawHistory, setHistory] = useStorage<ChatHistory<T> | undefined>(
|
|
CHAT_HISTORY_KEY,
|
|
undefined,
|
|
{ initializeWithValue: false, storageType: "indexedDB" },
|
|
);
|
|
|
|
const history = useMemo(() =>
|
|
rawHistory
|
|
? Object.fromEntries(
|
|
Object.entries(rawHistory)
|
|
.map(([k, v]) => [
|
|
k,
|
|
(v as UIMessage<unknown, any>[]).filter(acceptableMessagesFilter),
|
|
]),
|
|
)
|
|
: rawHistory, [rawHistory]);
|
|
|
|
const init = useRef(false);
|
|
const historyRef = useRef(history);
|
|
historyRef.current = history;
|
|
|
|
useEffect(() => {
|
|
if (history && !init.current) {
|
|
options?.onHistoryInit?.(history);
|
|
init.current = true;
|
|
}
|
|
return () => void (init.current = false);
|
|
}, [history]);
|
|
|
|
const saveChat = useCallback(
|
|
(id: string, messages: T[]) => {
|
|
setHistory((prevHistory) => ({
|
|
...prevHistory ?? {},
|
|
[id]: messages.filter(acceptableMessagesFilter),
|
|
}));
|
|
},
|
|
[setHistory],
|
|
);
|
|
|
|
const saveMessage = useCallback(
|
|
(id: string, message: T) => {
|
|
setHistory((prevHistory) => ({
|
|
...prevHistory ?? {},
|
|
[id]: [...prevHistory?.[id] ?? [], message].filter(
|
|
acceptableMessagesFilter,
|
|
),
|
|
}));
|
|
},
|
|
[setHistory],
|
|
);
|
|
|
|
const downloadChat = useCallback(
|
|
(id: string) => {
|
|
const messages = historyRef.current?.[id];
|
|
if (messages) {
|
|
downloadJSON(
|
|
messages,
|
|
`storytellurx-chat-${
|
|
getTitleFromMessages(messages ?? [])
|
|
?.replace(/\s/g, "_").slice(0, 40)
|
|
}-${id}.json`,
|
|
);
|
|
} else {
|
|
alert("Could not find chat in history to export.");
|
|
}
|
|
},
|
|
[setHistory],
|
|
);
|
|
|
|
const deleteChat = useCallback(
|
|
(id: string) => {
|
|
setHistory((prevHistory) => {
|
|
const newHistory = { ...prevHistory ?? {} };
|
|
delete newHistory[id];
|
|
return newHistory;
|
|
});
|
|
},
|
|
[setHistory],
|
|
);
|
|
|
|
const getChat = useCallback(
|
|
(id: string): T[] | undefined => {
|
|
return historyRef.current?.[id];
|
|
},
|
|
[],
|
|
);
|
|
|
|
return {
|
|
isHistoryLoading: history === undefined,
|
|
history,
|
|
saveChat,
|
|
deleteChat,
|
|
saveMessage,
|
|
getChat,
|
|
downloadChat,
|
|
};
|
|
}
|