| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import { useState, useRef, useCallback, useEffect } from 'react' |
|
|
| const WS_BASE = import.meta.env.VITE_WS_URL || |
| (window.location.protocol === 'https:' ? 'wss://' : 'ws://') + |
| window.location.host + '/api/ws' |
|
|
| export function useStreamingChat(sessionId) { |
| const [isStreaming, setIsStreaming] = useState(false) |
| const [currentTokens, setCurrentTokens] = useState('') |
| const [toolEvents, setToolEvents] = useState([]) |
| const [error, setError] = useState(null) |
|
|
| const ws = useRef(null) |
| const onDoneRef = useRef(null) |
|
|
| |
| const connect = useCallback(() => { |
| if (!sessionId) return |
| if (ws.current?.readyState === WebSocket.OPEN) return |
|
|
| const url = `${WS_BASE}/chat/${sessionId}` |
| ws.current = new WebSocket(url) |
|
|
| ws.current.onopen = () => { |
| setError(null) |
| } |
|
|
| ws.current.onmessage = (event) => { |
| try { |
| const msg = JSON.parse(event.data) |
|
|
| switch (msg.type) { |
| case 'thinking': |
| setCurrentTokens('...') |
| break |
|
|
| case 'token': |
| setCurrentTokens((prev) => prev === '...' ? msg.data : prev + msg.data) |
| break |
|
|
| case 'tool_use': |
| setToolEvents((prev) => [...prev, msg.data]) |
| break |
|
|
| case 'done': |
| setIsStreaming(false) |
| if (onDoneRef.current) { |
| onDoneRef.current({ |
| text: currentTokens, |
| used_tools: msg.data.used_tools || [], |
| memory_hits: msg.data.memory_hits || 0, |
| }) |
| } |
| setCurrentTokens('') |
| setToolEvents([]) |
| break |
|
|
| case 'error': |
| setError(msg.data) |
| setIsStreaming(false) |
| setCurrentTokens('') |
| break |
|
|
| default: |
| break |
| } |
| } catch (e) { |
| console.error('WS parse error:', e) |
| } |
| } |
|
|
| ws.current.onerror = () => { |
| setError('WebSocket connection error') |
| setIsStreaming(false) |
| } |
|
|
| ws.current.onclose = () => { |
| |
| setTimeout(() => { |
| if (sessionId) connect() |
| }, 2000) |
| } |
| }, [sessionId]) |
|
|
| |
| useEffect(() => { |
| connect() |
| return () => { |
| ws.current?.close() |
| } |
| }, [connect]) |
|
|
| |
| const sendMessage = useCallback( |
| (message, imageCaption = null, onDone = null) => { |
| if (!ws.current || ws.current.readyState !== WebSocket.OPEN) { |
| connect() |
| |
| setTimeout(() => sendMessage(message, imageCaption, onDone), 500) |
| return |
| } |
|
|
| setIsStreaming(true) |
| setCurrentTokens('') |
| setToolEvents([]) |
| setError(null) |
| onDoneRef.current = onDone |
|
|
| ws.current.send(JSON.stringify({ |
| message, |
| image_caption: imageCaption, |
| })) |
| }, |
| [connect], |
| ) |
|
|
| |
| const disconnect = useCallback(() => { |
| ws.current?.close() |
| }, []) |
|
|
| return { |
| sendMessage, |
| disconnect, |
| isStreaming, |
| currentTokens, |
| toolEvents, |
| error, |
| isConnected: ws.current?.readyState === WebSocket.OPEN, |
| } |
| } |
|
|
