import React from 'react' import { QueryClient, QueryClientProvider, useMutation, useQuery, useQueryClient } from '@tanstack/react-query' import { queryDocument, queryDocumentStream, ingestDocument, listDocuments, listChunks, deleteDocument, deleteChunk, listPromptProfiles, getPromptProfile, activatePromptProfile, updatePrompt, updateAllPrompts, resetPrompts, exportProfile, importProfile, listQueryHistory, getQueryHistoryDetail, deleteQueryHistory, clearQueryHistory, getHistoryStats, uploadVideo } from './api' import type { QueryRequest, QueryResponse, QueryStreamEvent, SourceMetadata, SubQuestionSources, IngestResponse, DocumentListResponse, ChunkInfo, DeleteResponse, PromptProfileListResponse, PromptSetResponse, PromptUpdateRequest, PromptBatchUpdateRequest, PromptActivateResponse, PromptStatusResponse, ProfileExportData, ProfileImportResponse, QueryHistoryList, QueryHistoryDetail, HistoryStats, HistoryDeleteResponse, VideoUploadResponse } from '../types' import { useState, useCallback, useRef } from 'react' export const queryClient = new QueryClient() export const useQueryDocument = () => { return useMutation({ mutationFn: queryDocument, }) } export interface QueryStreamState { extractedQuestions: string[] | null answer: string | null sources: SourceMetadata[] | null subQuestionSources: SubQuestionSources[] | null phase: 'idle' | 'decomposing' | 'retrieving' | 'filtering' | 'generating' | 'completed' | 'error' historyId: number | null error: Error | null } export const useQueryDocumentStream = () => { const [state, setState] = useState({ extractedQuestions: null, answer: null, sources: null, subQuestionSources: null, phase: 'idle', historyId: null, error: null, }) const abortRef = useRef(null) const mutate = useCallback(async (request: QueryRequest) => { setState({ extractedQuestions: null, answer: null, sources: null, subQuestionSources: null, phase: 'decomposing', historyId: null, error: null, }) abortRef.current = new AbortController() try { await queryDocumentStream(request, (event: QueryStreamEvent) => { switch (event.phase) { case 'decomposed': setState(prev => ({ ...prev, extractedQuestions: event.extracted_questions ?? null, phase: 'retrieving', })) break case 'retrieving': setState(prev => ({ ...prev, phase: 'retrieving' })) break case 'filtering': setState(prev => ({ ...prev, phase: 'filtering' })) break case 'generating': setState(prev => ({ ...prev, phase: 'generating' })) break case 'generating_subquestion': setState(prev => ({ ...prev, phase: 'generating' })) break case 'completed': setState(prev => ({ ...prev, answer: event.answer ?? null, sources: event.sources ?? null, subQuestionSources: event.sub_question_sources ?? null, phase: 'completed', historyId: (event as any).history_id ?? null, })) break case 'error': setState(prev => ({ ...prev, phase: 'error', error: new Error(event.message ?? 'Unknown error'), })) break } }, abortRef.current.signal) } catch (err) { if (err instanceof Error && err.name === 'AbortError') { setState(prev => ({ ...prev, phase: 'idle' })) return } setState(prev => ({ ...prev, phase: 'error', error: err instanceof Error ? err : new Error(String(err)), })) } }, []) const decomposeOnly = useCallback(async (request: QueryRequest) => { setState({ extractedQuestions: null, answer: null, sources: null, subQuestionSources: null, phase: 'decomposing', historyId: null, error: null, }) abortRef.current = new AbortController() try { await queryDocumentStream({ question: request.question, stop_after_decompose: true }, (event: QueryStreamEvent) => { switch (event.phase) { case 'decomposed': setState(prev => ({ ...prev, extractedQuestions: event.extracted_questions ?? null, phase: 'retrieving', })) break case 'retrieving': setState(prev => ({ ...prev, phase: 'retrieving' })) break case 'filtering': setState(prev => ({ ...prev, phase: 'filtering' })) break case 'generating': setState(prev => ({ ...prev, phase: 'generating' })) break case 'generating_subquestion': setState(prev => ({ ...prev, phase: 'generating' })) break case 'completed': setState(prev => ({ ...prev, answer: event.half_question ? null : (event.answer ?? null), sources: event.sources ?? null, subQuestionSources: event.sub_question_sources ?? null, phase: 'completed', historyId: (event as any).history_id ?? null, })) break case 'error': setState(prev => ({ ...prev, phase: 'error', error: new Error(event.message ?? 'Unknown error'), })) break } }, abortRef.current.signal) } catch (err) { if (err instanceof Error && err.name === 'AbortError') { setState(prev => ({ ...prev, phase: 'idle' })) return } setState(prev => ({ ...prev, phase: 'error', error: err instanceof Error ? err : new Error(String(err)), })) } }, []) const reset = useCallback(() => { abortRef.current?.abort() setState({ extractedQuestions: null, answer: null, sources: null, subQuestionSources: null, phase: 'idle', historyId: null, error: null, }) }, []) return { ...state, mutate, decomposeOnly, reset } } export const useIngestDocument = () => { return useMutation({ mutationFn: ingestDocument, }) } export const useDocuments = () => { return useQuery({ queryKey: ['documents'], queryFn: listDocuments, }) } export const useDocumentChunks = (documentId: string | null) => { return useQuery({ queryKey: ['documents', documentId, 'chunks'], queryFn: () => listChunks(documentId!), enabled: documentId !== null, }) } export const useDeleteDocument = () => { const queryClient = useQueryClient() return useMutation({ mutationFn: deleteDocument, onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['documents'] }) }, }) } export const useDeleteChunk = () => { const queryClient = useQueryClient() return useMutation({ mutationFn: deleteChunk, onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['documents'] }) }, }) } export const usePromptProfiles = () => { return useQuery({ queryKey: ['prompts', 'profiles'], queryFn: listPromptProfiles, }) } export const usePromptProfile = (name: string | null) => { return useQuery({ queryKey: ['prompts', 'profiles', name], queryFn: () => getPromptProfile(name!), enabled: name !== null, }) } export const useActivateProfile = () => { const queryClient = useQueryClient() return useMutation({ mutationFn: activatePromptProfile, onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['prompts'] }) }, }) } export const useUpdatePrompt = () => { const queryClient = useQueryClient() return useMutation({ mutationFn: ({ name, step, request }) => updatePrompt(name, step, request), onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['prompts'] }) }, }) } export const useUpdateAllPrompts = () => { const queryClient = useQueryClient() return useMutation({ mutationFn: ({ name, request }) => updateAllPrompts(name, request), onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['prompts'] }) }, }) } export const useResetPrompts = () => { const queryClient = useQueryClient() return useMutation({ mutationFn: ({ name, step }) => resetPrompts(name, step), onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['prompts'] }) }, }) } export const useImportProfile = () => { const queryClient = useQueryClient() return useMutation({ mutationFn: ({ name, data }) => importProfile(name, data), onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['prompts'] }) }, }) } export const useQueryHistoryList = (limit = 50, offset = 0) => { return useQuery({ queryKey: ['history', { limit, offset }], queryFn: () => listQueryHistory(limit, offset), }) } export const useQueryHistoryDetail = (id: number | null) => { return useQuery({ queryKey: ['history', id], queryFn: () => getQueryHistoryDetail(id!), enabled: id !== null, }) } export const useDeleteQueryHistory = () => { const queryClient = useQueryClient() return useMutation({ mutationFn: deleteQueryHistory, onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['history'] }) }, }) } export const useClearQueryHistory = () => { const queryClient = useQueryClient() return useMutation({ mutationFn: clearQueryHistory, onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['history'] }) }, }) } export const useHistoryStats = () => { return useQuery({ queryKey: ['history', 'stats'], queryFn: getHistoryStats, }) } export const AppQueryProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => { return {children} } export const useVideoUpload = () => { return useMutation void }>({ mutationFn: ({ file, onProgress }) => uploadVideo(file, onProgress), }) }