From b038e96a768601711784118513c1732357b94908 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Mon, 22 Dec 2025 23:17:48 +0200 Subject: [PATCH 1/7] add document mentions table --- .../components/assistant-ui/thread.tsx | 113 ++++ .../new-chat/DocumentsDataTable.tsx | 605 ++++++++++++++++++ 2 files changed, 718 insertions(+) create mode 100644 surfsense_web/components/new-chat/DocumentsDataTable.tsx 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)} +
+ +
+
+
+ ); +} From 7766161f84098804d7314a59fd646abad15fe4e4 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Mon, 22 Dec 2025 23:29:49 +0200 Subject: [PATCH 2/7] add document mentions table --- .../components/assistant-ui/thread.tsx | 30 +++++++++---------- .../new-chat/DocumentsDataTable.tsx | 24 +++++++-------- 2 files changed, 26 insertions(+), 28 deletions(-) diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index ba9baf2c5..60ae4b2cf 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -21,7 +21,9 @@ import { RefreshCwIcon, SquareIcon, } from "lucide-react"; +import { useParams } from "next/navigation"; import type { FC } from "react"; +import { useRef, useState } from "react"; import { ComposerAddAttachment, ComposerAttachments, @@ -30,14 +32,10 @@ import { import { MarkdownText } from "@/components/assistant-ui/markdown-text"; 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"; +import { Button } from "@/components/ui/button"; +import type { Document } from "@/contracts/types/document.types"; +import { cn } from "@/lib/utils"; export const Thread: FC = () => { return ( @@ -151,13 +149,9 @@ const Composer: FC = () => { 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 @@ -194,8 +188,8 @@ const Composer: FC = () => { const handleDocumentsMention = (documents: Document[]) => { // Add newly selected docs to allSelectedDocuments - setAllSelectedDocuments(prev => { - const toAdd = documents.filter(doc => !prev.find(p => p.id === doc.id)); + setAllSelectedDocuments((prev) => { + const toAdd = documents.filter((doc) => !prev.find((p) => p.id === doc.id)); return [...prev, ...toAdd]; }); let newValue = inputValue; @@ -217,7 +211,7 @@ const Composer: FC = () => { titlesMentioned.push(match[1]); } setMentionedDocuments( - allSelectedDocuments.filter(doc => titlesMentioned.includes(doc.title)) + allSelectedDocuments.filter((doc) => titlesMentioned.includes(doc.title)) ); }; @@ -256,10 +250,14 @@ const Composer: FC = () => {
)} {/* ---- Mention chips for selected/mentioned documents ---- */} - {mentionedDocuments.length > 0 && ( + {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 index 182543364..290417e53 100644 --- a/surfsense_web/components/new-chat/DocumentsDataTable.tsx +++ b/surfsense_web/components/new-chat/DocumentsDataTable.tsx @@ -182,7 +182,7 @@ export function DocumentsDataTable({ searchSpaceId, onSelectionChange, onDone, - initialSelectedDocuments = [] + initialSelectedDocuments = [], }: DocumentsDataTableProps) { const router = useRouter(); const [sorting, setSorting] = useState([]); @@ -545,18 +545,18 @@ export function DocumentsDataTable({

No documents found

-

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

+

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

- +
From ceb01dc544383eaedb82bb7a431c8aa102b2e408 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Tue, 23 Dec 2025 14:24:36 +0530 Subject: [PATCH 3/7] feat: enhance new chat functionality with document mentions support - Updated the new chat routes to include handling for mentioned document IDs, allowing users to reference specific documents in their chat. - Modified the NewChatRequest schema to accommodate optional document IDs. - Implemented document mention formatting in the chat streaming service for improved context. - Enhanced the frontend to manage document mentions, including a new atom for state management and UI updates for document selection. - Refactored the DocumentsDataTable component for better integration with the new mention functionality. --- .../app/routes/new_chat_routes.py | 3 +- surfsense_backend/app/schemas/new_chat.py | 3 + .../app/tasks/chat/stream_new_chat.py | 55 +- .../new-chat/[[...chat_id]]/page.tsx | 17 +- .../atoms/chat/mentioned-documents.atom.ts | 10 + .../components/assistant-ui/thread.tsx | 227 +++--- .../new-chat/DocumentsDataTable.tsx | 659 ++++-------------- 7 files changed, 346 insertions(+), 628 deletions(-) create mode 100644 surfsense_web/atoms/chat/mentioned-documents.atom.ts diff --git a/surfsense_backend/app/routes/new_chat_routes.py b/surfsense_backend/app/routes/new_chat_routes.py index 6cccdaa5b..81477ac6d 100644 --- a/surfsense_backend/app/routes/new_chat_routes.py +++ b/surfsense_backend/app/routes/new_chat_routes.py @@ -686,7 +686,7 @@ async def handle_new_chat( search_space = search_space_result.scalars().first() # TODO: Add new llm config arch then complete this - llm_config_id = -1 + llm_config_id = -4 # Return streaming response return StreamingResponse( @@ -698,6 +698,7 @@ async def handle_new_chat( llm_config_id=llm_config_id, messages=request.messages, attachments=request.attachments, + mentioned_document_ids=request.mentioned_document_ids, ), media_type="text/event-stream", headers={ diff --git a/surfsense_backend/app/schemas/new_chat.py b/surfsense_backend/app/schemas/new_chat.py index ffaf85554..78498cf04 100644 --- a/surfsense_backend/app/schemas/new_chat.py +++ b/surfsense_backend/app/schemas/new_chat.py @@ -160,3 +160,6 @@ class NewChatRequest(BaseModel): attachments: list[ChatAttachment] | None = ( None # Optional attachments with extracted content ) + mentioned_document_ids: list[int] | None = ( + None # Optional document IDs mentioned with @ in the chat + ) diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index de318a7d5..0f4ff26b8 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -9,6 +9,7 @@ import json from collections.abc import AsyncGenerator from langchain_core.messages import HumanMessage from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent from app.agents.new_chat.checkpointer import get_checkpointer @@ -16,6 +17,7 @@ from app.agents.new_chat.llm_config import ( create_chat_litellm_from_config, load_llm_config_from_yaml, ) +from app.db import Document from app.schemas.new_chat import ChatAttachment, ChatMessage from app.services.connector_service import ConnectorService from app.services.new_streaming_service import VercelStreamingService @@ -38,6 +40,27 @@ def format_attachments_as_context(attachments: list[ChatAttachment]) -> str: return "\n".join(context_parts) +def format_mentioned_documents_as_context(documents: list[Document]) -> str: + """Format mentioned documents as context for the agent.""" + if not documents: + return "" + + context_parts = [""] + context_parts.append( + "The user has explicitly mentioned the following documents from their knowledge base. " + "These documents are directly relevant to the query and should be prioritized as primary sources." + ) + for i, doc in enumerate(documents, 1): + context_parts.append( + f"" + ) + context_parts.append(f"") + context_parts.append("") + context_parts.append("") + + return "\n".join(context_parts) + + async def stream_new_chat( user_query: str, search_space_id: int, @@ -46,6 +69,7 @@ async def stream_new_chat( llm_config_id: int = -1, messages: list[ChatMessage] | None = None, attachments: list[ChatAttachment] | None = None, + mentioned_document_ids: list[int] | None = None, ) -> AsyncGenerator[str, None]: """ Stream chat responses from the new SurfSense deep agent. @@ -61,6 +85,8 @@ async def stream_new_chat( session: The database session llm_config_id: The LLM configuration ID (default: -1 for first global config) messages: Optional chat history from frontend (list of ChatMessage) + attachments: Optional attachments with extracted content + mentioned_document_ids: Optional list of document IDs mentioned with @ in the chat Yields: str: SSE formatted response strings @@ -105,13 +131,30 @@ async def stream_new_chat( # Build input with message history from frontend langchain_messages = [] - # Format the user query with attachment context if any - final_query = user_query - if attachments: - attachment_context = format_attachments_as_context(attachments) - final_query = ( - f"{attachment_context}\n\n{user_query}" + # Fetch mentioned documents if any + mentioned_documents: list[Document] = [] + if mentioned_document_ids: + result = await session.execute( + select(Document).filter( + Document.id.in_(mentioned_document_ids), + Document.search_space_id == search_space_id, + ) ) + mentioned_documents = list(result.scalars().all()) + + # Format the user query with context (attachments + mentioned documents) + final_query = user_query + context_parts = [] + + if attachments: + context_parts.append(format_attachments_as_context(attachments)) + + if mentioned_documents: + context_parts.append(format_mentioned_documents_as_context(mentioned_documents)) + + if context_parts: + context = "\n\n".join(context_parts) + final_query = f"{context}\n\n{user_query}" # if messages: # # Convert frontend messages to LangChain format diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index 9c7e3cb4a..339e1be3d 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -6,9 +6,11 @@ import { type ThreadMessageLike, useExternalStoreRuntime, } from "@assistant-ui/react"; +import { useAtomValue, useSetAtom } from "jotai"; import { useParams, useRouter } from "next/navigation"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; +import { mentionedDocumentIdsAtom } from "@/atoms/chat/mentioned-documents.atom"; import { Thread } from "@/components/assistant-ui/thread"; import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast"; import { LinkPreviewToolUI } from "@/components/tool-ui/link-preview"; @@ -107,6 +109,10 @@ export default function NewChatPage() { >(new Map()); const abortControllerRef = useRef(null); + // Get mentioned document IDs from the composer + const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom); + const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom); + // Create the attachment adapter for file processing const attachmentAdapter = useMemo(() => createAttachmentAdapter(), []); @@ -372,6 +378,14 @@ export default function NewChatPage() { // Extract attachment content to send with the request const attachments = extractAttachmentContent(messageAttachments); + // Get mentioned document IDs for context + const documentIds = mentionedDocumentIds.length > 0 ? [...mentionedDocumentIds] : undefined; + + // Clear mentioned documents after capturing them + if (mentionedDocumentIds.length > 0) { + setMentionedDocumentIds([]); + } + const response = await fetch(`${backendUrl}/api/v1/new_chat`, { method: "POST", headers: { @@ -384,6 +398,7 @@ export default function NewChatPage() { search_space_id: searchSpaceId, messages: messageHistory, attachments: attachments.length > 0 ? attachments : undefined, + mentioned_document_ids: documentIds, }), signal: controller.signal, }); @@ -546,7 +561,7 @@ export default function NewChatPage() { // Note: We no longer clear thinking steps - they persist with the message } }, - [threadId, searchSpaceId, messages] + [threadId, searchSpaceId, messages, mentionedDocumentIds, setMentionedDocumentIds] ); // Convert message (pass through since already in correct format) diff --git a/surfsense_web/atoms/chat/mentioned-documents.atom.ts b/surfsense_web/atoms/chat/mentioned-documents.atom.ts new file mode 100644 index 000000000..6fc0daf06 --- /dev/null +++ b/surfsense_web/atoms/chat/mentioned-documents.atom.ts @@ -0,0 +1,10 @@ +"use client"; + +import { atom } from "jotai"; + +/** + * Atom to store the IDs of documents mentioned in the current chat composer. + * This is used to pass document context to the backend when sending a message. + */ +export const mentionedDocumentIdsAtom = atom([]); + diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index ba12e626e..29b2df22d 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -7,7 +7,6 @@ import { MessagePrimitive, ThreadPrimitive, useAssistantState, - useMessage, } from "@assistant-ui/react"; import { ArrowDownIcon, @@ -27,18 +26,20 @@ import { Search, Sparkles, SquareIcon, + X, } from "lucide-react"; import { useParams } from "next/navigation"; import Link from "next/link"; -import { type FC, useState, useRef, useCallback, useEffect } from "react"; -import { useAtomValue } from "jotai"; +import { type FC, useState, useRef, useCallback, useEffect, createContext, useContext, useMemo } from "react"; +import { createPortal } from "react-dom"; +import { useAtomValue, useSetAtom } from "jotai"; +import { mentionedDocumentIdsAtom } from "@/atoms/chat/mentioned-documents.atom"; import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms"; import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms"; import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; import { getDocumentTypeLabel } from "@/app/dashboard/[search_space_id]/documents/(manage)/components/DocumentTypeIcon"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; -import { useRef, useState } from "react"; import { ComposerAddAttachment, ComposerAttachments, @@ -69,8 +70,6 @@ interface ThreadProps { } // Context to pass thinking steps to AssistantMessage -import { createContext, useContext } from "react"; - const ThinkingStepsContext = createContext>(new Map()); /** @@ -333,12 +332,15 @@ const getTimeBasedGreeting = (userEmail?: string): string => { const ThreadWelcome: FC = () => { const { data: user } = useAtomValue(currentUserAtom); + // Memoize greeting so it doesn't change on re-renders (only on user change) + const greeting = useMemo(() => getTimeBasedGreeting(user?.email), [user?.email]); + return (
{/* Greeting positioned above the composer - fixed position */}

- {getTimeBasedGreeting(user?.email)} + {greeting}

{/* Composer - top edge fixed, expands downward only */} @@ -351,122 +353,155 @@ const ThreadWelcome: FC = () => { const Composer: FC = () => { // ---- State for document mentions ---- - const [allSelectedDocuments, setAllSelectedDocuments] = useState([]); const [mentionedDocuments, setMentionedDocuments] = useState([]); const [showDocumentPopover, setShowDocumentPopover] = useState(false); - const [inputValue, setInputValue] = useState(""); const inputRef = useRef(null); const { search_space_id } = useParams(); + const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom); - const handleInputOrKeyUp = ( - e: React.FormEvent | React.KeyboardEvent - ) => { + // Sync mentioned document IDs to atom for use in chat request + useEffect(() => { + setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id)); + }, [mentionedDocuments, setMentionedDocumentIds]); + + const handleKeyUp = (e: 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]); + // Open document picker when user types '@' + if (e.key === "@" || (e.key === "2" && e.shiftKey)) { + setShowDocumentPopover(true); } - // 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 { + // Close popover if '@' is no longer in the input (user deleted it) + if (showDocumentPopover && !value.includes("@")) { 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]); + const handleKeyDown = (e: React.KeyboardEvent) => { + // Close popover on Escape + if (e.key === "Escape" && showDocumentPopover) { + e.preventDefault(); + setShowDocumentPopover(false); + return; } - setMentionedDocuments( - allSelectedDocuments.filter((doc) => titlesMentioned.includes(doc.title)) - ); + + // Remove last document chip when pressing backspace at the beginning of input + if (e.key === "Backspace" && mentionedDocuments.length > 0) { + const textarea = e.currentTarget; + const selectionStart = textarea.selectionStart; + const selectionEnd = textarea.selectionEnd; + + // Only remove chip if cursor is at position 0 and nothing is selected + if (selectionStart === 0 && selectionEnd === 0) { + e.preventDefault(); + // Remove the last document chip + setMentionedDocuments((prev) => prev.slice(0, -1)); + } + } + }; + + const handleDocumentsMention = (documents: Document[]) => { + // Update mentioned documents (merge with existing, avoid duplicates) + setMentionedDocuments((prev) => { + const existingIds = new Set(prev.map((d) => d.id)); + const newDocs = documents.filter((doc) => !existingIds.has(doc.id)); + return [...prev, ...newDocs]; + }); + + // Clean up the '@' trigger from input if present + if (inputRef.current) { + const input = inputRef.current; + const currentValue = input.value; + // Remove trailing @ if it exists + if (currentValue.endsWith("@")) { + // Use a native input event to properly update the controlled component + const nativeInputValueSetter = Object.getOwnPropertyDescriptor( + window.HTMLTextAreaElement.prototype, + "value" + )?.set; + if (nativeInputValueSetter) { + nativeInputValueSetter.call(input, currentValue.slice(0, -1)); + input.dispatchEvent(new Event("input", { bubbles: true })); + } + } + // Focus the input so user can continue typing + input.focus(); + } + }; + + const handleRemoveDocument = (docId: number) => { + setMentionedDocuments((prev) => prev.filter((doc) => doc.id !== docId)); }; return ( - {/* -------- Input field w/ refs and handlers -------- */} - + {/* -------- Input field with inline document chips -------- */} +
+ {/* Inline document chips */} + {mentionedDocuments.map((doc) => ( + + {doc.title} + + + ))} + {/* Text input */} + 0 ? "Ask about these documents..." : "Ask SurfSense (type @ to mention docs)"} + className="aui-composer-input flex-1 min-w-[120px] max-h-32 resize-none bg-transparent text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0 py-1" + rows={1} + autoFocus + aria-label="Message input" + /> +
- {/* -------- Document mention popover (simple version) -------- */} - {showDocumentPopover && ( -
-
+ {/* -------- Document mention popover (rendered via portal) -------- */} + {showDocumentPopover && typeof document !== "undefined" && createPortal( + <> + {/* Backdrop */} +
- )} - {/* ---- Mention chips for selected/mentioned documents ---- */} - {mentionedDocuments.length > 0 && ( -
- {mentionedDocuments.map((doc) => ( - - {doc.title} - - ))} -
+ , + document.body )} @@ -564,7 +599,7 @@ const ConnectorIndicator: FC = () => {
{/* Document types from the search space */} - {activeDocumentTypes.map(([docType, count]) => ( + {activeDocumentTypes.map(([docType]) => (
{ const thinkingStepsMap = useContext(ThinkingStepsContext); // Get the current message ID to look up thinking steps - const messageId = useMessage((m) => m.id); + const messageId = useAssistantState(({ message }) => message?.id); const thinkingSteps = thinkingStepsMap.get(messageId) || []; // Check if thread is still running (for stopping the spinner when cancelled) diff --git a/surfsense_web/components/new-chat/DocumentsDataTable.tsx b/surfsense_web/components/new-chat/DocumentsDataTable.tsx index 290417e53..d97096317 100644 --- a/surfsense_web/components/new-chat/DocumentsDataTable.tsx +++ b/surfsense_web/components/new-chat/DocumentsDataTable.tsx @@ -1,49 +1,21 @@ "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 { FileText, Search } from "lucide-react"; 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 { ScrollArea } from "@/components/ui/scroll-area"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; -import type { Document, DocumentTypeEnum } from "@/contracts/types/document.types"; +import type { Document } from "@/contracts/types/document.types"; import { documentsApiService } from "@/lib/apis/documents-api.service"; import { cacheKeys } from "@/lib/query-client/cache-keys"; +import { cn } from "@/lib/utils"; interface DocumentsDataTableProps { searchSpaceId: number; onSelectionChange: (documents: Document[]) => void; onDone: () => void; initialSelectedDocuments?: Document[]; - viewOnly?: boolean; } function useDebounced(value: T, delay = 300) { @@ -55,551 +27,190 @@ function useDebounced(value: T, delay = 300) { 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 [highlightedIndex, setHighlightedIndex] = useState(0); + const listRef = useRef(null); + const itemRefs = useRef>(new Map()); const fetchQueryParams = useMemo( () => ({ search_space_id: searchSpaceId, - page: pageIndex, - page_size: pageSize, - ...(documentTypeFilter.length > 0 && { document_types: documentTypeFilter }), + page: 0, + page_size: 20, }), - [searchSpaceId, pageIndex, pageSize, documentTypeFilter, debouncedSearch] + [searchSpaceId] ); const searchQueryParams = useMemo(() => { return { search_space_id: searchSpaceId, - page: pageIndex, - page_size: pageSize, - ...(documentTypeFilter.length > 0 && { document_types: documentTypeFilter }), + page: 0, + page_size: 20, title: debouncedSearch, }; - }, [debouncedSearch, searchSpaceId, pageIndex, pageSize, documentTypeFilter, debouncedSearch]); + }, [debouncedSearch, searchSpaceId]); // 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 + staleTime: 3 * 60 * 1000, enabled: !!searchSpaceId && !debouncedSearch.trim(), }); - // Seaching + // Searching const { data: searchedDocuments, isLoading: isSearchedDocumentsLoading } = useQuery({ queryKey: cacheKeys.documents.withQueryParams(searchQueryParams), queryFn: () => documentsApiService.searchDocuments({ queryParams: searchQueryParams }), - staleTime: 3 * 60 * 1000, // 3 minutes + staleTime: 3 * 60 * 1000, 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 + // Track already selected document IDs + const selectedIds = useMemo( + () => new Set(initialSelectedDocuments.map((d) => d.id)), + [initialSelectedDocuments] ); - // 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; - }); + // Filter out already selected documents for navigation + const selectableDocuments = useMemo( + () => actualDocuments.filter((doc) => !selectedIds.has(doc.id)), + [actualDocuments, selectedIds] + ); - // Track the last notified selection to avoid redundant parent calls - const lastNotifiedSelection = useRef(""); + const handleSelectDocument = useCallback((doc: Document) => { + onSelectionChange([...initialSelectedDocuments, doc]); + onDone(); + }, [initialSelectedDocuments, onSelectionChange, onDone]); - // Update row selection only when initialSelectedDocuments changes (not rowSelection itself) + // Scroll highlighted item into view 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; + const item = itemRefs.current.get(highlightedIndex); + if (item) { + item.scrollIntoView({ block: "nearest", behavior: "smooth" }); } + }, [highlightedIndex]); - setRowSelection(initialRowSelection); - }, [initialRowSelection]); // Remove rowSelection from dependencies to prevent loop + // Handle keyboard navigation + const handleKeyDown = useCallback((e: React.KeyboardEvent) => { + if (selectableDocuments.length === 0) return; - // 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; + switch (e.key) { + case "ArrowDown": + e.preventDefault(); + setHighlightedIndex((prev) => + prev < selectableDocuments.length - 1 ? prev + 1 : 0 + ); + break; + case "ArrowUp": + e.preventDefault(); + setHighlightedIndex((prev) => + prev > 0 ? prev - 1 : selectableDocuments.length - 1 + ); + break; + case "Enter": + e.preventDefault(); + if (selectableDocuments[highlightedIndex]) { + handleSelectDocument(selectableDocuments[highlightedIndex]); } - } - - // 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]); + break; + case "Escape": + e.preventDefault(); + onDone(); + break; + } + }, [selectableDocuments, highlightedIndex, handleSelectDocument, onDone]); return ( -
- {/* Header Controls */} -
- {/* Search and Filter Row */} -
-
- - { - setSearch(event.target.value); - setPageIndex(0); // Reset to first page on search - }} - className="pl-10 text-sm" - /> +
+ {/* Search */} +
+ + { + setSearch(e.target.value); + setHighlightedIndex(0); + }} + className="pl-8 h-8 text-sm border-0 focus-visible:ring-0 focus-visible:ring-offset-0" + autoFocus + /> +
+ + {/* Document List */} + + {actualLoading ? ( +
+
- - - - - -
-
Filter by Type
-
- {availableTypes.map((type) => ( -
- handleToggleType(type, !!checked)} - /> - -
- ))} -
- {documentTypeFilter.length > 0 && ( - - )} -
- - -
- - {/* Action Controls Row */} -
-
- - {selectedCount} selected {actualLoading && "· Loading..."} - -
-
- - - -
+ {/* Title */} + + {doc.title} + + + ); + })}
- -
-
- - {/* 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)} -
- -
-
+ )} +
); } From 8e3f4f4ed7fa1bd15e9f100fc60d69f543bd44d3 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Tue, 23 Dec 2025 15:13:03 +0530 Subject: [PATCH 4/7] feat: implement document mention extraction and management in new chat - Added functionality to extract and manage mentioned documents within the new chat interface. - Introduced new atoms for storing mentioned documents and their mappings to user messages. - Enhanced the message persistence logic to include mentioned documents, ensuring they are displayed correctly in the chat. - Updated the UI components to support document mentions, including a refined document selection interface. - Improved state management for document mentions to ensure a seamless user experience. --- .../new-chat/[[...chat_id]]/page.tsx | 77 +++- .../atoms/chat/mentioned-documents.atom.ts | 22 ++ .../components/assistant-ui/thread.tsx | 117 +++++- .../new-chat/DocumentsDataTable.tsx | 371 +++++++++--------- 4 files changed, 383 insertions(+), 204 deletions(-) diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index 339e1be3d..fa11e9ecf 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -10,7 +10,7 @@ import { useAtomValue, useSetAtom } from "jotai"; import { useParams, useRouter } from "next/navigation"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; -import { mentionedDocumentIdsAtom } from "@/atoms/chat/mentioned-documents.atom"; +import { mentionedDocumentIdsAtom, mentionedDocumentsAtom, messageDocumentsMapAtom, type MentionedDocumentInfo } from "@/atoms/chat/mentioned-documents.atom"; import { Thread } from "@/components/assistant-ui/thread"; import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast"; import { LinkPreviewToolUI } from "@/components/tool-ui/link-preview"; @@ -48,6 +48,23 @@ function extractThinkingSteps(content: unknown): ThinkingStep[] { return thinkingPart?.steps || []; } +/** + * Extract mentioned documents from message content + */ +function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] { + if (!Array.isArray(content)) return []; + + const docsPart = content.find( + (part: unknown) => + typeof part === "object" && + part !== null && + "type" in part && + (part as { type: string }).type === "mentioned-documents" + ) as { type: "mentioned-documents"; documents: MentionedDocumentInfo[] } | undefined; + + return docsPart?.documents || []; +} + /** * Convert backend message to assistant-ui ThreadMessageLike format * Filters out 'thinking-steps' part as it's handled separately @@ -58,13 +75,14 @@ function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike { if (typeof msg.content === "string") { content = [{ type: "text", text: msg.content }]; } else if (Array.isArray(msg.content)) { - // Filter out thinking-steps part - it's handled separately via messageThinkingSteps + // Filter out custom metadata parts - they're handled separately const filteredContent = msg.content.filter( - (part: unknown) => - !(typeof part === "object" && - part !== null && - "type" in part && - (part as { type: string }).type === "thinking-steps") + (part: unknown) => { + if (typeof part !== "object" || part === null || !("type" in part)) return true; + const partType = (part as { type: string }).type; + // Filter out thinking-steps and mentioned-documents + return partType !== "thinking-steps" && partType !== "mentioned-documents"; + } ); content = filteredContent.length > 0 ? (filteredContent as ThreadMessageLike["content"]) @@ -111,7 +129,10 @@ export default function NewChatPage() { // Get mentioned document IDs from the composer const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom); + const mentionedDocuments = useAtomValue(mentionedDocumentsAtom); const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom); + const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom); + const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom); // Create the attachment adapter for file processing const attachmentAdapter = useMemo(() => createAttachmentAdapter(), []); @@ -150,6 +171,9 @@ export default function NewChatPage() { // Extract and restore thinking steps from persisted messages const restoredThinkingSteps = new Map(); + // Extract and restore mentioned documents from persisted messages + const restoredDocsMap: Record = {}; + for (const msg of response.messages) { if (msg.role === "assistant") { const steps = extractThinkingSteps(msg.content); @@ -157,10 +181,19 @@ export default function NewChatPage() { restoredThinkingSteps.set(`msg-${msg.id}`, steps); } } + if (msg.role === "user") { + const docs = extractMentionedDocuments(msg.content); + if (docs.length > 0) { + restoredDocsMap[`msg-${msg.id}`] = docs; + } + } } if (restoredThinkingSteps.size > 0) { setMessageThinkingSteps(restoredThinkingSteps); } + if (Object.keys(restoredDocsMap).length > 0) { + setMessageDocumentsMap(restoredDocsMap); + } } } else { // Create new thread @@ -239,10 +272,33 @@ export default function NewChatPage() { }; setMessages((prev) => [...prev, userMessage]); - // Persist user message (don't await, fire and forget) + // Store mentioned documents with this message for display + if (mentionedDocuments.length > 0) { + const docsInfo: MentionedDocumentInfo[] = mentionedDocuments.map((doc) => ({ + id: doc.id, + title: doc.title, + document_type: doc.document_type, + })); + setMessageDocumentsMap((prev) => ({ + ...prev, + [userMsgId]: docsInfo, + })); + } + + // Persist user message with mentioned documents (don't await, fire and forget) + const persistContent = mentionedDocuments.length > 0 + ? [ + ...message.content, + { type: "mentioned-documents", documents: mentionedDocuments.map((doc) => ({ + id: doc.id, + title: doc.title, + document_type: doc.document_type, + })) }, + ] + : message.content; appendMessage(threadId, { role: "user", - content: message.content, + content: persistContent, }).catch((err) => console.error("Failed to persist user message:", err)); // Start streaming response @@ -384,6 +440,7 @@ export default function NewChatPage() { // Clear mentioned documents after capturing them if (mentionedDocumentIds.length > 0) { setMentionedDocumentIds([]); + setMentionedDocuments([]); } const response = await fetch(`${backendUrl}/api/v1/new_chat`, { @@ -561,7 +618,7 @@ export default function NewChatPage() { // Note: We no longer clear thinking steps - they persist with the message } }, - [threadId, searchSpaceId, messages, mentionedDocumentIds, setMentionedDocumentIds] + [threadId, searchSpaceId, messages, mentionedDocumentIds, mentionedDocuments, setMentionedDocumentIds, setMentionedDocuments, setMessageDocumentsMap] ); // Convert message (pass through since already in correct format) diff --git a/surfsense_web/atoms/chat/mentioned-documents.atom.ts b/surfsense_web/atoms/chat/mentioned-documents.atom.ts index 6fc0daf06..67ce10eee 100644 --- a/surfsense_web/atoms/chat/mentioned-documents.atom.ts +++ b/surfsense_web/atoms/chat/mentioned-documents.atom.ts @@ -1,6 +1,7 @@ "use client"; import { atom } from "jotai"; +import type { Document } from "@/contracts/types/document.types"; /** * Atom to store the IDs of documents mentioned in the current chat composer. @@ -8,3 +9,24 @@ import { atom } from "jotai"; */ export const mentionedDocumentIdsAtom = atom([]); +/** + * Atom to store the full document objects mentioned in the current chat composer. + * This persists across component remounts. + */ +export const mentionedDocumentsAtom = atom([]); + +/** + * Simplified document info for display purposes + */ +export interface MentionedDocumentInfo { + id: number; + title: string; + document_type: string; +} + +/** + * Atom to store mentioned documents per message ID. + * This allows displaying which documents were mentioned with each user message. + */ +export const messageDocumentsMapAtom = atom>({}); + diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index 29b2df22d..ea45c3cae 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -18,6 +18,7 @@ import { ChevronRightIcon, CopyIcon, DownloadIcon, + FileText, Loader2, PencilIcon, Plug2, @@ -32,8 +33,8 @@ import { useParams } from "next/navigation"; import Link from "next/link"; import { type FC, useState, useRef, useCallback, useEffect, createContext, useContext, useMemo } from "react"; import { createPortal } from "react-dom"; -import { useAtomValue, useSetAtom } from "jotai"; -import { mentionedDocumentIdsAtom } from "@/atoms/chat/mentioned-documents.atom"; +import { useAtom, useAtomValue, useSetAtom } from "jotai"; +import { mentionedDocumentIdsAtom, mentionedDocumentsAtom, messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom"; import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms"; import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms"; import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors"; @@ -55,7 +56,7 @@ import { ChainOfThoughtStep, ChainOfThoughtTrigger, } from "@/components/prompt-kit/chain-of-thought"; -import { DocumentsDataTable } from "@/components/new-chat/DocumentsDataTable"; +import { DocumentsDataTable, type DocumentsDataTableRef } from "@/components/new-chat/DocumentsDataTable"; import { Button } from "@/components/ui/button"; import type { Document } from "@/contracts/types/document.types"; import { cn } from "@/lib/utils"; @@ -352,10 +353,12 @@ const ThreadWelcome: FC = () => { }; const Composer: FC = () => { - // ---- State for document mentions ---- - const [mentionedDocuments, setMentionedDocuments] = useState([]); + // ---- State for document mentions (using atoms to persist across remounts) ---- + const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom); const [showDocumentPopover, setShowDocumentPopover] = useState(false); + const [mentionQuery, setMentionQuery] = useState(""); const inputRef = useRef(null); + const documentPickerRef = useRef(null); const { search_space_id } = useParams(); const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom); @@ -364,6 +367,13 @@ const Composer: FC = () => { setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id)); }, [mentionedDocuments, setMentionedDocumentIds]); + // Extract mention query (text after @) + const extractMentionQuery = useCallback((value: string): string => { + const atIndex = value.lastIndexOf("@"); + if (atIndex === -1) return ""; + return value.slice(atIndex + 1); + }, []); + const handleKeyUp = (e: React.KeyboardEvent) => { const textarea = e.currentTarget; const value = textarea.value; @@ -371,20 +381,60 @@ const Composer: FC = () => { // Open document picker when user types '@' if (e.key === "@" || (e.key === "2" && e.shiftKey)) { setShowDocumentPopover(true); + setMentionQuery(""); + return; } - // Close popover if '@' is no longer in the input (user deleted it) - if (showDocumentPopover && !value.includes("@")) { - setShowDocumentPopover(false); + // Check if value contains @ and extract query + if (value.includes("@")) { + const query = extractMentionQuery(value); + + // Close popup if query starts with space (user typed "@ ") + if (query.startsWith(" ")) { + setShowDocumentPopover(false); + setMentionQuery(""); + return; + } + + // Reopen popup if @ is present and query doesn't start with space + // (handles case where user deleted the space after @) + if (!showDocumentPopover) { + setShowDocumentPopover(true); + } + setMentionQuery(query); + } else { + // Close popover if '@' is no longer in the input (user deleted it) + if (showDocumentPopover) { + setShowDocumentPopover(false); + setMentionQuery(""); + } } }; const handleKeyDown = (e: React.KeyboardEvent) => { - // Close popover on Escape - if (e.key === "Escape" && showDocumentPopover) { - e.preventDefault(); - setShowDocumentPopover(false); - return; + // When popup is open, handle navigation keys + if (showDocumentPopover) { + if (e.key === "ArrowDown") { + e.preventDefault(); + documentPickerRef.current?.moveDown(); + return; + } + if (e.key === "ArrowUp") { + e.preventDefault(); + documentPickerRef.current?.moveUp(); + return; + } + if (e.key === "Enter") { + e.preventDefault(); + documentPickerRef.current?.selectHighlighted(); + return; + } + if (e.key === "Escape") { + e.preventDefault(); + setShowDocumentPopover(false); + setMentionQuery(""); + return; + } } // Remove last document chip when pressing backspace at the beginning of input @@ -410,25 +460,30 @@ const Composer: FC = () => { return [...prev, ...newDocs]; }); - // Clean up the '@' trigger from input if present + // Clean up the '@...' mention text from input if (inputRef.current) { const input = inputRef.current; const currentValue = input.value; - // Remove trailing @ if it exists - if (currentValue.endsWith("@")) { - // Use a native input event to properly update the controlled component + const atIndex = currentValue.lastIndexOf("@"); + + if (atIndex !== -1) { + // Remove @ and everything after it + const newValue = currentValue.slice(0, atIndex); const nativeInputValueSetter = Object.getOwnPropertyDescriptor( window.HTMLTextAreaElement.prototype, "value" )?.set; if (nativeInputValueSetter) { - nativeInputValueSetter.call(input, currentValue.slice(0, -1)); + nativeInputValueSetter.call(input, newValue); input.dispatchEvent(new Event("input", { bubbles: true })); } } // Focus the input so user can continue typing input.focus(); } + + // Reset mention query + setMentionQuery(""); }; const handleRemoveDocument = (docId: number) => { @@ -494,10 +549,15 @@ const Composer: FC = () => { }} > setShowDocumentPopover(false)} + onDone={() => { + setShowDocumentPopover(false); + setMentionQuery(""); + }} initialSelectedDocuments={mentionedDocuments} + externalSearch={mentionQuery} />
, @@ -819,6 +879,10 @@ const AssistantActionBar: FC = () => { }; const UserMessage: FC = () => { + const messageId = useAssistantState(({ message }) => message?.id); + const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom); + const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined; + return ( {
+ {/* Display mentioned documents as chips */} + {mentionedDocs && mentionedDocs.length > 0 && ( +
+ {mentionedDocs.map((doc) => ( + + + {doc.title} + + ))} +
+ )}
diff --git a/surfsense_web/components/new-chat/DocumentsDataTable.tsx b/surfsense_web/components/new-chat/DocumentsDataTable.tsx index d97096317..2c1ccf1cb 100644 --- a/surfsense_web/components/new-chat/DocumentsDataTable.tsx +++ b/surfsense_web/components/new-chat/DocumentsDataTable.tsx @@ -1,21 +1,26 @@ "use client"; import { useQuery } from "@tanstack/react-query"; -import { FileText, Search } from "lucide-react"; -import { useCallback, useEffect, useMemo, useRef, useState } from "react"; -import { Input } from "@/components/ui/input"; -import { ScrollArea } from "@/components/ui/scroll-area"; +import { FileText } from "lucide-react"; +import { forwardRef, useCallback, useEffect, useImperativeHandle, useMemo, useRef, useState } from "react"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; import type { Document } from "@/contracts/types/document.types"; import { documentsApiService } from "@/lib/apis/documents-api.service"; import { cacheKeys } from "@/lib/query-client/cache-keys"; import { cn } from "@/lib/utils"; +export interface DocumentsDataTableRef { + selectHighlighted: () => void; + moveUp: () => void; + moveDown: () => void; +} + interface DocumentsDataTableProps { searchSpaceId: number; onSelectionChange: (documents: Document[]) => void; onDone: () => void; initialSelectedDocuments?: Document[]; + externalSearch?: string; } function useDebounced(value: T, delay = 300) { @@ -27,190 +32,206 @@ function useDebounced(value: T, delay = 300) { return debounced; } -export function DocumentsDataTable({ - searchSpaceId, - onSelectionChange, - onDone, - initialSelectedDocuments = [], -}: DocumentsDataTableProps) { - const [search, setSearch] = useState(""); - const debouncedSearch = useDebounced(search, 300); - const [highlightedIndex, setHighlightedIndex] = useState(0); - const listRef = useRef(null); - const itemRefs = useRef>(new Map()); +export const DocumentsDataTable = forwardRef( + function DocumentsDataTable({ + searchSpaceId, + onSelectionChange, + onDone, + initialSelectedDocuments = [], + externalSearch = "", + }, ref) { + // Use external search + const search = externalSearch; + const debouncedSearch = useDebounced(search, 150); + const [highlightedIndex, setHighlightedIndex] = useState(0); + const itemRefs = useRef>(new Map()); - const fetchQueryParams = useMemo( - () => ({ - search_space_id: searchSpaceId, - page: 0, - page_size: 20, - }), - [searchSpaceId] - ); + const fetchQueryParams = useMemo( + () => ({ + search_space_id: searchSpaceId, + page: 0, + page_size: 20, + }), + [searchSpaceId] + ); - const searchQueryParams = useMemo(() => { - return { - search_space_id: searchSpaceId, - page: 0, - page_size: 20, - title: debouncedSearch, - }; - }, [debouncedSearch, searchSpaceId]); + const searchQueryParams = useMemo(() => { + return { + search_space_id: searchSpaceId, + page: 0, + page_size: 20, + title: debouncedSearch, + }; + }, [debouncedSearch, searchSpaceId]); - // 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, - enabled: !!searchSpaceId && !debouncedSearch.trim(), - }); + // 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, + enabled: !!searchSpaceId && !debouncedSearch.trim(), + }); - // Searching - const { data: searchedDocuments, isLoading: isSearchedDocumentsLoading } = useQuery({ - queryKey: cacheKeys.documents.withQueryParams(searchQueryParams), - queryFn: () => documentsApiService.searchDocuments({ queryParams: searchQueryParams }), - staleTime: 3 * 60 * 1000, - enabled: !!searchSpaceId && !!debouncedSearch.trim(), - }); + // Searching + const { data: searchedDocuments, isLoading: isSearchedDocumentsLoading } = useQuery({ + queryKey: cacheKeys.documents.withQueryParams(searchQueryParams), + queryFn: () => documentsApiService.searchDocuments({ queryParams: searchQueryParams }), + staleTime: 3 * 60 * 1000, + enabled: !!searchSpaceId && !!debouncedSearch.trim(), + }); - const actualDocuments = debouncedSearch.trim() - ? searchedDocuments?.items || [] - : documents?.items || []; - const actualLoading = debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading; + const actualDocuments = debouncedSearch.trim() + ? searchedDocuments?.items || [] + : documents?.items || []; + const actualLoading = debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading; - // Track already selected document IDs - const selectedIds = useMemo( - () => new Set(initialSelectedDocuments.map((d) => d.id)), - [initialSelectedDocuments] - ); + // Track already selected document IDs + const selectedIds = useMemo( + () => new Set(initialSelectedDocuments.map((d) => d.id)), + [initialSelectedDocuments] + ); - // Filter out already selected documents for navigation - const selectableDocuments = useMemo( - () => actualDocuments.filter((doc) => !selectedIds.has(doc.id)), - [actualDocuments, selectedIds] - ); + // Filter out already selected documents for navigation + const selectableDocuments = useMemo( + () => actualDocuments.filter((doc) => !selectedIds.has(doc.id)), + [actualDocuments, selectedIds] + ); - const handleSelectDocument = useCallback((doc: Document) => { - onSelectionChange([...initialSelectedDocuments, doc]); - onDone(); - }, [initialSelectedDocuments, onSelectionChange, onDone]); + const handleSelectDocument = useCallback((doc: Document) => { + onSelectionChange([...initialSelectedDocuments, doc]); + onDone(); + }, [initialSelectedDocuments, onSelectionChange, onDone]); - // Scroll highlighted item into view - useEffect(() => { - const item = itemRefs.current.get(highlightedIndex); - if (item) { - item.scrollIntoView({ block: "nearest", behavior: "smooth" }); + // Scroll highlighted item into view + useEffect(() => { + const item = itemRefs.current.get(highlightedIndex); + if (item) { + item.scrollIntoView({ block: "nearest", behavior: "smooth" }); + } + }, [highlightedIndex]); + + // Reset highlighted index when external search changes + const prevSearchRef = useRef(search); + if (prevSearchRef.current !== search) { + prevSearchRef.current = search; + if (highlightedIndex !== 0) { + setHighlightedIndex(0); + } } - }, [highlightedIndex]); - // Handle keyboard navigation - const handleKeyDown = useCallback((e: React.KeyboardEvent) => { - if (selectableDocuments.length === 0) return; - - switch (e.key) { - case "ArrowDown": - e.preventDefault(); - setHighlightedIndex((prev) => - prev < selectableDocuments.length - 1 ? prev + 1 : 0 - ); - break; - case "ArrowUp": - e.preventDefault(); - setHighlightedIndex((prev) => - prev > 0 ? prev - 1 : selectableDocuments.length - 1 - ); - break; - case "Enter": - e.preventDefault(); + // Expose methods to parent via ref + useImperativeHandle(ref, () => ({ + selectHighlighted: () => { if (selectableDocuments[highlightedIndex]) { handleSelectDocument(selectableDocuments[highlightedIndex]); } - break; - case "Escape": - e.preventDefault(); - onDone(); - break; - } - }, [selectableDocuments, highlightedIndex, handleSelectDocument, onDone]); + }, + moveUp: () => { + setHighlightedIndex((prev) => + prev > 0 ? prev - 1 : selectableDocuments.length - 1 + ); + }, + moveDown: () => { + setHighlightedIndex((prev) => + prev < selectableDocuments.length - 1 ? prev + 1 : 0 + ); + }, + }), [selectableDocuments, highlightedIndex, handleSelectDocument]); - return ( -
- {/* Search */} -
- - { - setSearch(e.target.value); - setHighlightedIndex(0); - }} - className="pl-8 h-8 text-sm border-0 focus-visible:ring-0 focus-visible:ring-offset-0" - autoFocus - /> + // Handle keyboard navigation + const handleKeyDown = useCallback((e: React.KeyboardEvent) => { + if (selectableDocuments.length === 0) return; + + switch (e.key) { + case "ArrowDown": + e.preventDefault(); + setHighlightedIndex((prev) => + prev < selectableDocuments.length - 1 ? prev + 1 : 0 + ); + break; + case "ArrowUp": + e.preventDefault(); + setHighlightedIndex((prev) => + prev > 0 ? prev - 1 : selectableDocuments.length - 1 + ); + break; + case "Enter": + e.preventDefault(); + if (selectableDocuments[highlightedIndex]) { + handleSelectDocument(selectableDocuments[highlightedIndex]); + } + break; + case "Escape": + e.preventDefault(); + onDone(); + break; + } + }, [selectableDocuments, highlightedIndex, handleSelectDocument, onDone]); + + return ( +
+ {/* Document List */} +
+ {actualLoading ? ( +
+
+
+ ) : actualDocuments.length === 0 ? ( +
+ +

No documents found

+
+ ) : ( +
+ {actualDocuments.map((doc) => { + const isAlreadySelected = selectedIds.has(doc.id); + const selectableIndex = selectableDocuments.findIndex((d) => d.id === doc.id); + const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex; + + return ( + + ); + })} +
+ )} +
- - {/* Document List */} - - {actualLoading ? ( -
-
-
- ) : actualDocuments.length === 0 ? ( -
- -

No documents found

-
- ) : ( -
- {actualDocuments.map((doc) => { - const isAlreadySelected = selectedIds.has(doc.id); - const selectableIndex = selectableDocuments.findIndex((d) => d.id === doc.id); - const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex; - - return ( - - ); - })} -
- )} - -
- ); -} + ); + } +); From 1d646d3b5fa3f7e7bd7cd80d3714cf828569f6e7 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Tue, 23 Dec 2025 16:04:39 +0530 Subject: [PATCH 5/7] refactor: improve message content filtering and update dependencies in new chat - Enhanced the message content filtering logic to exclude both "thinking-steps" and "mentioned-documents" for better clarity in chat messages. - Updated the dependency imports in the thread component for improved organization and consistency. - Adjusted the effect dependencies in the NewChatPage component to include the message documents map for better state management. --- .../new-chat/[[...chat_id]]/page.tsx | 17 ++++----- .../components/assistant-ui/thread.tsx | 36 ++++--------------- 2 files changed, 14 insertions(+), 39 deletions(-) diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index 6da678d7c..d50db6d9f 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -77,15 +77,12 @@ function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike { content = [{ type: "text", text: msg.content }]; } else if (Array.isArray(msg.content)) { // Filter out custom metadata parts - they're handled separately - const filteredContent = msg.content.filter( - (part: unknown) => - !( - typeof part === "object" && - part !== null && - "type" in part && - (part as { type: string }).type === "thinking-steps" - ) - ); + const filteredContent = msg.content.filter((part: unknown) => { + if (typeof part !== "object" || part === null || !("type" in part)) return true; + const partType = (part as { type: string }).type; + // Filter out thinking-steps and mentioned-documents + return partType !== "thinking-steps" && partType !== "mentioned-documents"; + }); content = filteredContent.length > 0 ? (filteredContent as ThreadMessageLike["content"]) @@ -219,7 +216,7 @@ export default function NewChatPage() { } finally { setIsInitializing(false); } - }, [urlChatId, searchSpaceId, router]); + }, [urlChatId, searchSpaceId, router, setMessageDocumentsMap]); // Initialize on mount useEffect(() => { diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index 98f5539a1..5c80842ef 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -9,7 +9,7 @@ import { useAssistantState, useThreadViewport, } from "@assistant-ui/react"; -import { useAtomValue } from "jotai"; +import { useAtom, useAtomValue, useSetAtom } from "jotai"; import { AlertCircle, ArrowDownIcon, @@ -32,17 +32,11 @@ import { SquareIcon, X, } from "lucide-react"; -import { useParams } from "next/navigation"; import Link from "next/link"; -import { type FC, useState, useRef, useCallback, useEffect, createContext, useContext, useMemo } from "react"; +import { useParams } from "next/navigation"; +import { type FC, createContext, useCallback, useContext, useEffect, useMemo, useRef, useState } from "react"; import { createPortal } from "react-dom"; -import { useAtom, useAtomValue, useSetAtom } from "jotai"; import { mentionedDocumentIdsAtom, mentionedDocumentsAtom, messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom"; -import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms"; -import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms"; -import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors"; -import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; -import { getDocumentTypeLabel } from "@/app/dashboard/[search_space_id]/documents/(manage)/components/DocumentTypeIcon"; import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms"; import { globalNewLLMConfigsAtom, @@ -59,6 +53,7 @@ import { import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { DocumentsDataTable, type DocumentsDataTableRef } from "@/components/new-chat/DocumentsDataTable"; import { ChainOfThought, ChainOfThoughtContent, @@ -66,12 +61,12 @@ import { ChainOfThoughtStep, ChainOfThoughtTrigger, } from "@/components/prompt-kit/chain-of-thought"; -import { DocumentsDataTable, type DocumentsDataTableRef } from "@/components/new-chat/DocumentsDataTable"; import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking"; import { Button } from "@/components/ui/button"; -import type { Document } from "@/contracts/types/document.types"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; +import type { Document } from "@/contracts/types/document.types"; +import { getDocumentTypeLabel } from "@/app/dashboard/[search_space_id]/documents/(manage)/components/DocumentTypeIcon"; import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors"; import { cn } from "@/lib/utils"; @@ -257,7 +252,7 @@ const ThinkingStepsScrollHandler: FC = () => { const scrollAttempt = () => { try { viewport.scrollToBottom(); - } catch (e) { + } catch { // Ignore errors - viewport might not be ready } }; @@ -531,23 +526,6 @@ const Composer: FC = () => { setMentionedDocuments((prev) => prev.filter((doc) => doc.id !== docId)); }; - // Check if a model is configured - needed to disable input - const { data: userConfigs } = useAtomValue(newLLMConfigsAtom); - const { data: globalConfigs } = useAtomValue(globalNewLLMConfigsAtom); - const { data: preferences } = useAtomValue(llmPreferencesAtom); - - const hasModelConfigured = useMemo(() => { - if (!preferences) return false; - const agentLlmId = preferences.agent_llm_id; - if (agentLlmId === null || agentLlmId === undefined) return false; - - // Check if the configured model actually exists - if (agentLlmId < 0) { - return globalConfigs?.some((c) => c.id === agentLlmId) ?? false; - } - return userConfigs?.some((c) => c.id === agentLlmId) ?? false; - }, [preferences, globalConfigs, userConfigs]); - return ( From 96bf469a1fe7cc26b8b4e32fdd809cb0e4bec94a Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 24 Dec 2025 06:46:58 +0200 Subject: [PATCH 6/7] align the user action edit button --- surfsense_web/components/assistant-ui/thread.tsx | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index 2e35bc295..d9275cebd 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -269,7 +269,6 @@ const ThinkingStepsScrollHandler: FC = () => { return null; // This component doesn't render anything }; -export const Thread: FC = ({ messageThinkingSteps = new Map() }) => { export const Thread: FC = ({ messageThinkingSteps = new Map(), header }) => { return ( @@ -989,7 +988,7 @@ const UserMessage: FC = () => {
-
+
From deec8c5c6c463ca813cf966f35b1034aa5929edd Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 24 Dec 2025 07:06:35 +0200 Subject: [PATCH 7/7] fix: formatting --- .../app/tasks/chat/stream_new_chat.py | 4 +- .../new-chat/[[...chat_id]]/page.tsx | 56 +++++--- .../atoms/chat/mentioned-documents.atom.ts | 1 - .../components/assistant-ui/thread.tsx | 116 +++++++++------ .../new-chat/DocumentsDataTable.tsx | 133 ++++++++++-------- 5 files changed, 183 insertions(+), 127 deletions(-) diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index daf7a20c7..2038e85dc 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -181,7 +181,9 @@ async def stream_new_chat( context_parts.append(format_attachments_as_context(attachments)) if mentioned_documents: - context_parts.append(format_mentioned_documents_as_context(mentioned_documents)) + context_parts.append( + format_mentioned_documents_as_context(mentioned_documents) + ) if context_parts: context = "\n\n".join(context_parts) diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index dc83d9204..a7fc23802 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -10,7 +10,12 @@ import { useAtomValue, useSetAtom } from "jotai"; import { useParams, useRouter } from "next/navigation"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; -import { mentionedDocumentIdsAtom, mentionedDocumentsAtom, messageDocumentsMapAtom, type MentionedDocumentInfo } from "@/atoms/chat/mentioned-documents.atom"; +import { + type MentionedDocumentInfo, + mentionedDocumentIdsAtom, + mentionedDocumentsAtom, + messageDocumentsMapAtom, +} from "@/atoms/chat/mentioned-documents.atom"; import { Thread } from "@/components/assistant-ui/thread"; import { ChatHeader } from "@/components/new-chat/chat-header"; import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking"; @@ -54,15 +59,15 @@ function extractThinkingSteps(content: unknown): ThinkingStep[] { */ function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] { if (!Array.isArray(content)) return []; - + const docsPart = content.find( - (part: unknown) => - typeof part === "object" && - part !== null && - "type" in part && + (part: unknown) => + typeof part === "object" && + part !== null && + "type" in part && (part as { type: string }).type === "mentioned-documents" ) as { type: "mentioned-documents"; documents: MentionedDocumentInfo[] } | undefined; - + return docsPart?.documents || []; } @@ -179,7 +184,7 @@ export default function NewChatPage() { const restoredThinkingSteps = new Map(); // Extract and restore mentioned documents from persisted messages const restoredDocsMap: Record = {}; - + for (const msg of response.messages) { if (msg.role === "assistant") { const steps = extractThinkingSteps(msg.content); @@ -292,16 +297,20 @@ export default function NewChatPage() { } // Persist user message with mentioned documents (don't await, fire and forget) - const persistContent = mentionedDocuments.length > 0 - ? [ - ...message.content, - { type: "mentioned-documents", documents: mentionedDocuments.map((doc) => ({ - id: doc.id, - title: doc.title, - document_type: doc.document_type, - })) }, - ] - : message.content; + const persistContent = + mentionedDocuments.length > 0 + ? [ + ...message.content, + { + type: "mentioned-documents", + documents: mentionedDocuments.map((doc) => ({ + id: doc.id, + title: doc.title, + document_type: doc.document_type, + })), + }, + ] + : message.content; appendMessage(threadId, { role: "user", content: persistContent, @@ -626,7 +635,16 @@ export default function NewChatPage() { // Note: We no longer clear thinking steps - they persist with the message } }, - [threadId, searchSpaceId, messages, mentionedDocumentIds, mentionedDocuments, setMentionedDocumentIds, setMentionedDocuments, setMessageDocumentsMap] + [ + threadId, + searchSpaceId, + messages, + mentionedDocumentIds, + mentionedDocuments, + setMentionedDocumentIds, + setMentionedDocuments, + setMessageDocumentsMap, + ] ); // Convert message (pass through since already in correct format) diff --git a/surfsense_web/atoms/chat/mentioned-documents.atom.ts b/surfsense_web/atoms/chat/mentioned-documents.atom.ts index 67ce10eee..79ea27d12 100644 --- a/surfsense_web/atoms/chat/mentioned-documents.atom.ts +++ b/surfsense_web/atoms/chat/mentioned-documents.atom.ts @@ -29,4 +29,3 @@ export interface MentionedDocumentInfo { * This allows displaying which documents were mentioned with each user message. */ export const messageDocumentsMapAtom = atom>({}); - diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index d9275cebd..191d60338 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -7,8 +7,8 @@ import { MessagePrimitive, ThreadPrimitive, useAssistantState, - useThreadViewport, useMessage, + useThreadViewport, } from "@assistant-ui/react"; import { useAtom, useAtomValue, useSetAtom } from "jotai"; import { @@ -35,9 +35,23 @@ import { } from "lucide-react"; import Link from "next/link"; import { useParams } from "next/navigation"; -import { type FC, createContext, useCallback, useContext, useEffect, useMemo, useRef, useState } from "react"; +import { + createContext, + type FC, + useCallback, + useContext, + useEffect, + useMemo, + useRef, + useState, +} from "react"; import { createPortal } from "react-dom"; -import { mentionedDocumentIdsAtom, mentionedDocumentsAtom, messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom"; +import { getDocumentTypeLabel } from "@/app/dashboard/[search_space_id]/documents/(manage)/components/DocumentTypeIcon"; +import { + mentionedDocumentIdsAtom, + mentionedDocumentsAtom, + messageDocumentsMapAtom, +} from "@/atoms/chat/mentioned-documents.atom"; import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms"; import { globalNewLLMConfigsAtom, @@ -54,7 +68,10 @@ import { import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; -import { DocumentsDataTable, type DocumentsDataTableRef } from "@/components/new-chat/DocumentsDataTable"; +import { + DocumentsDataTable, + type DocumentsDataTableRef, +} from "@/components/new-chat/DocumentsDataTable"; import { ChainOfThought, ChainOfThoughtContent, @@ -67,7 +84,6 @@ import { Button } from "@/components/ui/button"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; import type { Document } from "@/contracts/types/document.types"; -import { getDocumentTypeLabel } from "@/app/dashboard/[search_space_id]/documents/(manage)/components/DocumentTypeIcon"; import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors"; import { cn } from "@/lib/utils"; @@ -371,7 +387,7 @@ const getTimeBasedGreeting = (userEmail?: string): string => { const ThreadWelcome: FC = () => { const { data: user } = useAtomValue(currentUserAtom); - + // Memoize greeting so it doesn't change on re-renders (only on user change) const greeting = useMemo(() => getTimeBasedGreeting(user?.email), [user?.email]); @@ -427,14 +443,14 @@ const Composer: FC = () => { // Check if value contains @ and extract query if (value.includes("@")) { const query = extractMentionQuery(value); - + // Close popup if query starts with space (user typed "@ ") if (query.startsWith(" ")) { setShowDocumentPopover(false); setMentionQuery(""); return; } - + // Reopen popup if @ is present and query doesn't start with space // (handles case where user deleted the space after @) if (!showDocumentPopover) { @@ -504,7 +520,7 @@ const Composer: FC = () => { const input = inputRef.current; const currentValue = input.value; const atIndex = currentValue.lastIndexOf("@"); - + if (atIndex !== -1) { // Remove @ and everything after it const newValue = currentValue.slice(0, atIndex); @@ -520,7 +536,7 @@ const Composer: FC = () => { // Focus the input so user can continue typing input.focus(); } - + // Reset mention query setMentionQuery(""); }; @@ -558,7 +574,11 @@ const Composer: FC = () => { ref={inputRef} onKeyUp={handleKeyUp} onKeyDown={handleKeyDown} - placeholder={mentionedDocuments.length > 0 ? "Ask about these documents..." : "Ask SurfSense (type @ to mention docs)"} + placeholder={ + mentionedDocuments.length > 0 + ? "Ask about these documents..." + : "Ask SurfSense (type @ to mention docs)" + } className="aui-composer-input flex-1 min-w-[120px] max-h-32 resize-none bg-transparent text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0 py-1" rows={1} autoFocus @@ -567,41 +587,47 @@ const Composer: FC = () => {
{/* -------- Document mention popover (rendered via portal) -------- */} - {showDocumentPopover && typeof document !== "undefined" && createPortal( - <> - {/* Backdrop */} -
- , - document.body - )} + {/* Popover positioned above input */} +
+ { + setShowDocumentPopover(false); + setMentionQuery(""); + }} + initialSelectedDocuments={mentionedDocuments} + externalSearch={mentionQuery} + /> +
+ , + document.body + )} diff --git a/surfsense_web/components/new-chat/DocumentsDataTable.tsx b/surfsense_web/components/new-chat/DocumentsDataTable.tsx index 2c1ccf1cb..49aeff75c 100644 --- a/surfsense_web/components/new-chat/DocumentsDataTable.tsx +++ b/surfsense_web/components/new-chat/DocumentsDataTable.tsx @@ -2,7 +2,15 @@ import { useQuery } from "@tanstack/react-query"; import { FileText } from "lucide-react"; -import { forwardRef, useCallback, useEffect, useImperativeHandle, useMemo, useRef, useState } from "react"; +import { + forwardRef, + useCallback, + useEffect, + useImperativeHandle, + useMemo, + useRef, + useState, +} from "react"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; import type { Document } from "@/contracts/types/document.types"; import { documentsApiService } from "@/lib/apis/documents-api.service"; @@ -33,13 +41,16 @@ function useDebounced(value: T, delay = 300) { } export const DocumentsDataTable = forwardRef( - function DocumentsDataTable({ - searchSpaceId, - onSelectionChange, - onDone, - initialSelectedDocuments = [], - externalSearch = "", - }, ref) { + function DocumentsDataTable( + { + searchSpaceId, + onSelectionChange, + onDone, + initialSelectedDocuments = [], + externalSearch = "", + }, + ref + ) { // Use external search const search = externalSearch; const debouncedSearch = useDebounced(search, 150); @@ -97,10 +108,13 @@ export const DocumentsDataTable = forwardRef { - onSelectionChange([...initialSelectedDocuments, doc]); - onDone(); - }, [initialSelectedDocuments, onSelectionChange, onDone]); + const handleSelectDocument = useCallback( + (doc: Document) => { + onSelectionChange([...initialSelectedDocuments, doc]); + onDone(); + }, + [initialSelectedDocuments, onSelectionChange, onDone] + ); // Scroll highlighted item into view useEffect(() => { @@ -120,56 +134,55 @@ export const DocumentsDataTable = forwardRef ({ - selectHighlighted: () => { - if (selectableDocuments[highlightedIndex]) { - handleSelectDocument(selectableDocuments[highlightedIndex]); - } - }, - moveUp: () => { - setHighlightedIndex((prev) => - prev > 0 ? prev - 1 : selectableDocuments.length - 1 - ); - }, - moveDown: () => { - setHighlightedIndex((prev) => - prev < selectableDocuments.length - 1 ? prev + 1 : 0 - ); - }, - }), [selectableDocuments, highlightedIndex, handleSelectDocument]); - - // Handle keyboard navigation - const handleKeyDown = useCallback((e: React.KeyboardEvent) => { - if (selectableDocuments.length === 0) return; - - switch (e.key) { - case "ArrowDown": - e.preventDefault(); - setHighlightedIndex((prev) => - prev < selectableDocuments.length - 1 ? prev + 1 : 0 - ); - break; - case "ArrowUp": - e.preventDefault(); - setHighlightedIndex((prev) => - prev > 0 ? prev - 1 : selectableDocuments.length - 1 - ); - break; - case "Enter": - e.preventDefault(); + useImperativeHandle( + ref, + () => ({ + selectHighlighted: () => { if (selectableDocuments[highlightedIndex]) { handleSelectDocument(selectableDocuments[highlightedIndex]); } - break; - case "Escape": - e.preventDefault(); - onDone(); - break; - } - }, [selectableDocuments, highlightedIndex, handleSelectDocument, onDone]); + }, + moveUp: () => { + setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableDocuments.length - 1)); + }, + moveDown: () => { + setHighlightedIndex((prev) => (prev < selectableDocuments.length - 1 ? prev + 1 : 0)); + }, + }), + [selectableDocuments, highlightedIndex, handleSelectDocument] + ); + + // Handle keyboard navigation + const handleKeyDown = useCallback( + (e: React.KeyboardEvent) => { + if (selectableDocuments.length === 0) return; + + switch (e.key) { + case "ArrowDown": + e.preventDefault(); + setHighlightedIndex((prev) => (prev < selectableDocuments.length - 1 ? prev + 1 : 0)); + break; + case "ArrowUp": + e.preventDefault(); + setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableDocuments.length - 1)); + break; + case "Enter": + e.preventDefault(); + if (selectableDocuments[highlightedIndex]) { + handleSelectDocument(selectableDocuments[highlightedIndex]); + } + break; + case "Escape": + e.preventDefault(); + onDone(); + break; + } + }, + [selectableDocuments, highlightedIndex, handleSelectDocument, onDone] + ); return ( -
d.id === doc.id); const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex; - + return (