feat: add SurfSense docs mention support in chat

This commit is contained in:
CREDO23 2026-01-13 06:14:58 +02:00
parent 1b5f29afcc
commit cd3677b5fa
11 changed files with 226 additions and 141 deletions

View file

@ -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<Document, "id" | "title" | "document_type">[]) => {
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]

View file

@ -25,7 +25,7 @@ export interface InlineMentionEditorRef {
clear: () => void;
getText: () => string;
getMentionedDocuments: () => MentionedDocument[];
insertDocumentChip: (doc: Document) => void;
insertDocumentChip: (doc: Pick<Document, "id" | "title" | "document_type">) => 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<InlineMentionEditorRef, InlineMentionEditorProps>(
(
{
@ -84,15 +92,15 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
) => {
const editorRef = useRef<HTMLDivElement>(null);
const [isEmpty, setIsEmpty] = useState(true);
const [mentionedDocs, setMentionedDocs] = useState<Map<number, MentionedDocument>>(
() => new Map(initialDocuments.map((d) => [d.id, d]))
const [mentionedDocs, setMentionedDocs] = useState<Map<string, MentionedDocument>>(
() => 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<InlineMentionEditorRef, InlineMent
const chip = document.createElement("span");
chip.setAttribute(CHIP_DATA_ATTR, "true");
chip.setAttribute(CHIP_ID_ATTR, String(doc.id));
chip.setAttribute(CHIP_DOCTYPE_ATTR, doc.document_type ?? "UNKNOWN");
chip.contentEditable = "false";
chip.className =
"inline-flex items-center gap-0.5 mx-0.5 pl-1 pr-0.5 py-0.5 rounded bg-primary/10 text-xs font-bold text-primary border border-primary/10 select-none";
@ -175,13 +184,14 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
e.preventDefault();
e.stopPropagation();
chip.remove();
const docKey = `${doc.document_type ?? "UNKNOWN"}:${doc.id}`;
setMentionedDocs((prev) => {
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<InlineMentionEditorRef, InlineMent
// Insert a document chip at the current cursor position
const insertDocumentChip = useCallback(
(doc: Document) => {
(doc: Pick<Document, "id" | "title" | "document_type">) => {
if (!editorRef.current) return;
// Validate required fields for type safety
@ -210,8 +220,9 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
document_type: doc.document_type,
};
// Add to mentioned docs map
setMentionedDocs((prev) => 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<InlineMentionEditorRef, InlineMent
if (isChipElement(prevSibling)) {
e.preventDefault();
const chipId = getChipId(prevSibling);
const chipDocType = getChipDocType(prevSibling);
if (chipId !== null) {
prevSibling.remove();
const chipKey = `${chipDocType}:${chipId}`;
setMentionedDocs((prev) => {
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<InlineMentionEditorRef, InlineMent
if (isChipElement(prevChild)) {
e.preventDefault();
const chipId = getChipId(prevChild);
const chipDocType = getChipDocType(prevChild);
if (chipId !== null) {
prevChild.remove();
const chipKey = `${chipDocType}:${chipId}`;
setMentionedDocs((prev) => {
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);
}
}
}

View file

@ -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<Document, "id" | "title" | "document_type">[]) => {
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) => (
<span
key={doc.id}
key={`${doc.document_type}:${doc.id}`}
className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full bg-primary/10 text-xs font-medium text-primary border border-primary/20"
title={doc.title}
>

View file

@ -29,7 +29,7 @@ export const UserMessage: FC = () => {
{/* Mentioned documents as chips */}
{mentionedDocs?.map((doc) => (
<span
key={doc.id}
key={`${doc.document_type}:${doc.id}`}
className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full bg-primary/10 text-xs font-medium text-primary border border-primary/20"
title={doc.title}
>

View file

@ -25,9 +25,9 @@ export interface DocumentMentionPickerRef {
interface DocumentMentionPickerProps {
searchSpaceId: number;
onSelectionChange: (documents: Document[]) => void;
onSelectionChange: (documents: Pick<Document, "id" | "title" | "document_type">[]) => void;
onDone: () => void;
initialSelectedDocuments?: Document[];
initialSelectedDocuments?: Pick<Document, "id" | "title" | "document_type">[];
externalSearch?: string;
}
@ -57,7 +57,7 @@ export const DocumentMentionPicker = forwardRef<
const scrollContainerRef = useRef<HTMLDivElement>(null);
// State for pagination
const [accumulatedDocuments, setAccumulatedDocuments] = useState<Document[]>([]);
const [accumulatedDocuments, setAccumulatedDocuments] = useState<Pick<Document, "id" | "title" | "document_type">[]>([]);
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<Document, "id" | "title" | "document_type">[] = [];
// 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<Document, "id" | "title" | "document_type">) => {
onSelectionChange([...initialSelectedDocuments, doc]);
onDone();
},
@ -287,13 +317,16 @@ export const DocumentMentionPicker = forwardRef<
) : (
<div className="py-1">
{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 (
<button
key={doc.id}
key={docKey}
ref={(el) => {
if (el && selectableIndex >= 0) {
itemRefs.current.set(selectableIndex, el);