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/helpers/getRouteForAiStats.test.ts b/library/helpers/getRouteForAiStats.test.ts new file mode 100644 index 000000000..c026d68a6 --- /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 2cf8a51ec..441851ffc 100644 --- a/library/sinks/AwsSDKVersion3.ts +++ b/library/sinks/AwsSDKVersion3.ts @@ -2,6 +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 { isPlainObject } from "../helpers/isPlainObject"; type InvokeUsage = { @@ -78,6 +79,7 @@ export class AwsSDKVersion3 implements Wrapper { model: body.model, inputTokens: inputTokens, outputTokens: outputTokens, + route: getRouteForAiStats(), }); } } @@ -113,6 +115,7 @@ export class AwsSDKVersion3 implements Wrapper { model: modelId, inputTokens: inputTokens, outputTokens: outputTokens, + route: getRouteForAiStats(), }); } diff --git a/library/sinks/OpenAI.ts b/library/sinks/OpenAI.ts index 280a12a96..1f21f8b3c 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 { getRouteForAiStats } from "../helpers/getRouteForAiStats"; type Response = { model: string; @@ -64,6 +65,7 @@ export class OpenAI implements Wrapper { model: response.model ?? "", inputTokens: inputTokens, outputTokens: outputTokens, + route: getRouteForAiStats(), }); } @@ -93,6 +95,7 @@ export class OpenAI implements Wrapper { model: response.model ?? "", inputTokens: inputTokens, outputTokens: outputTokens, + route: getRouteForAiStats(), }); }