diff --git a/surfsense_web/app/dashboard/[search_space_id]/v2/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/v2/[[...chat_id]]/page.tsx index d3ea1a66f..cff02069c 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/v2/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/v2/[[...chat_id]]/page.tsx @@ -20,6 +20,7 @@ export default function ResearchChatPageV2() { isLoading, setIsLoading, searchMode, + setSearchMode, researchMode, setResearchMode, selectedConnectors, @@ -34,8 +35,6 @@ export default function ResearchChatPageV2() { const { fetchChatDetails, updateChat, createChat } = useChatAPI({ token, search_space_id: search_space_id as string, - researchMode, - selectedConnectors, }); // Memoize document IDs to prevent infinite re-renders @@ -43,31 +42,37 @@ export default function ResearchChatPageV2() { return selectedDocuments.map((doc) => doc.id); }, [selectedDocuments]); - // Helper functions for localStorage management - const getStorageKey = (searchSpaceId: string, chatId: string) => - `surfsense_selected_docs_${searchSpaceId}_${chatId}`; + // Unified localStorage management for chat state + interface ChatState { + selectedDocuments: Document[]; + searchMode: "DOCUMENTS" | "CHUNKS"; + researchMode: ResearchMode; + } - const storeSelectedDocuments = ( + const getChatStateStorageKey = (searchSpaceId: string, chatId: string) => + `surfsense_chat_state_${searchSpaceId}_${chatId}`; + + const storeChatState = ( searchSpaceId: string, chatId: string, - documents: Document[] + state: ChatState ) => { - const key = getStorageKey(searchSpaceId, chatId); - localStorage.setItem(key, JSON.stringify(documents)); + const key = getChatStateStorageKey(searchSpaceId, chatId); + localStorage.setItem(key, JSON.stringify(state)); }; - const restoreSelectedDocuments = ( + const restoreChatState = ( searchSpaceId: string, chatId: string - ): Document[] | null => { - const key = getStorageKey(searchSpaceId, chatId); + ): ChatState | null => { + const key = getChatStateStorageKey(searchSpaceId, chatId); const stored = localStorage.getItem(key); if (stored) { localStorage.removeItem(key); // Clean up after restoration try { return JSON.parse(stored); } catch (error) { - console.error("Error parsing stored documents:", error); + console.error("Error parsing stored chat state:", error); return null; } } @@ -99,14 +104,14 @@ export default function ResearchChatPageV2() { message: Message | CreateMessage, chatRequestOptions?: { data?: any } ) => { - const newChatId = await createChat(message.content); + const newChatId = await createChat(message.content, researchMode); if (newChatId) { - // Store selected documents before navigation - storeSelectedDocuments( - search_space_id as string, - newChatId, - selectedDocuments - ); + // Store chat state before navigation + storeChatState(search_space_id as string, newChatId, { + selectedDocuments, + searchMode, + researchMode, + }); router.replace(`/dashboard/${search_space_id}/v2/${newChatId}`); } return newChatId; @@ -119,18 +124,26 @@ export default function ResearchChatPageV2() { } }, [token, isNewChat, chatIdParam]); - // Restore selected documents from localStorage on page load + // Restore chat state from localStorage on page load useEffect(() => { if (chatIdParam && search_space_id) { - const restoredDocuments = restoreSelectedDocuments( + const restoredState = restoreChatState( search_space_id as string, chatIdParam ); - if (restoredDocuments && restoredDocuments.length > 0) { - setSelectedDocuments(restoredDocuments); + if (restoredState) { + setSelectedDocuments(restoredState.selectedDocuments); + setSearchMode(restoredState.searchMode); + setResearchMode(restoredState.researchMode); } } - }, [chatIdParam, search_space_id, setSelectedDocuments]); + }, [ + chatIdParam, + search_space_id, + setSelectedDocuments, + setSearchMode, + setResearchMode, + ]); const loadChatData = async (chatId: string) => { try { @@ -179,9 +192,9 @@ export default function ResearchChatPageV2() { handler.messages.length > 0 && handler.messages[handler.messages.length - 1]?.role === "assistant" ) { - updateChat(chatIdParam, handler.messages); + updateChat(chatIdParam, handler.messages, researchMode); } - }, [handler.messages, handler.status, chatIdParam, isNewChat, updateChat]); + }, [handler.messages, handler.status, chatIdParam, isNewChat]); if (isLoading) { return ( @@ -199,6 +212,10 @@ export default function ResearchChatPageV2() { }} onDocumentSelectionChange={setSelectedDocuments} selectedDocuments={selectedDocuments} + searchMode={searchMode} + onSearchModeChange={setSearchMode} + researchMode={researchMode} + onResearchModeChange={setResearchMode} /> ); } diff --git a/surfsense_web/components/chat_v2/ChatInterface.tsx b/surfsense_web/components/chat_v2/ChatInterface.tsx index cb2b3f3be..02aa77f86 100644 --- a/surfsense_web/components/chat_v2/ChatInterface.tsx +++ b/surfsense_web/components/chat_v2/ChatInterface.tsx @@ -17,16 +17,29 @@ import { DialogTitle, DialogTrigger, } from "../ui/dialog"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "../ui/select"; +import { Badge } from "../ui/badge"; import { Suspense, useState, useCallback } from "react"; import { useParams } from "next/navigation"; import { useDocuments, DocumentType, Document } from "@/hooks/use-documents"; import { DocumentsDataTable } from "./DocumentsDataTable"; +import { ResearchMode } from "@/components/chat"; import React from "react"; interface ChatInterfaceProps { handler: ChatHandler; onDocumentSelectionChange?: (documents: Document[]) => void; selectedDocuments?: Document[]; + searchMode?: "DOCUMENTS" | "CHUNKS"; + onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void; + researchMode?: ResearchMode; + onResearchModeChange?: (mode: ResearchMode) => void; } const DocumentSelector = React.memo( @@ -118,21 +131,127 @@ const DocumentSelector = React.memo( } ); +const SearchModeSelector = ({ + searchMode, + onSearchModeChange, +}: { + searchMode?: "DOCUMENTS" | "CHUNKS"; + onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void; +}) => { + return ( +
+ + Scope: + +
+ + +
+
+ ); +}; + +const ResearchModeSelector = ({ + researchMode, + onResearchModeChange, +}: { + researchMode?: ResearchMode; + onResearchModeChange?: (mode: ResearchMode) => void; +}) => { + const researchModeLabels: Record = { + QNA: "Q&A", + REPORT_GENERAL: "General Report", + REPORT_DEEP: "Deep Report", + REPORT_DEEPER: "Deeper Report", + }; + + const researchModeShortLabels: Record = { + QNA: "Q&A", + REPORT_GENERAL: "General", + REPORT_DEEP: "Deep", + REPORT_DEEPER: "Deeper", + }; + + return ( +
+ + Mode: + + +
+ ); +}; + const CustomChatInputOptions = ({ onDocumentSelectionChange, selectedDocuments, + searchMode, + onSearchModeChange, + researchMode, + onResearchModeChange, }: { onDocumentSelectionChange?: (documents: Document[]) => void; selectedDocuments?: Document[]; + searchMode?: "DOCUMENTS" | "CHUNKS"; + onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void; + researchMode?: ResearchMode; + onResearchModeChange?: (mode: ResearchMode) => void; }) => { return ( -
+
Loading...
}> + +
); }; @@ -140,9 +259,17 @@ const CustomChatInputOptions = ({ const CustomChatInput = ({ onDocumentSelectionChange, selectedDocuments, + searchMode, + onSearchModeChange, + researchMode, + onResearchModeChange, }: { onDocumentSelectionChange?: (documents: Document[]) => void; selectedDocuments?: Document[]; + searchMode?: "DOCUMENTS" | "CHUNKS"; + onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void; + researchMode?: ResearchMode; + onResearchModeChange?: (mode: ResearchMode) => void; }) => { return ( @@ -153,6 +280,10 @@ const CustomChatInput = ({ ); @@ -162,6 +293,10 @@ export default function ChatInterface({ handler, onDocumentSelectionChange, selectedDocuments = [], + searchMode, + onSearchModeChange, + researchMode, + onResearchModeChange, }: ChatInterfaceProps) { return ( @@ -177,6 +312,10 @@ export default function ChatInterface({ diff --git a/surfsense_web/hooks/useChat.ts b/surfsense_web/hooks/useChat.ts index 0bedcbfad..541654281 100644 --- a/surfsense_web/hooks/useChat.ts +++ b/surfsense_web/hooks/useChat.ts @@ -49,16 +49,9 @@ export function useChatState({ search_space_id, chat_id }: UseChatStateProps) { interface UseChatAPIProps { token: string | null; search_space_id: string; - researchMode: ResearchMode; - selectedConnectors: string[]; } -export function useChatAPI({ - token, - search_space_id, - researchMode, - selectedConnectors, -}: UseChatAPIProps) { +export function useChatAPI({ token, search_space_id }: UseChatAPIProps) { const fetchChatDetails = useCallback( async (chatId: string) => { if (!token) return null; @@ -93,7 +86,10 @@ export function useChatAPI({ ); const createChat = useCallback( - async (initialMessage: string): Promise => { + async ( + initialMessage: string, + researchMode: ResearchMode + ): Promise => { if (!token) { console.error("Authentication token not found"); return null; @@ -111,7 +107,7 @@ export function useChatAPI({ body: JSON.stringify({ type: researchMode, title: "Untitled Chat", - initial_connectors: selectedConnectors, + initial_connectors: [], messages: [ { role: "user", @@ -136,11 +132,15 @@ export function useChatAPI({ return null; } }, - [token, researchMode, selectedConnectors, search_space_id] + [token, search_space_id] ); const updateChat = useCallback( - async (chatId: string, messages: Message[]) => { + async ( + chatId: string, + messages: Message[], + researchMode: ResearchMode + ) => { if (!token) return; try { @@ -164,7 +164,7 @@ export function useChatAPI({ body: JSON.stringify({ type: researchMode, title: title, - initial_connectors: selectedConnectors, + initial_connectors: [], messages: messages, search_space_id: Number(search_space_id), }), @@ -180,7 +180,7 @@ export function useChatAPI({ console.error("Error updating chat:", err); } }, - [token, researchMode, selectedConnectors, search_space_id] + [token, search_space_id] ); return {