feat: AI Tracking to AI SDK

packages/sdk/ai/package.json

@@ -5,9 +5,9 @@
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "scripts": {
-    "build": "tsc",
-    "test": "jest",
-    "lint": "eslint . --ext .ts"
+    "build": "npx tsc",
+    "lint": "npx eslint . --ext .ts",
+    "lint:fix": "yarn run lint --fix"
   },
   "keywords": [
     "launchdarkly",

packages/sdk/ai/src/LDAIConfigTracker.ts

@@ -1,36 +1,97 @@
 import { LDClient, LDContext } from '@launchdarkly/node-server-sdk';
 
-import { BedrockTokenUsage, FeedbackKind, TokenUsage, UnderscoreTokenUsage } from './api/metrics';
-import { usageToTokenMetrics } from './trackUtils';
+import {
+  createBedrockTokenUsage,
+  FeedbackKind,
+  OpenAITokenUsage,
+  TokenMetrics,
+  TokenUsage,
+  UnderscoreTokenUsage,
+} from './api/metrics';
 
 export class LDAIConfigTracker {
   private ldClient: LDClient;
-  private configKey: string;
   private variationId: string;
+  private configKey: string;
   private context: LDContext;
 
   constructor(ldClient: LDClient, configKey: string, variationId: string, context: LDContext) {
     this.ldClient = ldClient;
-    this.configKey = configKey;
     this.variationId = variationId;
+    this.configKey = configKey;
     this.context = context;
   }
 
-  getTrackData() {
+  private getTrackData() {
     return {
-      configKey: this.configKey,
       variationId: this.variationId,
+      configKey: this.configKey,
     };
   }
 
   trackDuration(duration: number): void {
-    this.ldClient.track('$ld:ai:duration:total', this.context, this.variationId, duration);
+    this.ldClient.track('$ld:ai:duration:total', this.context, this.getTrackData(), duration);
+  }
+
+  async trackDurationOf(func: (...args: any[]) => Promise<any>, ...args: any[]): Promise<any> {
+    const startTime = Date.now();
+    const result = await func(...args);
+    const endTime = Date.now();
+    const duration = endTime - startTime; // duration in milliseconds
+    this.trackDuration(duration);
+    return result;
+  }
+
+  trackError(error: number): void {
+    this.ldClient.track('$ld:ai:error', this.context, this.getTrackData(), error);
   }
 
-  trackTokens(tokens: TokenUsage | UnderscoreTokenUsage | BedrockTokenUsage) {
-    console.log('tracking LLM tokens', tokens);
-    const tokenMetrics = usageToTokenMetrics(tokens);
-    console.log('token metrics', tokenMetrics);
+  trackFeedback(feedback: { kind: FeedbackKind }): void {
+    if (feedback.kind === FeedbackKind.Positive) {
+      this.ldClient.track('$ld:ai:feedback:user:positive', this.context, this.getTrackData(), 1);
+    } else if (feedback.kind === FeedbackKind.Negative) {
+      this.ldClient.track('$ld:ai:feedback:user:negative', this.context, this.getTrackData(), 1);
+    }
+  }
+
+  trackGeneration(generation: number): void {
+    this.ldClient.track('$ld:ai:generation', this.context, this.getTrackData(), generation);
+  }
+
+  async trackOpenAI(func: (...args: any[]) => Promise<any>, ...args: any[]): Promise<any> {
+    const result = await this.trackDurationOf(func, ...args);
+    this.trackGeneration(1);
+    if (result.usage) {
+      this.trackTokens(new OpenAITokenUsage(result.usage));
+    }
+    return result;
+  }
+
+  async trackBedrockConverse(res: {
+    $metadata?: { httpStatusCode: number };
+    metrics?: { latencyMs: number };
+    usage?: {
+      inputTokens: number;
+      outputTokens: number;
+      totalTokens: number;
+    };
+  }): Promise<any> {
+    if (res.$metadata?.httpStatusCode === 200) {
+      this.trackGeneration(1);
+    } else if (res.$metadata?.httpStatusCode && res.$metadata.httpStatusCode >= 400) {
+      this.trackError(res.$metadata.httpStatusCode);
+    }
+    if (res.metrics) {
+      this.trackDuration(res.metrics.latencyMs);
+    }
+    if (res.usage) {
+      this.trackTokens(createBedrockTokenUsage(res.usage));
+    }
+    return res;
+  }
+
+  trackTokens(tokens: TokenUsage | UnderscoreTokenUsage | { totalTokens: number; inputTokens: number; outputTokens: number }): void {
+    const tokenMetrics = toMetrics(tokens);
     if (tokenMetrics.total > 0) {
       this.ldClient.track(
         '$ld:ai:tokens:total',
@@ -40,7 +101,6 @@ export class LDAIConfigTracker {
       );
     }
     if (tokenMetrics.input > 0) {
-      console.log('tracking input tokens', tokenMetrics.input);
       this.ldClient.track(
         '$ld:ai:tokens:input',
         this.context,
@@ -49,7 +109,6 @@ export class LDAIConfigTracker {
       );
     }
     if (tokenMetrics.output > 0) {
-      console.log('tracking output tokens', tokenMetrics.output);
       this.ldClient.track(
         '$ld:ai:tokens:output',
         this.context,
@@ -58,20 +117,28 @@ export class LDAIConfigTracker {
       );
     }
   }
+}
 
-  trackError(error: number) {
-    this.ldClient.track('$ld:ai:error', this.context, this.getTrackData(), error);
-  }
-
-  trackGeneration(generation: number) {
-    this.ldClient.track('$ld:ai:generation', this.context, this.getTrackData(), generation);
+function toMetrics(
+  usage: TokenUsage | UnderscoreTokenUsage | { totalTokens: number; inputTokens: number; outputTokens: number },
+): TokenMetrics {
+  if ('inputTokens' in usage && 'outputTokens' in usage) {
+    // Bedrock usage
+    return {
+      total: usage.totalTokens,
+      input: usage.inputTokens,
+      output: usage.outputTokens,
+    };
   }
 
-  trackFeedback(feedback: { kind: FeedbackKind }) {
-    if (feedback.kind === FeedbackKind.Positive) {
-      this.ldClient.track('$ld:ai:feedback:user:positive', this.context, this.getTrackData(), 1);
-    } else if (feedback.kind === FeedbackKind.Negative) {
-      this.ldClient.track('$ld:ai:feedback:user:negative', this.context, this.getTrackData(), 1);
-    }
-  }
+  // OpenAI usage (both camelCase and snake_case)
+  return {
+    total: 'total_tokens' in usage ? usage.total_tokens! : (usage as TokenUsage).totalTokens ?? 0,
+    input:
+      'prompt_tokens' in usage ? usage.prompt_tokens! : (usage as TokenUsage).promptTokens ?? 0,
+    output:
+      'completion_tokens' in usage
+        ? usage.completion_tokens!
+        : (usage as TokenUsage).completionTokens ?? 0,
+  };
 }

packages/sdk/ai/src/api/config/LDAIConfig.ts

@@ -5,12 +5,17 @@ import { LDAIConfigTracker } from './LDAIConfigTracker';
  */
 export interface LDAIConfig {
   /**
-   * The result of the AI Config evaluation.
+   * The result of the AI Config customization.
    */
   config: unknown;
 
   /**
-   * A tracker which can be used to generate analytics for the migration.
+   * A tracker which can be used to generate analytics.
    */
   tracker: LDAIConfigTracker;
+
+  /**
+   * Whether the configuration is not found.
+   */
+  noConfiguration: boolean;
 }

packages/sdk/ai/src/api/config/LDAIConfigTracker.ts

@@ -1,9 +1,17 @@
-import { BedrockTokenUsage, FeedbackKind, TokenUsage, UnderscoreTokenUsage } from '../metrics';
+import { FeedbackKind, TokenUsage, UnderscoreTokenUsage } from '../metrics';
 
 export interface LDAIConfigTracker {
   trackDuration: (duration: number) => void;
-  trackTokens: (tokens: TokenUsage | UnderscoreTokenUsage | BedrockTokenUsage) => void;
+  trackTokens: (
+    tokens:
+      | TokenUsage
+      | UnderscoreTokenUsage
+      | { totalTokens: number; inputTokens: number; outputTokens: number },
+  ) => void;
   trackError: (error: number) => void;
   trackGeneration: (generation: number) => void;
   trackFeedback: (feedback: { kind: FeedbackKind }) => void;
+  trackDurationOf: (func: (...args: any[]) => Promise<any>, ...args: any[]) => Promise<any>;
+  trackOpenAI: (func: Function, ...args: any[]) => any;
+  trackBedrockConverse: (res: any) => any;
 }

packages/sdk/ai/src/api/config/index.ts

@@ -1,2 +1,2 @@
 export * from './LDAIConfig';
-export * from './LDAIConfigTracker';
+export { LDAIConfigTracker } from './LDAIConfigTracker';

packages/sdk/ai/src/api/metrics/BedrockTokenUsage.ts

@@ -1,5 +1,11 @@
-export interface BedrockTokenUsage {
+export function createBedrockTokenUsage(data: {
+  totalTokens: number;
   inputTokens: number;
   outputTokens: number;
-  totalTokens: number;
+}) {
+  return {
+    total: data.totalTokens || 0,
+    input: data.inputTokens || 0,
+    output: data.outputTokens || 0,
+  };
 }

packages/sdk/ai/src/api/metrics/OpenAITokenUsage.ts

@@ -0,0 +1,19 @@
+export class OpenAITokenUsage {
+  total_tokens: number;
+  prompt_tokens: number;
+  completion_tokens: number;
+
+  constructor(data: any) {
+    this.total_tokens = data.total_tokens;
+    this.prompt_tokens = data.prompt_tokens;
+    this.completion_tokens = data.completion_tokens;
+  }
+
+  toMetrics() {
+    return {
+      total: this.total_tokens,
+      input: this.prompt_tokens,
+      output: this.completion_tokens,
+    };
+  }
+}

packages/sdk/ai/src/api/metrics/TokenUsage.ts

@@ -1,5 +1,19 @@
-export interface TokenUsage {
-  completionTokens?: number;
-  promptTokens?: number;
-  totalTokens?: number;
+export class TokenUsage {
+  totalTokens: number;
+  promptTokens: number;
+  completionTokens: number;
+
+  constructor(data: any) {
+    this.totalTokens = data.total_tokens || 0;
+    this.promptTokens = data.prompt_tokens || 0;
+    this.completionTokens = data.completion_tokens || 0;
+  }
+
+  toMetrics() {
+    return {
+      total: this.totalTokens,
+      input: this.promptTokens,
+      output: this.completionTokens,
+    };
+  }
 }

packages/sdk/ai/src/api/metrics/UnderscoreTokenUsage.ts

@@ -1,5 +1,19 @@
-export interface UnderscoreTokenUsage {
-  completion_tokens?: number;
-  prompt_tokens?: number;
-  total_tokens?: number;
+export class UnderscoreTokenUsage {
+  total_tokens: number;
+  prompt_tokens: number;
+  completion_tokens: number;
+
+  constructor(data: any) {
+    this.total_tokens = data.total_tokens || 0;
+    this.prompt_tokens = data.prompt_tokens || 0;
+    this.completion_tokens = data.completion_tokens || 0;
+  }
+
+  toMetrics() {
+    return {
+      total: this.total_tokens,
+      input: this.prompt_tokens,
+      output: this.completion_tokens,
+    };
+  }
 }

packages/sdk/ai/src/api/metrics/index.ts

@@ -1,5 +1,6 @@
 export * from './BedrockTokenUsage';
 export * from './FeedbackKind';
+export * from './OpenAITokenUsage';
 export * from './TokenMetrics';
 export * from './TokenUsage';
 export * from './UnderScoreTokenUsage';

packages/sdk/ai/src/index.ts

@@ -2,8 +2,8 @@ import Mustache from 'mustache';
 
 import { LDClient, LDContext } from '@launchdarkly/node-server-sdk';
 
-import { LDAIConfigTracker } from './LDAIConfigTracker';
 import { LDAIConfig } from './api/config';
+import { LDAIConfigTracker } from './LDAIConfigTracker';
 
 export class AIClient {
   private ldClient: LDClient;
@@ -41,7 +41,6 @@ export class AIClient {
    * const defaultValue = {}};
    *
    * const result = modelConfig(key, context, defaultValue, variables);
-   * console.log(result);
    * // Output:
    * {
    * modelId: "gpt-4o",
@@ -67,7 +66,7 @@ export class AIClient {
     defaultValue: string,
     variables?: Record<string, unknown>,
   ): Promise<LDAIConfig> {
-    const detail = await this.ldClient.variationDetail(key, context, defaultValue);
+    const detail = await this.ldClient.variation(key, context, defaultValue);
 
     const allVariables = { ldctx: context, ...variables };
 
@@ -78,12 +77,9 @@ export class AIClient {
 
     return {
       config: detail.value,
-      tracker: new LDAIConfigTracker(
-        this.ldClient,
-        key,
-        detail.value["_ldMeta"]["variationId"],
-        context,
-      ),
+      // eslint-disable-next-line no-underscore-dangle
+      tracker: new LDAIConfigTracker(this.ldClient, key, detail.value._ldMeta.variationId, context),
+      noConfiguration: Object.keys(detail).length === 0,
     };
   }
 }
@@ -92,4 +88,5 @@ export function init(ldClient: LDClient): AIClient {
   return new AIClient(ldClient);
 }
 
-export { LDAIConfigTracker } from './LDAIConfigTracker';
+export * from './api/config/LDAIConfigTracker';
+export * from './api/metrics';

packages/sdk/ai/src/trackUtils.ts

@@ -1,4 +1,4 @@
-import { BedrockTokenUsage, TokenMetrics, TokenUsage, UnderscoreTokenUsage } from './types';
+import { BedrockTokenUsage, TokenMetrics, TokenUsage, UnderscoreTokenUsage } from './api/metrics';
 
 export function usageToTokenMetrics(
   usage: TokenUsage | UnderscoreTokenUsage | BedrockTokenUsage,

packages/sdk/ai/tsconfig.eslint.json

@@ -0,0 +1,5 @@
+{
+	"extends": "./tsconfig.json",
+	"include": ["/**/*.ts"],
+	"exclude": ["node_modules"]
+      }