Spaces:
Sleeping
Sleeping
| import { useState, useRef, useEffect } from 'react'; | |
| import type { Message } from '../types'; | |
| import { streamChat } from '../api/client'; | |
| // Helper to parse quota errors and generate user-friendly messages | |
| function parseQuotaError(content: string): { message: string; retryAfter?: number } | null { | |
| const rpmMatch = content.match(/\[QUOTA_ERROR:rpm_exhausted:try_after_(\d+)_seconds\]/); | |
| if (rpmMatch) { | |
| return { | |
| message: 'Rate limit exceeded (requests per minute). Please wait before sending another message.', | |
| retryAfter: parseInt(rpmMatch[1], 10), | |
| }; | |
| } | |
| const tpdMatch = content.match(/\[QUOTA_ERROR:tpd_exhausted:try_after_midnight\]/); | |
| if (tpdMatch) { | |
| return { | |
| message: 'Token limit exceeded for today. Please try again after 12:00 AM.', | |
| retryAfter: undefined, | |
| }; | |
| } | |
| const genericMatch = content.match(/\[QUOTA_ERROR:(\w+):([^\]]+)\]/); | |
| if (genericMatch) { | |
| return { | |
| message: `Rate limit exceeded. Please try again later.`, | |
| retryAfter: undefined, | |
| }; | |
| } | |
| return null; | |
| } | |
| export function useChat(selectedDocId: string | null) { | |
| const [messages, setMessages] = useState<Message[]>([]); | |
| const [inputText, setInputText] = useState(''); | |
| const [isStreaming, setIsStreaming] = useState(false); | |
| const messagesEndRef = useRef<HTMLDivElement>(null); | |
| // Scroll to bottom on new messages | |
| useEffect(() => { | |
| messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' }); | |
| }, [messages, isStreaming]); | |
| const handleSendMessage = async (e: React.FormEvent) => { | |
| e.preventDefault(); | |
| if (!inputText.trim() || isStreaming) return; | |
| const userMessage: Message = { role: 'user', content: inputText.trim() }; | |
| const updatedMessages = [...messages, userMessage]; | |
| setMessages(updatedMessages); | |
| setInputText(''); | |
| setIsStreaming(true); | |
| // Placeholder for streaming assistant response | |
| setMessages((prev) => [...prev, { role: 'assistant', content: '' }]); | |
| try { | |
| await streamChat( | |
| { messages: updatedMessages, document_id: selectedDocId }, | |
| { | |
| onChunk: (displayContent, sources) => { | |
| // Check for quota errors in the stream | |
| const quotaError = parseQuotaError(displayContent); | |
| if (quotaError) { | |
| setMessages((prev) => { | |
| const next = [...prev]; | |
| if (next.length > 0) { | |
| next[next.length - 1] = { | |
| role: 'assistant', | |
| content: quotaError.message, | |
| }; | |
| } | |
| return next; | |
| }); | |
| setIsStreaming(false); | |
| return; | |
| } | |
| setMessages((prev) => { | |
| const next = [...prev]; | |
| if (next.length > 0) { | |
| next[next.length - 1] = { | |
| role: 'assistant', | |
| content: displayContent, | |
| sources, | |
| }; | |
| } | |
| return next; | |
| }); | |
| }, | |
| onError: (message) => { | |
| setMessages((prev) => { | |
| const next = [...prev]; | |
| if (next.length > 0) { | |
| next[next.length - 1] = { role: 'assistant', content: `Error: ${message}` }; | |
| } | |
| return next; | |
| }); | |
| }, | |
| } | |
| ); | |
| } catch (err: unknown) { | |
| console.error(err); | |
| setMessages((prev) => { | |
| const next = [...prev]; | |
| if (next.length > 0) { | |
| next[next.length - 1] = { | |
| role: 'assistant', | |
| content: `Error: ${err instanceof Error ? err.message : 'An error occurred while streaming response.'}`, | |
| }; | |
| } | |
| return next; | |
| }); | |
| } finally { | |
| setIsStreaming(false); | |
| } | |
| }; | |
| const handleClearChat = () => setMessages([]); | |
| return { | |
| messages, | |
| inputText, | |
| setInputText, | |
| isStreaming, | |
| messagesEndRef, | |
| handleSendMessage, | |
| handleClearChat, | |
| }; | |
| } | |