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 ( - - ); - })} -
- )} - -
- ); -} + ); + } +);