Simplified chat creation logic. Created custom hooks and proper route handling for smooth chat handling

surfsense_web/app/dashboard/[search_space_id]/v2/[[...chat_id]]/page.tsx

@@ -0,0 +1,143 @@
+"use client";
+
+import { useChat, Message, CreateMessage } from "@ai-sdk/react";
+import { useParams, useRouter } from "next/navigation";
+import { useEffect } from "react";
+import ChatMain from "@/components/chat_v2/ChatMain";
+import { ResearchMode } from "@/components/chat";
+import { useChatState, useChatAPI } from "@/hooks/useChat";
+
+export default function ResearchChatPageV2() {
+    const { search_space_id, chat_id } = useParams();
+    const router = useRouter();
+
+    const chatIdParam = Array.isArray(chat_id) ? chat_id[0] : chat_id;
+    const isNewChat = !chatIdParam;
+
+    const {
+        token,
+        isLoading,
+        setIsLoading,
+        searchMode,
+        researchMode,
+        setResearchMode,
+        selectedConnectors,
+        setSelectedConnectors,
+    } = useChatState({
+        search_space_id: search_space_id as string,
+        chat_id: chatIdParam,
+    });
+
+    const { fetchChatDetails, updateChat, createChat } = useChatAPI({
+        token,
+        search_space_id: search_space_id as string,
+        researchMode,
+        selectedConnectors,
+    });
+
+    // Single useChat handler for both cases
+    const handler = useChat({
+        api: `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chat`,
+        streamProtocol: "data",
+        initialMessages: [],
+        headers: {
+            ...(token && { Authorization: `Bearer ${token}` }),
+        },
+        body: {
+            data: {
+                search_space_id: search_space_id,
+                selected_connectors: selectedConnectors,
+                research_mode: researchMode,
+                search_mode: searchMode,
+                document_ids_to_add_in_context: [],
+            },
+        },
+        onError: (error) => {
+            console.error("Chat error:", error);
+        },
+    });
+
+    const customHandlerAppend = async (
+        message: Message | CreateMessage,
+        chatRequestOptions?: { data?: any }
+    ) => {
+        const newChatId = await createChat(message.content);
+        router.replace(`/dashboard/${search_space_id}/v2/${newChatId}`);
+        return newChatId;
+    };
+
+    useEffect(() => {
+        if (token && !isNewChat && chatIdParam) {
+            setIsLoading(true);
+            loadChatData(chatIdParam);
+        }
+    }, [token, isNewChat, chatIdParam]);
+
+    const loadChatData = async (chatId: string) => {
+        try {
+            const chatData = await fetchChatDetails(chatId);
+            if (!chatData) return;
+
+            // Update configuration from chat data
+            if (chatData.type) {
+                setResearchMode(chatData.type as ResearchMode);
+            }
+
+            if (
+                chatData.initial_connectors &&
+                Array.isArray(chatData.initial_connectors)
+            ) {
+                setSelectedConnectors(chatData.initial_connectors);
+            }
+
+            // Load existing messages
+            if (chatData.messages && Array.isArray(chatData.messages)) {
+                if (
+                    chatData.messages.length === 1 &&
+                    chatData.messages[0].role === "user"
+                ) {
+                    // Single user message - append to trigger LLM response
+                    handler.append({
+                        role: "user",
+                        content: chatData.messages[0].content,
+                    });
+                } else if (chatData.messages.length > 1) {
+                    // Multiple messages - set them all
+                    handler.setMessages(chatData.messages);
+                }
+            }
+        } finally {
+            setIsLoading(false);
+        }
+    };
+
+    // Auto-update chat when messages change (only for existing chats)
+    useEffect(() => {
+        if (
+            !isNewChat &&
+            chatIdParam &&
+            handler.status === "ready" &&
+            handler.messages.length > 0 &&
+            handler.messages[handler.messages.length - 1]?.role === "assistant"
+        ) {
+            updateChat(chatIdParam, handler.messages);
+        }
+    }, [handler.messages, handler.status, chatIdParam, isNewChat, updateChat]);
+
+    if (isLoading) {
+        return (
+            <div className="flex items-center justify-center h-full">
+                <div>Loading...</div>
+            </div>
+        );
+    }
+
+    return (
+        <ChatMain
+            handler={{
+                ...handler,
+                append: isNewChat ? customHandlerAppend : handler.append,
+            }}
+        />
+    );
+}

