feat: feat: complete OpenAI tracer implementation, handling streaming and tool calls

Author: Stainless Bot

src/lib/integrations/openAiTracer.ts

@@ -9,27 +9,94 @@ export function traceOpenAI(openai: OpenAI): OpenAI {
     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;
+    const [params, options] = args;
+    const stream = params?.stream ?? false;
+
     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>
-        >;
+
+      // Call the original `create` function
+      let response = await createFunction.apply(this, args);
+
+      if (stream) {
+        // Handle streaming responses
+        const chunks: OpenAI.Chat.Completions.ChatCompletionChunk[] = [];
+        let collectedOutputData: any[] = [];
+        let firstTokenTime: number | undefined;
+        let completionTokens: number = 0;
+        if (isAsyncIterable(response)) {
+          async function* tracedOutputGenerator(): AsyncGenerator<
+            OpenAI.Chat.Completions.ChatCompletionChunk,
+            void,
+            unknown
+          > {
+            for await (const rawChunk of response as AsyncIterable<OpenAI.Chat.Completions.ChatCompletionChunk>) {
+              if (chunks.length === 0) {
+                firstTokenTime = performance.now();
+              }
+              chunks.push(rawChunk);
+              const delta = rawChunk.choices[0]?.delta;
+              if (delta?.content) {
+                collectedOutputData.push(delta?.content);
+              } else if (delta?.tool_calls) {
+                const tool_call = delta.tool_calls[0];
+                if (tool_call?.function?.name) {
+                  const functionName: string =
+                    '{\n    "name": ' + '"' + tool_call.function.name + '"' + '\n    "arguments": ';
+                  collectedOutputData.push(functionName);
+                } else if (tool_call?.function?.arguments) {
+                  collectedOutputData.push(tool_call.function.arguments);
+                }
+              }
+
+              if (rawChunk.choices[0]?.finish_reason === 'tool_calls') {
+                collectedOutputData.push('\n}');
+              }
+              completionTokens += 1;
+              yield rawChunk;
+            }
+            const endTime = performance.now();
+            const traceData = {
+              name: 'OpenAI Chat Completion',
+              inputs: { prompt: params.messages },
+              output: collectedOutputData.join(''),
+              latency: endTime - startTime,
+              model: chunks[0]?.model as string,
+              modelParameters: getModelParameters(args),
+              rawOutput: chunks.map((chunk) => JSON.stringify(chunk, null, 2)).join('\n'),
+              metadata: { timeToFistToken: firstTokenTime ? firstTokenTime - startTime : null },
+              provider: 'OpenAI',
+              completionTokens: completionTokens,
+              promptTokens: 0,
+              tokens: completionTokens,
+            };
+            addChatCompletionStepToTrace(traceData);
+          }
+          return tracedOutputGenerator() as unknown as Stream<OpenAI.Chat.Completions.ChatCompletionChunk>;
+        }
       } else {
-        const response = (await createFunction.apply(this, args)) as OpenAI.Chat.Completions.ChatCompletion;
+        // Handle non-streaming responses
+        response = response as OpenAI.Chat.Completions.ChatCompletion;
         const completion = response.choices[0];
         const endTime = performance.now();
+
+        let output: string = '';
+        if (completion?.message?.content) {
+          output = completion.message.content;
+        } else if (completion?.message.tool_calls) {
+          const tool_call = completion.message.tool_calls[0];
+          output = JSON.stringify(tool_call?.function, null, 2);
+        }
+
         const traceData = {
           name: 'OpenAI Chat Completion',
           inputs: { prompt: params.messages },
-          output: completion?.message.content,
+          output: output,
           latency: endTime - startTime,
-          tokens: response?.usage?.total_tokens ?? null,
-          promptTokens: response?.usage?.prompt_tokens ?? null,
-          completionTokens: response?.usage?.completion_tokens ?? null,
-          model: response?.model,
+          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: {},
@@ -42,6 +109,8 @@ export function traceOpenAI(openai: OpenAI): OpenAI {
       console.error('Failed to trace the create chat completion request with Openlayer', error);
       throw error;
     }
+    // Ensure a return statement is present
+    return undefined as any;
   } as typeof createFunction;
 
   return openai;
@@ -63,3 +132,6 @@ function getModelParameters(args: any): Record<string, any> {
     top_p: params?.topP ?? 1,
   };
 }
+
+const isAsyncIterable = (x: any) =>
+  x != null && typeof x === 'object' && typeof x[Symbol.asyncIterator] === 'function';

src/lib/tracing/tracer.ts

@@ -10,7 +10,7 @@ let currentTrace: Trace | null = null;
 const publish = process.env['OPENLAYER_DISABLE_PUBLISH'] != 'true';
 let client: Openlayer | null = null;
 if (publish) {
-  console.log('Publishing is enabled');
+  console.debug('Publishing is enabled');
   client = new Openlayer();
 }
 
@@ -45,13 +45,13 @@ function createStep(
   const isRootStep = parentStep === null;
 
   if (isRootStep) {
-    console.log('Starting a new trace...');
-    console.log(`Adding step ${name} as the root step`);
+    console.debug('Starting a new trace...');
+    console.debug(`Adding step ${name} as the root step`);
     const currentTrace = new Trace();
     setCurrentTrace(currentTrace);
     currentTrace.addStep(newStep);
   } else {
-    console.log(`Adding step ${name} as a nested step to ${parentStep!.name}`);
+    console.debug(`Adding step ${name} as a nested step to ${parentStep!.name}`);
     currentTrace = getCurrentTrace()!;
     parentStep!.addNestedStep(newStep);
   }
@@ -65,7 +65,7 @@ function createStep(
     stepStack.pop(); // Remove the current step from the stack
 
     if (isRootStep) {
-      console.log('Ending the trace...');
+      console.debug('Ending the trace...');
       const traceData = getCurrentTrace();
       // Post process trace and get the input variable names
       const { traceData: processedTraceData, inputVariableNames } = postProcessTrace(traceData!);
@@ -90,7 +90,7 @@ function createStep(
       setCurrentTrace(null);
       stepStack.length = 0; // Clear the step stack
     } else {
-      console.log(`Ending step ${name}`);
+      console.debug(`Ending step ${name}`);
     }
   };