From 888455b56a4ef55b600c08b5ed61707976775749 Mon Sep 17 00:00:00 2001
From: Vinlic <vinlic@vinlic.com>
Date: Thu, 21 Mar 2024 11:51:26 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E4=BB=A3=E7=A0=81=E8=B0=83?=
 =?UTF-8?q?=E7=94=A8=E8=BE=93=E5=87=BA=E6=95=88=E6=9E=9C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/api/controllers/chat.ts | 56 +++++++++++++++++++++++--------------
 1 file changed, 35 insertions(+), 21 deletions(-)

diff --git a/src/api/controllers/chat.ts b/src/api/controllers/chat.ts
index eb88a0e..4fc43fb 100644
--- a/src/api/controllers/chat.ts
+++ b/src/api/controllers/chat.ts
@@ -479,7 +479,8 @@ async function receiveStream(stream: any): Promise<any> {
       created: util.unixTimestamp()
     };
     let toolCall = false;
-    // let codeGenerating = false;
+    let codeGenerating = false;
+    let codeTemp = '';
     let lastExecutionOutput = '';
     let textOffset = 0;
     const parser = createParser(event => {
@@ -518,18 +519,24 @@ async function receiveStream(stream: any): Promise<any> {
                 toolCall = true;
                 return innerStr + imageText;
               }
-              // else if(type == 'code' && partStatus == 'init' && !codeGenerating) {
-              //   codeGenerating = true;
-              //   const label = '代码生成中...\n';
-              //   textOffset += label.length;
-              //   return innerStr + label;
-              // }
-              // else if(type == 'code' && partStatus == 'finish' && codeGenerating) {
-              //   codeGenerating = false;
-              //   const label = '代码执行中...\n';
-              //   textOffset += label.length;
-              //   return innerStr + label;
-              // }
+              else if(type == 'code' && partStatus == 'init') {
+                let codeHead = '';
+                if(!codeGenerating) {
+                  codeGenerating = true;
+                  codeHead = '```python\n';
+                }
+                const chunk = code.substring(codeTemp.length, code.length);
+                codeTemp += chunk;
+                textOffset += codeHead.length + chunk.length;
+                return innerStr + codeHead + chunk;
+              }
+              else if(type == 'code' && partStatus == 'finish' && codeGenerating) {
+                const codeFooter = '\n```';
+                codeGenerating = false;
+                codeTemp = '';
+                textOffset += codeFooter.length;
+                return innerStr + codeFooter;
+              }
               else if(type == 'execution_output' && _.isString(content) && partStatus == 'done' && lastExecutionOutput != content) {
                 lastExecutionOutput = content;
                 const _content = content.replace(/^\n/, '');
@@ -576,6 +583,7 @@ function createTransStream(stream: any, endCallback?: Function) {
   let content = '';
   let toolCall = false;
   let codeGenerating = false;
+  let codeTemp = '';
   let lastExecutionOutput = '';
   let textOffset = 0;
   !transStream.closed && transStream.write(`data: ${JSON.stringify({
@@ -621,17 +629,23 @@ function createTransStream(stream: any, endCallback?: Function) {
               toolCall = true;
               return innerStr + imageText;
             }
-            else if(type == 'code' && partStatus == 'init' && !codeGenerating) {
-              codeGenerating = true;
-              const label = '代码生成中...\n';
-              textOffset += label.length;
-              return innerStr + label;
+            else if(type == 'code' && partStatus == 'init') {
+              let codeHead = '';
+              if(!codeGenerating) {
+                codeGenerating = true;
+                codeHead = '```python\n';
+              }
+              const chunk = code.substring(codeTemp.length, code.length);
+              codeTemp += chunk;
+              textOffset += codeHead.length + chunk.length;
+              return innerStr + codeHead + chunk;
             }
             else if(type == 'code' && partStatus == 'finish' && codeGenerating) {
+              const codeFooter = '\n```';
               codeGenerating = false;
-              const label = '代码执行中...\n';
-              textOffset += label.length;
-              return innerStr + label;
+              codeTemp = '';
+              textOffset += codeFooter.length;
+              return innerStr + codeFooter;
             }
             else if(type == 'execution_output' && _.isString(content) && partStatus == 'done' && lastExecutionOutput != content) {
               lastExecutionOutput = content;