From d592429fff0c3b3ef8f3d2ecad8d317e84087ed4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6ssler?= Date: Tue, 17 Jun 2025 16:05:14 +0200 Subject: [PATCH 1/5] Re-add routes to AI stats Co-Authored-By: Hans Ott <3886384+hansott@users.noreply.github.com> --- library/agent/AIStatistics.test.ts | 261 +++++++++++++++++++++++++++++ library/agent/AIStatistics.ts | 71 ++++++++ library/agent/api/Event.ts | 10 ++ library/sinks/AwsSDKVersion3.ts | 13 ++ library/sinks/OpenAI.ts | 13 ++ 5 files changed, 368 insertions(+) diff --git a/library/agent/AIStatistics.test.ts b/library/agent/AIStatistics.test.ts index 49c1c1973..0b83485e0 100644 --- a/library/agent/AIStatistics.test.ts +++ b/library/agent/AIStatistics.test.ts @@ -16,6 +16,7 @@ t.test("it tracks basic AI calls", async () => { model: "gpt-4", inputTokens: 100, outputTokens: 50, + route: undefined, }); const result = stats.getStats(); @@ -29,6 +30,7 @@ t.test("it tracks basic AI calls", async () => { output: 50, total: 150, }, + routes: [], }); t.equal(stats.isEmpty(), false); @@ -42,6 +44,7 @@ t.test("it tracks multiple calls to the same provider/model", async () => { model: "gpt-4", inputTokens: 100, outputTokens: 50, + route: undefined, }); stats.onAICall({ @@ -49,6 +52,7 @@ t.test("it tracks multiple calls to the same provider/model", async () => { model: "gpt-4", inputTokens: 200, outputTokens: 75, + route: undefined, }); const result = stats.getStats(); @@ -62,6 +66,7 @@ t.test("it tracks multiple calls to the same provider/model", async () => { output: 125, total: 425, }, + routes: [], }); }); @@ -75,6 +80,7 @@ t.test( model: "gpt-4", inputTokens: 100, outputTokens: 50, + route: undefined, }); stats.onAICall({ @@ -82,6 +88,7 @@ t.test( model: "gpt-3.5-turbo", inputTokens: 80, outputTokens: 40, + route: undefined, }); stats.onAICall({ @@ -89,6 +96,7 @@ t.test( model: "claude-3", inputTokens: 120, outputTokens: 60, + route: undefined, }); const result = stats.getStats(); @@ -108,6 +116,7 @@ t.test( output: 60, total: 180, }, + routes: [], }); t.same(result[1], { @@ -119,6 +128,7 @@ t.test( output: 40, total: 120, }, + routes: [], }); t.same(result[2], { @@ -130,6 +140,7 @@ t.test( output: 50, total: 150, }, + routes: [], }); } ); @@ -142,6 +153,7 @@ t.test("it resets all statistics", async () => { model: "gpt-4", inputTokens: 100, outputTokens: 50, + route: undefined, }); stats.onAICall({ @@ -149,6 +161,7 @@ t.test("it resets all statistics", async () => { model: "claude-3", inputTokens: 120, outputTokens: 60, + route: undefined, }); t.equal(stats.isEmpty(), false); @@ -168,6 +181,7 @@ t.test("it handles zero token inputs", async () => { model: "gpt-4", inputTokens: 0, outputTokens: 0, + route: undefined, }); const result = stats.getStats(); @@ -187,6 +201,7 @@ t.test("called with empty provider", async () => { model: "gpt-4", inputTokens: 100, outputTokens: 50, + route: undefined, }); t.same(true, stats.isEmpty()); @@ -200,7 +215,253 @@ t.test("called with empty model", async () => { model: "", inputTokens: 100, outputTokens: 50, + route: undefined, }); t.same(true, stats.isEmpty()); }); + +t.test("it tracks route-specific statistics", async () => { + const stats = new AIStatistics(); + + stats.onAICall({ + provider: "openai", + model: "gpt-4", + route: { + path: "/api/chat", + method: "POST", + }, + 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, + }, + routes: [ + { + path: "/api/chat", + method: "POST", + calls: 1, + tokens: { + input: 100, + output: 50, + total: 150, + }, + }, + ], + }); +}); + +t.test( + "it tracks multiple route calls for the same provider/model", + async () => { + const stats = new AIStatistics(); + + // First call to /api/chat + stats.onAICall({ + provider: "openai", + model: "gpt-4", + route: { + path: "/api/chat", + method: "POST", + }, + inputTokens: 100, + outputTokens: 50, + }); + + // Second call to /api/chat + stats.onAICall({ + provider: "openai", + model: "gpt-4", + route: { + path: "/api/chat", + method: "POST", + }, + inputTokens: 120, + outputTokens: 60, + }); + + // Call to different route + stats.onAICall({ + provider: "openai", + model: "gpt-4", + route: { + path: "/api/summary", + method: "GET", + }, + inputTokens: 80, + outputTokens: 40, + }); + + const result = stats.getStats(); + t.equal(result.length, 1); + t.same(result[0].calls, 3); + t.same(result[0].tokens.total, 450); + t.same(result[0].routes.length, 2); + + t.same(result[0].routes[0], { + path: "/api/chat", + method: "POST", + calls: 2, + tokens: { + input: 220, + output: 110, + total: 330, + }, + }); + + t.same(result[0].routes[1], { + path: "/api/summary", + method: "GET", + calls: 1, + tokens: { + input: 80, + output: 40, + total: 120, + }, + }); + } +); + +t.test("it mixes calls with and without routes", async () => { + const stats = new AIStatistics(); + + // Call without route + stats.onAICall({ + provider: "openai", + model: "gpt-4", + inputTokens: 100, + outputTokens: 50, + route: undefined, + }); + + // Call with route + stats.onAICall({ + provider: "openai", + model: "gpt-4", + route: { + path: "/api/chat", + method: "POST", + }, + inputTokens: 120, + outputTokens: 60, + }); + + const result = stats.getStats(); + t.same(result.length, 1); + t.same(result[0].calls, 2); + t.same(result[0].tokens.total, 330); + t.same(result[0].routes.length, 1); + + t.same(result[0].routes[0], { + path: "/api/chat", + method: "POST", + calls: 1, + tokens: { + input: 120, + output: 60, + total: 180, + }, + }); +}); + +t.test("it respects LRU limit for routes", async () => { + const maxRoutes = 2; + const stats = new AIStatistics(maxRoutes); + + // Add three different routes to exceed the limit + stats.onAICall({ + provider: "openai", + model: "gpt-4", + route: { + path: "/api/route1", + method: "GET", + }, + inputTokens: 100, + outputTokens: 50, + }); + + stats.onAICall({ + provider: "openai", + model: "gpt-4", + route: { + path: "/api/route2", + method: "GET", + }, + inputTokens: 100, + outputTokens: 50, + }); + + stats.onAICall({ + provider: "openai", + model: "gpt-4", + route: { + path: "/api/route3", + method: "GET", + }, + inputTokens: 100, + outputTokens: 50, + }); + + const result = stats.getStats(); + t.equal(result.length, 1); + // All calls should be tracked in the provider stats + t.same(result[0].calls, 3); + t.same(result[0].tokens.total, 450); + + // But only the most recent routes should be kept (LRU eviction) + t.same(result[0].routes.length, 2); + + // The first route should have been evicted, keeping route2 and route3 + const routePaths = result[0].routes.map((r) => r.path); + t.notOk(routePaths.includes("/api/route1")); + t.ok(routePaths.includes("/api/route2")); + t.ok(routePaths.includes("/api/route3")); +}); + +t.test("called with empty path", async () => { + const stats = new AIStatistics(); + + stats.onAICall({ + provider: "openai", + model: "gpt-4", + route: { + path: "", + method: "POST", + }, + inputTokens: 100, + outputTokens: 50, + }); + + const result = stats.getStats(); + t.equal(result.length, 1); + t.same(result[0].routes.length, 0); +}); + +t.test("called with empty method", async () => { + const stats = new AIStatistics(); + + stats.onAICall({ + provider: "openai", + model: "gpt-4", + route: { + path: "/api/chat", + method: "", + }, + inputTokens: 100, + outputTokens: 50, + }); + + const result = stats.getStats(); + t.equal(result.length, 1); + t.same(result[0].routes.length, 0); +}); diff --git a/library/agent/AIStatistics.ts b/library/agent/AIStatistics.ts index 4cb01f605..788e8a65f 100644 --- a/library/agent/AIStatistics.ts +++ b/library/agent/AIStatistics.ts @@ -1,3 +1,16 @@ +import { LRUMap } from "../ratelimiting/LRUMap"; + +type AIRouteStats = { + path: string; + method: string; + calls: number; + tokens: { + input: number; + output: number; + total: number; + }; +}; + type AIProviderStats = { provider: string; model: string; @@ -7,10 +20,16 @@ type AIProviderStats = { output: number; total: number; }; + routesLRU: LRUMap; }; export class AIStatistics { private calls: Map = new Map(); + private maxRoutes: number; + + constructor(maxRoutes: number = 1000) { + this.maxRoutes = maxRoutes; + } private getProviderKey(provider: string, model: string): string { return `${provider}:${model}`; @@ -36,20 +55,54 @@ export class AIStatistics { output: 0, total: 0, }, + routesLRU: new LRUMap(this.maxRoutes), }); } return this.calls.get(key)!; } + private ensureRouteStats( + providerStats: AIProviderStats, + path: string, + method: string + ): AIRouteStats { + const routeKey = this.getRouteKey(path, method); + + let routeStats = providerStats.routesLRU.get(routeKey); + + if (!routeStats) { + routeStats = { + path, + method, + calls: 0, + tokens: { + input: 0, + output: 0, + total: 0, + }, + }; + providerStats.routesLRU.set(routeKey, routeStats); + } + + return routeStats; + } + onAICall({ provider, model, + route, inputTokens, outputTokens, }: { provider: string; model: string; + route: + | { + path: string; + method: string; + } + | undefined; inputTokens: number; outputTokens: number; }) { @@ -62,10 +115,27 @@ export class AIStatistics { providerStats.tokens.input += inputTokens; providerStats.tokens.output += outputTokens; providerStats.tokens.total += inputTokens + outputTokens; + + if (route && route.path && route.method) { + const routeStats = this.ensureRouteStats( + providerStats, + route.path, + route.method + ); + + routeStats.calls += 1; + routeStats.tokens.input += inputTokens; + routeStats.tokens.output += outputTokens; + routeStats.tokens.total += inputTokens + outputTokens; + } } getStats() { return Array.from(this.calls.values()).map((stats) => { + const routes = Array.from(stats.routesLRU.keys()).map( + (key) => stats.routesLRU.get(key) as AIRouteStats + ); + return { provider: stats.provider, model: stats.model, @@ -75,6 +145,7 @@ export class AIStatistics { output: stats.tokens.output, total: stats.tokens.total, }, + routes, }; }); } diff --git a/library/agent/api/Event.ts b/library/agent/api/Event.ts index 5758c6ab3..bc7208006 100644 --- a/library/agent/api/Event.ts +++ b/library/agent/api/Event.ts @@ -123,6 +123,16 @@ type Heartbeat = { output: number; total: number; }; + routes: { + path: string; + method: string; + calls: number; + tokens: { + input: number; + output: number; + total: number; + }; + }[]; }[]; packages: { name: string; diff --git a/library/sinks/AwsSDKVersion3.ts b/library/sinks/AwsSDKVersion3.ts index 2cf8a51ec..60a1892f1 100644 --- a/library/sinks/AwsSDKVersion3.ts +++ b/library/sinks/AwsSDKVersion3.ts @@ -1,4 +1,5 @@ import { Agent } from "../agent/Agent"; +import { getContext } from "../agent/Context"; import { Hooks } from "../agent/hooks/Hooks"; import { wrapExport } from "../agent/hooks/wrapExport"; import { Wrapper } from "../agent/Wrapper"; @@ -78,6 +79,7 @@ export class AwsSDKVersion3 implements Wrapper { model: body.model, inputTokens: inputTokens, outputTokens: outputTokens, + route: this.getRoute(), }); } } @@ -113,9 +115,20 @@ export class AwsSDKVersion3 implements Wrapper { model: modelId, inputTokens: inputTokens, outputTokens: outputTokens, + route: this.getRoute(), }); } + private getRoute() { + const context = getContext(); + + if (context && context.route && context.method) { + return { path: context.route, method: context.method }; + } + + return undefined; + } + wrap(hooks: Hooks) { hooks .addPackage("@aws-sdk/client-bedrock-runtime") diff --git a/library/sinks/OpenAI.ts b/library/sinks/OpenAI.ts index 280a12a96..6030cd611 100644 --- a/library/sinks/OpenAI.ts +++ b/library/sinks/OpenAI.ts @@ -3,6 +3,7 @@ import { Hooks } from "../agent/hooks/Hooks"; import { Wrapper } from "../agent/Wrapper"; import { wrapExport } from "../agent/hooks/wrapExport"; import { isPlainObject } from "../helpers/isPlainObject"; +import { getContext } from "../agent/Context"; type Response = { model: string; @@ -64,6 +65,7 @@ export class OpenAI implements Wrapper { model: response.model ?? "", inputTokens: inputTokens, outputTokens: outputTokens, + route: this.getRoute(), }); } @@ -93,6 +95,7 @@ export class OpenAI implements Wrapper { model: response.model ?? "", inputTokens: inputTokens, outputTokens: outputTokens, + route: this.getRoute(), }); } @@ -113,6 +116,16 @@ export class OpenAI implements Wrapper { return "openai"; } + private getRoute() { + const context = getContext(); + + if (context && context.route && context.method) { + return { path: context.route, method: context.method }; + } + + return undefined; + } + wrap(hooks: Hooks) { // Note: Streaming is not supported yet hooks From 32bc12686ebf95112008313aecaa191ecc2d6559 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6ssler?= Date: Wed, 18 Jun 2025 14:23:43 +0200 Subject: [PATCH 2/5] Extract getRoute to helper --- library/helpers/getRouteForAiStats.test.ts | 31 ++++++++++++++++++++++ library/helpers/getRouteForAiStats.ts | 11 ++++++++ library/sinks/AwsSDKVersion3.ts | 16 +++-------- library/sinks/OpenAI.ts | 16 +++-------- 4 files changed, 48 insertions(+), 26 deletions(-) create mode 100644 library/helpers/getRouteForAiStats.test.ts create mode 100644 library/helpers/getRouteForAiStats.ts diff --git a/library/helpers/getRouteForAiStats.test.ts b/library/helpers/getRouteForAiStats.test.ts new file mode 100644 index 000000000..2069a890f --- /dev/null +++ b/library/helpers/getRouteForAiStats.test.ts @@ -0,0 +1,31 @@ +import * as t from "tap"; +import { getRouteForAiStats } from "./getRouteForAIStats"; +import { runWithContext } from "../agent/Context"; + +const getTestContext = () => ({ + url: "/test/route", + method: "GET", + route: "/test/route", + query: {}, + body: undefined, + headers: {}, + routeParams: {}, + remoteAddress: "1.2.3.4", + source: "test", + cookies: {}, +}); + +t.test("it works", async (t) => { + t.same(getRouteForAiStats(), undefined); + + runWithContext(getTestContext(), () => { + t.same(getRouteForAiStats(), { path: "/test/route", method: "GET" }); + }); + + runWithContext( + { ...getTestContext(), route: undefined, method: undefined }, + () => { + t.same(getRouteForAiStats(), undefined); + } + ); +}); diff --git a/library/helpers/getRouteForAiStats.ts b/library/helpers/getRouteForAiStats.ts new file mode 100644 index 000000000..d3186098f --- /dev/null +++ b/library/helpers/getRouteForAiStats.ts @@ -0,0 +1,11 @@ +import { getContext } from "../agent/Context"; + +export function getRouteForAiStats() { + const context = getContext(); + + if (context && context.route && context.method) { + return { path: context.route, method: context.method }; + } + + return undefined; +} diff --git a/library/sinks/AwsSDKVersion3.ts b/library/sinks/AwsSDKVersion3.ts index 60a1892f1..a80e18836 100644 --- a/library/sinks/AwsSDKVersion3.ts +++ b/library/sinks/AwsSDKVersion3.ts @@ -1,8 +1,8 @@ import { Agent } from "../agent/Agent"; -import { getContext } from "../agent/Context"; import { Hooks } from "../agent/hooks/Hooks"; import { wrapExport } from "../agent/hooks/wrapExport"; import { Wrapper } from "../agent/Wrapper"; +import { getRouteForAiStats } from "../helpers/getRouteForAIStats"; import { isPlainObject } from "../helpers/isPlainObject"; type InvokeUsage = { @@ -79,7 +79,7 @@ export class AwsSDKVersion3 implements Wrapper { model: body.model, inputTokens: inputTokens, outputTokens: outputTokens, - route: this.getRoute(), + route: getRouteForAiStats(), }); } } @@ -115,20 +115,10 @@ export class AwsSDKVersion3 implements Wrapper { model: modelId, inputTokens: inputTokens, outputTokens: outputTokens, - route: this.getRoute(), + route: getRouteForAiStats(), }); } - private getRoute() { - const context = getContext(); - - if (context && context.route && context.method) { - return { path: context.route, method: context.method }; - } - - return undefined; - } - wrap(hooks: Hooks) { hooks .addPackage("@aws-sdk/client-bedrock-runtime") diff --git a/library/sinks/OpenAI.ts b/library/sinks/OpenAI.ts index 6030cd611..8c7487273 100644 --- a/library/sinks/OpenAI.ts +++ b/library/sinks/OpenAI.ts @@ -3,7 +3,7 @@ import { Hooks } from "../agent/hooks/Hooks"; import { Wrapper } from "../agent/Wrapper"; import { wrapExport } from "../agent/hooks/wrapExport"; import { isPlainObject } from "../helpers/isPlainObject"; -import { getContext } from "../agent/Context"; +import { getRouteForAiStats } from "../helpers/getRouteForAIStats"; type Response = { model: string; @@ -65,7 +65,7 @@ export class OpenAI implements Wrapper { model: response.model ?? "", inputTokens: inputTokens, outputTokens: outputTokens, - route: this.getRoute(), + route: getRouteForAiStats(), }); } @@ -95,7 +95,7 @@ export class OpenAI implements Wrapper { model: response.model ?? "", inputTokens: inputTokens, outputTokens: outputTokens, - route: this.getRoute(), + route: getRouteForAiStats(), }); } @@ -116,16 +116,6 @@ export class OpenAI implements Wrapper { return "openai"; } - private getRoute() { - const context = getContext(); - - if (context && context.route && context.method) { - return { path: context.route, method: context.method }; - } - - return undefined; - } - wrap(hooks: Hooks) { // Note: Streaming is not supported yet hooks From f00539d2950c7b2b2e800335e62bdb8f030f4c71 Mon Sep 17 00:00:00 2001 From: bitterpanda Date: Wed, 18 Jun 2025 12:29:23 +0000 Subject: [PATCH 3/5] Update library/sinks/OpenAI.ts --- library/sinks/OpenAI.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/sinks/OpenAI.ts b/library/sinks/OpenAI.ts index 8c7487273..1f21f8b3c 100644 --- a/library/sinks/OpenAI.ts +++ b/library/sinks/OpenAI.ts @@ -3,7 +3,7 @@ import { Hooks } from "../agent/hooks/Hooks"; import { Wrapper } from "../agent/Wrapper"; import { wrapExport } from "../agent/hooks/wrapExport"; import { isPlainObject } from "../helpers/isPlainObject"; -import { getRouteForAiStats } from "../helpers/getRouteForAIStats"; +import { getRouteForAiStats } from "../helpers/getRouteForAiStats"; type Response = { model: string; From 3b6d932b93a9f2f667e0eedf430508a2774cc329 Mon Sep 17 00:00:00 2001 From: bitterpanda Date: Wed, 18 Jun 2025 12:29:31 +0000 Subject: [PATCH 4/5] Update library/sinks/AwsSDKVersion3.ts --- library/sinks/AwsSDKVersion3.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/sinks/AwsSDKVersion3.ts b/library/sinks/AwsSDKVersion3.ts index a80e18836..441851ffc 100644 --- a/library/sinks/AwsSDKVersion3.ts +++ b/library/sinks/AwsSDKVersion3.ts @@ -2,7 +2,7 @@ import { Agent } from "../agent/Agent"; import { Hooks } from "../agent/hooks/Hooks"; import { wrapExport } from "../agent/hooks/wrapExport"; import { Wrapper } from "../agent/Wrapper"; -import { getRouteForAiStats } from "../helpers/getRouteForAIStats"; +import { getRouteForAiStats } from "../helpers/getRouteForAiStats"; import { isPlainObject } from "../helpers/isPlainObject"; type InvokeUsage = { From 02693bfbe30e3bc9113e27bfedd46f943962f768 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6ssler?= Date: Wed, 18 Jun 2025 14:31:07 +0200 Subject: [PATCH 5/5] Fix unit test --- library/helpers/getRouteForAiStats.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/helpers/getRouteForAiStats.test.ts b/library/helpers/getRouteForAiStats.test.ts index 2069a890f..c026d68a6 100644 --- a/library/helpers/getRouteForAiStats.test.ts +++ b/library/helpers/getRouteForAiStats.test.ts @@ -1,5 +1,5 @@ import * as t from "tap"; -import { getRouteForAiStats } from "./getRouteForAIStats"; +import { getRouteForAiStats } from "./getRouteForAiStats"; import { runWithContext } from "../agent/Context"; const getTestContext = () => ({