diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx new file mode 100644 index 000000000..62fbe0dd4 --- /dev/null +++ b/surfsense_web/components/assistant-ui/assistant-message.tsx @@ -0,0 +1,118 @@ +import { + ActionBarPrimitive, + AssistantIf, + ErrorPrimitive, + MessagePrimitive, + useAssistantState, +} from "@assistant-ui/react"; +import { CheckIcon, CopyIcon, DownloadIcon, RefreshCwIcon } from "lucide-react"; +import type { FC } from "react"; +import { useContext } from "react"; +import { MarkdownText } from "@/components/assistant-ui/markdown-text"; +import { ThinkingStepsContext, ThinkingStepsDisplay } from "@/components/assistant-ui/thinking-steps"; +import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; +import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { BranchPicker } from "@/components/assistant-ui/branch-picker"; + +export const MessageError: FC = () => { + return ( + + + + + + ); +}; + +/** + * Custom component to render thinking steps from Context + */ +const ThinkingStepsPart: FC = () => { + const thinkingStepsMap = useContext(ThinkingStepsContext); + + // Get the current message ID to look up thinking steps + const messageId = useAssistantState(({ message }) => message?.id); + const thinkingSteps = thinkingStepsMap.get(messageId) || []; + + // Check if this specific message is currently streaming + // A message is streaming if: thread is running AND this is the last assistant message + const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning); + const isLastMessage = useAssistantState(({ message }) => message?.isLast ?? false); + const isMessageStreaming = isThreadRunning && isLastMessage; + + if (thinkingSteps.length === 0) return null; + + return ( +
+ +
+ ); +}; + +const AssistantMessageInner: FC = () => { + return ( + <> + {/* Render thinking steps from message content - this ensures proper scroll tracking */} + + +
+ + +
+ +
+ + +
+ + ); +}; + +export const AssistantMessage: FC = () => { + return ( + + + + ); +}; + +const AssistantActionBar: FC = () => { + return ( + + + + message.isCopied}> + + + !message.isCopied}> + + + + + + + + + + + + + + + + ); +}; + diff --git a/surfsense_web/components/assistant-ui/branch-picker.tsx b/surfsense_web/components/assistant-ui/branch-picker.tsx new file mode 100644 index 000000000..1d9041309 --- /dev/null +++ b/surfsense_web/components/assistant-ui/branch-picker.tsx @@ -0,0 +1,33 @@ +import { BranchPickerPrimitive } from "@assistant-ui/react"; +import { ChevronLeftIcon, ChevronRightIcon } from "lucide-react"; +import type { FC } from "react"; +import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { cn } from "@/lib/utils"; + +export const BranchPicker: FC = ({ className, ...rest }) => { + return ( + + + + + + + + / + + + + + + + + ); +}; + diff --git a/surfsense_web/components/assistant-ui/composer-action.tsx b/surfsense_web/components/assistant-ui/composer-action.tsx new file mode 100644 index 000000000..9c5a95d88 --- /dev/null +++ b/surfsense_web/components/assistant-ui/composer-action.tsx @@ -0,0 +1,269 @@ +import { AssistantIf, ComposerPrimitive, useAssistantState } from "@assistant-ui/react"; +import { useAtomValue } from "jotai"; +import { AlertCircle, ArrowUpIcon, Loader2, Plus, Plug2, SquareIcon } from "lucide-react"; +import Link from "next/link"; +import type { FC } from "react"; +import { useCallback, useMemo, useRef, useState } from "react"; +import { getDocumentTypeLabel } from "@/app/dashboard/[search_space_id]/documents/(manage)/components/DocumentTypeIcon"; +import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms"; +import { + globalNewLLMConfigsAtom, + llmPreferencesAtom, + newLLMConfigsAtom, +} from "@/atoms/new-llm-config/new-llm-config-query.atoms"; +import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms"; +import { ComposerAddAttachment } from "@/components/assistant-ui/attachment"; +import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { Button } from "@/components/ui/button"; +import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; +import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; +import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors"; +import { cn } from "@/lib/utils"; +import { ChevronRightIcon } from "lucide-react"; + +const ConnectorIndicator: FC = () => { + const searchSpaceId = useAtomValue(activeSearchSpaceIdAtom); + const { connectors, isLoading: connectorsLoading } = useSearchSourceConnectors( + false, + searchSpaceId ? Number(searchSpaceId) : undefined + ); + const { data: documentTypeCounts, isLoading: documentTypesLoading } = + useAtomValue(documentTypeCountsAtom); + const [isOpen, setIsOpen] = useState(false); + const closeTimeoutRef = useRef(null); + + const isLoading = connectorsLoading || documentTypesLoading; + + // Get document types that have documents in the search space + const activeDocumentTypes = documentTypeCounts + ? Object.entries(documentTypeCounts).filter(([_, count]) => count > 0) + : []; + + const hasConnectors = connectors.length > 0; + const hasSources = hasConnectors || activeDocumentTypes.length > 0; + const totalSourceCount = connectors.length + activeDocumentTypes.length; + + const handleMouseEnter = useCallback(() => { + // Clear any pending close timeout + if (closeTimeoutRef.current) { + clearTimeout(closeTimeoutRef.current); + closeTimeoutRef.current = null; + } + setIsOpen(true); + }, []); + + const handleMouseLeave = useCallback(() => { + // Delay closing by 150ms for better UX + closeTimeoutRef.current = setTimeout(() => { + setIsOpen(false); + }, 150); + }, []); + + if (!searchSpaceId) return null; + + return ( + + + + + + {hasSources ? ( +
+
+

