feat: add search tool to ChatGLM-4

ai-shifu · Mar 16, 2024 · d10f71c · d10f71c
1 parent 478ff62
commit d10f71c
Showing 1 changed file with 28 additions and 4 deletions.
diff --git a/src/bots/zhipu/ChatGLM4Bot.js b/src/bots/zhipu/ChatGLM4Bot.js
@@ -51,20 +51,44 @@ export default class ChatGLM4Bot extends ChatGLMBot {
             withCredentials: true,
           },
         );
+
+        let beginning = "";
         source.addEventListener("message", (event) => {
+          let body = "";
+          let ending = "";
           let data = JSON.parse(event.data);
 
           if (!context.conversation_id && data.conversation_id) {
             context.conversation_id = data.conversation_id;
             this.setChatContext(context);
           }
 
-          const text = data.parts?.[0]?.content?.[0]?.text || "";
+          const response = data.parts?.[0];
+          if (!response || response.role !== "assistant") return;
+          const content = response.content[0];
+          if (!content) return;
+
+          // Parse search tool calls
+          if (content.type === "tool_calls" && response.status === "init") {
+            if (content.tool_calls.name === "browser") {
+              const info = content.tool_calls.arguments;
+              if (info.startsWith("search")) {
+                beginning += `> ${info}\n`;
+              }
+            }
+          } else if (content.type === "text") {
+            body = content.text;
+            response.meta_data?.citations?.forEach((citation) => {
+              ending += `> 1. [${citation.metadata.title}](${citation.metadata.url})\n`;
+            });
+          }
+
+          const done = data.status === "finish";
           onUpdateResponse(callbackParam, {
-            content: text,
-            done: data.status === "finish",
+            content: `${beginning}\n${body}\n${ending}`,
+            done,
           });
-          resolve();
+          done && resolve();
         });
         source.stream();
       } catch (err) {