fix: chat thinking render (#3544)

opensumi · Apr 16, 2024 · fcb8182 · fcb8182
1 parent 182e936
commit fcb8182
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 16 deletions.
diff --git a/packages/ai-native/src/browser/chat/chat.view.tsx b/packages/ai-native/src/browser/chat/chat.view.tsx
@@ -464,14 +464,19 @@ export const AIChatView = observer(() => {
 
   const handleReply = React.useCallback(
     (userInput: { type: AISerivceType; message: string }, replayCommandProps: IReplayComponentParam) => {
+      const { relationId } = replayCommandProps;
+
       if (chatRenderRegistry.chatAIRoleRender) {
         replayCommandProps.renderContent = (content: string, status: EMsgStreamStatus) =>
           chatRenderRegistry.chatAIRoleRender!({ content, status });
       }
 
+      aiChatService.setLatestSessionId(relationId);
+      aiChatService.messageWithStream(userInput.message, {}, relationId);
+
       const aiMessage = createMessageByAI({
         id: uuid(6),
-        relationId: replayCommandProps.relationId,
+        relationId,
         text: <StreamReplyRender prompt={userInput.message} params={replayCommandProps} />,
         className: styles.chat_with_more_actions,
       });
@@ -484,7 +489,7 @@ export const AIChatView = observer(() => {
       }
       setLoading(false);
     },
-    [messageListData, chatRenderRegistry, chatRenderRegistry.chatAIRoleRender],
+    [messageListData, aiChatService, chatRenderRegistry, chatRenderRegistry.chatAIRoleRender],
   );
 
   const handleClear = React.useCallback(() => {

diff --git a/packages/ai-native/src/browser/components/ChatThinking.tsx b/packages/ai-native/src/browser/components/ChatThinking.tsx
@@ -123,7 +123,7 @@ export const ChatThinkingResult = ({
     }
 
     return children;
-  }, [status, message, children]);
+  }, [status, message, hasMessage, children]);
 
   const isRenderRegenerate = useMemo(() => {
     if (isUndefined(showRegenerate)) {

diff --git a/packages/ai-native/src/browser/components/StreamMsg.tsx b/packages/ai-native/src/browser/components/StreamMsg.tsx
@@ -1,5 +1,5 @@
 import hljs from 'highlight.js';
-import React, { useCallback, useEffect } from 'react';
+import React, { useCallback, useEffect, useMemo } from 'react';
 
 import { DisposableCollection, useInjectable } from '@opensumi/ide-core-browser';
 import { IAIReporter, localize } from '@opensumi/ide-core-common';
@@ -20,6 +20,7 @@ interface IStreamMsgWrapperProps {
 
 export const StreamMsgWrapper = (props: IStreamMsgWrapperProps) => {
   const { sessionId, prompt, startTime = 0, onRegenerate, renderContent } = props;
+
   const [chunk, setChunk] = React.useState('');
   const [content, setContent] = React.useState<string>('');
   const [isError, setIsError] = React.useState<boolean>(false);
@@ -112,11 +113,17 @@ export const StreamMsgWrapper = (props: IStreamMsgWrapperProps) => {
     props.onStop?.();
   };
 
-  return status === EMsgStreamStatus.THINKING && msgStreamManager.currentSessionId === sessionId ? (
-    <ChatThinking status={status} message={content} onStop={onStop}>
-      {renderMsgList()}
-    </ChatThinking>
-  ) : (
+  const isThinking = useMemo(() => status === EMsgStreamStatus.THINKING && msgStreamManager.currentSessionId === sessionId, [status, sessionId, msgStreamManager.currentSessionId]);
+
+  if (isThinking) {
+    return (
+      <ChatThinking status={status} message={content} onStop={onStop}>
+        {renderMsgList()}
+      </ChatThinking>
+    );
+  }
+
+  return (
     <ChatThinkingResult status={status} message={content} onRegenerate={handleRegenerate} sessionId={sessionId}>
       {renderMsgList()}
     </ChatThinkingResult>

diff --git a/packages/ai-native/src/browser/components/StreamReplyRender.tsx b/packages/ai-native/src/browser/components/StreamReplyRender.tsx
@@ -17,23 +17,18 @@ export interface IReplayComponentParam {
 // 流式输出渲染组件
 export const StreamReplyRender = (props: { prompt: string; params: IReplayComponentParam }) => {
   const { prompt, params } = props;
+  const { relationId, renderContent } = params;
 
   const aiChatService = useInjectable<ChatInternalService>(IChatInternalService);
 
-  const { relationId, renderContent } = params;
-
   const send = React.useCallback(
     (isRetry = false) => {
       aiChatService.setLatestSessionId(relationId);
       aiChatService.messageWithStream(prompt, isRetry ? { enableGptCache: false } : {}, relationId);
     },
-    [aiChatService],
+    [aiChatService, relationId, prompt],
   );
 
-  React.useEffect(() => {
-    send();
-  }, []);
-
   return (
     <StreamMsgWrapper
       sessionId={relationId}