Connected Sources

+ + {totalSourceCount} + +
+
+ {/* Document types from the search space */} + {activeDocumentTypes.map(([docType]) => ( +
+ {getConnectorIcon(docType, "size-3.5")} + {getDocumentTypeLabel(docType)} +
+ ))} + {/* Search source connectors */} + {connectors.map((connector) => ( +
+ {getConnectorIcon(connector.connector_type, "size-3.5")} + {connector.name} +
+ ))} +
+
+ + + Add more sources + + +
+
+ ) : ( +
+

No sources yet

+

+ Add documents or connect data sources to enhance search results. +

+ + + Add Connector + +
+ )} + + + ); +}; + +export const ComposerAction: FC = () => { + // Check if any attachments are still being processed (running AND progress < 100) + // When progress is 100, processing is done but waiting for send() + const hasProcessingAttachments = useAssistantState(({ composer }) => + composer.attachments?.some((att) => { + const status = att.status; + if (status?.type !== "running") return false; + const progress = (status as { type: "running"; progress?: number }).progress; + return progress === undefined || progress < 100; + }) + ); + + // Check if composer text is empty + const isComposerEmpty = useAssistantState(({ composer }) => { + const text = composer.text?.trim() || ""; + return text.length === 0; + }); + + // Check if a model is configured + 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]); + + const isSendDisabled = hasProcessingAttachments || isComposerEmpty || !hasModelConfigured; + + return ( +
+
+ + +
+ + {/* Show processing indicator when attachments are being processed */} + {hasProcessingAttachments && ( +
+ + Processing... +
+ )} + + {/* Show warning when no model is configured */} + {!hasModelConfigured && !hasProcessingAttachments && ( +
+ + Select a model +
+ )} + + !thread.isRunning}> + + + + + + + + thread.isRunning}> + + + + +
+ ); +}; + diff --git a/surfsense_web/components/assistant-ui/composer.tsx b/surfsense_web/components/assistant-ui/composer.tsx new file mode 100644 index 000000000..1973726da --- /dev/null +++ b/surfsense_web/components/assistant-ui/composer.tsx @@ -0,0 +1,240 @@ +import { ComposerPrimitive, useAssistantState, useComposerRuntime } from "@assistant-ui/react"; +import { useAtom, useSetAtom } from "jotai"; +import { useParams } from "next/navigation"; +import type { FC } from "react"; +import { useCallback, useEffect, useRef, useState } from "react"; +import { createPortal } from "react-dom"; +import { + mentionedDocumentIdsAtom, + mentionedDocumentsAtom, +} from "@/atoms/chat/mentioned-documents.atom"; +import { + ComposerAddAttachment, + ComposerAttachments, +} from "@/components/assistant-ui/attachment"; +import { ComposerAction } from "@/components/assistant-ui/composer-action"; +import { + InlineMentionEditor, + type InlineMentionEditorRef, +} from "@/components/assistant-ui/inline-mention-editor"; +import { + DocumentMentionPicker, + type DocumentMentionPickerRef, +} from "@/components/new-chat/document-mention-picker"; +import type { Document } from "@/contracts/types/document.types"; + +export const Composer: FC = () => { + // ---- 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 editorRef = useRef(null); + const editorContainerRef = useRef(null); + const documentPickerRef = useRef(null); + const { search_space_id } = useParams(); + const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom); + const composerRuntime = useComposerRuntime(); + const hasAutoFocusedRef = useRef(false); + + // Check if thread is empty (new chat) + const isThreadEmpty = useAssistantState(({ thread }) => thread.isEmpty); + + // Check if thread is currently running (streaming response) + const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning); + + // Auto-focus editor when on new chat page + useEffect(() => { + if (isThreadEmpty && !hasAutoFocusedRef.current && editorRef.current) { + // Small delay to ensure the editor is fully mounted + const timeoutId = setTimeout(() => { + editorRef.current?.focus(); + hasAutoFocusedRef.current = true; + }, 100); + return () => clearTimeout(timeoutId); + } + }, [isThreadEmpty]); + + // Sync mentioned document IDs to atom for use in chat request + useEffect(() => { + setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id)); + }, [mentionedDocuments, setMentionedDocumentIds]); + + // Handle text change from inline editor - sync with assistant-ui composer + const handleEditorChange = useCallback( + (text: string) => { + composerRuntime.setText(text); + }, + [composerRuntime] + ); + + // Handle @ mention trigger from inline editor + const handleMentionTrigger = useCallback((query: string) => { + setShowDocumentPopover(true); + setMentionQuery(query); + }, []); + + // Handle mention close + const handleMentionClose = useCallback(() => { + if (showDocumentPopover) { + setShowDocumentPopover(false); + setMentionQuery(""); + } + }, [showDocumentPopover]); + + // Handle keyboard navigation when popover is open + const handleKeyDown = useCallback( + (e: React.KeyboardEvent) => { + 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; + } + } + }, + [showDocumentPopover] + ); + + // Handle submit from inline editor (Enter key) + const handleSubmit = useCallback(() => { + // Prevent sending while a response is still streaming + if (isThreadRunning) { + return; + } + if (!showDocumentPopover) { + composerRuntime.send(); + // Clear the editor after sending + editorRef.current?.clear(); + setMentionedDocuments([]); + setMentionedDocumentIds([]); + } + }, [ + showDocumentPopover, + isThreadRunning, + composerRuntime, + setMentionedDocuments, + setMentionedDocumentIds, + ]); + + // Handle document removal from inline editor + const handleDocumentRemove = useCallback( + (docId: number) => { + setMentionedDocuments((prev) => { + const updated = prev.filter((doc) => doc.id !== docId); + // Immediately sync document IDs to avoid race conditions + setMentionedDocumentIds(updated.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)); + + 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 updated = [...prev, ...uniqueNewDocs]; + // Immediately sync document IDs to avoid race conditions + setMentionedDocumentIds(updated.map((doc) => doc.id)); + return updated; + }); + + // Reset mention query but keep popover open for more selections + setMentionQuery(""); + }, + [mentionedDocuments, setMentionedDocuments, setMentionedDocumentIds] + ); + + return ( + + + + {/* -------- Inline Mention Editor -------- */} +
+ +
+ + {/* -------- Document mention popover (rendered via portal) -------- */} + {showDocumentPopover && + typeof document !== "undefined" && + createPortal( + <> + {/* Backdrop */} + + + + + + + + + ); +}; + diff --git a/surfsense_web/components/assistant-ui/thinking-steps.tsx b/surfsense_web/components/assistant-ui/thinking-steps.tsx new file mode 100644 index 000000000..f0cf4a7c1 --- /dev/null +++ b/surfsense_web/components/assistant-ui/thinking-steps.tsx @@ -0,0 +1,207 @@ +import { useAssistantState, useThreadViewport } from "@assistant-ui/react"; +import type { FC } from "react"; +import { createContext, useCallback, useContext, useEffect, useRef, useState } from "react"; +import { ChevronRightIcon } from "lucide-react"; +import { ChainOfThoughtItem } from "@/components/prompt-kit/chain-of-thought"; +import { TextShimmerLoader } from "@/components/prompt-kit/loader"; +import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking"; +import { cn } from "@/lib/utils"; + +// Context to pass thinking steps to AssistantMessage +export const ThinkingStepsContext = createContext>(new Map()); + +/** + * Chain of thought display component - single collapsible dropdown design + */ +export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolean }> = ({ + steps, + isThreadRunning = true, +}) => { + const [isOpen, setIsOpen] = useState(true); + + // Derive effective status for each step + const getEffectiveStatus = useCallback( + (step: ThinkingStep): "pending" | "in_progress" | "completed" => { + if (step.status === "in_progress" && !isThreadRunning) { + return "completed"; + } + return step.status; + }, + [isThreadRunning] + ); + + // Calculate summary info + const completedSteps = steps.filter((s) => getEffectiveStatus(s) === "completed").length; + const inProgressStep = steps.find((s) => getEffectiveStatus(s) === "in_progress"); + const allCompleted = completedSteps === steps.length && steps.length > 0 && !isThreadRunning; + const isProcessing = isThreadRunning && !allCompleted; + + // Auto-collapse when all tasks are completed + useEffect(() => { + if (allCompleted) { + setIsOpen(false); + } + }, [allCompleted]); + + if (steps.length === 0) return null; + + // Generate header text + const getHeaderText = () => { + if (allCompleted) { + return `Reviewed ${completedSteps} ${completedSteps === 1 ? "step" : "steps"}`; + } + if (inProgressStep) { + return inProgressStep.title; + } + if (isProcessing) { + return `Processing ${completedSteps}/${steps.length} steps`; + } + return `Reviewed ${completedSteps} ${completedSteps === 1 ? "step" : "steps"}`; + }; + + return ( +
+
+ {/* Main collapsible header */} + + + {/* Collapsible content with CSS grid animation */} +
+
+
+ {steps.map((step, index) => { + const effectiveStatus = getEffectiveStatus(step); + const isLast = index === steps.length - 1; + + return ( +
+ {/* Dot and line column */} +
+ {/* Vertical connection line - extends to next dot */} + {!isLast && ( +
+ )} + {/* Step dot - on top of line */} +
+ {effectiveStatus === "in_progress" ? ( + + ) : ( + + )} +
+
+ + {/* Step content */} +
+ {/* Step title */} +
+ {effectiveStatus === "in_progress" ? ( + + ) : ( + step.title + )} +
+ + {/* Step items (sub-content) */} + {step.items && step.items.length > 0 && ( +
+ {step.items.map((item, idx) => ( + + {item} + + ))} +
+ )} +
+
+ ); + })} +
+
+
+
+
+ ); +}; + +/** + * Component that handles auto-scroll when thinking steps update. + * Uses useThreadViewport to scroll to bottom when thinking steps change, + * ensuring the user always sees the latest content during streaming. + */ +export const ThinkingStepsScrollHandler: FC = () => { + const thinkingStepsMap = useContext(ThinkingStepsContext); + const viewport = useThreadViewport(); + const isRunning = useAssistantState(({ thread }) => thread.isRunning); + // Track the serialized state to detect any changes + const prevStateRef = useRef(""); + + useEffect(() => { + // Only act during streaming + if (!isRunning) { + prevStateRef.current = ""; + return; + } + + // Serialize the thinking steps state to detect any changes + // This catches new steps, status changes, and item additions + let stateString = ""; + thinkingStepsMap.forEach((steps, msgId) => { + steps.forEach((step) => { + stateString += `${msgId}:${step.id}:${step.status}:${step.items?.length || 0};`; + }); + }); + + // If state changed at all during streaming, scroll + if (stateString !== prevStateRef.current && stateString !== "") { + prevStateRef.current = stateString; + + // Multiple attempts to ensure scroll happens after DOM updates + const scrollAttempt = () => { + try { + viewport.scrollToBottom(); + } catch { + // Ignore errors - viewport might not be ready + } + }; + + // Delayed attempts to handle async DOM updates + requestAnimationFrame(scrollAttempt); + setTimeout(scrollAttempt, 100); + } + }, [thinkingStepsMap, viewport, isRunning]); + + return null; // This component doesn't render anything +}; + diff --git a/surfsense_web/components/assistant-ui/thread-scroll-to-bottom.tsx b/surfsense_web/components/assistant-ui/thread-scroll-to-bottom.tsx new file mode 100644 index 000000000..6f641615e --- /dev/null +++ b/surfsense_web/components/assistant-ui/thread-scroll-to-bottom.tsx @@ -0,0 +1,19 @@ +import { ThreadPrimitive } from "@assistant-ui/react"; +import { ArrowDownIcon } from "lucide-react"; +import type { FC } from "react"; +import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; + +export const ThreadScrollToBottom: FC = () => { + return ( + + + + + + ); +}; + diff --git a/surfsense_web/components/assistant-ui/thread-welcome.tsx b/surfsense_web/components/assistant-ui/thread-welcome.tsx new file mode 100644 index 000000000..b5e4bbac0 --- /dev/null +++ b/surfsense_web/components/assistant-ui/thread-welcome.tsx @@ -0,0 +1,72 @@ +import { useAtomValue } from "jotai"; +import type { FC } from "react"; +import { useMemo } from "react"; +import { currentUserAtom } from "@/atoms/user/user-query.atoms"; +import { Composer } from "@/components/assistant-ui/composer"; + +const getTimeBasedGreeting = (userEmail?: string): string => { + const hour = new Date().getHours(); + + // Extract first name from email if available + const firstName = userEmail + ? userEmail.split("@")[0].split(".")[0].charAt(0).toUpperCase() + + userEmail.split("@")[0].split(".")[0].slice(1) + : null; + + // Array of greeting variations for each time period + const morningGreetings = ["Good morning", "Fresh start today", "Morning", "Hey there"]; + + const afternoonGreetings = ["Good afternoon", "Afternoon", "Hey there", "Hi there"]; + + const eveningGreetings = ["Good evening", "Evening", "Hey there", "Hi there"]; + + const nightGreetings = ["Good night", "Evening", "Hey there", "Winding down"]; + + const lateNightGreetings = ["Still up", "Night owl mode", "Up past bedtime", "Hi there"]; + + // Select a random greeting based on time + let greeting: string; + if (hour < 5) { + // Late night: midnight to 5 AM + greeting = lateNightGreetings[Math.floor(Math.random() * lateNightGreetings.length)]; + } else if (hour < 12) { + greeting = morningGreetings[Math.floor(Math.random() * morningGreetings.length)]; + } else if (hour < 18) { + greeting = afternoonGreetings[Math.floor(Math.random() * afternoonGreetings.length)]; + } else if (hour < 22) { + greeting = eveningGreetings[Math.floor(Math.random() * eveningGreetings.length)]; + } else { + // Night: 10 PM to midnight + greeting = nightGreetings[Math.floor(Math.random() * nightGreetings.length)]; + } + + // Add personalization with first name if available + if (firstName) { + return `${greeting}, ${firstName}!`; + } + + return `${greeting}!`; +}; + +export 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 */} +
+

+ {greeting} +

+
+ {/* Composer - top edge fixed, expands downward only */} +
+ +
+
+ ); +}; + diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index cb01e7605..a354414f0 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -1,85 +1,13 @@ -import { - ActionBarPrimitive, - AssistantIf, - BranchPickerPrimitive, - ComposerPrimitive, - ErrorPrimitive, - MessagePrimitive, - ThreadPrimitive, - useAssistantState, - useComposerRuntime, - useThreadViewport, -} from "@assistant-ui/react"; -import { useAtom, useAtomValue, useSetAtom } from "jotai"; -import { - AlertCircle, - ArrowDownIcon, - ArrowUpIcon, - CheckIcon, - ChevronLeftIcon, - ChevronRightIcon, - CopyIcon, - DownloadIcon, - FileText, - Loader2, - PencilIcon, - Plug2, - Plus, - RefreshCwIcon, - SquareIcon, -} from "lucide-react"; -import Link from "next/link"; -import { useParams } from "next/navigation"; -import { - createContext, - type FC, - useCallback, - useContext, - useEffect, - useMemo, - useRef, - useState, -} from "react"; -import { createPortal } from "react-dom"; -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, - llmPreferencesAtom, - newLLMConfigsAtom, -} from "@/atoms/new-llm-config/new-llm-config-query.atoms"; -import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms"; -import { currentUserAtom } from "@/atoms/user/user-query.atoms"; -import { - ComposerAddAttachment, - ComposerAttachments, - UserMessageAttachments, -} from "@/components/assistant-ui/attachment"; -import { - InlineMentionEditor, - type InlineMentionEditorRef, -} from "@/components/assistant-ui/inline-mention-editor"; -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 { - DocumentMentionPicker, - type DocumentMentionPickerRef, -} from "@/components/new-chat/document-mention-picker"; -import { ChainOfThoughtItem } from "@/components/prompt-kit/chain-of-thought"; -import { TextShimmerLoader } from "@/components/prompt-kit/loader"; +import { AssistantIf, ThreadPrimitive } from "@assistant-ui/react"; +import type { FC } from "react"; import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking"; -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 { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors"; -import { cn } from "@/lib/utils"; +import { ThinkingStepsContext } from "@/components/assistant-ui/thinking-steps"; +import { ThreadWelcome } from "@/components/assistant-ui/thread-welcome"; +import { Composer } from "@/components/assistant-ui/composer"; +import { ThreadScrollToBottom } from "@/components/assistant-ui/thread-scroll-to-bottom"; +import { AssistantMessage } from "@/components/assistant-ui/assistant-message"; +import { UserMessage } from "@/components/assistant-ui/user-message"; +import { EditComposer } from "@/components/assistant-ui/edit-composer"; /** * Props for the Thread component @@ -90,204 +18,6 @@ interface ThreadProps { header?: React.ReactNode; } -// Context to pass thinking steps to AssistantMessage -const ThinkingStepsContext = createContext>(new Map()); - -/** - * Chain of thought display component - single collapsible dropdown design - */ -const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolean }> = ({ - steps, - isThreadRunning = true, -}) => { - const [isOpen, setIsOpen] = useState(true); - - // Derive effective status for each step - const getEffectiveStatus = useCallback( - (step: ThinkingStep): "pending" | "in_progress" | "completed" => { - if (step.status === "in_progress" && !isThreadRunning) { - return "completed"; - } - return step.status; - }, - [isThreadRunning] - ); - - // Calculate summary info - const completedSteps = steps.filter((s) => getEffectiveStatus(s) === "completed").length; - const inProgressStep = steps.find((s) => getEffectiveStatus(s) === "in_progress"); - const allCompleted = completedSteps === steps.length && steps.length > 0 && !isThreadRunning; - const isProcessing = isThreadRunning && !allCompleted; - - // Auto-collapse when all tasks are completed - useEffect(() => { - if (allCompleted) { - setIsOpen(false); - } - }, [allCompleted]); - - if (steps.length === 0) return null; - - // Generate header text - const getHeaderText = () => { - if (allCompleted) { - return `Reviewed ${completedSteps} ${completedSteps === 1 ? "step" : "steps"}`; - } - if (inProgressStep) { - return inProgressStep.title; - } - if (isProcessing) { - return `Processing ${completedSteps}/${steps.length} steps`; - } - return `Reviewed ${completedSteps} ${completedSteps === 1 ? "step" : "steps"}`; - }; - - return ( -
-
- {/* Main collapsible header */} - - - {/* Collapsible content with CSS grid animation */} -
-
-
- {steps.map((step, index) => { - const effectiveStatus = getEffectiveStatus(step); - const isLast = index === steps.length - 1; - - return ( -
- {/* Dot and line column */} -
- {/* Vertical connection line - extends to next dot */} - {!isLast && ( -
- )} - {/* Step dot - on top of line */} -
- {effectiveStatus === "in_progress" ? ( - - ) : ( - - )} -
-
- - {/* Step content */} -
- {/* Step title */} -
- {effectiveStatus === "in_progress" ? ( - - ) : ( - step.title - )} -
- - {/* Step items (sub-content) */} - {step.items && step.items.length > 0 && ( -
- {step.items.map((item, idx) => ( - - {item} - - ))} -
- )} -
-
- ); - })} -
-
-
-
-
- ); -}; - -/** - * Component that handles auto-scroll when thinking steps update. - * Uses useThreadViewport to scroll to bottom when thinking steps change, - * ensuring the user always sees the latest content during streaming. - */ -const _ThinkingStepsScrollHandler: FC = () => { - const thinkingStepsMap = useContext(ThinkingStepsContext); - const viewport = useThreadViewport(); - const isRunning = useAssistantState(({ thread }) => thread.isRunning); - // Track the serialized state to detect any changes - const prevStateRef = useRef(""); - - useEffect(() => { - // Only act during streaming - if (!isRunning) { - prevStateRef.current = ""; - return; - } - - // Serialize the thinking steps state to detect any changes - // This catches new steps, status changes, and item additions - let stateString = ""; - thinkingStepsMap.forEach((steps, msgId) => { - steps.forEach((step) => { - stateString += `${msgId}:${step.id}:${step.status}:${step.items?.length || 0};`; - }); - }); - - // If state changed at all during streaming, scroll - if (stateString !== prevStateRef.current && stateString !== "") { - prevStateRef.current = stateString; - - // Multiple attempts to ensure scroll happens after DOM updates - const scrollAttempt = () => { - try { - viewport.scrollToBottom(); - } catch { - // Ignore errors - viewport might not be ready - } - }; - - // Delayed attempts to handle async DOM updates - requestAnimationFrame(scrollAttempt); - setTimeout(scrollAttempt, 100); - } - }, [thinkingStepsMap, viewport, isRunning]); - - return null; // This component doesn't render anything -}; - export const Thread: FC = ({ messageThinkingSteps = new Map(), header }) => { return ( @@ -329,760 +59,3 @@ export const Thread: FC = ({ messageThinkingSteps = new Map(), head ); }; - -const ThreadScrollToBottom: FC = () => { - return ( - - - - - - ); -}; - -const getTimeBasedGreeting = (userEmail?: string): string => { - const hour = new Date().getHours(); - - // Extract first name from email if available - const firstName = userEmail - ? userEmail.split("@")[0].split(".")[0].charAt(0).toUpperCase() + - userEmail.split("@")[0].split(".")[0].slice(1) - : null; - - // Array of greeting variations for each time period - const morningGreetings = ["Good morning", "Fresh start today", "Morning", "Hey there"]; - - const afternoonGreetings = ["Good afternoon", "Afternoon", "Hey there", "Hi there"]; - - const eveningGreetings = ["Good evening", "Evening", "Hey there", "Hi there"]; - - const nightGreetings = ["Good night", "Evening", "Hey there", "Winding down"]; - - const lateNightGreetings = ["Still up", "Night owl mode", "Up past bedtime", "Hi there"]; - - // Select a random greeting based on time - let greeting: string; - if (hour < 5) { - // Late night: midnight to 5 AM - greeting = lateNightGreetings[Math.floor(Math.random() * lateNightGreetings.length)]; - } else if (hour < 12) { - greeting = morningGreetings[Math.floor(Math.random() * morningGreetings.length)]; - } else if (hour < 18) { - greeting = afternoonGreetings[Math.floor(Math.random() * afternoonGreetings.length)]; - } else if (hour < 22) { - greeting = eveningGreetings[Math.floor(Math.random() * eveningGreetings.length)]; - } else { - // Night: 10 PM to midnight - greeting = nightGreetings[Math.floor(Math.random() * nightGreetings.length)]; - } - - // Add personalization with first name if available - if (firstName) { - return `${greeting}, ${firstName}!`; - } - - return `${greeting}!`; -}; - -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 */} -
-

- {greeting} -

-
- {/* Composer - top edge fixed, expands downward only */} -
- -
-
- ); -}; - -const Composer: FC = () => { - // ---- 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 editorRef = useRef(null); - const editorContainerRef = useRef(null); - const documentPickerRef = useRef(null); - const { search_space_id } = useParams(); - const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom); - const composerRuntime = useComposerRuntime(); - const hasAutoFocusedRef = useRef(false); - - // Check if thread is empty (new chat) - const isThreadEmpty = useAssistantState(({ thread }) => thread.isEmpty); - - // Check if thread is currently running (streaming response) - const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning); - - // Auto-focus editor when on new chat page - useEffect(() => { - if (isThreadEmpty && !hasAutoFocusedRef.current && editorRef.current) { - // Small delay to ensure the editor is fully mounted - const timeoutId = setTimeout(() => { - editorRef.current?.focus(); - hasAutoFocusedRef.current = true; - }, 100); - return () => clearTimeout(timeoutId); - } - }, [isThreadEmpty]); - - // Sync mentioned document IDs to atom for use in chat request - useEffect(() => { - setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id)); - }, [mentionedDocuments, setMentionedDocumentIds]); - - // Handle text change from inline editor - sync with assistant-ui composer - const handleEditorChange = useCallback( - (text: string) => { - composerRuntime.setText(text); - }, - [composerRuntime] - ); - - // Handle @ mention trigger from inline editor - const handleMentionTrigger = useCallback((query: string) => { - setShowDocumentPopover(true); - setMentionQuery(query); - }, []); - - // Handle mention close - const handleMentionClose = useCallback(() => { - if (showDocumentPopover) { - setShowDocumentPopover(false); - setMentionQuery(""); - } - }, [showDocumentPopover]); - - // Handle keyboard navigation when popover is open - const handleKeyDown = useCallback( - (e: React.KeyboardEvent) => { - 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; - } - } - }, - [showDocumentPopover] - ); - - // Handle submit from inline editor (Enter key) - const handleSubmit = useCallback(() => { - // Prevent sending while a response is still streaming - if (isThreadRunning) { - return; - } - if (!showDocumentPopover) { - composerRuntime.send(); - // Clear the editor after sending - editorRef.current?.clear(); - setMentionedDocuments([]); - setMentionedDocumentIds([]); - } - }, [ - showDocumentPopover, - isThreadRunning, - composerRuntime, - setMentionedDocuments, - setMentionedDocumentIds, - ]); - - // Handle document removal from inline editor - const handleDocumentRemove = useCallback( - (docId: number) => { - setMentionedDocuments((prev) => { - const updated = prev.filter((doc) => doc.id !== docId); - // Immediately sync document IDs to avoid race conditions - setMentionedDocumentIds(updated.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)); - - 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 updated = [...prev, ...uniqueNewDocs]; - // Immediately sync document IDs to avoid race conditions - setMentionedDocumentIds(updated.map((doc) => doc.id)); - return updated; - }); - - // Reset mention query but keep popover open for more selections - setMentionQuery(""); - }, - [mentionedDocuments, setMentionedDocuments, setMentionedDocumentIds] - ); - - return ( - - - - {/* -------- Inline Mention Editor -------- */} -
- -
- - {/* -------- Document mention popover (rendered via portal) -------- */} - {showDocumentPopover && - typeof document !== "undefined" && - createPortal( - <> - {/* Backdrop */} - - - - {hasSources ? ( -
-
-

Connected Sources

- - {totalSourceCount} - -
-
- {/* Document types from the search space */} - {activeDocumentTypes.map(([docType]) => ( -
- {getConnectorIcon(docType, "size-3.5")} - {getDocumentTypeLabel(docType)} -
- ))} - {/* Search source connectors */} - {connectors.map((connector) => ( -
- {getConnectorIcon(connector.connector_type, "size-3.5")} - {connector.name} -
- ))} -
-
- - - Add more sources - - -
-
- ) : ( -
-

No sources yet

-

- Add documents or connect data sources to enhance search results. -

- - - Add Connector - -
- )} - - - ); -}; - -const ComposerAction: FC = () => { - // Check if any attachments are still being processed (running AND progress < 100) - // When progress is 100, processing is done but waiting for send() - const hasProcessingAttachments = useAssistantState(({ composer }) => - composer.attachments?.some((att) => { - const status = att.status; - if (status?.type !== "running") return false; - const progress = (status as { type: "running"; progress?: number }).progress; - return progress === undefined || progress < 100; - }) - ); - - // Check if composer text is empty - const isComposerEmpty = useAssistantState(({ composer }) => { - const text = composer.text?.trim() || ""; - return text.length === 0; - }); - - // Check if a model is configured - 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]); - - const isSendDisabled = hasProcessingAttachments || isComposerEmpty || !hasModelConfigured; - - return ( -
-
- - -
- - {/* Show processing indicator when attachments are being processed */} - {hasProcessingAttachments && ( -
- - Processing... -
- )} - - {/* Show warning when no model is configured */} - {!hasModelConfigured && !hasProcessingAttachments && ( -
- - Select a model -
- )} - - !thread.isRunning}> - - - - - - - - thread.isRunning}> - - - - -
- ); -}; - -const MessageError: FC = () => { - return ( - - - - - - ); -}; - -/** - * Custom component to render thinking steps from Context - */ -const ThinkingStepsPart: FC = () => { - const thinkingStepsMap = useContext(ThinkingStepsContext); - - // Get the current message ID to look up thinking steps - const messageId = useAssistantState(({ message }) => message?.id); - const thinkingSteps = thinkingStepsMap.get(messageId) || []; - - // Check if this specific message is currently streaming - // A message is streaming if: thread is running AND this is the last assistant message - const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning); - const isLastMessage = useAssistantState(({ message }) => message?.isLast ?? false); - const isMessageStreaming = isThreadRunning && isLastMessage; - - if (thinkingSteps.length === 0) return null; - - return ( -
- -
- ); -}; - -const AssistantMessageInner: FC = () => { - return ( - <> - {/* Render thinking steps from message content - this ensures proper scroll tracking */} - - -
- - -
- -
- - -
- - ); -}; - -const AssistantMessage: FC = () => { - return ( - - - - ); -}; - -const AssistantActionBar: FC = () => { - return ( - - - - message.isCopied}> - - - !message.isCopied}> - - - - - - - - - - - - - - - - ); -}; - -const UserMessage: FC = () => { - const messageId = useAssistantState(({ message }) => message?.id); - const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom); - const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined; - const hasAttachments = useAssistantState( - ({ message }) => message?.attachments && message.attachments.length > 0 - ); - - return ( - -
- {/* Display attachments and mentioned documents */} - {(hasAttachments || (mentionedDocs && mentionedDocs.length > 0)) && ( -
- {/* Attachments (images show as thumbnails, documents as chips) */} - - {/* Mentioned documents as chips */} - {mentionedDocs?.map((doc) => ( - - - {doc.title} - - ))} -
- )} - {/* Message bubble with action bar positioned relative to it */} -
-
- -
-
- -
-
-
- - - - ); -}; - -const UserActionBar: FC = () => { - return ( - - - - - - - - ); -}; - -const EditComposer: FC = () => { - return ( - - - -
- - - - - - -
- - - ); -}; - -const BranchPicker: FC = ({ className, ...rest }) => { - return ( - - - - - - - - / - - - - - - - - ); -}; diff --git a/surfsense_web/components/assistant-ui/user-message.tsx b/surfsense_web/components/assistant-ui/user-message.tsx new file mode 100644 index 000000000..fbbcf42bf --- /dev/null +++ b/surfsense_web/components/assistant-ui/user-message.tsx @@ -0,0 +1,73 @@ +import { ActionBarPrimitive, MessagePrimitive, useAssistantState } from "@assistant-ui/react"; +import { useAtomValue } from "jotai"; +import { FileText, PencilIcon } from "lucide-react"; +import type { FC } from "react"; +import { messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom"; +import { UserMessageAttachments } from "@/components/assistant-ui/attachment"; +import { BranchPicker } from "@/components/assistant-ui/branch-picker"; +import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; + +export const UserMessage: FC = () => { + const messageId = useAssistantState(({ message }) => message?.id); + const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom); + const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined; + const hasAttachments = useAssistantState( + ({ message }) => message?.attachments && message.attachments.length > 0 + ); + + return ( + +
+ {/* Display attachments and mentioned documents */} + {(hasAttachments || (mentionedDocs && mentionedDocs.length > 0)) && ( +
+ {/* Attachments (images show as thumbnails, documents as chips) */} + + {/* Mentioned documents as chips */} + {mentionedDocs?.map((doc) => ( + + + {doc.title} + + ))} +
+ )} + {/* Message bubble with action bar positioned relative to it */} +
+
+ +
+
+ +
+
+
+ + + + ); +}; + +const UserActionBar: FC = () => { + return ( + + + + + + + + ); +}; +