diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index 25cec6959..45899c2ef 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -76,17 +76,42 @@ def format_mentioned_surfsense_docs_as_context( if not documents: return "" + import json + context_parts = [""] context_parts.append( "The user has explicitly mentioned the following SurfSense documentation pages. " "These are official documentation about how to use SurfSense and should be used to answer questions about the application." ) - for i, doc in enumerate(documents, 1): - context_parts.append( - f"" - ) - context_parts.append(f"") - context_parts.append("") + + for doc in documents: + metadata_json = json.dumps({"source": doc.source}, ensure_ascii=False) + + context_parts.append("") + context_parts.append("") + context_parts.append(f" doc-{doc.id}") + context_parts.append(" SURFSENSE_DOCS") + context_parts.append(f" <![CDATA[{doc.title}]]>") + context_parts.append(f" ") + context_parts.append(f" ") + context_parts.append("") + context_parts.append("") + context_parts.append("") + + if hasattr(doc, 'chunks') and doc.chunks: + for chunk in doc.chunks: + context_parts.append( + f" " + ) + else: + context_parts.append( + f" " + ) + + context_parts.append("") + context_parts.append("") + context_parts.append("") + context_parts.append("") return "\n".join(context_parts) @@ -236,8 +261,11 @@ async def stream_new_chat( # Fetch mentioned SurfSense docs if any mentioned_surfsense_docs: list[SurfsenseDocsDocument] = [] if mentioned_surfsense_doc_ids: + from sqlalchemy.orm import selectinload result = await session.execute( - select(SurfsenseDocsDocument).filter( + select(SurfsenseDocsDocument) + .options(selectinload(SurfsenseDocsDocument.chunks)) + .filter( SurfsenseDocsDocument.id.in_(mentioned_surfsense_doc_ids), ) ) 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 b1abd647f..489e17b20 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 @@ -265,7 +265,10 @@ export default function NewChatPage() { setMessages([]); setThreadId(null); setMessageThinkingSteps(new Map()); - setMentionedDocumentIds([]); + setMentionedDocumentIds({ + surfsense_doc_ids: [], + document_ids: [], + }); setMentionedDocuments([]); setMessageDocumentsMap({}); clearPlanOwnerRegistry(); // Reset plan ownership for new chat @@ -429,7 +432,7 @@ export default function NewChatPage() { // Track message sent trackChatMessageSent(searchSpaceId, currentThreadId, { hasAttachments: messageAttachments.length > 0, - hasMentionedDocuments: mentionedDocumentIds.length > 0, + hasMentionedDocuments: mentionedDocumentIds.surfsense_doc_ids.length > 0 || mentionedDocumentIds.document_ids.length > 0, messageLength: userQuery.length, }); @@ -627,12 +630,16 @@ 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; + // Get mentioned document IDs for context (separate fields for backend) + const hasDocumentIds = mentionedDocumentIds.document_ids.length > 0; + const hasSurfsenseDocIds = mentionedDocumentIds.surfsense_doc_ids.length > 0; // Clear mentioned documents after capturing them - if (mentionedDocumentIds.length > 0) { - setMentionedDocumentIds([]); + if (hasDocumentIds || hasSurfsenseDocIds) { + setMentionedDocumentIds({ + surfsense_doc_ids: [], + document_ids: [], + }); setMentionedDocuments([]); } @@ -648,7 +655,8 @@ export default function NewChatPage() { search_space_id: searchSpaceId, messages: messageHistory, attachments: attachments.length > 0 ? attachments : undefined, - mentioned_document_ids: documentIds, + mentioned_document_ids: hasDocumentIds ? mentionedDocumentIds.document_ids : undefined, + mentioned_surfsense_doc_ids: hasSurfsenseDocIds ? mentionedDocumentIds.surfsense_doc_ids : undefined, }), signal: controller.signal, }); diff --git a/surfsense_web/atoms/chat/mentioned-documents.atom.ts b/surfsense_web/atoms/chat/mentioned-documents.atom.ts index 79ea27d12..17ae38616 100644 --- a/surfsense_web/atoms/chat/mentioned-documents.atom.ts +++ b/surfsense_web/atoms/chat/mentioned-documents.atom.ts @@ -1,19 +1,25 @@ "use client"; import { atom } from "jotai"; -import type { Document } from "@/contracts/types/document.types"; +import type { Document, SurfsenseDocsDocument } from "@/contracts/types/document.types"; /** * 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([]); +export const mentionedDocumentIdsAtom = atom<{ + surfsense_doc_ids: number[]; + document_ids: number[]; +}>({ + surfsense_doc_ids: [], + document_ids: [], +}); /** * Atom to store the full document objects mentioned in the current chat composer. * This persists across component remounts. */ -export const mentionedDocumentsAtom = atom([]); +export const mentionedDocumentsAtom = atom<(Pick)[]>([]); /** * Simplified document info for display purposes diff --git a/surfsense_web/atoms/chat/mentioned-surfsense-docs.atom.ts b/surfsense_web/atoms/chat/mentioned-surfsense-docs.atom.ts deleted file mode 100644 index aa4a84e05..000000000 --- a/surfsense_web/atoms/chat/mentioned-surfsense-docs.atom.ts +++ /dev/null @@ -1,32 +0,0 @@ -"use client"; - -import { atom } from "jotai"; -import type { SurfsenseDocsDocument } from "@/contracts/types/document.types"; - -/** - * Atom to store the IDs of SurfSense docs mentioned in the current chat composer. - * This is used to pass documentation context to the backend when sending a message. - */ -export const mentionedSurfsenseDocIdsAtom = atom([]); - -/** - * Atom to store the full SurfSense doc objects mentioned in the current chat composer. - * This persists across component remounts. - */ -export const mentionedSurfsenseDocsAtom = atom([]); - -/** - * Simplified SurfSense doc info for display purposes - */ -export interface MentionedSurfsenseDocInfo { - id: number; - title: string; - source: string; -} - -/** - * Atom to store mentioned SurfSense docs per message ID. - * This allows displaying which docs were mentioned with each user message. - */ -export const messageSurfsenseDocsMapAtom = atom>({}); - diff --git a/surfsense_web/components/assistant-ui/composer.tsx b/surfsense_web/components/assistant-ui/composer.tsx index 8f8ee5e0b..417f7c70f 100644 --- a/surfsense_web/components/assistant-ui/composer.tsx +++ b/surfsense_web/components/assistant-ui/composer.tsx @@ -53,7 +53,10 @@ export const Composer: FC = () => { // Sync mentioned document IDs to atom for use in chat request useEffect(() => { - setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id)); + setMentionedDocumentIds({ + surfsense_doc_ids: mentionedDocuments.filter((doc) => doc.document_type === "SURFSENSE_DOCS").map((doc) => doc.id), + document_ids: mentionedDocuments.filter((doc) => doc.document_type !== "SURFSENSE_DOCS").map((doc) => doc.id), + }); }, [mentionedDocuments, setMentionedDocumentIds]); // Handle text change from inline editor - sync with assistant-ui composer @@ -119,7 +122,10 @@ export const Composer: FC = () => { // Clear the editor after sending editorRef.current?.clear(); setMentionedDocuments([]); - setMentionedDocumentIds([]); + setMentionedDocumentIds({ + surfsense_doc_ids: [], + document_ids: [], + }); } }, [ showDocumentPopover, @@ -129,41 +135,48 @@ export const Composer: FC = () => { setMentionedDocumentIds, ]); - // Handle document removal from inline editor const handleDocumentRemove = useCallback( - (docId: number) => { + (docId: number, docType?: string) => { setMentionedDocuments((prev) => { - const updated = prev.filter((doc) => doc.id !== docId); - // Immediately sync document IDs to avoid race conditions - setMentionedDocumentIds(updated.map((doc) => doc.id)); + const updated = prev.filter( + (doc) => !(doc.id === docId && doc.document_type === docType) + ); + setMentionedDocumentIds({ + surfsense_doc_ids: updated.filter((doc) => doc.document_type === "SURFSENSE_DOCS").map((doc) => doc.id), + document_ids: updated.filter((doc) => doc.document_type !== "SURFSENSE_DOCS").map((doc) => doc.id), + }); return updated; }); }, [setMentionedDocuments, setMentionedDocumentIds] ); - // Handle document selection from picker const handleDocumentsMention = useCallback( - (documents: Document[]) => { - // Insert chips into the inline editor for each new document - const existingIds = new Set(mentionedDocuments.map((d) => d.id)); - const newDocs = documents.filter((doc) => !existingIds.has(doc.id)); + (documents: Pick[]) => { + const existingKeys = new Set( + mentionedDocuments.map((d) => `${d.document_type}:${d.id}`) + ); + const newDocs = documents.filter( + (doc) => !existingKeys.has(`${doc.document_type}:${doc.id}`) + ); for (const doc of newDocs) { editorRef.current?.insertDocumentChip(doc); } - // Update mentioned documents state setMentionedDocuments((prev) => { - const existingIdSet = new Set(prev.map((d) => d.id)); - const uniqueNewDocs = documents.filter((doc) => !existingIdSet.has(doc.id)); + const existingKeySet = new Set(prev.map((d) => `${d.document_type}:${d.id}`)); + const uniqueNewDocs = documents.filter( + (doc) => !existingKeySet.has(`${doc.document_type}:${doc.id}`) + ); const updated = [...prev, ...uniqueNewDocs]; - // Immediately sync document IDs to avoid race conditions - setMentionedDocumentIds(updated.map((doc) => doc.id)); + setMentionedDocumentIds({ + surfsense_doc_ids: updated.filter((doc) => doc.document_type === "SURFSENSE_DOCS").map((doc) => doc.id), + document_ids: updated.filter((doc) => doc.document_type !== "SURFSENSE_DOCS").map((doc) => doc.id), + }); return updated; }); - // Reset mention query but keep popover open for more selections setMentionQuery(""); }, [mentionedDocuments, setMentionedDocuments, setMentionedDocumentIds] diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx index 12a8f895f..4fa847a95 100644 --- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx +++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx @@ -25,7 +25,7 @@ export interface InlineMentionEditorRef { clear: () => void; getText: () => string; getMentionedDocuments: () => MentionedDocument[]; - insertDocumentChip: (doc: Document) => void; + insertDocumentChip: (doc: Pick) => void; } interface InlineMentionEditorProps { @@ -34,7 +34,7 @@ interface InlineMentionEditorProps { onMentionClose?: () => void; onSubmit?: () => void; onChange?: (text: string, docs: MentionedDocument[]) => void; - onDocumentRemove?: (docId: number) => void; + onDocumentRemove?: (docId: number, docType?: string) => void; onKeyDown?: (e: React.KeyboardEvent) => void; disabled?: boolean; className?: string; @@ -44,6 +44,7 @@ interface InlineMentionEditorProps { // Unique data attribute to identify chip elements const CHIP_DATA_ATTR = "data-mention-chip"; const CHIP_ID_ATTR = "data-mention-id"; +const CHIP_DOCTYPE_ATTR = "data-mention-doctype"; /** * Type guard to check if a node is a chip element @@ -66,6 +67,13 @@ function getChipId(element: Element): number | null { return Number.isNaN(id) ? null : id; } +/** + * Get chip document type from element attribute + */ +function getChipDocType(element: Element): string { + return element.getAttribute(CHIP_DOCTYPE_ATTR) ?? "UNKNOWN"; +} + export const InlineMentionEditor = forwardRef( ( { @@ -84,15 +92,15 @@ export const InlineMentionEditor = forwardRef { const editorRef = useRef(null); const [isEmpty, setIsEmpty] = useState(true); - const [mentionedDocs, setMentionedDocs] = useState>( - () => new Map(initialDocuments.map((d) => [d.id, d])) + const [mentionedDocs, setMentionedDocs] = useState>( + () => new Map(initialDocuments.map((d) => [`${d.document_type ?? "UNKNOWN"}:${d.id}`, d])) ); const isComposingRef = useRef(false); // Sync initial documents useEffect(() => { if (initialDocuments.length > 0) { - setMentionedDocs(new Map(initialDocuments.map((d) => [d.id, d]))); + setMentionedDocs(new Map(initialDocuments.map((d) => [`${d.document_type ?? "UNKNOWN"}:${d.id}`, d]))); } }, [initialDocuments]); @@ -153,6 +161,7 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); - next.delete(doc.id); + next.delete(docKey); return next; }); // Notify parent that a document was removed - onDocumentRemove?.(doc.id); + onDocumentRemove?.(doc.id, doc.document_type); focusAtEnd(); }; @@ -195,7 +205,7 @@ export const InlineMentionEditor = forwardRef { + (doc: Pick) => { if (!editorRef.current) return; // Validate required fields for type safety @@ -210,8 +220,9 @@ export const InlineMentionEditor = forwardRef new Map(prev).set(doc.id, mentionDoc)); + // Add to mentioned docs map using unique key + const docKey = `${doc.document_type ?? "UNKNOWN"}:${doc.id}`; + setMentionedDocs((prev) => new Map(prev).set(docKey, mentionDoc)); // Find and remove the @query text const selection = window.getSelection(); @@ -413,15 +424,17 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); - next.delete(chipId); + next.delete(chipKey); return next; }); // Notify parent that a document was removed - onDocumentRemove?.(chipId); + onDocumentRemove?.(chipId, chipDocType); } return; } @@ -448,15 +461,17 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); - next.delete(chipId); + next.delete(chipKey); return next; }); // Notify parent that a document was removed - onDocumentRemove?.(chipId); + onDocumentRemove?.(chipId, chipDocType); } } } diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index 6dc13fddf..92b8ad786 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -229,7 +229,10 @@ const Composer: FC = () => { // Sync mentioned document IDs to atom for use in chat request useEffect(() => { - setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id)); + setMentionedDocumentIds({ + surfsense_doc_ids: mentionedDocuments.filter((doc) => doc.document_type === "SURFSENSE_DOCS").map((doc) => doc.id), + document_ids: mentionedDocuments.filter((doc) => doc.document_type !== "SURFSENSE_DOCS").map((doc) => doc.id), + }); }, [mentionedDocuments, setMentionedDocumentIds]); // Handle text change from inline editor - sync with assistant-ui composer @@ -295,7 +298,10 @@ const Composer: FC = () => { // Clear the editor after sending editorRef.current?.clear(); setMentionedDocuments([]); - setMentionedDocumentIds([]); + setMentionedDocumentIds({ + surfsense_doc_ids: [], + document_ids: [], + }); } }, [ showDocumentPopover, @@ -305,41 +311,48 @@ const Composer: FC = () => { setMentionedDocumentIds, ]); - // Handle document removal from inline editor const handleDocumentRemove = useCallback( - (docId: number) => { + (docId: number, docType?: string) => { setMentionedDocuments((prev) => { - const updated = prev.filter((doc) => doc.id !== docId); - // Immediately sync document IDs to avoid race conditions - setMentionedDocumentIds(updated.map((doc) => doc.id)); + const updated = prev.filter( + (doc) => !(doc.id === docId && doc.document_type === docType) + ); + setMentionedDocumentIds({ + surfsense_doc_ids: updated.filter((doc) => doc.document_type === "SURFSENSE_DOCS").map((doc) => doc.id), + document_ids: updated.filter((doc) => doc.document_type !== "SURFSENSE_DOCS").map((doc) => doc.id), + }); return updated; }); }, [setMentionedDocuments, setMentionedDocumentIds] ); - // Handle document selection from picker const handleDocumentsMention = useCallback( - (documents: Document[]) => { - // Insert chips into the inline editor for each new document - const existingIds = new Set(mentionedDocuments.map((d) => d.id)); - const newDocs = documents.filter((doc) => !existingIds.has(doc.id)); + (documents: Pick[]) => { + const existingKeys = new Set( + mentionedDocuments.map((d) => `${d.document_type}:${d.id}`) + ); + const newDocs = documents.filter( + (doc) => !existingKeys.has(`${doc.document_type}:${doc.id}`) + ); for (const doc of newDocs) { editorRef.current?.insertDocumentChip(doc); } - // Update mentioned documents state setMentionedDocuments((prev) => { - const existingIdSet = new Set(prev.map((d) => d.id)); - const uniqueNewDocs = documents.filter((doc) => !existingIdSet.has(doc.id)); + const existingKeySet = new Set(prev.map((d) => `${d.document_type}:${d.id}`)); + const uniqueNewDocs = documents.filter( + (doc) => !existingKeySet.has(`${doc.document_type}:${doc.id}`) + ); const updated = [...prev, ...uniqueNewDocs]; - // Immediately sync document IDs to avoid race conditions - setMentionedDocumentIds(updated.map((doc) => doc.id)); + setMentionedDocumentIds({ + surfsense_doc_ids: updated.filter((doc) => doc.document_type === "SURFSENSE_DOCS").map((doc) => doc.id), + document_ids: updated.filter((doc) => doc.document_type !== "SURFSENSE_DOCS").map((doc) => doc.id), + }); return updated; }); - // Reset mention query but keep popover open for more selections setMentionQuery(""); }, [mentionedDocuments, setMentionedDocuments, setMentionedDocumentIds] @@ -640,7 +653,7 @@ const UserMessage: FC = () => { {/* Mentioned documents as chips */} {mentionedDocs?.map((doc) => ( diff --git a/surfsense_web/components/assistant-ui/user-message.tsx b/surfsense_web/components/assistant-ui/user-message.tsx index dcf626461..745542304 100644 --- a/surfsense_web/components/assistant-ui/user-message.tsx +++ b/surfsense_web/components/assistant-ui/user-message.tsx @@ -29,7 +29,7 @@ export const UserMessage: FC = () => { {/* Mentioned documents as chips */} {mentionedDocs?.map((doc) => ( diff --git a/surfsense_web/components/new-chat/document-mention-picker.tsx b/surfsense_web/components/new-chat/document-mention-picker.tsx index 7a9e7aaa5..90515d52d 100644 --- a/surfsense_web/components/new-chat/document-mention-picker.tsx +++ b/surfsense_web/components/new-chat/document-mention-picker.tsx @@ -25,9 +25,9 @@ export interface DocumentMentionPickerRef { interface DocumentMentionPickerProps { searchSpaceId: number; - onSelectionChange: (documents: Document[]) => void; + onSelectionChange: (documents: Pick[]) => void; onDone: () => void; - initialSelectedDocuments?: Document[]; + initialSelectedDocuments?: Pick[]; externalSearch?: string; } @@ -57,7 +57,7 @@ export const DocumentMentionPicker = forwardRef< const scrollContainerRef = useRef(null); // State for pagination - const [accumulatedDocuments, setAccumulatedDocuments] = useState([]); + const [accumulatedDocuments, setAccumulatedDocuments] = useState[]>([]); const [currentPage, setCurrentPage] = useState(0); const [hasMore, setHasMore] = useState(false); const [isLoadingMore, setIsLoadingMore] = useState(false); @@ -90,6 +90,13 @@ export const DocumentMentionPicker = forwardRef< }; }, [debouncedSearch, searchSpaceId]); + const surfsenseDocsQueryParams = useMemo(() => { + return { + page: 0, + page_size: PAGE_SIZE, + }; + }, []); + // Use query for fetching first page of documents const { data: documents, isLoading: isDocumentsLoading } = useQuery({ queryKey: cacheKeys.documents.withQueryParams(fetchQueryParams), @@ -106,22 +113,45 @@ export const DocumentMentionPicker = forwardRef< enabled: !!searchSpaceId && !!debouncedSearch.trim() && currentPage === 0, }); - // Update accumulated documents when first page loads + // Use query for fetching first page of SurfSense docs + const { data: surfsenseDocs, isLoading: isSurfsenseDocsLoading } = useQuery({ + queryKey: ["surfsense-docs-mention", surfsenseDocsQueryParams], + queryFn: () => documentsApiService.getSurfsenseDocs({ queryParams: surfsenseDocsQueryParams }), + staleTime: 3 * 60 * 1000, + }); + + // Update accumulated documents when first page loads - combine both sources useEffect(() => { if (currentPage === 0) { + const combinedDocs: Pick[] = []; + + // Add SurfSense docs first (they appear at top) + if (surfsenseDocs?.items) { + for (const doc of surfsenseDocs.items) { + combinedDocs.push({ + id: doc.id, + title: doc.title, + document_type: "SURFSENSE_DOCS", + }); + } + } + + // Add regular documents if (debouncedSearch.trim()) { - if (searchedDocuments) { - setAccumulatedDocuments(searchedDocuments.items); + if (searchedDocuments?.items) { + combinedDocs.push(...searchedDocuments.items); setHasMore(searchedDocuments.has_more); } } else { - if (documents) { - setAccumulatedDocuments(documents.items); + if (documents?.items) { + combinedDocs.push(...documents.items); setHasMore(documents.has_more); } } + + setAccumulatedDocuments(combinedDocs); } - }, [documents, searchedDocuments, debouncedSearch, currentPage]); + }, [documents, searchedDocuments, surfsenseDocs, debouncedSearch, currentPage]); // Function to load next page const loadNextPage = useCallback(async () => { @@ -175,22 +205,22 @@ export const DocumentMentionPicker = forwardRef< const actualDocuments = accumulatedDocuments; const actualLoading = - (debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading) && currentPage === 0; + ((debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading) || isSurfsenseDocsLoading) && currentPage === 0; - // Track already selected document IDs - const selectedIds = useMemo( - () => new Set(initialSelectedDocuments.map((d) => d.id)), + // Track already selected documents using unique key (document_type:id) to avoid ID collisions + const selectedKeys = useMemo( + () => new Set(initialSelectedDocuments.map((d) => `${d.document_type}:${d.id}`)), [initialSelectedDocuments] ); // Filter out already selected documents for navigation const selectableDocuments = useMemo( - () => actualDocuments.filter((doc) => !selectedIds.has(doc.id)), - [actualDocuments, selectedIds] + () => actualDocuments.filter((doc) => !selectedKeys.has(`${doc.document_type}:${doc.id}`)), + [actualDocuments, selectedKeys] ); const handleSelectDocument = useCallback( - (doc: Document) => { + (doc: Pick) => { onSelectionChange([...initialSelectedDocuments, doc]); onDone(); }, @@ -287,13 +317,16 @@ export const DocumentMentionPicker = forwardRef< ) : (
{actualDocuments.map((doc) => { - const isAlreadySelected = selectedIds.has(doc.id); - const selectableIndex = selectableDocuments.findIndex((d) => d.id === doc.id); + const docKey = `${doc.document_type}:${doc.id}`; + const isAlreadySelected = selectedKeys.has(docKey); + const selectableIndex = selectableDocuments.findIndex( + (d) => d.document_type === doc.document_type && d.id === doc.id + ); const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex; return (