fix(sdk): properly initialize token enrich value for instrumentations (…

…#384)
traceloop · Jul 28, 2024 · 143bc66 · 143bc66
1 parent 1e404b7
commit 143bc66
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 12 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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);