feat: add initial working version of the OpenAI tracer

Author: gustavocidornelas

src/lib/integrations/openAiTracer.ts

@@ -0,0 +1,65 @@
+import OpenAI from 'openai';
+import { Stream } from 'openai/src/streaming';
+import { addChatCompletionStepToTrace } from '../tracing/tracer';
+
+export function traceOpenAI(openai: OpenAI): OpenAI {
+  const createFunction = openai.chat.completions.create;
+
+  openai.chat.completions.create = async function (
+    this: typeof openai.chat.completions,
+    ...args: Parameters<typeof createFunction>
+  ): Promise<Stream<OpenAI.Chat.Completions.ChatCompletionChunk> | OpenAI.Chat.Completions.ChatCompletion> {
+    const [params, options = { stream: false }] = args;
+    try {
+      const startTime = performance.now();
+      if (options.stream) {
+        console.log('streaming not implemented yet');
+        return createFunction.apply(this, args) as unknown as Promise<
+          Stream<OpenAI.Chat.Completions.ChatCompletionChunk>
+        >;
+      } else {
+        const response = (await createFunction.apply(this, args)) as OpenAI.Chat.Completions.ChatCompletion;
+        const completion = response.choices[0];
+        const endTime = performance.now();
+        const traceData = {
+          name: 'OpenAI Chat Completion',
+          inputs: { prompt: params.messages },
+          output: completion?.message.content,
+          latency: endTime - startTime,
+          tokens: response?.usage?.total_tokens ?? null,
+          promptTokens: response?.usage?.prompt_tokens ?? null,
+          completionTokens: response?.usage?.completion_tokens ?? null,
+          model: response?.model,
+          modelParameters: getModelParameters(args),
+          rawOutput: JSON.stringify(response, null, 2),
+          metadata: {},
+          provider: 'OpenAI',
+        };
+        addChatCompletionStepToTrace(traceData);
+        return response;
+      }
+    } catch (error) {
+      console.error('Failed to trace the create chat completion request with Openlayer', error);
+      throw error;
+    }
+  } as typeof createFunction;
+
+  return openai;
+}
+
+function getModelParameters(args: any): Record<string, any> {
+  const params = args[0];
+  return {
+    frequency_penalty: params?.frequencyPenalty ?? 0,
+    logit_bias: params?.logitBias ?? null,
+    logprobs: params?.logprobs ?? false,
+    top_logprobs: params?.topLogprobs ?? null,
+    max_tokens: params?.maxTokens ?? null,
+    n: params?.n ?? 1,
+    presence_penalty: params?.presencePenalty ?? 0,
+    seed: params?.seed ?? null,
+    stop: params?.stop ?? null,
+    temperature: params?.temperature ?? 1,
+    top_p: params?.topP ?? 1,
+  };
+}

src/lib/tracing/tracer.ts

@@ -69,8 +69,6 @@ function createStep(
       const traceData = getCurrentTrace();
       // Post process trace and get the input variable names
       const { traceData: processedTraceData, inputVariableNames } = postProcessTrace(traceData!);
-      console.log('Processed trace data:', JSON.stringify(processedTraceData, null, 2));
-      console.log('Input variable names:', inputVariableNames);
 
       if (publish && process.env['OPENLAYER_INFERENCE_PIPELINE_ID']) {
         client!.inferencePipelines.data.stream(process.env['OPENLAYER_INFERENCE_PIPELINE_ID'], {
@@ -87,7 +85,6 @@ function createStep(
           rows: [processedTraceData],
         });
       }
-      console.log('Trace data ready for upload:', JSON.stringify(traceData, null, 2));
 
       // Reset the entire trace state
       setCurrentTrace(null);
@@ -134,15 +131,47 @@ function trace(fn: Function, stepType: StepType = StepType.USER_CALL, stepName?:
   };
 }
 
-// Example usage of specialized function to add a chat completion step
-export function addChatCompletionStepToTrace(
-  name: string,
-  inputs: any,
-  output: any,
-  metadata?: Record<string, any>,
-) {
+export function addChatCompletionStepToTrace({
+  name,
+  inputs,
+  output,
+  latency,
+  tokens = null,
+  promptTokens = null,
+  completionTokens = null,
+  model = null,
+  modelParameters = null,
+  rawOutput = null,
+  metadata = {},
+  provider = 'OpenAI',
+}: {
+  name: string;
+  inputs: any;
+  output: any;
+  latency: number;
+  tokens?: number | null;
+  promptTokens?: number | null;
+  completionTokens?: number | null;
+  model?: string | null;
+  modelParameters?: Record<string, any> | null;
+  rawOutput?: string | null;
+  metadata?: Record<string, any>;
+  provider?: string;
+}) {
   const [step, endStep] = createStep(name, StepType.CHAT_COMPLETION, inputs, output, metadata);
-  step.log({ inputs, output });
+
+  if (step instanceof ChatCompletionStep) {
+    step.provider = provider;
+    step.promptTokens = promptTokens;
+    step.completionTokens = completionTokens;
+    step.tokens = tokens;
+    step.model = model;
+    step.modelParameters = modelParameters;
+    step.rawOutput = rawOutput;
+    step.latency = latency;
+  }
+
+  step.log({ inputs, output, metadata });
   endStep();
 }