legco_ai_assistant/frontend/src/lib/queries.tsx

58 lines
1.8 KiB
TypeScript

import React from 'react'
import { QueryClient, QueryClientProvider, useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
import { queryDocument, ingestDocument, listDocuments, listChunks, deleteDocument, deleteChunk } from './api'
import type { QueryRequest, QueryResponse, IngestResponse, DocumentListResponse, ChunkInfo, DeleteResponse } from '../types'
export const queryClient = new QueryClient()
export const useQueryDocument = () => {
return useMutation<QueryResponse, Error, QueryRequest>({
mutationFn: queryDocument,
})
}
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 AppQueryProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
return <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
}