feat(openai): add tags & metadata at the call level

src/api.ts

@@ -423,7 +423,6 @@ export class API {
 
       return response.data;
     } catch (e) {
-      console.error(e);
       if (e instanceof AxiosError) {
         throw new Error(JSON.stringify(e.response?.data.errors));
       } else {
@@ -451,7 +450,6 @@ export class API {
 
       return response.data;
     } catch (e) {
-      console.error(e);
       if (e instanceof AxiosError) {
         throw new Error(JSON.stringify(e.response?.data));
       } else {

src/instrumentation/index.ts

@@ -1,15 +1,16 @@
-import { LiteralClient } from '..';
+import { LiteralClient, Maybe } from '..';
 import { LiteralCallbackHandler } from './langchain';
 import { instrumentLlamaIndex, withThread } from './llamaindex';
 import instrumentOpenAI from './openai';
-import { InstrumentOpenAIOptions } from './openai';
 import { makeInstrumentVercelSDK } from './vercel-sdk';
 
-export type { InstrumentOpenAIOptions } from './openai';
+export type OpenAIGlobalOptions = {
+  tags?: Maybe<string[]>;
+  metadata?: Maybe<Record<string, any>>;
+};
 
 export default (client: LiteralClient) => ({
-  openai: (options?: InstrumentOpenAIOptions) =>
-    instrumentOpenAI(client, options),
+  openai: (options?: OpenAIGlobalOptions) => instrumentOpenAI(client, options),
   langchain: {
     literalCallback: (threadId?: string) => {
       try {

src/instrumentation/openai.ts

@@ -1,4 +1,5 @@
 import OpenAI from 'openai';
+import { RequestOptions } from 'openai/core';
 import type {
   ChatCompletion,
   ChatCompletionChunk,
@@ -13,40 +14,70 @@ import {
   IGenerationMessage,
   LiteralClient,
   Maybe,
+  OpenAIGlobalOptions,
   StepConstructor
 } from '..';
 
 // Define a generic type for the original function to be wrapped
-type OriginalFunction<T extends any[], R> = (...args: T) => Promise<R>;
+type OriginalFunction<Output> = (
+  body: any,
+  options?: RequestOptions
+) => Promise<Output>;
+
+type OpenAICallOptions = {
+  literalaiTags?: Maybe<string[]>;
+  literalaiMetadata?: Maybe<Record<string, any>>;
+};
+
+function cleanOpenAIArgs(
+  body: any,
+  callOptions?: RequestOptions & OpenAICallOptions
+): [any, RequestOptions?] {
+  if (!callOptions) {
+    return [body];
+  }
+
+  const { literalaiTags, literalaiMetadata, ...restCallOptions } = callOptions;
+  return [body, restCallOptions];
+}
 
 // Utility function to wrap a method
-function wrapFunction<T extends any[], R>(
-  originalFunction: OriginalFunction<T, R>,
+function wrapFunction<Output>(
+  originalFunction: OriginalFunction<Output>,
   client: LiteralClient,
-  options: InstrumentOpenAIOptions = {}
-): OriginalFunction<T, R> {
-  return async function (this: any, ...args: T): Promise<R> {
+  globalOptions: OpenAIGlobalOptions = {}
+) {
+  return async function (
+    this: any,
+    body: any,
+    callOptions?: RequestOptions & OpenAICallOptions
+  ): Promise<Output> {
     const start = Date.now();
 
     // Call the original function
-    const result = await originalFunction.apply(this, args);
+    const result = await originalFunction.apply(
+      this,
+      cleanOpenAIArgs(body, callOptions)
+    );
 
     if (result instanceof Stream) {
       const streamResult = result;
       const [returnedResult, processedResult] = streamResult.tee();
 
       await processOpenAIOutput(client, processedResult, {
-        ...options,
+        globalOptions,
+        callOptions,
         start,
-        inputs: args[0]
+        inputs: body
       });
 
-      return returnedResult as R;
+      return returnedResult as Output;
     } else {
       await processOpenAIOutput(client, result as ChatCompletion | Completion, {
-        ...options,
+        globalOptions,
+        callOptions,
         start,
-        inputs: args[0]
+        inputs: body
       });
 
       return result;
@@ -56,32 +87,51 @@ function wrapFunction<T extends any[], R>(
 
 function instrumentOpenAI(
   client: LiteralClient,
-  options: InstrumentOpenAIOptions = {}
+  options: OpenAIGlobalOptions = {}
 ) {
   // Patching the chat.completions.create function
   const originalChatCompletionsCreate =
     OpenAI.Chat.Completions.prototype.create;
-  OpenAI.Chat.Completions.prototype.create = wrapFunction(
+  const wrappedChatCompletionsCreate = wrapFunction(
     originalChatCompletionsCreate,
     client,
     options
-  ) as any;
+  );
 
   // Patching the completions.create function
   const originalCompletionsCreate = OpenAI.Completions.prototype.create;
-  OpenAI.Completions.prototype.create = wrapFunction(
+  const wrappedCompletionsCreate = wrapFunction(
     originalCompletionsCreate,
     client,
     options
-  ) as any;
+  );
 
   // Patching the images.generate function
   const originalImagesGenerate = OpenAI.Images.prototype.generate;
-  OpenAI.Images.prototype.generate = wrapFunction(
+  const wrappedImagesGenerate = wrapFunction(
     originalImagesGenerate,
     client,
     options
-  ) as any;
+  );
+
+  OpenAI.Chat.Completions.prototype.create =
+    wrappedChatCompletionsCreate as any;
+  OpenAI.Completions.prototype.create = wrappedCompletionsCreate as any;
+  OpenAI.Images.prototype.generate = wrappedImagesGenerate as any;
+
+  return {
+    chat: {
+      completions: {
+        create: wrappedChatCompletionsCreate
+      }
+    },
+    completions: {
+      create: wrappedCompletionsCreate
+    },
+    images: {
+      generate: wrappedImagesGenerate
+    }
+  };
 }
 
 function processChatDelta(
@@ -277,11 +327,9 @@ export type OpenAIOutput =
   | Stream<ChatCompletionChunk>
   | ImagesResponse;
 
-export interface InstrumentOpenAIOptions {
-  tags?: Maybe<string[]>;
-}
-
-export interface ProcessOpenAIOutput extends InstrumentOpenAIOptions {
+export interface ProcessOpenAIOutput {
+  globalOptions: Maybe<OpenAIGlobalOptions>;
+  callOptions: Maybe<OpenAICallOptions>;
   start: number;
   inputs: Record<string, any>;
 }
@@ -299,13 +347,22 @@ function isStream(obj: any): boolean {
 const processOpenAIOutput = async (
   client: LiteralClient,
   output: OpenAIOutput,
-  { start, tags, inputs }: ProcessOpenAIOutput
+  { start, globalOptions, callOptions, inputs }: ProcessOpenAIOutput
 ) => {
+  const tags = [
+    ...(globalOptions?.tags ?? []),
+    ...(callOptions?.literalaiTags ?? [])
+  ];
+  const metadata = {
+    ...globalOptions?.metadata,
+    ...callOptions?.literalaiMetadata
+  };
+
   const baseGeneration = {
     provider: 'openai',
     model: inputs.model,
     settings: getSettings(inputs),
-    tags: tags
+    tags
   };
 
   const threadFromStore = client._currentThread();
@@ -322,7 +379,8 @@ const processOpenAIOutput = async (
       output: output,
       startTime: new Date(start).toISOString(),
       endTime: new Date().toISOString(),
-      tags: tags
+      tags,
+      metadata
     };
 
     const step = parent
@@ -366,7 +424,9 @@ const processOpenAIOutput = async (
           generation,
           output: messageCompletion,
           startTime: new Date(start).toISOString(),
-          endTime: new Date(start + metrics.duration).toISOString()
+          endTime: new Date(start + metrics.duration).toISOString(),
+          tags,
+          metadata
         });
         await step.send();
       } else {
@@ -392,7 +452,9 @@ const processOpenAIOutput = async (
           generation,
           output: { content: completion },
           startTime: new Date(start).toISOString(),
-          endTime: new Date(start + metrics.duration).toISOString()
+          endTime: new Date(start + metrics.duration).toISOString(),
+          tags,
+          metadata
         });
         await step.send();
       } else {
@@ -430,7 +492,9 @@ const processOpenAIOutput = async (
           generation,
           output: messageCompletion,
           startTime: new Date(start).toISOString(),
-          endTime: new Date().toISOString()
+          endTime: new Date().toISOString(),
+          tags,
+          metadata
         });
         await step.send();
       } else {
@@ -454,7 +518,9 @@ const processOpenAIOutput = async (
           generation,
           output: { content: completion },
           startTime: new Date(start).toISOString(),
-          endTime: new Date().toISOString()
+          endTime: new Date().toISOString(),
+          tags,
+          metadata
         });
         await step.send();
       } else {