surfsense_web/app/dashboard/[search_space_id]/v2/[chat_id]/page.tsx

@@ -1,193 +0,0 @@
-"use client";
-
-import { Message, useChat } from "@ai-sdk/react";
-import { useParams } from "next/navigation";
-import { useEffect, useState } from "react";
-import ChatMain from "@/components/chat_v2/ChatMain";
-import { ResearchMode } from "@/components/chat";
-
-export default function ResearcherChatPageV2() {
-    const { search_space_id, chat_id } = useParams();
-
-    const [token, setToken] = useState<string | null>(null);
-    const [isLoading, setIsLoading] = useState(false);
-
-    // const [initialMessages, setInitialMessages] = useState<any[]>([]);
-
-    const [searchMode, setSearchMode] = useState<"DOCUMENTS" | "CHUNKS">(
-        "DOCUMENTS"
-    );
-    const [researchMode, setResearchMode] = useState<ResearchMode>("QNA");
-    const [selectedConnectors, setSelectedConnectors] = useState<string[]>([]);
-
-    const handler = useChat({
-        api: `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chat`,
-        streamProtocol: "data",
-        initialMessages: [],
-        headers: {
-            ...(token && { Authorization: `Bearer ${token}` }),
-        },
-        body: {
-            data: {
-                search_space_id: search_space_id,
-                selected_connectors: selectedConnectors,
-                research_mode: researchMode,
-                search_mode: searchMode,
-                document_ids_to_add_in_context: [],
-            },
-        },
-        onError: (error) => {
-            console.error("Chat error:", error);
-            // You can add additional error handling here if needed
-        },
-    });
-
-    useEffect(() => {
-        setIsLoading(true);
-        let token = localStorage.getItem("surfsense_bearer_token");
-        if (token) {
-            setToken(token);
-            fetchChatDetails(token);
-            setIsLoading(false);
-        }
-    }, [chat_id]);
-
-    const fetchChatDetails = async (token: string) => {
-        try {
-            if (!token) return;
-
-            // console.log('Fetching chat details for chat ID:', chat_id);
-
-            const response = await fetch(
-                `${
-                    process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL
-                }/api/v1/chats/${Number(chat_id)}`,
-                {
-                    method: "GET",
-                    headers: {
-                        "Content-Type": "application/json",
-                        Authorization: `Bearer ${token}`,
-                    },
-                }
-            );
-
-            if (!response.ok) {
-                throw new Error(
-                    `Failed to fetch chat details: ${response.statusText}`
-                );
-            }
-
-            const chatData = await response.json();
-            // console.log('Chat details fetched:', chatData);
-
-            // Set research mode from chat data
-            if (chatData.type) {
-                setResearchMode(chatData.type as ResearchMode);
-            }
-
-            // Set connectors from chat data
-            if (
-                chatData.initial_connectors &&
-                Array.isArray(chatData.initial_connectors)
-            ) {
-                setSelectedConnectors(chatData.initial_connectors);
-            }
-
-            if (chatData.messages && Array.isArray(chatData.messages)) {
-                console.log("chatData.messages", chatData.messages);
-
-                if (
-                    chatData.messages.length === 1 &&
-                    chatData.messages[0].role === "user"
-                ) {
-                    console.log("appending");
-                    handler.append({
-                        role: "user",
-                        content: chatData.messages[0].content,
-                    });
-                } else {
-                    console.log("setting");
-                    handler.setMessages(chatData.messages);
-                }
-            }
-        } catch (err) {
-            console.error("Error fetching chat details:", err);
-        }
-    };
-
-    const updateChat = async (messages: Message[]) => {
-        try {
-            const token = localStorage.getItem("surfsense_bearer_token");
-            console.log("updating chat", messages, token);
-            if (!token) return;
-
-            // Find the first user message to use as title
-            const userMessages = handler.messages.filter(
-                (msg: any) => msg.role === "user"
-            );
-
-            console.log("userMessages", userMessages);
-            console.log("handler.messages", handler.messages);
-
-            if (userMessages.length === 0) return;
-
-            // Use the first user message as the title
-            const title = userMessages[0].content;
-
-            // console.log('Updating chat with title:', title);
-
-            // Update the chat
-            console.log("messages", messages);
-            const response = await fetch(
-                `${
-                    process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL
-                }/api/v1/chats/${Number(chat_id)}`,
-                {
-                    method: "PUT",
-                    headers: {
-                        "Content-Type": "application/json",
-                        Authorization: `Bearer ${token}`,
-                    },
-                    body: JSON.stringify({
-                        type: researchMode,
-                        title: title,
-                        initial_connectors: selectedConnectors,
-                        messages: messages,
-                        search_space_id: Number(search_space_id),
-                    }),
-                }
-            );
-
-            if (!response.ok) {
-                throw new Error(
-                    `Failed to update chat: ${response.statusText}`
-                );
-            }
-
-            // console.log('Chat updated successfully');
-        } catch (err) {
-            console.error("Error updating chat:", err);
-        }
-    };
-
-    useEffect(() => {
-        console.log("handler.messages", handler.messages, handler.status);
-        if (
-            handler.status === "ready" &&
-            handler.messages.length > 0 &&
-            handler.messages[handler.messages.length - 1]?.role === "assistant"
-        ) {
-            updateChat(handler.messages);
-        }
-    }, [handler.messages, handler.status]);
-
-    const handleQuerySubmit = (input: string, handleSubmit: () => void) => {
-        handleSubmit();
-    };
-
-    if (isLoading) {
-        return <div>Loading...</div>;
-    }
-
-    return <ChatMain handler={handler} handleQuerySubmit={handleQuerySubmit} />;
-}

