Fix real-world streaming

Author: lukemelia

packages/ai-bot/lib/responder.ts

@@ -75,10 +75,10 @@ export class Responder {
   ) {
     const toolCallsSnapshot = snapshot.choices[0].message.tool_calls;
     if (toolCallsSnapshot?.length) {
-      if (
-        JSON.stringify(this.toolCalls) !== JSON.stringify(toolCallsSnapshot)
-      ) {
+      let latestToolCallsJson = JSON.stringify(toolCallsSnapshot);
+      if (this.toolCallsJson !== latestToolCallsJson) {
         this.toolCalls = toolCallsSnapshot;
+        this.toolCallsJson = latestToolCallsJson;
         await this.sendMessageEventWithThrottling();
       }
     }
@@ -134,7 +134,14 @@ export class Responder {
       result['name'] = f.name;
     }
     if (f.arguments) {
-      result['arguments'] = JSON.parse(f.arguments);
+      try {
+        result['arguments'] = JSON.parse(f.arguments);
+      } catch (error) {
+        // If the arguments are not valid JSON, we'll just return an empty object
+        // This will happen during streaming, when the tool call is not yet complete
+        // and the arguments are not yet available
+        result['arguments'] = {};
+      }
     }
     return result;
   }
@@ -162,6 +169,10 @@ export class Responder {
   }
 
   async finalize() {
+    if (!this.isStreamingFinished) {
+      this.isStreamingFinished = true;
+      await this.sendMessageEventWithThrottling();
+    }
     await this.flush();
   }
 }

packages/ai-bot/main.ts

@@ -212,7 +212,7 @@ Common issues are:
         let eventList = (initial!.messages?.chunk ||
           []) as DiscreteMatrixEvent[];
         try {
-          promptParts = getPromptParts(eventList, aiBotUserId);
+          promptParts = await getPromptParts(eventList, aiBotUserId);
           if (!promptParts.shouldRespond) {
             return;
           }

packages/ai-bot/package.json

@@ -11,7 +11,7 @@
     "@types/stream-json": "^1.7.3",
     "lodash": "~4.17.21",
     "matrix-js-sdk": "^31.0.0",
-    "openai": "4.47.1",
+    "openai": "4.86.1",
     "qunit": "^2.18.0",
     "stream-chain": "^2.2.5",
     "stream-json": "^1.8.0",

packages/ai-bot/tests/responding-test.ts

@@ -354,7 +354,7 @@ module('Responding', (hooks) => {
     let sentEvents = fakeMatrixClient.getSentEvents();
     assert.equal(
       sentEvents.length,
-      4,
+      5,
       'Thinking message, and event with content, event with partial tool call, and event with full tool call should be sent',
     );
     assert.equal(
@@ -424,6 +424,16 @@ module('Responding', (hooks) => {
       },
       'The replacement event with the tool call event should replace the original message',
     );
+    assert.deepEqual(
+      sentEvents[3].content.isStreamingFinished,
+      false,
+      'The tool call event should not be sent with isStreamingFinished set to true',
+    );
+    assert.deepEqual(
+      sentEvents[4].content.isStreamingFinished,
+      true,
+      'The final event should be sent with isStreamingFinished set to true',
+    );
   });
 
   test('Handles multiple tool calls', async () => {