diff --git a/surfsense_backend/app/services/streaming_service.py b/surfsense_backend/app/services/streaming_service.py index cbd0e6d..2df442e 100644 --- a/surfsense_backend/app/services/streaming_service.py +++ b/surfsense_backend/app/services/streaming_service.py @@ -60,7 +60,7 @@ class StreamingService: self.message_annotations[1]["content"] = sources # Return only the delta annotation - annotation = {"type": "SOURCES", "content": sources} + annotation = {"type": "SOURCES", "data": sources} return f"8:[{json.dumps(annotation)}]\n" def format_answer_delta(self, answer_chunk: str) -> str: diff --git a/surfsense_web/components/chat_v2/ChatInputGroup.tsx b/surfsense_web/components/chat_v2/ChatInputGroup.tsx index f20496a..a2bb91b 100644 --- a/surfsense_web/components/chat_v2/ChatInputGroup.tsx +++ b/surfsense_web/components/chat_v2/ChatInputGroup.tsx @@ -138,7 +138,7 @@ const ConnectorSelector = React.memo( const [isOpen, setIsOpen] = useState(false); const { connectorSourceItems, isLoading, isLoaded, fetchConnectors } = - useSearchSourceConnectors(); + useSearchSourceConnectors(true); const handleOpenChange = useCallback( (open: boolean) => { diff --git a/surfsense_web/components/chat_v2/ChatInterface.tsx b/surfsense_web/components/chat_v2/ChatInterface.tsx index 6a02ca7..d8e66bf 100644 --- a/surfsense_web/components/chat_v2/ChatInterface.tsx +++ b/surfsense_web/components/chat_v2/ChatInterface.tsx @@ -1,5 +1,6 @@ "use client"; +import React from "react"; import { ChatSection, ChatHandler, @@ -12,8 +13,8 @@ import { import { Document } from "@/hooks/use-documents"; import { CustomChatInput } from "@/components/chat_v2/ChatInputGroup"; import { ResearchMode } from "@/components/chat"; -import React from "react"; import TerminalDisplay from "@/components/chat_v2/ChatTerminal"; +import ChatSourcesDisplay from "@/components/chat_v2/ChatSources"; interface ChatInterfaceProps { handler: ChatHandler; @@ -43,6 +44,7 @@ function ChatMessageDisplay({ {message.role === "assistant" ? (
+ diff --git a/surfsense_web/components/chat_v2/ChatSources.tsx b/surfsense_web/components/chat_v2/ChatSources.tsx new file mode 100644 index 0000000..997ee64 --- /dev/null +++ b/surfsense_web/components/chat_v2/ChatSources.tsx @@ -0,0 +1,172 @@ +"use client"; + +import { useState } from "react"; +import { getAnnotationData, Message } from "@llamaindex/chat-ui"; +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogHeader, + DialogTitle, + DialogTrigger, +} from "@/components/ui/dialog"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { + Card, + CardContent, + CardDescription, + CardHeader, + CardTitle, +} from "@/components/ui/card"; +import { Badge } from "@/components/ui/badge"; +import { ExternalLink, FileText, Github, Globe } from "lucide-react"; + +interface Source { + id: number; + title: string; + description: string; + url: string; +} + +interface SourceGroup { + id: number; + name: string; + type: string; + sources: Source[]; +} + +function getSourceIcon(type: string) { + switch (type) { + case "GITHUB_CONNECTOR": + return ; + case "NOTION_CONNECTOR": + return ; + case "FILE": + case "USER_SELECTED_FILE": + return ; + default: + return ; + } +} + +function SourceCard({ source }: { source: Source }) { + const hasUrl = source.url && source.url.trim() !== ""; + + return ( + + +
+ + {source.title} + + {hasUrl && ( + + )} +
+
+ + + {source.description} + + +
+ ); +} + +export default function ChatSourcesDisplay({ message }: { message: Message }) { + const [open, setOpen] = useState(false); + const annotations = getAnnotationData(message, "SOURCES"); + + // Flatten the nested array structure and ensure we have source groups + const sourceGroups: SourceGroup[] = + Array.isArray(annotations) && annotations.length > 0 + ? annotations + .flat() + .filter( + (group): group is SourceGroup => + group !== null && + group !== undefined && + typeof group === "object" && + "sources" in group && + Array.isArray(group.sources) + ) + : []; + + if (sourceGroups.length === 0) { + return null; + } + + const totalSources = sourceGroups.reduce( + (acc, group) => acc + group.sources.length, + 0 + ); + + return ( + + + + + + + Sources + + + + {sourceGroups.map((group) => ( + + {getSourceIcon(group.type)} + {group.name} + + {group.sources.length} + + + ))} + + {sourceGroups.map((group) => ( + +
+
+ {group.sources.map((source) => ( + + ))} +
+
+
+ ))} +
+
+
+ ); +} diff --git a/surfsense_web/hooks/useSearchSourceConnectors.ts b/surfsense_web/hooks/useSearchSourceConnectors.ts index ddc16f9..5204c74 100644 --- a/surfsense_web/hooks/useSearchSourceConnectors.ts +++ b/surfsense_web/hooks/useSearchSourceConnectors.ts @@ -21,10 +21,10 @@ export interface ConnectorSourceItem { /** * Hook to fetch search source connectors from the API */ -export const useSearchSourceConnectors = () => { +export const useSearchSourceConnectors = (lazy: boolean = false) => { const [connectors, setConnectors] = useState([]); - const [isLoading, setIsLoading] = useState(false); - const [isLoaded, setIsLoaded] = useState(false); + const [isLoading, setIsLoading] = useState(!lazy); // Don't show loading initially for lazy mode + const [isLoaded, setIsLoaded] = useState(false); // Memoization flag const [error, setError] = useState(null); const [connectorSourceItems, setConnectorSourceItems] = useState< ConnectorSourceItem[] @@ -56,7 +56,7 @@ export const useSearchSourceConnectors = () => { ]); const fetchConnectors = useCallback(async () => { - if (isLoaded) return; // Don't fetch if already loaded + if (isLoaded && lazy) return; // Avoid redundant calls in lazy mode try { setIsLoading(true); @@ -100,7 +100,19 @@ export const useSearchSourceConnectors = () => { } finally { setIsLoading(false); } - }, [isLoaded]); + }, [isLoaded, lazy]); + + useEffect(() => { + if (!lazy) { + fetchConnectors(); + } + }, [lazy, fetchConnectors]); + + // Function to refresh the connectors list + const refreshConnectors = useCallback(async () => { + setIsLoaded(false); // Reset memoization flag to allow refetch + await fetchConnectors(); + }, [fetchConnectors]); // Update connector source items when connectors change const updateConnectorSourceItems = ( @@ -363,5 +375,6 @@ export const useSearchSourceConnectors = () => { indexConnector, getConnectorSourceItems, connectorSourceItems, + refreshConnectors, }; };