Merge branch 'main' into prevent-double-require

Author: timokoessler

README.md

@@ -98,6 +98,14 @@ 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
+
+Zen instruments the following AI SDKs to track which models are used and how many tokens are consumed, allowing you to monitor your AI usage and costs:
+
+* ✅ [`openai`](https://www.npmjs.com/package/openai) 4.x
+* ✅ [`@aws-sdk/client-bedrock-runtime`](https://www.npmjs.com/package/@aws-sdk/client-bedrock-runtime) 3.x
+
+_Note: Prompt injection attacks are currently not covered by Zen._
 
 ## Installation
 

docs/esbuild.md

@@ -4,25 +4,24 @@ Note: Zen runs only on the server side, it does not run in the browser.
 
 Note: If `bundle` is set to `false` in the esbuild configuration, Zen will work without any additional configuration.
 
-Modify your esbuild configuration to include the external option using this utility:
+Modify your esbuild configuration to set the `packages` option to `external`:
 
 ```js
 const { build } = require("esbuild");
-const { externals } = require("@aikidosec/firewall/bundler"); // <-- Add this line
 
 build({
   entryPoints: ["./app.js"],
   bundle: true,
   platform: "node",
   target: "node18",
   outfile: "./dist/app.js",
-  external: externals(), // <-- Add this line
+  packages: "external", // <-- Add this line
 });
 ```
 
-This tells esbuild to exclude @aikidosec/firewall and any packages that Zen hooks into from the bundle.
+This tells esbuild to load packages (including @aikidosec/firewall and any packages that Zen hooks into) from the `node_modules` directory, while still bundling your application code.
 
-⚠️ Don't forget to copy the node_modules directory to the output directory.
+⚠️ Don't forget to copy the node_modules directory to the output directory (only `dependencies`, not `devDependencies`).
 
 ## Why do I need to do this?
 

library/agent/AIStatistics.test.ts

@@ -0,0 +1,206 @@
+import * as t from "tap";
+import { AIStatistics } from "./AIStatistics";
+
+t.test("it initializes with empty state", async () => {
+  const stats = new AIStatistics();
+
+  t.same(stats.getStats(), []);
+  t.equal(stats.isEmpty(), true);
+});
+
+t.test("it tracks basic AI calls", async () => {
+  const stats = new AIStatistics();
+
+  stats.onAICall({
+    provider: "openai",
+    model: "gpt-4",
+    inputTokens: 100,
+    outputTokens: 50,
+  });
+
+  const result = stats.getStats();
+  t.equal(result.length, 1);
+  t.same(result[0], {
+    provider: "openai",
+    model: "gpt-4",
+    calls: 1,
+    tokens: {
+      input: 100,
+      output: 50,
+      total: 150,
+    },
+  });
+
+  t.equal(stats.isEmpty(), false);
+});
+
+t.test("it tracks multiple calls to the same provider/model", async () => {
+  const stats = new AIStatistics();
+
+  stats.onAICall({
+    provider: "openai",
+    model: "gpt-4",
+    inputTokens: 100,
+    outputTokens: 50,
+  });
+
+  stats.onAICall({
+    provider: "openai",
+    model: "gpt-4",
+    inputTokens: 200,
+    outputTokens: 75,
+  });
+
+  const result = stats.getStats();
+  t.same(result.length, 1);
+  t.same(result[0], {
+    provider: "openai",
+    model: "gpt-4",
+    calls: 2,
+    tokens: {
+      input: 300,
+      output: 125,
+      total: 425,
+    },
+  });
+});
+
+t.test(
+  "it tracks different provider/model combinations separately",
+  async () => {
+    const stats = new AIStatistics();
+
+    stats.onAICall({
+      provider: "openai",
+      model: "gpt-4",
+      inputTokens: 100,
+      outputTokens: 50,
+    });
+
+    stats.onAICall({
+      provider: "openai",
+      model: "gpt-3.5-turbo",
+      inputTokens: 80,
+      outputTokens: 40,
+    });
+
+    stats.onAICall({
+      provider: "anthropic",
+      model: "claude-3",
+      inputTokens: 120,
+      outputTokens: 60,
+    });
+
+    const result = stats.getStats();
+    t.equal(result.length, 3);
+
+    // Sort by provider:model for consistent testing
+    result.sort((a, b) =>
+      `${a.provider}:${a.model}`.localeCompare(`${b.provider}:${b.model}`)
+    );
+
+    t.same(result[0], {
+      provider: "anthropic",
+      model: "claude-3",
+      calls: 1,
+      tokens: {
+        input: 120,
+        output: 60,
+        total: 180,
+      },
+    });
+
+    t.same(result[1], {
+      provider: "openai",
+      model: "gpt-3.5-turbo",
+      calls: 1,
+      tokens: {
+        input: 80,
+        output: 40,
+        total: 120,
+      },
+    });
+
+    t.same(result[2], {
+      provider: "openai",
+      model: "gpt-4",
+      calls: 1,
+      tokens: {
+        input: 100,
+        output: 50,
+        total: 150,
+      },
+    });
+  }
+);
+
+t.test("it resets all statistics", async () => {
+  const stats = new AIStatistics();
+
+  stats.onAICall({
+    provider: "openai",
+    model: "gpt-4",
+    inputTokens: 100,
+    outputTokens: 50,
+  });
+
+  stats.onAICall({
+    provider: "anthropic",
+    model: "claude-3",
+    inputTokens: 120,
+    outputTokens: 60,
+  });
+
+  t.equal(stats.isEmpty(), false);
+  t.equal(stats.getStats().length, 2);
+
+  stats.reset();
+
+  t.equal(stats.isEmpty(), true);
+  t.same(stats.getStats(), []);
+});
+
+t.test("it handles zero token inputs", async () => {
+  const stats = new AIStatistics();
+
+  stats.onAICall({
+    provider: "openai",
+    model: "gpt-4",
+    inputTokens: 0,
+    outputTokens: 0,
+  });
+
+  const result = stats.getStats();
+  t.equal(result.length, 1);
+  t.same(result[0].tokens, {
+    input: 0,
+    output: 0,
+    total: 0,
+  });
+});
+
+t.test("called with empty provider", async () => {
+  const stats = new AIStatistics();
+
+  stats.onAICall({
+    provider: "",
+    model: "gpt-4",
+    inputTokens: 100,
+    outputTokens: 50,
+  });
+
+  t.same(true, stats.isEmpty());
+});
+
+t.test("called with empty model", async () => {
+  const stats = new AIStatistics();
+
+  stats.onAICall({
+    provider: "openai",
+    model: "",
+    inputTokens: 100,
+    outputTokens: 50,
+  });
+
+  t.same(true, stats.isEmpty());
+});

library/agent/AIStatistics.ts

@@ -0,0 +1,89 @@
+type AIProviderStats = {
+  provider: string;
+  model: string;
+  calls: number;
+  tokens: {
+    input: number;
+    output: number;
+    total: number;
+  };
+};
+
+export class AIStatistics {
+  private calls: Map<string, AIProviderStats> = new Map();
+
+  private getProviderKey(provider: string, model: string): string {
+    return `${provider}:${model}`;
+  }
+
+  private getRouteKey(path: string, method: string): string {
+    return `${method}:${path}`;
+  }
+
+  private ensureProviderStats(
+    provider: string,
+    model: string
+  ): AIProviderStats {
+    const key = this.getProviderKey(provider, model);
+
+    if (!this.calls.has(key)) {
+      this.calls.set(key, {
+        provider,
+        model,
+        calls: 0,
+        tokens: {
+          input: 0,
+          output: 0,
+          total: 0,
+        },
+      });
+    }
+
+    return this.calls.get(key)!;
+  }
+
+  onAICall({
+    provider,
+    model,
+    inputTokens,
+    outputTokens,
+  }: {
+    provider: string;
+    model: string;
+    inputTokens: number;
+    outputTokens: number;
+  }) {
+    if (!provider || !model) {
+      return;
+    }
+
+    const providerStats = this.ensureProviderStats(provider, model);
+    providerStats.calls += 1;
+    providerStats.tokens.input += inputTokens;
+    providerStats.tokens.output += outputTokens;
+    providerStats.tokens.total += inputTokens + outputTokens;
+  }
+
+  getStats() {
+    return Array.from(this.calls.values()).map((stats) => {
+      return {
+        provider: stats.provider,
+        model: stats.model,
+        calls: stats.calls,
+        tokens: {
+          input: stats.tokens.input,
+          output: stats.tokens.output,
+          total: stats.tokens.total,
+        },
+      };
+    });
+  }
+
+  reset() {
+    this.calls.clear();
+  }
+
+  isEmpty(): boolean {
+    return this.calls.size === 0;
+  }
+}