File size: 2,871 Bytes
38c6487 585d8da 40d9188 585d8da 259cdfe 38c6487 259cdfe 585d8da 259cdfe 491d80b 259cdfe 491d80b 259cdfe 38c6487 585d8da 259cdfe 585d8da 38c6487 259cdfe 38c6487 585d8da 38c6487 585d8da 38c6487 585d8da 259cdfe 585d8da 259cdfe 585d8da 130dce0 585d8da 259cdfe 585d8da 259cdfe 585d8da |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import { useState, useEffect, useRef, useCallback } from "react";
import { IMessage } from "../interfaces/conversation";
import { conversationWebSocket } from "../services/websockets/conversation";
interface StreamingMessage extends IMessage {
isComplete?: boolean;
}
const useWebSocket = () => {
const [messages, setMessages] = useState<StreamingMessage[]>([]);
const [streamingMessage, setStreamingMessage] = useState("");
const [isStreaming, setIsStreaming] = useState(false);
const wsRef = useRef<WebSocket | null>(null);
const simulateStreaming = (completeMessage: string) => {
setIsStreaming(true);
setStreamingMessage("");
const cleanMessage = completeMessage.replace(/<userStyle>.*<\/userStyle>/, "").trim();
const chars = cleanMessage.split("");
let currentIndex = -1;
setStreamingMessage("");
const streamInterval = setInterval(() => {
if (currentIndex < chars.length) {
currentIndex++;
setStreamingMessage((prev) => prev + chars[currentIndex]);
} else {
clearInterval(streamInterval);
setMessages((prev) => [
...prev,
{
text: cleanMessage,
sender: "bot",
isComplete: true,
},
]);
setIsStreaming(false);
setStreamingMessage("");
}
}, 20);
return () => clearInterval(streamInterval);
};
const connectWebSocket = useCallback((conversationId: string) => {
const ws = conversationWebSocket({ conversationId, modality: "text" });
wsRef.current = ws;
ws.onmessage = (event) => {
const data = JSON.parse(event.data);
if (!data.ai_message) return;
simulateStreaming(data.ai_message);
};
ws.onclose = () => {
wsRef.current = null;
};
ws.onerror = (error) => {
console.error("WebSocket error:", error);
wsRef.current = null;
};
}, []);
useEffect(() => {
return () => {
wsRef.current?.close();
wsRef.current = null;
};
}, []);
const sendMessage = useCallback((message: string) => {
if (message.trim()) {
setMessages((prev) => [
...prev,
{
text: message,
sender: "user",
isComplete: true,
},
]);
if (wsRef.current?.readyState === WebSocket.OPEN) {
wsRef.current.send(JSON.stringify({ user_message: message }));
}
}
}, []);
const getAllMessages = useCallback(() => {
const displayMessages = [...messages];
if (isStreaming) {
displayMessages.push({
text: streamingMessage,
sender: "bot",
isComplete: false,
});
}
return displayMessages;
}, [messages, isStreaming, streamingMessage]);
return {
messages: getAllMessages(),
sendMessage,
connectWebSocket,
isStreaming,
};
};
export default useWebSocket;
|