131 lines
3.9 KiB
TypeScript
131 lines
3.9 KiB
TypeScript
import { useCallback } from "react";
|
|
import type { Message } from "@/ipc/ipc_types";
|
|
import { useAtom, useSetAtom } from "jotai";
|
|
import {
|
|
chatErrorAtom,
|
|
chatMessagesAtom,
|
|
chatStreamCountAtom,
|
|
isStreamingAtom,
|
|
} from "@/atoms/chatAtoms";
|
|
import { IpcClient } from "@/ipc/ipc_client";
|
|
import { isPreviewOpenAtom } from "@/atoms/viewAtoms";
|
|
import type { ChatResponseEnd } from "@/ipc/ipc_types";
|
|
import { useChats } from "./useChats";
|
|
import { useLoadApp } from "./useLoadApp";
|
|
import { selectedAppIdAtom } from "@/atoms/appAtoms";
|
|
import { useVersions } from "./useVersions";
|
|
import { showUncommittedFilesWarning } from "@/lib/toast";
|
|
import { useProposal } from "./useProposal";
|
|
import { useSearch } from "@tanstack/react-router";
|
|
import { useRunApp } from "./useRunApp";
|
|
import { useCountTokens } from "./useCountTokens";
|
|
|
|
export function getRandomNumberId() {
|
|
return Math.floor(Math.random() * 1_000_000_000_000_000);
|
|
}
|
|
|
|
export function useStreamChat({
|
|
hasChatId = true,
|
|
}: { hasChatId?: boolean } = {}) {
|
|
const [, setMessages] = useAtom(chatMessagesAtom);
|
|
const [isStreaming, setIsStreaming] = useAtom(isStreamingAtom);
|
|
const [error, setError] = useAtom(chatErrorAtom);
|
|
const setIsPreviewOpen = useSetAtom(isPreviewOpenAtom);
|
|
const [selectedAppId] = useAtom(selectedAppIdAtom);
|
|
const { refreshChats } = useChats(selectedAppId);
|
|
const { refreshApp } = useLoadApp(selectedAppId);
|
|
const setStreamCount = useSetAtom(chatStreamCountAtom);
|
|
const { refreshVersions } = useVersions(selectedAppId);
|
|
const { refreshAppIframe } = useRunApp();
|
|
const { countTokens } = useCountTokens();
|
|
|
|
let chatId: number | undefined;
|
|
|
|
if (hasChatId) {
|
|
const { id } = useSearch({ from: "/chat" });
|
|
chatId = id;
|
|
}
|
|
let { refreshProposal } = hasChatId ? useProposal(chatId) : useProposal();
|
|
|
|
const streamMessage = useCallback(
|
|
async ({
|
|
prompt,
|
|
chatId,
|
|
redo,
|
|
attachments,
|
|
}: {
|
|
prompt: string;
|
|
chatId: number;
|
|
redo?: boolean;
|
|
attachments?: File[];
|
|
}) => {
|
|
if (
|
|
(!prompt.trim() && (!attachments || attachments.length === 0)) ||
|
|
!chatId
|
|
) {
|
|
return;
|
|
}
|
|
|
|
setError(null);
|
|
setIsStreaming(true);
|
|
let hasIncrementedStreamCount = false;
|
|
try {
|
|
IpcClient.getInstance().streamMessage(prompt, {
|
|
chatId,
|
|
redo,
|
|
attachments,
|
|
onUpdate: (updatedMessages: Message[]) => {
|
|
if (!hasIncrementedStreamCount) {
|
|
setStreamCount((streamCount) => streamCount + 1);
|
|
hasIncrementedStreamCount = true;
|
|
}
|
|
|
|
setMessages(updatedMessages);
|
|
},
|
|
onEnd: (response: ChatResponseEnd) => {
|
|
if (response.updatedFiles) {
|
|
setIsPreviewOpen(true);
|
|
refreshAppIframe();
|
|
}
|
|
if (response.uncommittedFiles) {
|
|
showUncommittedFilesWarning(response.uncommittedFiles);
|
|
}
|
|
refreshProposal(chatId);
|
|
|
|
// Keep the same as below
|
|
setIsStreaming(false);
|
|
refreshChats();
|
|
refreshApp();
|
|
refreshVersions();
|
|
countTokens(chatId, "");
|
|
},
|
|
onError: (errorMessage: string) => {
|
|
console.error(`[CHAT] Stream error for ${chatId}:`, errorMessage);
|
|
setError(errorMessage);
|
|
|
|
// Keep the same as above
|
|
setIsStreaming(false);
|
|
refreshChats();
|
|
refreshApp();
|
|
refreshVersions();
|
|
countTokens(chatId, "");
|
|
},
|
|
});
|
|
} catch (error) {
|
|
console.error("[CHAT] Exception during streaming setup:", error);
|
|
setIsStreaming(false);
|
|
setError(error instanceof Error ? error.message : String(error));
|
|
}
|
|
},
|
|
[setMessages, setIsStreaming, setIsPreviewOpen],
|
|
);
|
|
|
|
return {
|
|
streamMessage,
|
|
isStreaming,
|
|
error,
|
|
setError,
|
|
setIsStreaming,
|
|
};
|
|
}
|