legco_ai_assistant/frontend/src/lib/queries.tsx

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 } from './api'
import type { QueryRequest, QueryResponse, QueryStreamEvent, SourceMetadata, IngestResponse, DocumentListResponse, ChunkInfo, DeleteResponse, PromptProfileListResponse, PromptSetResponse, PromptUpdateRequest, PromptBatchUpdateRequest, PromptActivateResponse, PromptStatusResponse } from '../types'
import { useState, useCallback, useRef } from 'react'

export const queryClient = new QueryClient()

export const useQueryDocument = () => {
  return useMutation<QueryResponse, Error, QueryRequest>({
    mutationFn: queryDocument,
  })
}

export interface QueryStreamState {
  extractedQuestions: string[] | null
  answer: string | null
  sources: SourceMetadata[] | null
  phase: 'idle' | 'decomposing' | 'retrieving' | 'filtering' | 'generating' | 'completed' | 'error'
  error: Error | null
}

export const useQueryDocumentStream = () => {
  const [state, setState] = useState<QueryStreamState>({
    extractedQuestions: null,
    answer: null,
    sources: null,
    phase: 'idle',
    error: null,
  })
  const abortRef = useRef<AbortController | null>(null)

  const mutate = useCallback(async (request: QueryRequest) => {
    setState({
      extractedQuestions: null,
      answer: null,
      sources: null,
      phase: 'decomposing',
      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 'completed':
            setState(prev => ({
              ...prev,
              answer: event.answer ?? null,
              sources: event.sources ?? null,
              phase: 'completed',
            }))
            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,
      phase: 'idle',
      error: null,
    })
  }, [])

  return { ...state, mutate, reset }
}

export const useIngestDocument = () => {
  return useMutation<IngestResponse, Error, File>({
    mutationFn: ingestDocument,
  })
}

export const useDocuments = () => {
  return useQuery<DocumentListResponse, Error>({
    queryKey: ['documents'],
    queryFn: listDocuments,
  })
}

export const useDocumentChunks = (documentId: string | null) => {
  return useQuery<ChunkInfo[], Error>({
    queryKey: ['documents', documentId, 'chunks'],
    queryFn: () => listChunks(documentId!),
    enabled: documentId !== null,
  })
}

export const useDeleteDocument = () => {
  const queryClient = useQueryClient()
  return useMutation<DeleteResponse, Error, string>({
    mutationFn: deleteDocument,
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['documents'] })
    },
  })
}

export const useDeleteChunk = () => {
  const queryClient = useQueryClient()
  return useMutation<DeleteResponse, Error, string>({
    mutationFn: deleteChunk,
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['documents'] })
    },
  })
}

export const usePromptProfiles = () => {
  return useQuery<PromptProfileListResponse, Error>({
    queryKey: ['prompts', 'profiles'],
    queryFn: listPromptProfiles,
  })
}

export const usePromptProfile = (name: string | null) => {
  return useQuery<PromptSetResponse, Error>({
    queryKey: ['prompts', 'profiles', name],
    queryFn: () => getPromptProfile(name!),
    enabled: name !== null,
  })
}

export const useActivateProfile = () => {
  const queryClient = useQueryClient()
  return useMutation<PromptActivateResponse, Error, string>({
    mutationFn: activatePromptProfile,
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['prompts'] })
    },
  })
}

export const useUpdatePrompt = () => {
  const queryClient = useQueryClient()
  return useMutation<PromptStatusResponse, Error, { name: string; step: string; request: PromptUpdateRequest }>({
    mutationFn: ({ name, step, request }) => updatePrompt(name, step, request),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['prompts'] })
    },
  })
}

export const useUpdateAllPrompts = () => {
  const queryClient = useQueryClient()
  return useMutation<PromptStatusResponse, Error, { name: string; request: PromptBatchUpdateRequest }>({
    mutationFn: ({ name, request }) => updateAllPrompts(name, request),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['prompts'] })
    },
  })
}

export const useResetPrompts = () => {
  const queryClient = useQueryClient()
  return useMutation<PromptStatusResponse, Error, { name: string; step?: string }>({
    mutationFn: ({ name, step }) => resetPrompts(name, step),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['prompts'] })
    },
  })
}

export const AppQueryProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
  return <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
}