From 47e3cf82e09279efc498325abdab7d42db8b0bd5 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Sat, 5 Oct 2024 19:08:57 +0300
Subject: [PATCH] Don't call tools on stopped streams. Emit events

---
 public/script.js               | 12 ++++++++----
 public/scripts/tool-calling.js |  7 +++++--
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/public/script.js b/public/script.js
index d2e96c756..f2176fe39 100644
--- a/public/script.js
+++ b/public/script.js
@@ -467,6 +467,8 @@ export const event_types = {
     ONLINE_STATUS_CHANGED: 'online_status_changed',
     IMAGE_SWIPED: 'image_swiped',
     CONNECTION_PROFILE_LOADED: 'connection_profile_loaded',
+    TOOL_CALLS_PERFORMED: 'tool_calls_performed',
+    TOOL_CALLS_RENDERED: 'tool_calls_rendered',
 };
 
 export const eventSource = new EventEmitter();
@@ -4421,7 +4423,9 @@ export async function Generate(type, { automatic_trigger, force_name2, quiet_pro
                 getMessage = continue_mag + getMessage;
             }
 
-            if (canPerformToolCalls && Array.isArray(streamingProcessor.toolCalls) && streamingProcessor.toolCalls.length) {
+            const isStreamFinished = streamingProcessor && !streamingProcessor.isStopped && streamingProcessor.isFinished;
+            const isStreamWithToolCalls = streamingProcessor && Array.isArray(streamingProcessor.toolCalls) && streamingProcessor.toolCalls.length;
+            if (canPerformToolCalls && isStreamFinished && isStreamWithToolCalls) {
                 const lastMessage = chat[chat.length - 1];
                 const hasToolCalls = ToolManager.hasToolCalls(streamingProcessor.toolCalls);
                 const shouldDeleteMessage = ['', '...'].includes(lastMessage?.mes) && ['', '...'].includes(streamingProcessor?.result);
@@ -4437,12 +4441,12 @@ export async function Generate(type, { automatic_trigger, force_name2, quiet_pro
                     }
 
                     streamingProcessor = null;
-                    ToolManager.saveFunctionToolInvocations(invocationResult.invocations);
+                    await ToolManager.saveFunctionToolInvocations(invocationResult.invocations);
                     return Generate(type, { automatic_trigger, force_name2, quiet_prompt, quietToLoud, skipWIAN, force_chid, signal, quietImage, quietName }, dryRun);
                 }
             }
 
-            if (streamingProcessor && !streamingProcessor.isStopped && streamingProcessor.isFinished) {
+            if (isStreamFinished) {
                 await streamingProcessor.onFinishStreaming(streamingProcessor.messageId, getMessage);
                 streamingProcessor = null;
                 triggerAutoContinue(messageChunk, isImpersonate);
@@ -4529,7 +4533,7 @@ export async function Generate(type, { automatic_trigger, force_name2, quiet_pro
                     return;
                 }
 
-                ToolManager.saveFunctionToolInvocations(invocationResult.invocations);
+                await ToolManager.saveFunctionToolInvocations(invocationResult.invocations);
                 return Generate(type, { automatic_trigger, force_name2, quiet_prompt, quietToLoud, skipWIAN, force_chid, signal, quietImage, quietName }, dryRun);
             }
         }
diff --git a/public/scripts/tool-calling.js b/public/scripts/tool-calling.js
index a8c026242..ffef21071 100644
--- a/public/scripts/tool-calling.js
+++ b/public/scripts/tool-calling.js
@@ -1,4 +1,4 @@
-import { addOneMessage, chat, main_api, system_avatar, systemUserName } from '../script.js';
+import { addOneMessage, chat, event_types, eventSource, main_api, saveChatConditional, system_avatar, systemUserName } from '../script.js';
 import { chat_completion_sources, oai_settings } from './openai.js';
 import { Popup } from './popup.js';
 
@@ -542,7 +542,7 @@ export class ToolManager {
      * Saves function tool invocations to the last user chat message extra metadata.
      * @param {ToolInvocation[]} invocations Successful tool invocations
      */
-    static saveFunctionToolInvocations(invocations) {
+    static async saveFunctionToolInvocations(invocations) {
         if (!Array.isArray(invocations) || invocations.length === 0) {
             return;
         }
@@ -558,7 +558,10 @@ export class ToolManager {
             },
         };
         chat.push(message);
+        await eventSource.emit(event_types.TOOL_CALLS_PERFORMED, invocations);
         addOneMessage(message);
+        await eventSource.emit(event_types.TOOL_CALLS_RENDERED, invocations);
+        await saveChatConditional();
     }
 
     /**