From 3f051b0a194021b00712ef9a13a6cab3d274e12c Mon Sep 17 00:00:00 2001 From: Utkarsh-Patel-13 Date: Tue, 22 Jul 2025 14:06:10 -0700 Subject: [PATCH] Added document selection and document selection dialog to new chat interface --- .../v2/[[...chat_id]]/page.tsx | 73 ++- .../components/chat_v2/ChatInterface.tsx | 187 ++++++ surfsense_web/components/chat_v2/ChatMain.tsx | 11 - .../components/chat_v2/DocumentsDataTable.tsx | 530 ++++++++++++++++++ surfsense_web/components/ui/table.tsx | 183 +++--- surfsense_web/hooks/use-documents.ts | 222 ++++---- surfsense_web/hooks/useChat.ts | 209 ++++--- 7 files changed, 1119 insertions(+), 296 deletions(-) create mode 100644 surfsense_web/components/chat_v2/ChatInterface.tsx delete mode 100644 surfsense_web/components/chat_v2/ChatMain.tsx create mode 100644 surfsense_web/components/chat_v2/DocumentsDataTable.tsx diff --git a/surfsense_web/app/dashboard/[search_space_id]/v2/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/v2/[[...chat_id]]/page.tsx index 52ce7d8..d3ea1a6 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/v2/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/v2/[[...chat_id]]/page.tsx @@ -2,10 +2,11 @@ import { useChat, Message, CreateMessage } from "@ai-sdk/react"; import { useParams, useRouter } from "next/navigation"; -import { useEffect } from "react"; -import ChatMain from "@/components/chat_v2/ChatMain"; +import { useEffect, useMemo } from "react"; +import ChatInterface from "@/components/chat_v2/ChatInterface"; import { ResearchMode } from "@/components/chat"; import { useChatState, useChatAPI } from "@/hooks/useChat"; +import { Document } from "@/hooks/use-documents"; export default function ResearchChatPageV2() { const { search_space_id, chat_id } = useParams(); @@ -23,6 +24,8 @@ export default function ResearchChatPageV2() { setResearchMode, selectedConnectors, setSelectedConnectors, + selectedDocuments, + setSelectedDocuments, } = useChatState({ search_space_id: search_space_id as string, chat_id: chatIdParam, @@ -35,7 +38,42 @@ export default function ResearchChatPageV2() { selectedConnectors, }); - // Single useChat handler for both cases + // Memoize document IDs to prevent infinite re-renders + const documentIds = useMemo(() => { + return selectedDocuments.map((doc) => doc.id); + }, [selectedDocuments]); + + // Helper functions for localStorage management + const getStorageKey = (searchSpaceId: string, chatId: string) => + `surfsense_selected_docs_${searchSpaceId}_${chatId}`; + + const storeSelectedDocuments = ( + searchSpaceId: string, + chatId: string, + documents: Document[] + ) => { + const key = getStorageKey(searchSpaceId, chatId); + localStorage.setItem(key, JSON.stringify(documents)); + }; + + const restoreSelectedDocuments = ( + searchSpaceId: string, + chatId: string + ): Document[] | null => { + const key = getStorageKey(searchSpaceId, chatId); + const stored = localStorage.getItem(key); + if (stored) { + localStorage.removeItem(key); // Clean up after restoration + try { + return JSON.parse(stored); + } catch (error) { + console.error("Error parsing stored documents:", error); + return null; + } + } + return null; + }; + const handler = useChat({ api: `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chat`, streamProtocol: "data", @@ -49,7 +87,7 @@ export default function ResearchChatPageV2() { selected_connectors: selectedConnectors, research_mode: researchMode, search_mode: searchMode, - document_ids_to_add_in_context: [], + document_ids_to_add_in_context: documentIds, }, }, onError: (error) => { @@ -62,7 +100,15 @@ export default function ResearchChatPageV2() { chatRequestOptions?: { data?: any } ) => { const newChatId = await createChat(message.content); - router.replace(`/dashboard/${search_space_id}/v2/${newChatId}`); + if (newChatId) { + // Store selected documents before navigation + storeSelectedDocuments( + search_space_id as string, + newChatId, + selectedDocuments + ); + router.replace(`/dashboard/${search_space_id}/v2/${newChatId}`); + } return newChatId; }; @@ -73,6 +119,19 @@ export default function ResearchChatPageV2() { } }, [token, isNewChat, chatIdParam]); + // Restore selected documents from localStorage on page load + useEffect(() => { + if (chatIdParam && search_space_id) { + const restoredDocuments = restoreSelectedDocuments( + search_space_id as string, + chatIdParam + ); + if (restoredDocuments && restoredDocuments.length > 0) { + setSelectedDocuments(restoredDocuments); + } + } + }, [chatIdParam, search_space_id, setSelectedDocuments]); + const loadChatData = async (chatId: string) => { try { const chatData = await fetchChatDetails(chatId); @@ -133,11 +192,13 @@ export default function ResearchChatPageV2() { } return ( - ); } diff --git a/surfsense_web/components/chat_v2/ChatInterface.tsx b/surfsense_web/components/chat_v2/ChatInterface.tsx new file mode 100644 index 0000000..cb2b3f3 --- /dev/null +++ b/surfsense_web/components/chat_v2/ChatInterface.tsx @@ -0,0 +1,187 @@ +"use client"; + +import { + ChatSection, + ChatHandler, + ChatInput, + ChatCanvas, + ChatMessages, +} from "@llamaindex/chat-ui"; +import { Button } from "../ui/button"; +import { FolderOpen } from "lucide-react"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogHeader, + DialogTitle, + DialogTrigger, +} from "../ui/dialog"; +import { Suspense, useState, useCallback } from "react"; +import { useParams } from "next/navigation"; +import { useDocuments, DocumentType, Document } from "@/hooks/use-documents"; +import { DocumentsDataTable } from "./DocumentsDataTable"; +import React from "react"; + +interface ChatInterfaceProps { + handler: ChatHandler; + onDocumentSelectionChange?: (documents: Document[]) => void; + selectedDocuments?: Document[]; +} + +const DocumentSelector = React.memo( + ({ + onSelectionChange, + selectedDocuments = [], + }: { + onSelectionChange?: (documents: Document[]) => void; + selectedDocuments?: Document[]; + }) => { + const { search_space_id } = useParams(); + const [isOpen, setIsOpen] = useState(false); + + const { documents, loading, isLoaded, fetchDocuments } = useDocuments( + Number(search_space_id) + ); + + const handleOpenChange = useCallback( + (open: boolean) => { + setIsOpen(open); + if (open && !isLoaded) { + fetchDocuments(); + } + }, + [fetchDocuments, isLoaded] + ); + + const handleSelectionChange = useCallback( + (documents: Document[]) => { + onSelectionChange?.(documents); + }, + [onSelectionChange] + ); + + const handleDone = useCallback(() => { + setIsOpen(false); + }, [selectedDocuments]); + + const selectedCount = selectedDocuments.length; + + return ( + + + + + + +
+
+ + Select Documents + + + Choose documents to include in your research + context + +
+ +
+ {loading ? ( +
+
+
+

+ Loading documents... +

+
+
+ ) : isLoaded ? ( + + ) : null} +
+
+ +
+ ); + } +); + +const CustomChatInputOptions = ({ + onDocumentSelectionChange, + selectedDocuments, +}: { + onDocumentSelectionChange?: (documents: Document[]) => void; + selectedDocuments?: Document[]; +}) => { + return ( +
+ Loading...
}> + + + + ); +}; + +const CustomChatInput = ({ + onDocumentSelectionChange, + selectedDocuments, +}: { + onDocumentSelectionChange?: (documents: Document[]) => void; + selectedDocuments?: Document[]; +}) => { + return ( + + + + + + + + ); +}; + +export default function ChatInterface({ + handler, + onDocumentSelectionChange, + selectedDocuments = [], +}: ChatInterfaceProps) { + return ( + +
+ + + {/* Custom message rendering */} + + + + +
+ +
+
+ + + + ); +} diff --git a/surfsense_web/components/chat_v2/ChatMain.tsx b/surfsense_web/components/chat_v2/ChatMain.tsx deleted file mode 100644 index 5e969ba..0000000 --- a/surfsense_web/components/chat_v2/ChatMain.tsx +++ /dev/null @@ -1,11 +0,0 @@ -"use client"; - -import { ChatSection, ChatHandler } from "@llamaindex/chat-ui"; - -interface ChatMainProps { - handler: ChatHandler; -} - -export default function ChatMain({ handler }: ChatMainProps) { - return ; -} diff --git a/surfsense_web/components/chat_v2/DocumentsDataTable.tsx b/surfsense_web/components/chat_v2/DocumentsDataTable.tsx new file mode 100644 index 0000000..8e6e16b --- /dev/null +++ b/surfsense_web/components/chat_v2/DocumentsDataTable.tsx @@ -0,0 +1,530 @@ +"use client"; + +import * as React from "react"; +import { + ColumnDef, + ColumnFiltersState, + flexRender, + getCoreRowModel, + getFilteredRowModel, + getPaginationRowModel, + getSortedRowModel, + SortingState, + useReactTable, + VisibilityState, +} from "@tanstack/react-table"; +import { ArrowUpDown, Calendar, FileText, Search } from "lucide-react"; + +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { Input } from "@/components/ui/input"; +import { + Table, + TableBody, + TableCell, + TableHead, + TableHeader, + TableRow, +} from "@/components/ui/table"; +import { Badge } from "@/components/ui/badge"; +import { Document, DocumentType } from "@/hooks/use-documents"; + +interface DocumentsDataTableProps { + documents: Document[]; + onSelectionChange: (documents: Document[]) => void; + onDone: () => void; + initialSelectedDocuments?: Document[]; +} + +const DOCUMENT_TYPES: (DocumentType | "ALL")[] = [ + "ALL", + "FILE", + "EXTENSION", + "CRAWLED_URL", + "YOUTUBE_VIDEO", + "SLACK_CONNECTOR", + "NOTION_CONNECTOR", + "GITHUB_CONNECTOR", + "LINEAR_CONNECTOR", + "DISCORD_CONNECTOR", +]; + +const getDocumentTypeColor = (type: DocumentType) => { + const colors = { + FILE: "bg-blue-50 text-blue-700 border-blue-200", + EXTENSION: "bg-green-50 text-green-700 border-green-200", + CRAWLED_URL: "bg-purple-50 text-purple-700 border-purple-200", + YOUTUBE_VIDEO: "bg-red-50 text-red-700 border-red-200", + SLACK_CONNECTOR: "bg-yellow-50 text-yellow-700 border-yellow-200", + NOTION_CONNECTOR: "bg-indigo-50 text-indigo-700 border-indigo-200", + GITHUB_CONNECTOR: "bg-gray-50 text-gray-700 border-gray-200", + LINEAR_CONNECTOR: "bg-pink-50 text-pink-700 border-pink-200", + DISCORD_CONNECTOR: "bg-violet-50 text-violet-700 border-violet-200", + }; + return colors[type] || "bg-gray-50 text-gray-700 border-gray-200"; +}; + +const columns: ColumnDef[] = [ + { + id: "select", + header: ({ table }) => ( + + table.toggleAllPageRowsSelected(!!value) + } + aria-label="Select all" + /> + ), + cell: ({ row }) => ( + row.toggleSelected(!!value)} + aria-label="Select row" + /> + ), + enableSorting: false, + enableHiding: false, + size: 40, + }, + { + accessorKey: "title", + header: ({ column }) => ( + + ), + cell: ({ row }) => { + const title = row.getValue("title") as string; + return ( +
+ {title} +
+ ); + }, + }, + { + accessorKey: "document_type", + header: "Type", + cell: ({ row }) => { + const type = row.getValue("document_type") as DocumentType; + return ( + + + {type.replace(/_/g, " ")} + + {type.split("_")[0]} + + ); + }, + size: 80, + meta: { + className: "hidden sm:table-cell", + }, + }, + { + accessorKey: "content", + header: "Preview", + cell: ({ row }) => { + const content = row.getValue("content") as string; + return ( +
+ + {content.substring(0, 30)}... + + + {content.substring(0, 100)}... + +
+ ); + }, + enableSorting: false, + meta: { + className: "hidden md:table-cell", + }, + }, + { + accessorKey: "created_at", + header: ({ column }) => ( + + ), + cell: ({ row }) => { + const date = new Date(row.getValue("created_at")); + return ( +
+ + {date.toLocaleDateString("en-US", { + month: "short", + day: "numeric", + year: "numeric", + })} + + + {date.toLocaleDateString("en-US", { + month: "numeric", + day: "numeric", + })} + +
+ ); + }, + size: 80, + }, +]; + +export function DocumentsDataTable({ + documents, + onSelectionChange, + onDone, + initialSelectedDocuments = [], +}: DocumentsDataTableProps) { + const [sorting, setSorting] = React.useState([]); + const [columnFilters, setColumnFilters] = + React.useState([]); + const [columnVisibility, setColumnVisibility] = + React.useState({}); + const [documentTypeFilter, setDocumentTypeFilter] = React.useState< + DocumentType | "ALL" + >("ALL"); + + // Memoize initial row selection to prevent infinite loops + const initialRowSelection = React.useMemo(() => { + if (!documents.length || !initialSelectedDocuments.length) return {}; + + const selection: Record = {}; + initialSelectedDocuments.forEach((selectedDoc) => { + const docIndex = documents.findIndex( + (doc) => doc.id === selectedDoc.id + ); + if (docIndex !== -1) { + selection[docIndex.toString()] = true; + } + }); + return selection; + }, [documents, initialSelectedDocuments]); + + const [rowSelection, setRowSelection] = React.useState< + Record + >({}); + + // Only update row selection when initialRowSelection actually changes and is not empty + React.useEffect(() => { + const hasChanges = + JSON.stringify(rowSelection) !== + JSON.stringify(initialRowSelection); + if (hasChanges && Object.keys(initialRowSelection).length > 0) { + setRowSelection(initialRowSelection); + } + }, [initialRowSelection]); + + // Initialize row selection on mount + React.useEffect(() => { + if ( + Object.keys(rowSelection).length === 0 && + Object.keys(initialRowSelection).length > 0 + ) { + setRowSelection(initialRowSelection); + } + }, []); + + const filteredDocuments = React.useMemo(() => { + if (documentTypeFilter === "ALL") return documents; + return documents.filter( + (doc) => doc.document_type === documentTypeFilter + ); + }, [documents, documentTypeFilter]); + + const table = useReactTable({ + data: filteredDocuments, + columns, + onSortingChange: setSorting, + onColumnFiltersChange: setColumnFilters, + getCoreRowModel: getCoreRowModel(), + getPaginationRowModel: getPaginationRowModel(), + getSortedRowModel: getSortedRowModel(), + getFilteredRowModel: getFilteredRowModel(), + onColumnVisibilityChange: setColumnVisibility, + onRowSelectionChange: setRowSelection, + initialState: { pagination: { pageSize: 10 } }, + state: { sorting, columnFilters, columnVisibility, rowSelection }, + }); + + React.useEffect(() => { + const selectedRows = table.getFilteredSelectedRowModel().rows; + const selectedDocuments = selectedRows.map((row) => row.original); + onSelectionChange(selectedDocuments); + }, [rowSelection, onSelectionChange, table]); + + const handleClearAll = () => setRowSelection({}); + + const handleSelectPage = () => { + const currentPageRows = table.getRowModel().rows; + const newSelection = { ...rowSelection }; + currentPageRows.forEach((row) => { + newSelection[row.id] = true; + }); + setRowSelection(newSelection); + }; + + const handleSelectAllFiltered = () => { + const allFilteredRows = table.getFilteredRowModel().rows; + const newSelection: Record = {}; + allFilteredRows.forEach((row) => { + newSelection[row.id] = true; + }); + setRowSelection(newSelection); + }; + + const selectedCount = table.getFilteredSelectedRowModel().rows.length; + const totalFiltered = table.getFilteredRowModel().rows.length; + + return ( +
+ {/* Header Controls */} +
+ {/* Search and Filter Row */} +
+
+ + + table + .getColumn("title") + ?.setFilterValue(event.target.value) + } + className="pl-10 text-sm" + /> +
+ +
+ + {/* Action Controls Row */} +
+
+ + {selectedCount} of {totalFiltered} selected + +
+
+ + + + +
+
+ +
+
+ + {/* Table Container */} +
+
+ + + {table.getHeaderGroups().map((headerGroup) => ( + + {headerGroup.headers.map((header) => ( + + {header.isPlaceholder + ? null + : flexRender( + header.column.columnDef + .header, + header.getContext() + )} + + ))} + + ))} + + + {table.getRowModel().rows?.length ? ( + table.getRowModel().rows.map((row) => ( + + {row.getVisibleCells().map((cell) => ( + + {flexRender( + cell.column.columnDef.cell, + cell.getContext() + )} + + ))} + + )) + ) : ( + + + No documents found. + + + )} + +
+
+
+ + {/* Footer Pagination */} +
+
+ Showing{" "} + {table.getState().pagination.pageIndex * + table.getState().pagination.pageSize + + 1}{" "} + to{" "} + {Math.min( + (table.getState().pagination.pageIndex + 1) * + table.getState().pagination.pageSize, + table.getFilteredRowModel().rows.length + )}{" "} + of {table.getFilteredRowModel().rows.length} documents +
+
+ +
+ Page + + {table.getState().pagination.pageIndex + 1} + + of + {table.getPageCount()} +
+ +
+
+
+ ); +} diff --git a/surfsense_web/components/ui/table.tsx b/surfsense_web/components/ui/table.tsx index 27697c0..51b74dd 100644 --- a/surfsense_web/components/ui/table.tsx +++ b/surfsense_web/components/ui/table.tsx @@ -1,95 +1,116 @@ -import * as React from "react"; +"use client" -import { cn } from "@/lib/utils"; +import * as React from "react" -const Table = React.forwardRef>( - ({ className, ...props }, ref) => ( -
- +import { cn } from "@/lib/utils" + +function Table({ className, ...props }: React.ComponentProps<"table">) { + return ( +
+
- ), -); -Table.displayName = "Table"; + ) +} -const TableHeader = React.forwardRef< - HTMLTableSectionElement, - React.HTMLAttributes ->(({ className, ...props }, ref) => ); -TableHeader.displayName = "TableHeader"; +function TableHeader({ className, ...props }: React.ComponentProps<"thead">) { + return ( + + ) +} -const TableBody = React.forwardRef< - HTMLTableSectionElement, - React.HTMLAttributes ->(({ className, ...props }, ref) => ( - -)); -TableBody.displayName = "TableBody"; +function TableBody({ className, ...props }: React.ComponentProps<"tbody">) { + return ( + + ) +} -const TableFooter = React.forwardRef< - HTMLTableSectionElement, - React.HTMLAttributes ->(({ className, ...props }, ref) => ( - tr]:last:border-b-0", - className, - )} - {...props} - /> -)); -TableFooter.displayName = "TableFooter"; - -const TableRow = React.forwardRef>( - ({ className, ...props }, ref) => ( - ) { + return ( + tr]:last:border-b-0", + className )} {...props} /> - ), -); -TableRow.displayName = "TableRow"; + ) +} -const TableHead = React.forwardRef< - HTMLTableCellElement, - React.ThHTMLAttributes ->(({ className, ...props }, ref) => ( - + ) +} -const TableCell = React.forwardRef< - HTMLTableCellElement, - React.TdHTMLAttributes ->(({ className, ...props }, ref) => ( -
[role=checkbox]]:translate-y-0.5", - className, - )} - {...props} - /> -)); -TableHead.displayName = "TableHead"; +function TableRow({ className, ...props }: React.ComponentProps<"tr">) { + return ( +
[role=checkbox]]:translate-y-0.5", - className, - )} - {...props} - /> -)); -TableCell.displayName = "TableCell"; +function TableHead({ className, ...props }: React.ComponentProps<"th">) { + return ( + [role=checkbox]]:translate-y-[2px]", + className + )} + {...props} + /> + ) +} -const TableCaption = React.forwardRef< - HTMLTableCaptionElement, - React.HTMLAttributes ->(({ className, ...props }, ref) => ( -
-)); -TableCaption.displayName = "TableCaption"; +function TableCell({ className, ...props }: React.ComponentProps<"td">) { + return ( +
[role=checkbox]]:translate-y-[2px]", + className + )} + {...props} + /> + ) +} -export { Table, TableBody, TableCaption, TableCell, TableFooter, TableHead, TableHeader, TableRow }; +function TableCaption({ + className, + ...props +}: React.ComponentProps<"caption">) { + return ( +
+ ) +} + +export { + Table, + TableHeader, + TableBody, + TableFooter, + TableHead, + TableRow, + TableCell, + TableCaption, +} diff --git a/surfsense_web/hooks/use-documents.ts b/surfsense_web/hooks/use-documents.ts index cfe2b05..d38b374 100644 --- a/surfsense_web/hooks/use-documents.ts +++ b/surfsense_web/hooks/use-documents.ts @@ -1,115 +1,121 @@ -"use client" -import { useState, useEffect } from 'react'; -import { toast } from 'sonner'; +"use client"; +import { useState, useEffect, useCallback } from "react"; +import { toast } from "sonner"; export interface Document { - id: number; - title: string; - document_type: "EXTENSION" | "CRAWLED_URL" | "SLACK_CONNECTOR" | "NOTION_CONNECTOR" | "FILE"; - document_metadata: any; - content: string; - created_at: string; - search_space_id: number; + id: number; + title: string; + document_type: DocumentType; + document_metadata: any; + content: string; + created_at: string; + search_space_id: number; } -export function useDocuments(searchSpaceId: number) { - const [documents, setDocuments] = useState([]); - const [loading, setLoading] = useState(true); - const [error, setError] = useState(null); +export type DocumentType = + | "EXTENSION" + | "CRAWLED_URL" + | "SLACK_CONNECTOR" + | "NOTION_CONNECTOR" + | "FILE" + | "YOUTUBE_VIDEO" + | "GITHUB_CONNECTOR" + | "LINEAR_CONNECTOR" + | "DISCORD_CONNECTOR"; - useEffect(() => { - const fetchDocuments = async () => { - try { - setLoading(true); - const response = await fetch( - `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/documents?search_space_id=${searchSpaceId}`, - { - headers: { - Authorization: `Bearer ${localStorage.getItem('surfsense_bearer_token')}`, - }, - method: "GET", - } - ); - - if (!response.ok) { - toast.error("Failed to fetch documents"); - throw new Error("Failed to fetch documents"); +export function useDocuments(searchSpaceId: number, lazy: boolean = true) { + const [documents, setDocuments] = useState([]); + const [loading, setLoading] = useState(!lazy); // Don't show loading initially for lazy mode + const [error, setError] = useState(null); + const [isLoaded, setIsLoaded] = useState(false); // Memoization flag + + const fetchDocuments = useCallback(async () => { + if (isLoaded && lazy) return; // Avoid redundant calls in lazy mode + + try { + setLoading(true); + const response = await fetch( + `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/documents?search_space_id=${searchSpaceId}`, + { + headers: { + Authorization: `Bearer ${localStorage.getItem( + "surfsense_bearer_token" + )}`, + }, + method: "GET", + } + ); + + if (!response.ok) { + toast.error("Failed to fetch documents"); + throw new Error("Failed to fetch documents"); + } + + const data = await response.json(); + setDocuments(data); + setError(null); + setIsLoaded(true); + } catch (err: any) { + setError(err.message || "Failed to fetch documents"); + console.error("Error fetching documents:", err); + } finally { + setLoading(false); } - - const data = await response.json(); - setDocuments(data); - setError(null); - } catch (err: any) { - setError(err.message || 'Failed to fetch documents'); - console.error('Error fetching documents:', err); - } finally { - setLoading(false); - } + }, [searchSpaceId, isLoaded, lazy]); + + useEffect(() => { + if (!lazy && searchSpaceId) { + fetchDocuments(); + } + }, [searchSpaceId, lazy, fetchDocuments]); + + // Function to refresh the documents list + const refreshDocuments = useCallback(async () => { + setIsLoaded(false); // Reset memoization flag to allow refetch + await fetchDocuments(); + }, [fetchDocuments]); + + // Function to delete a document + const deleteDocument = useCallback( + async (documentId: number) => { + try { + const response = await fetch( + `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/documents/${documentId}`, + { + headers: { + Authorization: `Bearer ${localStorage.getItem( + "surfsense_bearer_token" + )}`, + }, + method: "DELETE", + } + ); + + if (!response.ok) { + toast.error("Failed to delete document"); + throw new Error("Failed to delete document"); + } + + toast.success("Document deleted successfully"); + // Update the local state after successful deletion + setDocuments(documents.filter((doc) => doc.id !== documentId)); + return true; + } catch (err: any) { + toast.error(err.message || "Failed to delete document"); + console.error("Error deleting document:", err); + return false; + } + }, + [documents] + ); + + return { + documents, + loading, + error, + isLoaded, + fetchDocuments, // Manual fetch function for lazy mode + refreshDocuments, + deleteDocument, }; - - if (searchSpaceId) { - fetchDocuments(); - } - }, [searchSpaceId]); - - // Function to refresh the documents list - const refreshDocuments = async () => { - setLoading(true); - try { - const response = await fetch( - `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/documents?search_space_id=${searchSpaceId}`, - { - headers: { - Authorization: `Bearer ${localStorage.getItem('surfsense_bearer_token')}`, - }, - method: "GET", - } - ); - - if (!response.ok) { - toast.error("Failed to fetch documents"); - throw new Error("Failed to fetch documents"); - } - - const data = await response.json(); - setDocuments(data); - setError(null); - } catch (err: any) { - setError(err.message || 'Failed to fetch documents'); - console.error('Error fetching documents:', err); - } finally { - setLoading(false); - } - }; - - // Function to delete a document - const deleteDocument = async (documentId: number) => { - try { - const response = await fetch( - `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/documents/${documentId}`, - { - headers: { - Authorization: `Bearer ${localStorage.getItem('surfsense_bearer_token')}`, - }, - method: "DELETE", - } - ); - - if (!response.ok) { - toast.error("Failed to delete document"); - throw new Error("Failed to delete document"); - } - - toast.success("Document deleted successfully"); - // Update the local state after successful deletion - setDocuments(documents.filter(doc => doc.id !== documentId)); - return true; - } catch (err: any) { - toast.error(err.message || 'Failed to delete document'); - console.error('Error deleting document:', err); - return false; - } - }; - - return { documents, loading, error, refreshDocuments, deleteDocument }; -} \ No newline at end of file +} diff --git a/surfsense_web/hooks/useChat.ts b/surfsense_web/hooks/useChat.ts index 023ec4b..0bedcbf 100644 --- a/surfsense_web/hooks/useChat.ts +++ b/surfsense_web/hooks/useChat.ts @@ -1,6 +1,7 @@ import { useState, useEffect, useCallback } from "react"; import { Message } from "@ai-sdk/react"; import { ResearchMode } from "@/components/chat"; +import { Document } from "@/hooks/use-documents"; interface UseChatStateProps { search_space_id: string; @@ -10,12 +11,17 @@ interface UseChatStateProps { export function useChatState({ search_space_id, chat_id }: UseChatStateProps) { const [token, setToken] = useState(null); const [isLoading, setIsLoading] = useState(false); - const [currentChatId, setCurrentChatId] = useState(chat_id || null); - + const [currentChatId, setCurrentChatId] = useState( + chat_id || null + ); + // Chat configuration state - const [searchMode, setSearchMode] = useState<"DOCUMENTS" | "CHUNKS">("DOCUMENTS"); + const [searchMode, setSearchMode] = useState<"DOCUMENTS" | "CHUNKS">( + "DOCUMENTS" + ); const [researchMode, setResearchMode] = useState("QNA"); const [selectedConnectors, setSelectedConnectors] = useState([]); + const [selectedDocuments, setSelectedDocuments] = useState([]); useEffect(() => { const bearerToken = localStorage.getItem("surfsense_bearer_token"); @@ -35,6 +41,8 @@ export function useChatState({ search_space_id, chat_id }: UseChatStateProps) { setResearchMode, selectedConnectors, setSelectedConnectors, + selectedDocuments, + setSelectedDocuments, }; } @@ -51,112 +59,133 @@ export function useChatAPI({ researchMode, selectedConnectors, }: UseChatAPIProps) { - const fetchChatDetails = useCallback(async (chatId: string) => { - if (!token) return null; + const fetchChatDetails = useCallback( + async (chatId: string) => { + if (!token) return null; - try { - const response = await fetch( - `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chats/${Number(chatId)}`, - { - method: "GET", - headers: { - "Content-Type": "application/json", - Authorization: `Bearer ${token}`, - }, + try { + const response = await fetch( + `${ + process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL + }/api/v1/chats/${Number(chatId)}`, + { + method: "GET", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${token}`, + }, + } + ); + + if (!response.ok) { + throw new Error( + `Failed to fetch chat details: ${response.statusText}` + ); } - ); - if (!response.ok) { - throw new Error(`Failed to fetch chat details: ${response.statusText}`); + return await response.json(); + } catch (err) { + console.error("Error fetching chat details:", err); + return null; + } + }, + [token] + ); + + const createChat = useCallback( + async (initialMessage: string): Promise => { + if (!token) { + console.error("Authentication token not found"); + return null; } - return await response.json(); - } catch (err) { - console.error("Error fetching chat details:", err); - return null; - } - }, [token]); + try { + const response = await fetch( + `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chats/`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${token}`, + }, + body: JSON.stringify({ + type: researchMode, + title: "Untitled Chat", + initial_connectors: selectedConnectors, + messages: [ + { + role: "user", + content: initialMessage, + }, + ], + search_space_id: Number(search_space_id), + }), + } + ); - const createChat = useCallback(async (initialMessage: string): Promise => { - if (!token) { - console.error("Authentication token not found"); - return null; - } - - try { - const response = await fetch( - `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chats/`, - { - method: "POST", - headers: { - "Content-Type": "application/json", - Authorization: `Bearer ${token}`, - }, - body: JSON.stringify({ - type: researchMode, - title: "Untitled Chat", - initial_connectors: selectedConnectors, - messages: [ - { - role: "user", - content: initialMessage, - }, - ], - search_space_id: Number(search_space_id), - }), + if (!response.ok) { + throw new Error( + `Failed to create chat: ${response.statusText}` + ); } - ); - if (!response.ok) { - throw new Error(`Failed to create chat: ${response.statusText}`); + const data = await response.json(); + return data.id; + } catch (err) { + console.error("Error creating chat:", err); + return null; } + }, + [token, researchMode, selectedConnectors, search_space_id] + ); - const data = await response.json(); - return data.id; - } catch (err) { - console.error("Error creating chat:", err); - return null; - } - }, [token, researchMode, selectedConnectors, search_space_id]); + const updateChat = useCallback( + async (chatId: string, messages: Message[]) => { + if (!token) return; - const updateChat = useCallback(async (chatId: string, messages: Message[]) => { - if (!token) return; + try { + const userMessages = messages.filter( + (msg) => msg.role === "user" + ); + if (userMessages.length === 0) return; - try { - const userMessages = messages.filter(msg => msg.role === "user"); - if (userMessages.length === 0) return; + const title = userMessages[0].content; - const title = userMessages[0].content; + const response = await fetch( + `${ + process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL + }/api/v1/chats/${Number(chatId)}`, + { + method: "PUT", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${token}`, + }, + body: JSON.stringify({ + type: researchMode, + title: title, + initial_connectors: selectedConnectors, + messages: messages, + search_space_id: Number(search_space_id), + }), + } + ); - const response = await fetch( - `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chats/${Number(chatId)}`, - { - method: "PUT", - headers: { - "Content-Type": "application/json", - Authorization: `Bearer ${token}`, - }, - body: JSON.stringify({ - type: researchMode, - title: title, - initial_connectors: selectedConnectors, - messages: messages, - search_space_id: Number(search_space_id), - }), + if (!response.ok) { + throw new Error( + `Failed to update chat: ${response.statusText}` + ); } - ); - - if (!response.ok) { - throw new Error(`Failed to update chat: ${response.statusText}`); + } catch (err) { + console.error("Error updating chat:", err); } - } catch (err) { - console.error("Error updating chat:", err); - } - }, [token, researchMode, selectedConnectors, search_space_id]); + }, + [token, researchMode, selectedConnectors, search_space_id] + ); return { fetchChatDetails, createChat, updateChat, }; -} \ No newline at end of file +}