Function: Finish Chat Branching

Author: ShenKSPZ

src/components/chat/ChatMessageArea.tsx

@@ -1,8 +1,9 @@
 import React, { useState, FormEvent, useRef, useEffect } from 'react';
-import { Conversation } from '../../types/chat';
+import { Conversation, Message } from '../../types/chat';
 import { Send, Square, Copy, RotateCcw, Share2, Pencil, Loader2 } from 'lucide-react';
 import MarkdownContent from './MarkdownContent';
 import MessageToolboxMenu, { ToolboxAction } from '../ui/MessageToolboxMenu';
+import { MessageHelper } from '../../services/message-helper';
 
 interface ChatMessageAreaProps {
   activeConversation: Conversation | null;
@@ -30,12 +31,19 @@ export const ChatMessageArea: React.FC<ChatMessageAreaProps> = ({
   const [hoveredMessageId, setHoveredMessageId] = useState<string | null>(null);
   const [editingMessageId, setEditingMessageId] = useState<string | null>(null);
   const [editingContent, setEditingContent] = useState('');
+  const [messagesList, setMessagesList] = useState<Message[]>([]);
 
   // Scroll to bottom when messages change
   useEffect(() => {
     setTimeout(() => {
       messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
     }, 50);
+
+    console.log(activeConversation);
+
+    if(activeConversation) {
+      setMessagesList(MessageHelper.mapMessagesTreeToList(activeConversation));
+    }
   }, [activeConversation?.messages]);
 
   const handleSubmit = (e: FormEvent) => {
@@ -70,23 +78,16 @@ export const ChatMessageArea: React.FC<ChatMessageAreaProps> = ({
     setEditingContent(content);
   };
 
-  // Handle save edit
-  const handleSaveEdit = () => {
-    if (editingMessageId && onEditMessage && editingContent.trim()) {
-      onEditMessage(editingMessageId, editingContent);
-      setEditingMessageId(null);
-      setEditingContent('');
-    }
-  };
-
   // Handle send edited message
   const handleSendEditedMessage = () => {
-    if (editingMessageId && editingContent.trim()) {
+    if (editingMessageId && onEditMessage && editingContent.trim()) {
+      const newContent = editingContent;
+
       // First cancel the edit mode
       setEditingMessageId(null);
       setEditingContent('');
 
-      onSendMessage(input);
+      onEditMessage(editingMessageId, newContent);
     }
   };
 
@@ -109,6 +110,10 @@ export const ChatMessageArea: React.FC<ChatMessageAreaProps> = ({
     console.error(`Function not implemented yet: ${action}`);
   };
 
+  const getMessagesList = () => {
+    return messagesList;
+  }
+
   // If no active conversation is selected
   if (!activeConversation) {
     return (
@@ -125,7 +130,7 @@ export const ChatMessageArea: React.FC<ChatMessageAreaProps> = ({
     <div className="flex flex-col h-full">
       {/* Messages area */}
       <div className="flex-1 p-4 space-y-4 overflow-y-auto">
-        {activeConversation.messages.filter(m => m.role !== 'system').map((message) => {
+        {getMessagesList().map((message) => {
           const isUserMessage = message.role === 'user';
           const isEditing = editingMessageId === message.messageId;
 

src/pages/ChatPage.tsx

@@ -171,11 +171,11 @@ export const ChatPage: React.FC<ChatPageProps> = ({
     } catch (err) {
       console.error('Error sending streaming message:', err);
 
-      // If streaming fails, we'll try to fall back to regular mode
-      const error = err as Error;
-      if (error.message && error.message.includes('does not support streaming')) {
-        await handleSendMessage(content);
-      }
+      // // If streaming fails, we'll try to fall back to regular mode
+      // const error = err as Error;
+      // if (error.message && error.message.includes('does not support streaming')) {
+      //   await handleSendMessage(content);
+      // }
     }
   };
 
@@ -264,37 +264,33 @@ export const ChatPage: React.FC<ChatPageProps> = ({
 
   // Handle editing a message
   const handleEditMessage = async (messageId: string, newContent: string) => {
-    if (!isServiceInitialized || !chatServiceRef.current || !activeConversation) return;
+    if (!activeConversationId || !isServiceInitialized || !chatServiceRef.current) return;
 
     try {
-      // Find the message being edited
-      const message = activeConversation.messages.find(m => m.messageId === messageId);
-      
-      if (message && message.role === 'user') {
-        // Delete this message and all subsequent messages
-        const messageIndex = activeConversation.messages.findIndex(m => m.messageId === messageId);
-        
-        // If this is not the last message, we need to delete all subsequent messages
-        if (messageIndex < activeConversation.messages.length - 1) {
-          // Delete from last to first to avoid index shifting issues
-          for (let i = activeConversation.messages.length - 1; i > messageIndex; i--) {
-            const msgToDelete = activeConversation.messages[i];
-            await chatServiceRef.current.deleteMessage(msgToDelete.messageId, (updatedConversation) => {
-              setConversations(updatedConversation);
-            });
-          }
-        }
-        
-        // Then delete the edited message itself
-        await chatServiceRef.current.deleteMessage(messageId, (updatedConversation) => {
+      const chatService = chatServiceRef.current;
+      const selectedModel = SettingsService.getInstance().getSelectedModel();
+      const selectedProvider = SettingsService.getInstance().getSelectedProvider();
+      console.log('Using streaming with provider:', selectedProvider);
+      console.log('Using streaming with model:', selectedModel);
+
+      // Send user message with streaming
+      await chatService.editMessage(
+        messageId, 
+        activeConversationId,
+        newContent,
+        true,
+        (updatedConversation) => {
           setConversations(updatedConversation);
-        });
-        
-        // Finally, send the new message
-        await handleSendMessage(newContent);
-      }
-    } catch (error) {
-      console.error('Error editing message:', error);
+        }
+      );
+      
+    } catch (err) {
+      console.error('Error editing message:', err);
+      // // If streaming fails, we'll try to fall back to regular mode
+      // const error = err as Error;
+      // if (error.message && error.message.includes('does not support streaming')) {
+      //   await handleSendMessage(content);
+      // }
     }
   };
 

src/services/ai-service.ts

@@ -71,7 +71,9 @@ export class AIService {
     const settings = settingsService.getSettings();
     for (const provider of Object.keys(settings.providers)) {
       const providerSettings = settings.providers[provider];
-      console.log('Provider: ', provider, ' Provider settings: ', providerSettings);
+
+      // console.log('Provider: ', provider, ' Provider settings: ', providerSettings);
+
       if (providerSettings && providerSettings.apiKey && providerSettings.apiKey.length > 0) {
         const providerInstance = ProviderFactory.getNewProvider(provider as AIProvider);
         if (providerInstance) {