diff --git a/surfsense_web/components/assistant-ui/chat-viewport.tsx b/surfsense_web/components/assistant-ui/chat-viewport.tsx index f7f1ac188..c0684407e 100644 --- a/surfsense_web/components/assistant-ui/chat-viewport.tsx +++ b/surfsense_web/components/assistant-ui/chat-viewport.tsx @@ -39,7 +39,7 @@ export const ChatViewport: FC = ({ children, footer }) => ( {children} {footer ? (
diff --git a/surfsense_web/components/assistant-ui/user-message.tsx b/surfsense_web/components/assistant-ui/user-message.tsx index fb7212119..145ac2d7e 100644 --- a/surfsense_web/components/assistant-ui/user-message.tsx +++ b/surfsense_web/components/assistant-ui/user-message.tsx @@ -1,4 +1,10 @@ -import { ActionBarPrimitive, AuiIf, MessagePrimitive, useAuiState } from "@assistant-ui/react"; +import { + ActionBarPrimitive, + AuiIf, + MessagePrimitive, + useAuiState, + useMessagePartText, +} from "@assistant-ui/react"; import { useAtomValue } from "jotai"; import { CheckIcon, CopyIcon, Pencil } from "lucide-react"; import Image from "next/image"; @@ -7,6 +13,8 @@ import { currentThreadAtom } from "@/atoms/chat/current-thread.atom"; import { messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; +import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; +import { parseMentionSegments } from "@/lib/chat/parse-mention-segments"; interface AuthorMetadata { displayName: string | null; @@ -47,23 +55,40 @@ const UserAvatar: FC = ({ displayName, avatarUrl }) => { ); }; -export const UserMessage: FC = () => { +const UserTextPart: FC = () => { const messageId = useAuiState(({ message }) => message?.id); - const messageText = useAuiState(({ message }) => - (message?.content ?? []) - .map((part) => - typeof part === "object" && - part !== null && - "type" in part && - (part as { type?: string }).type === "text" && - "text" in part - ? String((part as { text?: string }).text ?? "") - : "" - ) - .join("") - ); + const part = useMessagePartText(); + const text = (part as { text?: string }).text ?? ""; const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom); - const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined; + const mentionedDocs = (messageId ? messageDocumentsMap[messageId] : undefined) ?? []; + + const segments = parseMentionSegments(text, mentionedDocs); + + return ( +

+ {segments.map((segment) => + segment.type === "text" ? ( + {segment.value} + ) : ( + + + {getConnectorIcon(segment.doc.document_type ?? "UNKNOWN", "h-3 w-3")} + + {segment.doc.title} + + ) + )} +

+ ); +}; + +const userMessageParts = { Text: UserTextPart }; + +export const UserMessage: FC = () => { const metadata = useAuiState(({ message }) => message?.metadata); const author = metadata?.custom?.author as AuthorMetadata | undefined; const isSharedChat = useAtomValue(currentThreadAtom).visibility === "SEARCH_SPACE"; @@ -78,11 +103,7 @@ export const UserMessage: FC = () => {
- {mentionedDocs && mentionedDocs.length > 0 ? ( - - ) : ( - - )} +
@@ -99,64 +120,6 @@ export const UserMessage: FC = () => { ); }; -const UserMessageWithMentionChips: FC<{ - text: string; - mentionedDocs: { id: number; title: string; document_type: string }[]; -}> = ({ text, mentionedDocs }) => { - type Segment = - | { type: "text"; value: string; start: number } - | { type: "mention"; doc: { id: number; title: string; document_type: string }; start: number }; - - const tokens = mentionedDocs - .map((doc) => ({ doc, token: `@${doc.title}` })) - .sort((a, b) => b.token.length - a.token.length); - - const segments: Segment[] = []; - let i = 0; - let buffer = ""; - let bufferStart = 0; - while (i < text.length) { - const tokenMatch = tokens.find(({ token }) => text.startsWith(token, i)); - if (tokenMatch) { - if (buffer) { - segments.push({ type: "text", value: buffer, start: bufferStart }); - buffer = ""; - } - segments.push({ type: "mention", doc: tokenMatch.doc, start: i }); - i += tokenMatch.token.length; - bufferStart = i; - continue; - } - if (!buffer) bufferStart = i; - buffer += text[i]; - i += 1; - } - if (buffer) { - segments.push({ type: "text", value: buffer, start: bufferStart }); - } - - return ( - - {segments.map((segment) => - segment.type === "text" ? ( - {segment.value} - ) : ( - - - {getConnectorIcon(segment.doc.document_type ?? "UNKNOWN", "h-3 w-3")} - - {segment.doc.title} - - ) - )} - - ); -}; - const UserActionBar: FC = () => { const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); diff --git a/surfsense_web/lib/chat/parse-mention-segments.ts b/surfsense_web/lib/chat/parse-mention-segments.ts new file mode 100644 index 000000000..b9cf59792 --- /dev/null +++ b/surfsense_web/lib/chat/parse-mention-segments.ts @@ -0,0 +1,54 @@ +import type { MentionedDocumentInfo } from "@/atoms/chat/mentioned-documents.atom"; + +export type MentionSegment = + | { type: "text"; value: string; start: number } + | { type: "mention"; doc: MentionedDocumentInfo; start: number }; + +/** + * Tokenizes a user message into text and `@mention` segments. + * + * Pure: no React, no DOM, no side effects. Safe to unit-test and reuse. + * + * Mentions are matched greedily by longest title first so that a longer title + * (e.g. `@Project Roadmap`) is never shadowed by a shorter prefix + * (e.g. `@Project`). + */ +export function parseMentionSegments( + text: string, + docs: ReadonlyArray +): MentionSegment[] { + if (text.length === 0) return []; + if (docs.length === 0) return [{ type: "text", value: text, start: 0 }]; + + const tokens = docs + .map((doc) => ({ doc, token: `@${doc.title}` })) + .sort((a, b) => b.token.length - a.token.length); + + const segments: MentionSegment[] = []; + let i = 0; + let buffer = ""; + let bufferStart = 0; + + while (i < text.length) { + const tokenMatch = tokens.find(({ token }) => text.startsWith(token, i)); + if (tokenMatch) { + if (buffer) { + segments.push({ type: "text", value: buffer, start: bufferStart }); + buffer = ""; + } + segments.push({ type: "mention", doc: tokenMatch.doc, start: i }); + i += tokenMatch.token.length; + bufferStart = i; + continue; + } + if (!buffer) bufferStart = i; + buffer += text[i]; + i += 1; + } + + if (buffer) { + segments.push({ type: "text", value: buffer, start: bufferStart }); + } + + return segments; +}