Better OpenAI types.

Author: kinyoklion

packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts

@@ -28,9 +28,9 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
     this._ldClient.track('$ld:ai:duration:total', this._context, this._getTrackData(), duration);
   }
 
-  async trackDurationOf(func: (...args: any[]) => Promise<any>, ...args: any[]): Promise<any> {
+  async trackDurationOf<TRes>(func: () => Promise<TRes>): Promise<TRes> {
     const startTime = Date.now();
-    const result = await func(...args);
+    const result = await func();
     const endTime = Date.now();
     const duration = endTime - startTime; // duration in milliseconds
     this.trackDuration(duration);
@@ -49,8 +49,16 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
     this._ldClient.track('$ld:ai:generation', this._context, this._getTrackData(), 1);
   }
 
-  async trackOpenAI<TRes>(func: (...args: any[]) => Promise<TRes>, ...args: any[]): Promise<TRes> {
-    const result = await this.trackDurationOf(func, ...args);
+  async trackOpenAI<
+    TRes extends {
+      usage?: {
+        total_tokens?: number;
+        prompt_token?: number;
+        completion_token?: number;
+      };
+    },
+  >(func: () => Promise<TRes>): Promise<TRes> {
+    const result = await this.trackDurationOf(func);
     this.trackSuccess();
     if (result.usage) {
       this.trackTokens(createOpenAiUsage(result.usage));

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

@@ -35,19 +35,27 @@ export interface LDAIConfigTracker {
   /**
    * Track the duration of execution of the provided function.
    * @param func The function to track the duration of.
-   * @param args Arguments for the function.
    * @returns The result of the function.
    */
-  trackDurationOf(func: (...args: any[]) => Promise<any>, ...args: any[]): Promise<any>;
+  trackDurationOf(func: () => Promise<any>): Promise<any>;
 
   /**
    * Track an OpenAI operation.
    *
    * @param func Function which executes the operation.
-   * @param args Arguments for the operation.
    * @returns The result of the operation.
    */
-  trackOpenAI<TRes>(func: (...args: any[]) => Promise<TRes>, ...args: any[]): Promise<TRes>;
+  trackOpenAI<
+    TRes extends {
+      usage?: {
+        total_tokens?: number;
+        prompt_token?: number;
+        completion_token?: number;
+      };
+    },
+  >(
+    func: () => Promise<TRes>,
+  ): Promise<TRes>;
 
   /**
    * Track an operation which uses Bedrock.

packages/sdk/server-ai/src/api/metrics/OpenAiUsage.ts

@@ -1,6 +1,10 @@
 import { LDTokenUsage } from './LDTokenUsage';
 
-export function createOpenAiUsage(data: any): LDTokenUsage {
+export function createOpenAiUsage(data: {
+  total_tokens?: number;
+  prompt_token?: number;
+  completion_token?: number;
+}): LDTokenUsage {
   return {
     total: data.total_tokens ?? 0,
     input: data.prompt_token ?? 0,