diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index a41d2e143..ba9baf2c5 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -32,6 +32,12 @@ import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { Button } from "@/components/ui/button"; import { cn } from "@/lib/utils"; +import { useRef, useState, useMemo } from "react"; +import { useAtomValue } from "jotai"; +import { useParams } from "next/navigation"; +import type { Document } from "@/contracts/types/document.types"; +import { documentsAtom } from "@/atoms/documents/document-query.atoms"; +import { DocumentsDataTable } from "@/components/new-chat/DocumentsDataTable"; export const Thread: FC = () => { return ( @@ -141,17 +147,124 @@ const ThreadSuggestions: FC = () => { }; const Composer: FC = () => { + // ---- State for document mentions ---- + const [allSelectedDocuments, setAllSelectedDocuments] = useState([]); + const [mentionedDocuments, setMentionedDocuments] = useState([]); + const [showDocumentPopover, setShowDocumentPopover] = useState(false); + const [mentionModalSelectedDocs, setMentionModalSelectedDocs] = useState([]); + const [inputValue, setInputValue] = useState(""); + const inputRef = useRef(null); + const documentsResult = useAtomValue(documentsAtom); // Atom fetches all docs for this workspace + const allDocuments = documentsResult?.data?.items || []; + const { search_space_id } = useParams(); + const searchSpaceId = typeof search_space_id === "string" ? parseInt(search_space_id, 10) : 0; + + const handleInputOrKeyUp = ( + e: React.FormEvent | React.KeyboardEvent + ) => { + const textarea = e.currentTarget; + const value = textarea.value; + setInputValue(value); + + // Regex: finds all [title] occurrences + const mentionRegex = /\[([^\]]+)\]/g; + const titlesMentioned: string[] = []; + let match; + while ((match = mentionRegex.exec(value)) !== null) { + titlesMentioned.push(match[1]); + } + + // Use allSelectedDocuments to filter down for current chips + setMentionedDocuments( + allSelectedDocuments.filter((doc) => titlesMentioned.includes(doc.title)) + ); + + const selectionStart = textarea.selectionStart; + // Only open if the last character before the caret is exactly '@' + if ( + selectionStart !== null && + value[selectionStart - 1] === "@" && + value.length === selectionStart + ) { + setShowDocumentPopover(true); + } else { + setShowDocumentPopover(false); + } + }; + + const handleDocumentsMention = (documents: Document[]) => { + // Add newly selected docs to allSelectedDocuments + setAllSelectedDocuments(prev => { + const toAdd = documents.filter(doc => !prev.find(p => p.id === doc.id)); + return [...prev, ...toAdd]; + }); + let newValue = inputValue; + documents.forEach((doc) => { + const refString = `[${doc.title}]`; + if (!newValue.includes(refString)) { + if (newValue.trim() !== "" && !newValue.endsWith(" ")) { + newValue += " "; + } + newValue += refString; + } + }); + setInputValue(newValue); + // Run the chip update as well right after change + const mentionRegex = /\[([^\]]+)\]/g; + const titlesMentioned: string[] = []; + let match; + while ((match = mentionRegex.exec(newValue)) !== null) { + titlesMentioned.push(match[1]); + } + setMentionedDocuments( + allSelectedDocuments.filter(doc => titlesMentioned.includes(doc.title)) + ); + }; + return ( + {/* -------- Input field w/ refs and handlers -------- */} + + {/* -------- Document mention popover (simple version) -------- */} + {showDocumentPopover && ( +
+
+ setShowDocumentPopover(false)} + initialSelectedDocuments={mentionedDocuments} + viewOnly={true} + /> +
+
+ )} + {/* ---- Mention chips for selected/mentioned documents ---- */} + {mentionedDocuments.length > 0 && ( +
+ {mentionedDocuments.map((doc) => ( + + {doc.title} + + ))} +
+ )} diff --git a/surfsense_web/components/new-chat/DocumentsDataTable.tsx b/surfsense_web/components/new-chat/DocumentsDataTable.tsx new file mode 100644 index 000000000..182543364 --- /dev/null +++ b/surfsense_web/components/new-chat/DocumentsDataTable.tsx @@ -0,0 +1,605 @@ +"use client"; + +import { useQuery } from "@tanstack/react-query"; +import { + type ColumnDef, + flexRender, + getCoreRowModel, + type SortingState, + useReactTable, +} from "@tanstack/react-table"; +import { useAtomValue } from "jotai"; +import { ArrowUpDown, Calendar, FileText, Filter, Plus, Search } from "lucide-react"; +import { useRouter } from "next/navigation"; +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { + Table, + TableBody, + TableCell, + TableHead, + TableHeader, + TableRow, +} from "@/components/ui/table"; +import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; +import type { Document, DocumentTypeEnum } from "@/contracts/types/document.types"; +import { documentsApiService } from "@/lib/apis/documents-api.service"; +import { cacheKeys } from "@/lib/query-client/cache-keys"; + +interface DocumentsDataTableProps { + searchSpaceId: number; + onSelectionChange: (documents: Document[]) => void; + onDone: () => void; + initialSelectedDocuments?: Document[]; + viewOnly?: boolean; +} + +function useDebounced(value: T, delay = 300) { + const [debounced, setDebounced] = useState(value); + useEffect(() => { + const t = setTimeout(() => setDebounced(value), delay); + return () => clearTimeout(t); + }, [value, delay]); + return debounced; +} + +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 ( +
+ {getConnectorIcon(String(type))} +
+ ); + }, + 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({ + searchSpaceId, + onSelectionChange, + onDone, + initialSelectedDocuments = [] +}: DocumentsDataTableProps) { + const router = useRouter(); + const [sorting, setSorting] = useState([]); + const [search, setSearch] = useState(""); + const debouncedSearch = useDebounced(search, 300); + const [documentTypeFilter, setDocumentTypeFilter] = useState([]); + const [pageIndex, setPageIndex] = useState(0); + const [pageSize, setPageSize] = useState(10); + const { data: typeCounts } = useAtomValue(documentTypeCountsAtom); + + const fetchQueryParams = useMemo( + () => ({ + search_space_id: searchSpaceId, + page: pageIndex, + page_size: pageSize, + ...(documentTypeFilter.length > 0 && { document_types: documentTypeFilter }), + }), + [searchSpaceId, pageIndex, pageSize, documentTypeFilter, debouncedSearch] + ); + + const searchQueryParams = useMemo(() => { + return { + search_space_id: searchSpaceId, + page: pageIndex, + page_size: pageSize, + ...(documentTypeFilter.length > 0 && { document_types: documentTypeFilter }), + title: debouncedSearch, + }; + }, [debouncedSearch, searchSpaceId, pageIndex, pageSize, documentTypeFilter, debouncedSearch]); + + // Use query for fetching documents + const { data: documents, isLoading: isDocumentsLoading } = useQuery({ + queryKey: cacheKeys.documents.withQueryParams(fetchQueryParams), + queryFn: () => documentsApiService.getDocuments({ queryParams: fetchQueryParams }), + staleTime: 3 * 60 * 1000, // 3 minutes + enabled: !!searchSpaceId && !debouncedSearch.trim(), + }); + + // Seaching + const { data: searchedDocuments, isLoading: isSearchedDocumentsLoading } = useQuery({ + queryKey: cacheKeys.documents.withQueryParams(searchQueryParams), + queryFn: () => documentsApiService.searchDocuments({ queryParams: searchQueryParams }), + staleTime: 3 * 60 * 1000, // 3 minutes + enabled: !!searchSpaceId && !!debouncedSearch.trim(), + }); + + // Use query data when not searching, otherwise use hook data + const actualDocuments = debouncedSearch.trim() + ? searchedDocuments?.items || [] + : documents?.items || []; + const actualTotal = debouncedSearch.trim() + ? searchedDocuments?.total || 0 + : documents?.total || 0; + const actualLoading = debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading; + + // Memoize initial row selection to prevent infinite loops + const initialRowSelection = useMemo(() => { + if (!initialSelectedDocuments.length) return {}; + + const selection: Record = {}; + initialSelectedDocuments.forEach((selectedDoc) => { + selection[selectedDoc.id] = true; + }); + return selection; + }, [initialSelectedDocuments]); + + const [rowSelection, setRowSelection] = useState>( + () => initialRowSelection + ); + + // Maintain a separate state for actually selected documents (across all pages) + const [selectedDocumentsMap, setSelectedDocumentsMap] = useState>(() => { + const map = new Map(); + initialSelectedDocuments.forEach((doc) => map.set(doc.id, doc)); + return map; + }); + + // Track the last notified selection to avoid redundant parent calls + const lastNotifiedSelection = useRef(""); + + // Update row selection only when initialSelectedDocuments changes (not rowSelection itself) + useEffect(() => { + const initialKeys = Object.keys(initialRowSelection); + if (initialKeys.length === 0) return; + + const currentKeys = Object.keys(rowSelection); + // Quick length check before expensive comparison + if (currentKeys.length === initialKeys.length) { + // Check if all keys match (order doesn't matter for Sets) + const hasAllKeys = initialKeys.every((key) => rowSelection[key]); + if (hasAllKeys) return; + } + + setRowSelection(initialRowSelection); + }, [initialRowSelection]); // Remove rowSelection from dependencies to prevent loop + + // Update the selected documents map when row selection changes + useEffect(() => { + if (!actualDocuments || actualDocuments.length === 0) return; + + setSelectedDocumentsMap((prev) => { + const newMap = new Map(prev); + let hasChanges = false; + + // Process only current page documents + for (const doc of actualDocuments) { + const docId = doc.id; + const isSelected = rowSelection[docId.toString()]; + const wasInMap = newMap.has(docId); + + if (isSelected && !wasInMap) { + newMap.set(docId, doc); + hasChanges = true; + } else if (!isSelected && wasInMap) { + newMap.delete(docId); + hasChanges = true; + } + } + + // Return same reference if no changes to avoid unnecessary re-renders + return hasChanges ? newMap : prev; + }); + }, [rowSelection, documents]); + + // Memoize selected documents array + const selectedDocumentsArray = useMemo(() => { + return Array.from(selectedDocumentsMap.values()); + }, [selectedDocumentsMap]); + + // Notify parent of selection changes only when content actually changes + useEffect(() => { + // Create a stable string representation for comparison + const selectionKey = selectedDocumentsArray + .map((d) => d.id) + .sort() + .join(","); + + // Skip if selection hasn't actually changed + if (selectionKey === lastNotifiedSelection.current) return; + + lastNotifiedSelection.current = selectionKey; + onSelectionChange(selectedDocumentsArray); + }, [selectedDocumentsArray, onSelectionChange]); + + const table = useReactTable({ + data: actualDocuments || [], + columns, + getRowId: (row) => row.id.toString(), + onSortingChange: setSorting, + getCoreRowModel: getCoreRowModel(), + onRowSelectionChange: setRowSelection, + manualPagination: true, + pageCount: Math.ceil(actualTotal / pageSize), + state: { sorting, rowSelection, pagination: { pageIndex, pageSize } }, + }); + + const handleClearAll = useCallback(() => { + setRowSelection({}); + setSelectedDocumentsMap(new Map()); + }, []); + + const handleSelectPage = useCallback(() => { + const currentPageRows = table.getRowModel().rows; + const newSelection = { ...rowSelection }; + currentPageRows.forEach((row) => { + newSelection[row.id] = true; + }); + setRowSelection(newSelection); + }, [table, rowSelection]); + + const handleToggleType = useCallback((type: DocumentTypeEnum, checked: boolean) => { + setDocumentTypeFilter((prev) => { + if (checked) { + return [...prev, type]; + } + return prev.filter((t) => t !== type); + }); + setPageIndex(0); // Reset to first page when filter changes + }, []); + + const selectedCount = selectedDocumentsMap.size; + + // Get available document types from type counts (memoized) + const availableTypes = useMemo(() => { + const types = typeCounts ? (Object.keys(typeCounts) as DocumentTypeEnum[]) : []; + return types.length > 0 ? types.sort() : []; + }, [typeCounts]); + + return ( +
+ {/* Header Controls */} +
+ {/* Search and Filter Row */} +
+
+ + { + setSearch(event.target.value); + setPageIndex(0); // Reset to first page on search + }} + className="pl-10 text-sm" + /> +
+ + + + + +
+
Filter by Type
+
+ {availableTypes.map((type) => ( +
+ handleToggleType(type, !!checked)} + /> + +
+ ))} +
+ {documentTypeFilter.length > 0 && ( + + )} +
+ + +
+ + {/* Action Controls Row */} +
+
+ + {selectedCount} selected {actualLoading && "ยท Loading..."} + +
+
+ + + +
+
+ +
+
+ + {/* Table Container */} +
+
+ {actualLoading ? ( +
+
+
+

Loading documents...

+
+
+ ) : ( + + + {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

+

+ Get started by adding your first data source to build your knowledge + base. +

+
+ +
+ + + )} + +
+ )} +
+
+ + {/* Footer Pagination */} +
+
+ Showing {pageIndex * pageSize + 1} to {Math.min((pageIndex + 1) * pageSize, actualTotal)}{" "} + of {actualTotal} documents +
+
+ +
+ Page + {pageIndex + 1} + of + {Math.ceil(actualTotal / pageSize)} +
+ +
+
+
+ ); +}