surfsense_web/app/dashboard/[search_space_id]/v2/page.tsx

@@ -1,94 +0,0 @@
-"use client";
-
-import { useChat } from "@ai-sdk/react";
-
-import { useParams } from "next/navigation";
-import { useEffect, useState } from "react";
-import { useRouter } from "next/navigation";
-import ChatMain from "@/components/chat_v2/ChatMain";
-
-export default function ResearcherPageV2() {
-    const { search_space_id, chat_id } = useParams();
-    const router = useRouter();
-
-    const [token, setToken] = useState<string | null>(null);
-
-    useEffect(() => {
-        setToken(localStorage.getItem("surfsense_bearer_token"));
-    }, []);
-
-    const handleQuerySubmit = (input: string, handleSubmit: () => void) => {
-        const createChat = async () => {
-            try {
-                if (!token) {
-                    console.error("Authentication token not found");
-                    return;
-                }
-
-                // Create a new chat
-                const response = await fetch(
-                    `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chats/`,
-                    {
-                        method: "POST",
-                        headers: {
-                            "Content-Type": "application/json",
-                            Authorization: `Bearer ${token}`,
-                        },
-                        body: JSON.stringify({
-                            type: "QNA",
-                            title: "Untitled Chat", // Empty title initially
-                            initial_connectors: [], // No default connectors
-                            messages: [
-                                {
-                                    role: "user",
-                                    content: input,
-                                },
-                            ],
-                            search_space_id: Number(search_space_id),
-                        }),
-                    }
-                );
-
-                if (!response.ok) {
-                    throw new Error(
-                        `Failed to create chat: ${response.statusText}`
-                    );
-                }
-
-                const data = await response.json();
-
-                router.replace(`/dashboard/${search_space_id}/v2/${data.id}`);
-            } catch (err) {
-                console.error("Error creating chat:", err);
-            }
-        };
-
-        if (!chat_id) {
-            createChat();
-            return;
-        }
-    };
-
-    const handler = useChat({
-        api: `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chat`,
-        streamProtocol: "data",
-        headers: {
-            ...(token && { Authorization: `Bearer ${token}` }),
-        },
-        body: {
-            data: {
-                search_space_id: search_space_id,
-                selected_connectors: [],
-                research_mode: "QNA",
-                search_mode: "DOCUMENTS",
-                document_ids_to_add_in_context: [],
-            },
-        },
-        onError: (error) => {
-            console.error("Chat error:", error);
-            // You can add additional error handling here if needed
-        },
-    });
-
-    return <ChatMain handler={handler} handleQuerySubmit={handleQuerySubmit} />;
-}

