Skip to content

Commit

Permalink
fix: some bugs in chat (#163)
Browse files Browse the repository at this point in the history
  • Loading branch information
@RainyNight9
RainyNight9 authored Feb 20, 2025
1 parent c520ef8 commit dd32ec5
Show file tree
Hide file tree
Showing 7 changed files with 144 additions and 269 deletions.
194 changes: 110 additions & 84 deletions src/components/Assistant/Chat.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,15 @@ import {
useImperativeHandle,
useRef,
useState,
useMemo,
} from "react";
import { MessageSquarePlus, PanelLeft } from "lucide-react";
import { isTauri } from "@tauri-apps/api/core";
import { useTranslation } from "react-i18next";
import { debounce } from "lodash-es";

import { ChatMessage } from "./ChatMessage";
import type { Chat, Message } from "./types";
import type { Chat } from "./types";
import { tauriFetch } from "@/api/tauriFetchClient";
import { useWebSocket } from "@/hooks/useWebSocket";
import { useChatStore } from "@/stores/chatStore";
Expand All @@ -22,22 +23,32 @@ import { clientEnv } from "@/utils/env";

interface ChatAIProps {
isTransitioned: boolean;
changeInput: (val: string) => void;
isSearchActive?: boolean;
isDeepThinkActive?: boolean;
isChatPage?: boolean;
activeChatProp?: Chat;
changeInput?: (val: string) => void;
}

export interface ChatAIRef {
init: (value: string) => void;
cancelChat: () => void;
connected: boolean;
reconnect: () => void;
handleSendMessage: (value: string) => void;
}

const ChatAI = memo(
forwardRef<ChatAIRef, ChatAIProps>(
(
{ isTransitioned, changeInput, isSearchActive, isDeepThinkActive },
{
isTransitioned,
changeInput,
isSearchActive,
isDeepThinkActive,
isChatPage = false,
activeChatProp,
},
ref
) => {
const { t } = useTranslation();
Expand All @@ -46,6 +57,7 @@ const ChatAI = memo(
cancelChat: cancelChat,
connected: connected,
reconnect: reconnect,
handleSendMessage: handleSendMessage,
}));

const { createWin } = useWindows();
Expand All @@ -57,17 +69,18 @@ const ChatAI = memo(
const [timedoutShow, setTimedoutShow] = useState(false);
const messagesEndRef = useRef<HTMLDivElement>(null);

const [websocketId, setWebsocketId] = useState("");
const [curMessage, setCurMessage] = useState("");
const [curId, setCurId] = useState("");

const websocketIdRef = useRef("");

const curChatEndRef = useRef(curChatEnd);
curChatEndRef.current = curChatEnd;

const curIdRef = useRef(curId);
curIdRef.current = curId;
const curIdRef = useRef("");

useEffect(() => {
activeChatProp && setActiveChat(activeChatProp);
}, [activeChatProp]);

const handleMessageChunk = useCallback((chunk: string) => {
setCurMessage((prev) => prev + chunk);
Expand All @@ -83,7 +96,6 @@ const ChatAI = memo(

if (msg.includes("websocket-session-id")) {
const array = msg.split(" ");
setWebsocketId(array[2]);
websocketIdRef.current = array[2];
return "";
} else if (msg.includes("PRIVATE")) {
Expand All @@ -105,7 +117,7 @@ const ChatAI = memo(
} else {
const cleanedData = msg.replace(/^PRIVATE /, "");
try {
console.log("cleanedData", cleanedData);
// console.log("cleanedData", cleanedData);
const chunkData = JSON.parse(cleanedData);
if (chunkData.reply_to_message === curIdRef.current) {
handleMessageChunk(chunkData.message_chunk);
Expand All @@ -117,37 +129,41 @@ const ChatAI = memo(
}
}
}
}, []);
}, [curChatEnd, isTyping]);

const { messages, setMessages, connected, reconnect } = useWebSocket(
clientEnv.COCO_WEBSOCKET_URL,
dealMsg
);

const simulateAssistantResponse = useCallback(() => {
if (!activeChat?._id) return;

// console.log("curMessage", curMessage);

const assistantMessage: Message = {
const assistantMessage = useMemo(() => {
if (!activeChat?._id || (!curMessage && !messages)) return null;
return {
_id: activeChat._id,
_source: {
type: "assistant",
message: curMessage || messages,
},
};
}, [activeChat?._id, curMessage, messages]);

const updatedChat = {
const updatedChat = useMemo(() => {
if (!activeChat?._id || !assistantMessage) return null;
return {
...activeChat,
messages: [...(activeChat.messages || []), assistantMessage],
};
}, [activeChat, assistantMessage]);

const simulateAssistantResponse = useCallback(() => {
if (!updatedChat) return;

// console.log("updatedChat:", updatedChat);
setActiveChat(updatedChat);
setMessages("");
setCurMessage("");
setIsTyping(false);
}, [activeChat?._id, curMessage, messages]);
}, [updatedChat]);

useEffect(() => {
if (curChatEnd) {
Expand Down Expand Up @@ -182,7 +198,7 @@ const ChatAI = memo(
});
console.log("_new", response);
const newChat: Chat = response.data;

setActiveChat(newChat);
handleSendMessage(value, newChat);
} catch (error) {
Expand All @@ -199,36 +215,42 @@ const ChatAI = memo(
}
};

const handleSendMessage = useCallback(async (content: string, newChat?: Chat) => {
newChat = newChat || activeChat;
if (!newChat?._id || !content) return;
setTimedoutShow(false);
try {
const response = await tauriFetch({
url: `/chat/${newChat?._id}/_send?search=${isSearchActive}&deep_thinking=${isDeepThinkActive}`,
method: "POST",
headers: {
"WEBSOCKET-SESSION-ID": websocketIdRef.current || websocketId
},
body: JSON.stringify({ message: content }),
});
console.log("_send", response, websocketId);
setCurId(response.data[0]?._id);

const updatedChat: Chat = {
...newChat,
messages: [...(newChat?.messages || []), ...(response.data || [])],
};

changeInput("");
// console.log("updatedChat2", updatedChat);
setActiveChat(updatedChat);
setIsTyping(true);
setCurChatEnd(false);
} catch (error) {
console.error("Failed to fetch user data:", error);
}
}, [JSON.stringify(activeChat), websocketId]);
const handleSendMessage = useCallback(
async (content: string, newChat?: Chat) => {
newChat = newChat || activeChat;
if (!newChat?._id || !content) return;
setTimedoutShow(false);
try {
const response = await tauriFetch({
url: `/chat/${newChat?._id}/_send?search=${isSearchActive}&deep_thinking=${isDeepThinkActive}`,
method: "POST",
headers: {
"WEBSOCKET-SESSION-ID": websocketIdRef.current,
},
body: JSON.stringify({ message: content }),
});
console.log("_send", response, websocketIdRef.current);
curIdRef.current = response.data[0]?._id;

const updatedChat: Chat = {
...newChat,
messages: [
...(newChat?.messages || []),
...(response.data || []),
],
};

changeInput && changeInput("");
// console.log("updatedChat2", updatedChat);
setActiveChat(updatedChat);
setIsTyping(true);
setCurChatEnd(false);
} catch (error) {
console.error("Failed to fetch user data:", error);
}
},
[activeChat?._id]
);

const chatClose = async () => {
if (!activeChat?._id) return;
Expand Down Expand Up @@ -302,39 +324,41 @@ const ChatAI = memo(
return (
<div
data-tauri-drag-region
className={`h-[500px] flex flex-col rounded-xl overflow-hidden`}
className={`h-full flex flex-col rounded-xl overflow-hidden`}
>
<header
data-tauri-drag-region
className={`flex items-center justify-between py-2 px-1`}
>
<button
onClick={() => openChatAI()}
className={`p-2 rounded-lg transition-colors text-[#333] dark:text-[#d8d8d8]`}
{isChatPage ? null : (
<header
data-tauri-drag-region
className={`flex items-center justify-between py-2 px-1`}
>
<PanelLeft className="h-4 w-4" />
</button>

<button
onClick={() => {
createNewChat();
}}
className={`p-2 rounded-lg transition-colors text-[#333] dark:text-[#d8d8d8]`}
>
<MessageSquarePlus className="h-4 w-4" />
</button>
</header>
<button
onClick={() => openChatAI()}
className={`p-2 rounded-lg transition-colors text-[#333] dark:text-[#d8d8d8]`}
>
<PanelLeft className="h-4 w-4" />
</button>

<button
onClick={() => {
createNewChat();
}}
className={`p-2 rounded-lg transition-colors text-[#333] dark:text-[#d8d8d8]`}
>
<MessageSquarePlus className="h-4 w-4" />
</button>
</header>
)}

{/* Chat messages */}
<div className="w-full overflow-x-hidden overflow-y-auto border-t border-[rgba(0,0,0,0.1)] dark:border-[rgba(255,255,255,0.15)] custom-scrollbar relative">
<ChatMessage
key={"greetings"}
message={{
_id: 'greetings',
_id: "greetings",
_source: {
type: 'assistant',
message: t("assistant.chat.greetings")
}
type: "assistant",
message: t("assistant.chat.greetings"),
},
}}
isTyping={false}
/>
Expand Down Expand Up @@ -365,17 +389,19 @@ const ChatAI = memo(
/>
) : null}

{timedoutShow ? <ChatMessage
key={"timedout"}
message={{
_id: 'timedout',
_source: {
type: 'assistant',
message: t("assistant.chat.timedout")
}
}}
isTyping={false}
/>: null}
{timedoutShow ? (
<ChatMessage
key={"timedout"}
message={{
_id: "timedout",
_source: {
type: "assistant",
message: t("assistant.chat.timedout"),
},
}}
isTyping={false}
/>
) : null}

<div ref={messagesEndRef} />
</div>
Expand Down
28 changes: 7 additions & 21 deletions src/components/Assistant/ChatMessage.tsx
View file
Original file line number Diff line number Diff line change
@@ -1,37 +1,25 @@
import { Brain, ChevronDown, ChevronUp } from "lucide-react";
import { useState, useEffect, useRef } from "react";
import { useState, memo } from "react";
import { useTranslation } from "react-i18next";

import type { Message } from "./types";
import Markdown from "./Markdown";
import { formatThinkingMessage } from "@/utils/index";
import logoImg from "@/assets/icon.svg";
import { SourceResult } from "./SourceResult";
import { useTranslation } from "react-i18next";

interface ChatMessageProps {
message: Message;
isTyping?: boolean;
}

export function ChatMessage({ message, isTyping }: ChatMessageProps) {
export const ChatMessage = memo(function ChatMessage({ message, isTyping }: ChatMessageProps) {
const { t } = useTranslation();
const [isThinkingExpanded, setIsThinkingExpanded] = useState(true);
const [responseTime, setResponseTime] = useState(0);
const startTimeRef = useRef<number | null>(null);
const hasStartedRef = useRef(false);

const isAssistant = message._source?.type === "assistant";
const segments = formatThinkingMessage(message._source.message);

useEffect(() => {
if (isTyping && !hasStartedRef.current) {
startTimeRef.current = Date.now();
hasStartedRef.current = true;
} else if (!isTyping && hasStartedRef.current && startTimeRef.current) {
const duration = (Date.now() - startTimeRef.current) / 1000;
setResponseTime(duration);
hasStartedRef.current = false;
}
}, [isTyping]);
const segments = formatThinkingMessage(message._source.message);

return (
<div
Expand Down Expand Up @@ -85,9 +73,7 @@ export function ChatMessage({ message, isTyping }: ChatMessageProps) {
<>
<Brain className="w-4 h-4 text-[#999999]" />
<span className="text-xs text-[#999999]">
{t("assistant.message.thoughtTime", {
time: responseTime.toFixed(1),
})}
{t("assistant.message.thoughtTime")}
</span>
</>
)}
Expand Down Expand Up @@ -140,4 +126,4 @@ export function ChatMessage({ message, isTyping }: ChatMessageProps) {
</div>
</div>
);
}
})
Loading

0 comments on commit dd32ec5

Please sign in to comment.