fix(sdk): propagate association properties within a workflow (#318)

traceloop · Jun 12, 2024 · 3e530bc · 3e530bc
1 parent 1314a24
commit 3e530bc
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 6 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/sample-app/package.json b/packages/sample-app/package.json
@@ -43,7 +43,7 @@
     "cohere-ai": "^7.7.5",
     "langchain": "^0.1.7",
     "llamaindex": "^0.1.16",
-    "openai": "^4.12.4",
+    "openai": "^4.28.4",
     "chromadb": "^1.8.1"
   },
   "private": true,

diff --git a/packages/traceloop-sdk/src/lib/tracing/decorators.ts b/packages/traceloop-sdk/src/lib/tracing/decorators.ts
@@ -1,11 +1,14 @@
 import { Span, context } from "@opentelemetry/api";
-import { getTracer, WORKFLOW_NAME_KEY } from "./tracing";
+import {
+  ASSOCATION_PROPERTIES_KEY,
+  getTracer,
+  WORKFLOW_NAME_KEY,
+} from "./tracing";
 import {
   CONTEXT_KEY_ALLOW_TRACE_CONTENT,
   SpanAttributes,
   TraceloopSpanKindValues,
 } from "@traceloop/ai-semantic-conventions";
-import { withAssociationProperties } from "./association";
 import { shouldSendTraces } from ".";
 import { Telemetry } from "../telemetry/telemetry";
 
@@ -44,8 +47,14 @@ function withEntity<
       overrideTraceContent,
     );
   }
+  if (associationProperties) {
+    entityContext = entityContext.setValue(
+      ASSOCATION_PROPERTIES_KEY,
+      associationProperties,
+    );
+  }
 
-  return withAssociationProperties(associationProperties ?? {}, () =>
+  return context.with(entityContext, () =>
     getTracer().startActiveSpan(
       `${name}.${type}`,
       {},

diff --git a/packages/traceloop-sdk/test/decorators.test.ts b/packages/traceloop-sdk/test/decorators.test.ts
@@ -76,7 +76,7 @@ describe("Test SDK Decorators", () => {
   it("should create spans for workflows using withWorkflow syntax", async () => {
     const jokeSubject = "OpenTelemetry";
     const result = await traceloop.withWorkflow(
-      { name: "sample_chat" },
+      { name: "sample_chat", associationProperties: { userId: "123" } },
       async () => {
         const chatCompletion = await openai.chat.completions.create({
           messages: [
@@ -110,6 +110,12 @@ describe("Test SDK Decorators", () => {
       workflowSpan.attributes[`${SpanAttributes.TRACELOOP_ENTITY_NAME}`],
       "sample_chat",
     );
+    assert.strictEqual(
+      workflowSpan.attributes[
+        `${SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.userId`
+      ],
+      "123",
+    );
     assert.strictEqual(
       workflowSpan.attributes[`${SpanAttributes.TRACELOOP_ENTITY_INPUT}`],
       JSON.stringify({ args: [], kwargs: { jokeSubject: "OpenTelemetry" } }),
@@ -119,6 +125,12 @@ describe("Test SDK Decorators", () => {
       JSON.stringify(result),
     );
     assert.ok(chatSpan);
+    assert.strictEqual(
+      chatSpan.attributes[
+        `${SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.userId`
+      ],
+      "123",
+    );
     assert.strictEqual(
       chatSpan.attributes[`${SpanAttributes.TRACELOOP_WORKFLOW_NAME}`],
       "sample_chat",