surfsense_web/components/chat_v2/ChatMain.tsx

@@ -1,75 +1,11 @@
 "use client";
 
-import {
-    ChatCanvas,
-    ChatMessages,
-    ChatSection,
-    useChatUI,
-    ChatHandler,
-} from "@llamaindex/chat-ui";
-import { Textarea } from "@/components/ui/textarea";
-import { Button } from "@/components/ui/button";
-import { Loader2 } from "lucide-react";
-import { useEffect } from "react";
+import { ChatSection, ChatHandler } from "@llamaindex/chat-ui";
 
 interface ChatMainProps {
     handler: ChatHandler;
-    handleQuerySubmit: (input: string, handleSubmit: () => void) => void;
 }
 
-const ChatInput = (props: {
-    handleQuerySubmit: (input: string, handleSubmit: () => void) => void;
-}) => {
-    const { input, setInput, handleSubmit } = useChatUI();
-    const { handleQuerySubmit } = props;
-
-    const handleFormSubmit = (e: React.FormEvent<HTMLFormElement>) => {
-        e.preventDefault();
-        if (!input.trim()) return;
-
-        handleQuerySubmit(input, handleSubmit);
-    };
-
-    return (
-        <form
-            className="flex flex-row items-center justify-between gap-2"
-            onSubmit={handleFormSubmit}
-        >
-            <Textarea
-                value={input}
-                onChange={(e) => setInput(e.target.value)}
-                placeholder="Type your message here..."
-                rows={4}
-                className="max-h-[150px] overflow-y-auto"
-            />
-
-            <Button type="submit">Send</Button>
-        </form>
-    );
-};
-
-export default function ChatMain({
-    handler,
-    handleQuerySubmit,
-}: ChatMainProps) {
-    return (
-        <ChatSection handler={handler} className="flex h-full">
-            <div className="flex flex-1 flex-col">
-                <ChatMessages className="flex-1">
-                    <ChatMessages.List className="p-4">
-                        {/* Custom message rendering */}
-                    </ChatMessages.List>
-                    <ChatMessages.Loading>
-                        <Loader2 className="animate-spin" />
-                    </ChatMessages.Loading>
-                </ChatMessages>
-
-                <div className="border-t p-4">
-                    <ChatInput handleQuerySubmit={handleQuerySubmit} />
-                </div>
-            </div>
-
-            <ChatCanvas className="w-1/2 border-l" />
-        </ChatSection>
-    );
+export default function ChatMain({ handler }: ChatMainProps) {
+    return <ChatSection handler={handler} className="flex h-full" />;
 }

surfsense_web/hooks/useChat.ts

@@ -0,0 +1,162 @@
+import { useState, useEffect, useCallback } from "react";
+import { Message } from "@ai-sdk/react";
+import { ResearchMode } from "@/components/chat";
+
+interface UseChatStateProps {
+    search_space_id: string;
+    chat_id?: string;
+}
+
+export function useChatState({ search_space_id, chat_id }: UseChatStateProps) {
+    const [token, setToken] = useState<string | null>(null);
+    const [isLoading, setIsLoading] = useState(false);
+    const [currentChatId, setCurrentChatId] = useState<string | null>(chat_id || null);
+    
+    // Chat configuration state
+    const [searchMode, setSearchMode] = useState<"DOCUMENTS" | "CHUNKS">("DOCUMENTS");
+    const [researchMode, setResearchMode] = useState<ResearchMode>("QNA");
+    const [selectedConnectors, setSelectedConnectors] = useState<string[]>([]);
+
+    useEffect(() => {
+        const bearerToken = localStorage.getItem("surfsense_bearer_token");
+        setToken(bearerToken);
+    }, []);
+
+    return {
+        token,
+        setToken,
+        isLoading,
+        setIsLoading,
+        currentChatId,
+        setCurrentChatId,
+        searchMode,
+        setSearchMode,
+        researchMode,
+        setResearchMode,
+        selectedConnectors,
+        setSelectedConnectors,
+    };
+}
+
+interface UseChatAPIProps {
+    token: string | null;
+    search_space_id: string;
+    researchMode: ResearchMode;
+    selectedConnectors: string[];
+}
+
+export function useChatAPI({
+    token,
+    search_space_id,
+    researchMode,
+    selectedConnectors,
+}: UseChatAPIProps) {
+    const fetchChatDetails = useCallback(async (chatId: string) => {
+        if (!token) return null;
+
+        try {
+            const response = await fetch(
+                `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chats/${Number(chatId)}`,
+                {
+                    method: "GET",
+                    headers: {
+                        "Content-Type": "application/json",
+                        Authorization: `Bearer ${token}`,
+                    },
+                }
+            );
+
+            if (!response.ok) {
+                throw new Error(`Failed to fetch chat details: ${response.statusText}`);
+            }
+
+            return await response.json();
+        } catch (err) {
+            console.error("Error fetching chat details:", err);
+            return null;
+        }
+    }, [token]);
+
+    const createChat = useCallback(async (initialMessage: string): Promise<string | null> => {
+        if (!token) {
+            console.error("Authentication token not found");
+            return null;
+        }
+
+        try {
+            const response = await fetch(
+                `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chats/`,
+                {
+                    method: "POST",
+                    headers: {
+                        "Content-Type": "application/json",
+                        Authorization: `Bearer ${token}`,
+                    },
+                    body: JSON.stringify({
+                        type: researchMode,
+                        title: "Untitled Chat",
+                        initial_connectors: selectedConnectors,
+                        messages: [
+                            {
+                                role: "user",
+                                content: initialMessage,
+                            },
+                        ],
+                        search_space_id: Number(search_space_id),
+                    }),
+                }
+            );
+
+            if (!response.ok) {
+                throw new Error(`Failed to create chat: ${response.statusText}`);
+            }
+
+            const data = await response.json();
+            return data.id;
+        } catch (err) {
+            console.error("Error creating chat:", err);
+            return null;
+        }
+    }, [token, researchMode, selectedConnectors, search_space_id]);
+
+    const updateChat = useCallback(async (chatId: string, messages: Message[]) => {
+        if (!token) return;
+
+        try {
+            const userMessages = messages.filter(msg => msg.role === "user");
+            if (userMessages.length === 0) return;
+
+            const title = userMessages[0].content;
+
+            const response = await fetch(
+                `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/chats/${Number(chatId)}`,
+                {
+                    method: "PUT",
+                    headers: {
+                        "Content-Type": "application/json",
+                        Authorization: `Bearer ${token}`,
+                    },
+                    body: JSON.stringify({
+                        type: researchMode,
+                        title: title,
+                        initial_connectors: selectedConnectors,
+                        messages: messages,
+                        search_space_id: Number(search_space_id),
+                    }),
+                }
+            );
+
+            if (!response.ok) {
+                throw new Error(`Failed to update chat: ${response.statusText}`);
+            }
+        } catch (err) {
+            console.error("Error updating chat:", err);
+        }
+    }, [token, researchMode, selectedConnectors, search_space_id]);
+
+    return {
+        fetchChatDetails,
+        createChat,
+        updateChat,
+    };
+}