Added search type and search mode selections for chcat input

This commit is contained in:
Utkarsh-Patel-13 2025-07-22 15:11:33 -07:00
parent 3f051b0a19
commit ac99613b24
3 changed files with 198 additions and 42 deletions

View file

@ -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}
/>
);
}

View file

@ -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 (
<div className="flex items-center gap-1 sm:gap-2">
<span className="text-xs text-muted-foreground hidden sm:block">
Scope:
</span>
<div className="flex rounded-md border">
<Button
variant={searchMode === "DOCUMENTS" ? "default" : "ghost"}
size="sm"
className="rounded-r-none border-r h-8 px-2 sm:px-3 text-xs"
onClick={() => onSearchModeChange?.("DOCUMENTS")}
>
<span className="hidden sm:inline">Documents</span>
<span className="sm:hidden">Docs</span>
</Button>
<Button
variant={searchMode === "CHUNKS" ? "default" : "ghost"}
size="sm"
className="rounded-l-none h-8 px-2 sm:px-3 text-xs"
onClick={() => onSearchModeChange?.("CHUNKS")}
>
Chunks
</Button>
</div>
</div>
);
};
const ResearchModeSelector = ({
researchMode,
onResearchModeChange,
}: {
researchMode?: ResearchMode;
onResearchModeChange?: (mode: ResearchMode) => void;
}) => {
const researchModeLabels: Record<ResearchMode, string> = {
QNA: "Q&A",
REPORT_GENERAL: "General Report",
REPORT_DEEP: "Deep Report",
REPORT_DEEPER: "Deeper Report",
};
const researchModeShortLabels: Record<ResearchMode, string> = {
QNA: "Q&A",
REPORT_GENERAL: "General",
REPORT_DEEP: "Deep",
REPORT_DEEPER: "Deeper",
};
return (
<div className="flex items-center gap-1 sm:gap-2">
<span className="text-xs text-muted-foreground hidden sm:block">
Mode:
</span>
<Select
value={researchMode}
onValueChange={(value) =>
onResearchModeChange?.(value as ResearchMode)
}
>
<SelectTrigger className="w-auto min-w-[80px] sm:min-w-[120px] h-8 text-xs">
<SelectValue placeholder="Mode" />
</SelectTrigger>
<SelectContent>
<SelectItem value="QNA">Q&A</SelectItem>
<SelectItem value="REPORT_GENERAL">
<span className="hidden sm:inline">General Report</span>
<span className="sm:hidden">General</span>
</SelectItem>
<SelectItem value="REPORT_DEEP">
<span className="hidden sm:inline">Deep Report</span>
<span className="sm:hidden">Deep</span>
</SelectItem>
<SelectItem value="REPORT_DEEPER">
<span className="hidden sm:inline">Deeper Report</span>
<span className="sm:hidden">Deeper</span>
</SelectItem>
</SelectContent>
</Select>
</div>
);
};
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 (
<div className="flex flex-row gap-2">
<div className="flex flex-wrap gap-2 sm:gap-3 items-center justify-start">
<Suspense fallback={<div>Loading...</div>}>
<DocumentSelector
onSelectionChange={onDocumentSelectionChange}
selectedDocuments={selectedDocuments}
/>
</Suspense>
<SearchModeSelector
searchMode={searchMode}
onSearchModeChange={onSearchModeChange}
/>
<ResearchModeSelector
researchMode={researchMode}
onResearchModeChange={onResearchModeChange}
/>
</div>
);
};
@ -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 (
<ChatInput>
@ -153,6 +280,10 @@ const CustomChatInput = ({
<CustomChatInputOptions
onDocumentSelectionChange={onDocumentSelectionChange}
selectedDocuments={selectedDocuments}
searchMode={searchMode}
onSearchModeChange={onSearchModeChange}
researchMode={researchMode}
onResearchModeChange={onResearchModeChange}
/>
</ChatInput>
);
@ -162,6 +293,10 @@ export default function ChatInterface({
handler,
onDocumentSelectionChange,
selectedDocuments = [],
searchMode,
onSearchModeChange,
researchMode,
onResearchModeChange,
}: ChatInterfaceProps) {
return (
<ChatSection handler={handler} className="flex h-full">
@ -177,6 +312,10 @@ export default function ChatInterface({
<CustomChatInput
onDocumentSelectionChange={onDocumentSelectionChange}
selectedDocuments={selectedDocuments}
searchMode={searchMode}
onSearchModeChange={onSearchModeChange}
researchMode={researchMode}
onResearchModeChange={onResearchModeChange}
/>
</div>
</div>

View file

@ -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<string | null> => {
async (
initialMessage: string,
researchMode: ResearchMode
): Promise<string | null> => {
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 {