Update readme and sources, fix tests

Author: timokoessler

.github/workflows/unit-test.yml

@@ -80,6 +80,8 @@ jobs:
       - run: npm run install-lib-only
       - run: npm run build
       - run: npm run test:ci
+        env:
+          GOOGLE_GENERATIVE_AI_API_KEY: ${{ secrets.GOOGLE_GENERATIVE_AI_API_KEY }}
       - name: "Upload coverage"
         uses: codecov/codecov-action@0565863a31f2c772f9f0395002a31e3f06189574 # v5
         with:

README.md

@@ -98,6 +98,8 @@ See list above for supported database drivers.
 
 * ✅ [`@koa/router`](https://www.npmjs.com/package/@koa/router) 13.x, 12.x, 11.x and 10.x
 
+### AI SDKs
+* ✅ [`ai`](https://www.npmjs.com/package/ai) 4.x
 
 ## Installation
 

library/agent/Source.ts

@@ -9,6 +9,7 @@ export const SOURCES = [
   "subdomains",
   "markUnsafe",
   "url",
+  "aiToolParams",
 ] as const;
 
 export type Source = (typeof SOURCES)[number];

library/agent/protect.ts

@@ -48,6 +48,7 @@ import { Fastify } from "../sources/Fastify";
 import { Koa } from "../sources/Koa";
 import { ClickHouse } from "../sinks/ClickHouse";
 import { Prisma } from "../sinks/Prisma";
+import { AiSDK } from "../sources/AiSDK";
 
 function getLogger(): Logger {
   if (isDebugging()) {
@@ -141,6 +142,7 @@ export function getWrappers() {
     new ClickHouse(),
     new Prisma(),
     // new Function(), Disabled because functionName.constructor === Function is false after patching global
+    new AiSDK(),
   ];
 }
 

library/sources/AiSDK.test.ts

@@ -2,88 +2,98 @@ import * as t from "tap";
 import { startTestAgent } from "../helpers/startTestAgent";
 import { AiSDK } from "./AiSDK";
 import { getContext, runWithContext, type Context } from "../agent/Context";
+import { getMajorNodeVersion } from "../helpers/getNodeVersion";
 
-t.test("It works with agentic", async (t) => {
-  startTestAgent({
-    wrappers: [new AiSDK()],
-    rewrite: {},
-  });
-
-  const getTestContext = (message: string): Context => {
-    return {
-      remoteAddress: "::1",
-      method: "POST",
-      url: "http://localhost:4000",
-      query: {
-        message,
-      },
-      body: undefined,
-      headers: {},
-      cookies: {},
-      routeParams: {},
-      source: "express",
-      route: "/posts/:id",
-    };
-  };
-
-  const { google } =
-    require("@ai-sdk/google") as typeof import("@ai-sdk/google");
-  const { generateText, tool } = require("ai") as typeof import("ai");
-  const { z } = require("zod") as typeof import("zod");
-
-  const callWithPrompt = async (prompt: string) => {
-    return await generateText({
-      model: google("models/gemini-2.0-flash"),
-      tools: {
-        weather: tool({
-          description: "Get the weather in a location",
-          parameters: z.object({
-            location: z
-              .string()
-              .describe("The location to get the weather for"),
-          }),
-          execute: async ({ location }) => {
-            const temperature = location === "Norway" ? 5 : 24;
-            return {
-              temperature,
-              context: getContext(),
-            };
-          },
-        }),
-      },
-      prompt: prompt,
+t.test(
+  "It works with AI tool execution",
+  {
+    skip:
+      !process.env.GOOGLE_GENERATIVE_AI_API_KEY || getMajorNodeVersion() < 20
+        ? "Google API key not set or Node version < 20"
+        : undefined,
+  },
+  async (t) => {
+    startTestAgent({
+      wrappers: [new AiSDK()],
+      rewrite: {},
     });
-  };
-
-  await runWithContext(
-    getTestContext("What is the weather in San Francisco?"),
-    async () => {
-      const result = await callWithPrompt(
-        "What is the weather in San Francisco?"
-      );
 
-      t.same(result.toolResults.length, 1);
-      t.same(result.toolResults[0].toolName, "weather");
-      t.same(result.toolResults[0].result.temperature, 24);
-      t.match(result.toolResults[0].result.context, {
+    const getTestContext = (message: string): Context => {
+      return {
         remoteAddress: "::1",
         method: "POST",
         url: "http://localhost:4000",
         query: {
-          message: "What is the weather in San Francisco?",
+          message,
         },
         body: undefined,
         headers: {},
         cookies: {},
         routeParams: {},
         source: "express",
         route: "/posts/:id",
-        aiToolParams: [
-          {
-            location: "San Francisco",
-          },
-        ],
+      };
+    };
+
+    const { google } =
+      require("@ai-sdk/google") as typeof import("@ai-sdk/google");
+    const { generateText, tool } = require("ai") as typeof import("ai");
+    const { z } = require("zod") as typeof import("zod");
+
+    const callWithPrompt = async (prompt: string) => {
+      return await generateText({
+        model: google("models/gemini-2.0-flash-lite"),
+        tools: {
+          weather: tool({
+            description: "Get the weather in a location",
+            parameters: z.object({
+              location: z
+                .string()
+                .describe("The location to get the weather for"),
+            }),
+            execute: async ({ location }) => {
+              const temperature = location === "Norway" ? 5 : 24;
+              return {
+                temperature,
+                context: getContext(),
+              };
+            },
+          }),
+        },
+        prompt: prompt,
       });
-    }
-  );
-});
+    };
+
+    await runWithContext(
+      getTestContext("What is the weather in San Francisco?"),
+      async () => {
+        const result = await callWithPrompt(
+          "What is the weather in San Francisco?"
+        );
+
+        t.same(result.toolResults.length, 1);
+        t.same(result.toolResults[0].toolName, "weather");
+        t.same(result.toolResults[0].result.temperature, 24);
+        t.match(result.toolResults[0].result.context, {
+          remoteAddress: "::1",
+          method: "POST",
+          url: "http://localhost:4000",
+          query: {
+            message: "What is the weather in San Francisco?",
+          },
+          body: undefined,
+          headers: {},
+          cookies: {},
+          routeParams: {},
+          source: "express",
+          route: "/posts/:id",
+          aiToolParams: [
+            {
+              location: "San Francisco",
+            },
+          ],
+        });
+      }
+    );
+  }
+);