Add support for streams

Author: timokoessler

library/sinks/AiSDK.test.ts

@@ -5,6 +5,7 @@ import { runWithContext, type Context } from "../agent/Context";
 import { getMajorNodeVersion } from "../helpers/getNodeVersion";
 import { getInstance } from "../agent/AgentSingleton";
 import { z } from "zod";
+import { setTimeout } from "timers/promises";
 
 t.test(
   "It works",
@@ -37,7 +38,7 @@ t.test(
 
     const { google } =
       require("@ai-sdk/google") as typeof import("@ai-sdk/google");
-    const { generateText, generateObject } =
+    const { generateText, generateObject, streamText, streamObject } =
       require("ai") as typeof import("ai");
 
     await runWithContext(getTestContext(), async () => {
@@ -73,6 +74,8 @@ t.test(
         "Total tokens should match input + output"
       );
 
+      await setTimeout(400);
+
       const resultObj = await generateObject({
         model: google("models/gemini-2.0-flash-lite"),
         prompt: "Return numbers one to five",
@@ -91,6 +94,74 @@ t.test(
           },
         },
       ]);
+
+      await setTimeout(400);
+
+      const stream = streamText({
+        model: google("models/gemini-2.0-flash"),
+        prompt: "What is Zen by Aikido Security? Return one sentence.",
+      });
+
+      let streamedText = "";
+      for await (const chunk of stream.textStream) {
+        streamedText += chunk;
+      }
+
+      t.ok(streamedText.length > 0, "Streamed text should not be empty");
+
+      t.match(agent.getAIStatistics().getStats(), [
+        {
+          provider: "gemini",
+          model: "gemini-2.0-flash-lite",
+          calls: 2,
+          tokens: {
+            input: 23,
+          },
+        },
+        {
+          provider: "gemini",
+          model: "gemini-2.0-flash",
+          calls: 1,
+          tokens: {
+            input: 12,
+          },
+        },
+      ]);
+
+      await setTimeout(400);
+
+      const objectStream = streamObject({
+        model: google("models/gemini-2.0-flash"),
+        prompt: "Return numbers one to five",
+        output: "array",
+        schema: z.array(z.number()),
+      });
+
+      const streamedObject = [];
+      for await (const chunk of objectStream.elementStream) {
+        streamedObject.push(chunk);
+      }
+
+      t.same(streamedObject, [[1], [2], [3], [4], [5]]);
+
+      t.match(agent.getAIStatistics().getStats(), [
+        {
+          provider: "gemini",
+          model: "gemini-2.0-flash-lite",
+          calls: 2,
+          tokens: {
+            input: 23,
+          },
+        },
+        {
+          provider: "gemini",
+          model: "gemini-2.0-flash",
+          calls: 2,
+          tokens: {
+            input: 23,
+          },
+        },
+      ]);
     });
   }
 );

library/sinks/AiSDK.ts

@@ -111,8 +111,8 @@ export class AiSDK implements Wrapper {
     return modelName;
   }
 
-  wrap(hooks: Hooks) {
-    const interceptors: InterceptorObject = {
+  private getInterceptors(): InterceptorObject {
+    return {
       kind: "ai_op",
       modifyReturnValue: (args, returnValue, agent) => {
         if (returnValue instanceof Promise) {
@@ -128,14 +128,54 @@ export class AiSDK implements Wrapper {
         return returnValue;
       },
     };
+  }
+
+  private getStreamInterceptors(): InterceptorObject {
+    return {
+      kind: "ai_op",
+      modifyReturnValue: (args, returnValue, agent) => {
+        if (
+          !returnValue ||
+          typeof returnValue !== "object" ||
+          !("response" in returnValue) ||
+          !(returnValue.response instanceof Promise) ||
+          !("usage" in returnValue) ||
+          !(returnValue.usage instanceof Promise)
+        ) {
+          return returnValue;
+        }
+
+        Promise.all([returnValue.response, returnValue.usage]).then(
+          ([response, usage]) => {
+            try {
+              this.inspectAiCall(agent, args, {
+                response,
+                usage,
+              });
+            } catch {
+              // If we don't catch these errors, it will result in an unhandled promise rejection!
+            }
+          }
+        );
+
+        return returnValue;
+      },
+    };
+  }
 
+  wrap(hooks: Hooks) {
     hooks
       .addPackage("ai")
       .withVersion("^4.0.0")
       .onRequire((exports, pkgInfo) => {
         // Can't wrap it directly because it's a readonly proxy
         const generateTextFunc = exports.generateText;
         const generateObjectFunc = exports.generateObject;
+        const streamTextFunc = exports.streamText;
+        const streamObjectFunc = exports.streamObject;
+
+        const interceptors = this.getInterceptors();
+        const streamInterceptors = this.getStreamInterceptors();
 
         return {
           ...exports,
@@ -151,6 +191,18 @@ export class AiSDK implements Wrapper {
             pkgInfo,
             interceptors
           ),
+          streamText: wrapExport(
+            streamTextFunc,
+            undefined,
+            pkgInfo,
+            streamInterceptors
+          ),
+          streamObject: wrapExport(
+            streamObjectFunc,
+            undefined,
+            pkgInfo,
+            streamInterceptors
+          ),
         };
       });
   }