From 0a6637cac5ab031380d7c25a75ecb3d489d76df2 Mon Sep 17 00:00:00 2001
From: Nir Gazit <nirga@users.noreply.github.com>
Date: Sun, 28 Jul 2024 19:49:22 +0300
Subject: [PATCH] fix(sdk): properly initialize token enrich value for
 instrumentations

---
 packages/sample-app/package.json              |  1 +
 .../sample-app/src/sample_openai_streaming.ts | 32 +++++++++++++++++++
 .../src/lib/configuration/index.ts            |  4 ---
 .../initialize-options.interface.ts           |  6 ----
 .../traceloop-sdk/src/lib/tracing/index.ts    |  8 +++--
 5 files changed, 39 insertions(+), 12 deletions(-)
 create mode 100644 packages/sample-app/src/sample_openai_streaming.ts

diff --git a/packages/sample-app/package.json b/packages/sample-app/package.json
index 9ec07b7d..c4b229cf 100644
--- a/packages/sample-app/package.json
+++ b/packages/sample-app/package.json
@@ -19,6 +19,7 @@
     "run:prompt_mgmt": "npm run build && node dist/src/sample_prompt_mgmt.js",
     "run:sample_vision": "npm run build && node dist/src/sample_vision_prompt.js",
     "run:sample_azure": "npm run build && node dist/src/sample_azure.js",
+    "run:openai_streaming": "npm run build && node dist/src/sample_openai_streaming.js",
     "run:sampler": "npm run build && node dist/src/sample_sampler.js",
     "run:llamaindex": "npm run build && node dist/src/sample_llamaindex.js",
     "run:pinecone": "npm run build && node dist/src/sample_pinecone.js",
diff --git a/packages/sample-app/src/sample_openai_streaming.ts b/packages/sample-app/src/sample_openai_streaming.ts
new file mode 100644
index 00000000..99c32f06
--- /dev/null
+++ b/packages/sample-app/src/sample_openai_streaming.ts
@@ -0,0 +1,32 @@
+import * as traceloop from "@traceloop/node-server-sdk";
+import OpenAI from "openai";
+
+traceloop.initialize({
+  appName: "sample_openai_streaming",
+  apiKey: process.env.TRACELOOP_API_KEY,
+  disableBatch: true,
+});
+const openai = new OpenAI();
+
+async function create_joke() {
+  const responseStream = await traceloop.withTask(
+    { name: "joke_creation" },
+    () => {
+      return openai.chat.completions.create({
+        model: "gpt-3.5-turbo",
+        messages: [
+          { role: "user", content: "Tell me a joke about opentelemetry" },
+        ],
+        stream: true,
+      });
+    },
+  );
+  let result = "";
+  for await (const chunk of responseStream) {
+    result += chunk.choices[0]?.delta?.content || "";
+  }
+  console.log(result);
+  return result;
+}
+
+create_joke();
diff --git a/packages/traceloop-sdk/src/lib/configuration/index.ts b/packages/traceloop-sdk/src/lib/configuration/index.ts
index 184d8df3..fccb66a3 100644
--- a/packages/traceloop-sdk/src/lib/configuration/index.ts
+++ b/packages/traceloop-sdk/src/lib/configuration/index.ts
@@ -54,10 +54,6 @@ export const initialize = (options: InitializeOptions) => {
       options.traceloopSyncDevPollingInterval =
         Number(process.env.TRACELOOP_SYNC_DEV_POLLING_INTERVAL) || 5;
     }
-
-    if (options.shouldEnrichMetrics === undefined) {
-      options.shouldEnrichMetrics = true;
-    }
   }
 
   validateConfiguration(options);
diff --git a/packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts b/packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts
index 3c080415..e80b7d6f 100644
--- a/packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts
+++ b/packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts
@@ -49,12 +49,6 @@ export interface InitializeOptions {
    */
   logLevel?: "debug" | "info" | "warn" | "error";
 
-  /**
-   * Whether to enrich metrics with additional data like OpenAI token usage for streaming requests. Optional.
-   * Defaults to true.
-   */
-  shouldEnrichMetrics?: boolean;
-
   /**
    * Whether to log prompts, completions and embeddings on traces. Optional.
    * Defaults to true.
diff --git a/packages/traceloop-sdk/src/lib/tracing/index.ts b/packages/traceloop-sdk/src/lib/tracing/index.ts
index c22a572f..4d072f74 100644
--- a/packages/traceloop-sdk/src/lib/tracing/index.ts
+++ b/packages/traceloop-sdk/src/lib/tracing/index.ts
@@ -56,9 +56,11 @@ const instrumentations: Instrumentation[] = [];
 
 export const initInstrumentations = () => {
   const exceptionLogger = (e: Error) => Telemetry.getInstance().logException(e);
+  const enrichTokens =
+    (process.env.TRACELOOP_ENRICH_TOKENS || "true").toLowerCase() === "true";
 
   openAIInstrumentation = new OpenAIInstrumentation({
-    enrichTokens: _configuration?.shouldEnrichMetrics,
+    enrichTokens,
     exceptionLogger,
   });
   instrumentations.push(openAIInstrumentation);
@@ -109,13 +111,15 @@ export const manuallyInitInstrumentations = (
   instrumentModules: InitializeOptions["instrumentModules"],
 ) => {
   const exceptionLogger = (e: Error) => Telemetry.getInstance().logException(e);
+  const enrichTokens =
+    (process.env.TRACELOOP_ENRICH_TOKENS || "true").toLowerCase() === "true";
 
   // Clear the instrumentations array that was initialized by default
   instrumentations.length = 0;
 
   if (instrumentModules?.openAI) {
     openAIInstrumentation = new OpenAIInstrumentation({
-      enrichTokens: _configuration?.shouldEnrichMetrics,
+      enrichTokens,
       exceptionLogger,
     });
     instrumentations.push(openAIInstrumentation);