'use client'; import { useState, useCallback, useRef } from 'react'; import type { Message, Citation, Widget, StreamEvent, ChatAttachmentInfo } from '../types'; import { streamChat, generateId } from '../api'; import { getLanguageFromStorage } from '../contexts/LanguageContext'; import type { SendOptions, ChatAttachment } from '@/components/ChatInput'; function fileToDataUrl(file: File): Promise { return new Promise((resolve, reject) => { const reader = new FileReader(); reader.onload = () => resolve(reader.result as string); reader.onerror = reject; reader.readAsDataURL(file); }); } interface UseChatOptions { onError?: (error: Error) => void; } export function useChat(options: UseChatOptions = {}) { const [messages, setMessages] = useState([]); const [isLoading, setIsLoading] = useState(false); const [citations, setCitations] = useState([]); const [widgets, setWidgets] = useState([]); const abortControllerRef = useRef(null); const sendMessage = useCallback(async (content: string, sendOptions?: SendOptions | 'speed' | 'balanced' | 'quality') => { const mode = typeof sendOptions === 'string' ? sendOptions : sendOptions?.mode ?? 'balanced'; const webSearch = typeof sendOptions === 'object' ? sendOptions.webSearch : false; const attachments = typeof sendOptions === 'object' ? sendOptions.attachments : []; const locale = getLanguageFromStorage(); if (!content.trim() || isLoading) return; const userMessage: Message = { id: generateId(), role: 'user', content: content.trim(), createdAt: new Date(), }; const assistantMessage: Message = { id: generateId(), role: 'assistant', content: '', citations: [], widgets: [], isStreaming: true, createdAt: new Date(), }; setMessages((prev) => [...prev, userMessage, assistantMessage]); setIsLoading(true); setCitations([]); setWidgets([]); const history: [string, string][] = messages .filter((m) => !m.isStreaming) .reduce((acc, m, i, arr) => { if (m.role === 'user' && arr[i + 1]?.role === 'assistant') { acc.push([m.content, arr[i + 1].content]); } return acc; }, [] as [string, string][]); const attachmentInfos = await Promise.all( attachments.map(async (att) => { let dataUrl: string | undefined; if (att.type === 'image' || att.file.type.startsWith('image/')) { dataUrl = await fileToDataUrl(att.file); } return { name: att.file.name, type: att.file.type, size: att.file.size, dataUrl, }; }) ); try { const stream = streamChat({ message: { messageId: assistantMessage.id, chatId: userMessage.id, content: content.trim(), }, optimizationMode: mode, history, locale, webSearch, attachments: attachmentInfos.length > 0 ? attachmentInfos : undefined, }); let fullContent = ''; const collectedCitations: Citation[] = []; const collectedWidgets: Widget[] = []; for await (const event of stream) { switch (event.type) { case 'messageStart': break; case 'message': if (event.data && typeof event.data === 'object' && 'content' in event.data) { const chunk = (event.data as { content: string }).content; fullContent += chunk; setMessages((prev) => prev.map((m) => m.id === assistantMessage.id ? { ...m, content: fullContent } : m ) ); } break; case 'block': // Handle new block format from agent-svc if (event.block && typeof event.block === 'object') { const block = event.block as { type: string; data: unknown }; if (block.type === 'text' && typeof block.data === 'string') { fullContent = block.data; setMessages((prev) => prev.map((m) => m.id === assistantMessage.id ? { ...m, content: fullContent } : m ) ); } else if (block.type === 'widget' && block.data) { const widget = block.data as Widget; collectedWidgets.push(widget); setWidgets([...collectedWidgets]); } } break; case 'textChunk': // Handle streaming text chunks if (event.chunk && typeof event.chunk === 'string') { fullContent += event.chunk; setMessages((prev) => prev.map((m) => m.id === assistantMessage.id ? { ...m, content: fullContent } : m ) ); } break; case 'updateBlock': // Handle block updates (final text) if (event.patch && Array.isArray(event.patch)) { for (const p of event.patch) { if (p.op === 'replace' && p.path === '/data' && typeof p.value === 'string') { fullContent = p.value; setMessages((prev) => prev.map((m) => m.id === assistantMessage.id ? { ...m, content: fullContent } : m ) ); } } } break; case 'sources': if (event.data && Array.isArray(event.data)) { const sources = event.data as Array<{ url: string; title?: string; metadata?: { title?: string }; }>; sources.forEach((source, idx) => { const citation: Citation = { index: idx + 1, url: source.url, title: source.title || source.metadata?.title || new URL(source.url).hostname, domain: new URL(source.url).hostname, }; collectedCitations.push(citation); }); setCitations([...collectedCitations]); } break; case 'widget': if (event.data && typeof event.data === 'object') { const widget = event.data as Widget; collectedWidgets.push(widget); setWidgets([...collectedWidgets]); } break; case 'researchComplete': // Research phase complete, text generation starting break; case 'messageEnd': setMessages((prev) => prev.map((m) => m.id === assistantMessage.id ? { ...m, content: fullContent, citations: collectedCitations, widgets: collectedWidgets, isStreaming: false, } : m ) ); break; case 'error': throw new Error( event.data && typeof event.data === 'object' && 'message' in event.data ? String((event.data as { message: string }).message) : 'Unknown error' ); } } } catch (error) { const errorMessage = error instanceof Error ? error.message : 'Ошибка при отправке сообщения'; setMessages((prev) => prev.map((m) => m.id === assistantMessage.id ? { ...m, content: `Ошибка: ${errorMessage}`, isStreaming: false } : m ) ); options.onError?.(error instanceof Error ? error : new Error(errorMessage)); } finally { setIsLoading(false); } }, [isLoading, messages, options]); const stopGeneration = useCallback(() => { abortControllerRef.current?.abort(); setIsLoading(false); setMessages((prev) => prev.map((m) => (m.isStreaming ? { ...m, isStreaming: false } : m)) ); }, []); const clearMessages = useCallback(() => { setMessages([]); setCitations([]); setWidgets([]); }, []); return { messages, isLoading, citations, widgets, sendMessage, stopGeneration, clearMessages, }; }