refactor: migrate document types fetch to jotai + tanstack query

- Replace useDocumentTypes hook with documentTypeCountsAtom in ChatInputGroup
- Replace useDocumentTypes hook with documentTypeCountsAtom in researcher page
- Delete obsolete use-document-types.ts hook
- Transform atom response to maintain backward compatibility
This commit is contained in:
CREDO23 2025-12-04 13:40:08 +00:00
parent e5d1598270
commit bccbd65333
3 changed files with 32 additions and 100 deletions

View file

@ -1,15 +1,15 @@
"use client";
import { type CreateMessage, type Message, useChat } from "@ai-sdk/react";
import { useAtomValue } from "jotai";
import { useAtom, useAtomValue } from "jotai";
import { useParams, useRouter } from "next/navigation";
import { useEffect, useMemo, useRef } from "react";
import { createChatMutationAtom, updateChatMutationAtom } from "@/atoms/chats/chat-mutation.atoms";
import { activeChatAtom } from "@/atoms/chats/chat-query.atoms";
import { activeChatIdAtom } from "@/atoms/chats/ui.atoms";
import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms";
import ChatInterface from "@/components/chat/ChatInterface";
import { useChatState } from "@/hooks/use-chat";
import { useDocumentTypes } from "@/hooks/use-document-types";
import type { Document } from "@/hooks/use-documents";
import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors";
@ -46,7 +46,19 @@ export default function ResearcherPage() {
});
// Fetch all available sources (document types + live search connectors)
const { documentTypes } = useDocumentTypes(Number(search_space_id));
// Use the documentTypeCountsAtom for fetching document types
const [documentTypeCountsQuery] = useAtom(documentTypeCountsAtom);
const { data: documentTypeCountsData } = documentTypeCountsQuery;
// Transform the response into the expected format
const documentTypes = useMemo(() => {
if (!documentTypeCountsData) return [];
return Object.entries(documentTypeCountsData).map(([type, count]) => ({
type,
count,
}));
}, [documentTypeCountsData]);
const { connectors: searchConnectors } = useSearchSourceConnectors(
false,
Number(search_space_id)

View file

@ -3,7 +3,9 @@
import { ChatInput } from "@llamaindex/chat-ui";
import { Brain, Check, FolderOpen, Minus, Plus, PlusCircle, Zap } from "lucide-react";
import { useParams, useRouter } from "next/navigation";
import React, { Suspense, useCallback, useState } from "react";
import React, { Suspense, useCallback, useState, useMemo } from "react";
import { useAtom } from "jotai";
import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms";
import { DocumentsDataTable } from "@/components/chat/DocumentsDataTable";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
@ -25,7 +27,6 @@ import {
} from "@/components/ui/select";
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
import { useDocumentTypes } from "@/hooks/use-document-types";
import type { Document } from "@/hooks/use-documents";
import { useGlobalLLMConfigs, useLLMConfigs, useLLMPreferences } from "@/hooks/use-llm-configs";
import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors";
@ -118,11 +119,20 @@ const ConnectorSelector = React.memo(
const router = useRouter();
const [isOpen, setIsOpen] = useState(false);
// Fetch immediately (not lazy) so the button can show the correct count
const { documentTypes, isLoading, isLoaded, fetchDocumentTypes } = useDocumentTypes(
Number(search_space_id),
false
);
// Use the documentTypeCountsAtom for fetching document types
const [documentTypeCountsQuery] = useAtom(documentTypeCountsAtom);
const { data: documentTypeCountsData, isLoading, refetch: fetchDocumentTypes } = documentTypeCountsQuery;
// Transform the response into the expected format
const documentTypes = useMemo(() => {
if (!documentTypeCountsData) return [];
return Object.entries(documentTypeCountsData).map(([type, count]) => ({
type,
count,
}));
}, [documentTypeCountsData]);
const isLoaded = !!documentTypeCountsData;
// Fetch live search connectors immediately (non-indexable)
const {

View file

@ -1,90 +0,0 @@
import { useCallback, useEffect, useState } from "react";
import { authenticatedFetch } from "@/lib/auth-utils";
export interface DocumentTypeCount {
type: string;
count: number;
}
/**
* Hook to fetch document type counts from the API
* @param searchSpaceId - The search space ID to filter document types
* @param lazy - If true, types won't be fetched on mount
*/
export const useDocumentTypes = (searchSpaceId?: number, lazy: boolean = false) => {
const [documentTypes, setDocumentTypes] = useState<DocumentTypeCount[]>([]);
const [isLoading, setIsLoading] = useState(!lazy);
const [isLoaded, setIsLoaded] = useState(false);
const [error, setError] = useState<Error | null>(null);
const fetchDocumentTypes = useCallback(
async (spaceId?: number) => {
if (isLoaded && lazy) return;
try {
setIsLoading(true);
setError(null);
// Build URL with optional search_space_id query parameter
const url = new URL(
`${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/documents/type-counts`
);
if (spaceId !== undefined) {
url.searchParams.append("search_space_id", spaceId.toString());
}
const response = await authenticatedFetch(url.toString(), {
method: "GET",
headers: { "Content-Type": "application/json" },
});
if (!response.ok) {
throw new Error(`Failed to fetch document types: ${response.statusText}`);
}
const data = await response.json();
// Convert the object to an array of DocumentTypeCount
const typeCounts: DocumentTypeCount[] = Object.entries(data).map(([type, count]) => ({
type,
count: count as number,
}));
setDocumentTypes(typeCounts);
setIsLoaded(true);
return typeCounts;
} catch (err) {
setError(err instanceof Error ? err : new Error("An unknown error occurred"));
console.error("Error fetching document types:", err);
} finally {
setIsLoading(false);
}
},
[isLoaded, lazy]
);
useEffect(() => {
if (!lazy) {
fetchDocumentTypes(searchSpaceId);
}
}, [lazy, fetchDocumentTypes, searchSpaceId]);
// Function to refresh the document types
const refreshDocumentTypes = useCallback(
async (spaceId?: number) => {
setIsLoaded(false);
await fetchDocumentTypes(spaceId !== undefined ? spaceId : searchSpaceId);
},
[fetchDocumentTypes, searchSpaceId]
);
return {
documentTypes,
isLoading,
isLoaded,
error,
fetchDocumentTypes,
refreshDocumentTypes,
};
};