fix(sdk-js) Chained entity name on Task & Tool spans (#383)

OzBenSimhonTraceloop · web-flow · commit a191c9d97975 · 2024-07-24T15:39:29.000+03:00
diff --git a/packages/traceloop-sdk/recordings/Test-SDK-Decorators_847855269/should-create-workflow-and-tasks-spans-with-chained-entity-names_971051426/recording.har b/packages/traceloop-sdk/recordings/Test-SDK-Decorators_847855269/should-create-workflow-and-tasks-spans-with-chained-entity-names_971051426/recording.har
@@ -0,0 +1,236 @@
+{
+  "log": {
+    "_recordingName": "Test SDK Decorators/should create workflow and tasks spans with chained entity names",
+    "creator": {
+      "comment": "persister:fs",
+      "name": "Polly.JS",
+      "version": "6.0.6"
+    },
+    "entries": [
+      {
+        "_id": "1f5e5f3a0f77648b2315477cbbf33879",
+        "_order": 0,
+        "cache": {},
+        "request": {
+          "bodySize": 133,
+          "cookies": [],
+          "headers": [
+            {
+              "_fromType": "array",
+              "name": "content-length",
+              "value": "133"
+            },
+            {
+              "_fromType": "array",
+              "name": "accept",
+              "value": "application/json"
+            },
+            {
+              "_fromType": "array",
+              "name": "content-type",
+              "value": "application/json"
+            },
+            {
+              "_fromType": "array",
+              "name": "user-agent",
+              "value": "OpenAI/JS 4.51.0"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-lang",
+              "value": "js"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-package-version",
+              "value": "4.51.0"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-os",
+              "value": "MacOS"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-arch",
+              "value": "arm64"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-runtime",
+              "value": "node"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-runtime-version",
+              "value": "v22.0.0"
+            },
+            {
+              "_fromType": "array",
+              "name": "accept-encoding",
+              "value": "gzip,deflate"
+            },
+            {
+              "name": "host",
+              "value": "api.openai.com"
+            }
+          ],
+          "headersSize": 469,
+          "httpVersion": "HTTP/1.1",
+          "method": "POST",
+          "postData": {
+            "mimeType": "application/json",
+            "params": [],
+            "text": "{\n  \"messages\": [\n    {\n      \"role\": \"user\",\n      \"content\": \"Tell me a joke about pirates\"\n    }\n  ],\n  \"model\": \"gpt-3.5-turbo\"\n}"
+          },
+          "queryString": [],
+          "url": "https://api.openai.com/v1/chat/completions"
+        },
+        "response": {
+          "bodySize": 451,
+          "content": {
+            "encoding": "base64",
+            "mimeType": "application/json",
+            "size": 451,
+            "text": "[\"H4sIAAAAAAAAA1SQwU7DMBBE7/mKxecGNalCaS4IJJA4cQDUIoQqx9kmBsfr2luJUvXfkdO0hcseZjyzb71LAISuRQlCtZJV50w6o/v8+SFU6/l8dfWq5k/ju8efFnGxvl28iVFMUPWJio+pS0WdM8ia7MFWHiVjbM2meTYtppNZ3hsd1WhirHGcTi6LlDe+onSc5cWQbEkrDKKE9wQAYNfPyGhr/BYljEdHpcMQZIOiPD0CEJ5MVIQMQQeWlsXobCqyjLbHnrdbqHUN\",\"3CI47SUjNARMvVBJfwMvBA0ytDpAQAkGm3Ahhq79CcJQ4zxVEdhujDnpK211aJceZSAbFwYmd4jvE4CP/tjNP37hPHWOl0xfaGNhNjnUifP3/jGLwWRiac56fp0MfCJsA2O3XGnboHde95dHymSf/AIAAP//AwC01euq+AEAAA==\"]"
+          },
+          "cookies": [
+            {
+              "domain": ".api.openai.com",
+              "expires": "2024-07-23T18:26:32.000Z",
+              "httpOnly": true,
+              "name": "__cf_bm",
+              "path": "/",
+              "sameSite": "None",
+              "secure": true,
+              "value": "H91GR4kfVHvrZNLG6C4p7cTy0Rje3q95RXaNRra.E7w-1721757392-1.0.1.1-THuuXeeZjJ_kAnzn8lv.3KorJ7IgwtB9_uFiWMgJmsAZgN16Q6Y3ilq2JBkbQtS8SyVJjCY9Dup01866tcpTYQ"
+            },
+            {
+              "domain": ".api.openai.com",
+              "httpOnly": true,
+              "name": "_cfuvid",
+              "path": "/",
+              "sameSite": "None",
+              "secure": true,
+              "value": ".sNQn3QGP675mXVWZPtdu1NGO9R3ZBNRvsCPus1F6mY-1721757392468-0.0.1.1-604800000"
+            }
+          ],
+          "headers": [
+            {
+              "name": "date",
+              "value": "Tue, 23 Jul 2024 17:56:32 GMT"
+            },
+            {
+              "name": "content-type",
+              "value": "application/json"
+            },
+            {
+              "name": "transfer-encoding",
+              "value": "chunked"
+            },
+            {
+              "name": "connection",
+              "value": "keep-alive"
+            },
+            {
+              "name": "openai-organization",
+              "value": "traceloop"
+            },
+            {
+              "name": "openai-processing-ms",
+              "value": "311"
+            },
+            {
+              "name": "openai-version",
+              "value": "2020-10-01"
+            },
+            {
+              "name": "strict-transport-security",
+              "value": "max-age=15552000; includeSubDomains; preload"
+            },
+            {
+              "name": "x-ratelimit-limit-requests",
+              "value": "5000"
+            },
+            {
+              "name": "x-ratelimit-limit-tokens",
+              "value": "4000000"
+            },
+            {
+              "name": "x-ratelimit-remaining-requests",
+              "value": "4999"
+            },
+            {
+              "name": "x-ratelimit-remaining-tokens",
+              "value": "3999976"
+            },
+            {
+              "name": "x-ratelimit-reset-requests",
+              "value": "12ms"
+            },
+            {
+              "name": "x-ratelimit-reset-tokens",
+              "value": "0s"
+            },
+            {
+              "name": "x-request-id",
+              "value": "req_992d6f84a8255b5722ef98d224d589a3"
+            },
+            {
+              "name": "cf-cache-status",
+              "value": "DYNAMIC"
+            },
+            {
+              "_fromType": "array",
+              "name": "set-cookie",
+              "value": "__cf_bm=H91GR4kfVHvrZNLG6C4p7cTy0Rje3q95RXaNRra.E7w-1721757392-1.0.1.1-THuuXeeZjJ_kAnzn8lv.3KorJ7IgwtB9_uFiWMgJmsAZgN16Q6Y3ilq2JBkbQtS8SyVJjCY9Dup01866tcpTYQ; path=/; expires=Tue, 23-Jul-24 18:26:32 GMT; domain=.api.openai.com; HttpOnly; Secure; SameSite=None"
+            },
+            {
+              "_fromType": "array",
+              "name": "set-cookie",
+              "value": "_cfuvid=.sNQn3QGP675mXVWZPtdu1NGO9R3ZBNRvsCPus1F6mY-1721757392468-0.0.1.1-604800000; path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None"
+            },
+            {
+              "name": "x-content-type-options",
+              "value": "nosniff"
+            },
+            {
+              "name": "server",
+              "value": "cloudflare"
+            },
+            {
+              "name": "cf-ray",
+              "value": "8a7d8c339925b077-ATL"
+            },
+            {
+              "name": "content-encoding",
+              "value": "gzip"
+            },
+            {
+              "name": "alt-svc",
+              "value": "h3=\":443\"; ma=86400"
+            }
+          ],
+          "headersSize": 1143,
+          "httpVersion": "HTTP/1.1",
+          "redirectURL": "",
+          "status": 200,
+          "statusText": "OK"
+        },
+        "startedDateTime": "2024-07-23T17:56:31.166Z",
+        "time": 1192,
+        "timings": {
+          "blocked": -1,
+          "connect": -1,
+          "dns": -1,
+          "receive": 0,
+          "send": 0,
+          "ssl": -1,
+          "wait": 1192
+        }
+      }
+    ],
+    "pages": [],
+    "version": "1.2"
+  }
+}
diff --git a/packages/traceloop-sdk/src/lib/tracing/decorators.ts b/packages/traceloop-sdk/src/lib/tracing/decorators.ts
@@ -1,6 +1,8 @@
 import { Span, context } from "@opentelemetry/api";
 import {
   ASSOCATION_PROPERTIES_KEY,
+  ENTITY_NAME_KEY,
+  getChainedEntityName,
   getTracer,
   WORKFLOW_NAME_KEY,
 } from "./tracing";
@@ -43,6 +45,16 @@ function withEntity<
   ) {
     entityContext = entityContext.setValue(WORKFLOW_NAME_KEY, name);
   }
+
+  let entityName = name;
+  if (
+    type === TraceloopSpanKindValues.TOOL ||
+    type === TraceloopSpanKindValues.TASK
+  ) {
+    entityName = getChainedEntityName(entityContext, name);
+    entityContext = entityContext.setValue(ENTITY_NAME_KEY, entityName);
+  }
+
   if (overrideTraceContent != undefined) {
     entityContext = entityContext.setValue(
       CONTEXT_KEY_ALLOW_TRACE_CONTENT,
@@ -68,8 +80,8 @@ function withEntity<
         ) {
           span.setAttribute(SpanAttributes.TRACELOOP_WORKFLOW_NAME, name);
         }
+        span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_NAME, entityName);
         span.setAttribute(SpanAttributes.TRACELOOP_SPAN_KIND, type);
-        span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name);
 
         if (version) {
           span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_VERSION, version);
diff --git a/packages/traceloop-sdk/src/lib/tracing/index.ts b/packages/traceloop-sdk/src/lib/tracing/index.ts
@@ -11,7 +11,11 @@ import { Resource } from "@opentelemetry/resources";
 import { SEMRESATTRS_SERVICE_NAME } from "@opentelemetry/semantic-conventions";
 import { Instrumentation } from "@opentelemetry/instrumentation";
 import { InitializeOptions } from "../interfaces";
-import { ASSOCATION_PROPERTIES_KEY, WORKFLOW_NAME_KEY } from "./tracing";
+import {
+  ASSOCATION_PROPERTIES_KEY,
+  ENTITY_NAME_KEY,
+  WORKFLOW_NAME_KEY,
+} from "./tracing";
 import { Telemetry } from "../telemetry/telemetry";
 import { _configuration } from "../configuration";
 import {
@@ -257,6 +261,14 @@ export const startTracing = (options: InitializeOptions) => {
       );
     }
 
+    const entityName = context.active().getValue(ENTITY_NAME_KEY);
+    if (entityName) {
+      span.setAttribute(
+        SpanAttributes.TRACELOOP_ENTITY_NAME,
+        entityName as string,
+      );
+    }
+
     const associationProperties = context
       .active()
       .getValue(ASSOCATION_PROPERTIES_KEY);
diff --git a/packages/traceloop-sdk/src/lib/tracing/tracing.ts b/packages/traceloop-sdk/src/lib/tracing/tracing.ts
@@ -1,11 +1,22 @@
 import { trace, createContextKey } from "@opentelemetry/api";
+import { Context } from "@opentelemetry/api/build/src/context/types";
 
 const TRACER_NAME = "traceloop.tracer";
 export const WORKFLOW_NAME_KEY = createContextKey("workflow_name");
+export const ENTITY_NAME_KEY = createContextKey("entity_name");
 export const ASSOCATION_PROPERTIES_KEY = createContextKey(
   "association_properties",
 );
 
 export const getTracer = () => {
   return trace.getTracer(TRACER_NAME);
 };
+
+export const getChainedEntityName = (
+  entityContext: Context,
+  entityName: string,
+): string => {
+  const parentEntityName = entityContext.getValue(ENTITY_NAME_KEY);
+
+  return parentEntityName ? `${parentEntityName}.${entityName}` : entityName;
+};
diff --git a/packages/traceloop-sdk/test/decorators.test.ts b/packages/traceloop-sdk/test/decorators.test.ts
@@ -537,4 +537,62 @@ describe("Test SDK Decorators", () => {
       "456",
     );
   });
+
+  it("should create workflow and tasks spans with chained entity names", async () => {
+    class TestOpenAI {
+      @traceloop.workflow({ name: "joke_creation_chat", version: 3 })
+      async chat() {
+        return await this.jokeCreationTaskWrapper();
+      }
+
+      @traceloop.task({ name: "joke_creation_task_wrapper", version: 2 })
+      async jokeCreationTaskWrapper() {
+        return await this.jokeCreation();
+      }
+
+      @traceloop.task({ name: "joke_creation", version: 2 })
+      async jokeCreation() {
+        const chatCompletion = await openai.chat.completions.create({
+          messages: [{ role: "user", content: "Tell me a joke about pirates" }],
+          model: "gpt-3.5-turbo",
+        });
+
+        return chatCompletion.choices[0].message.content;
+      }
+    }
+
+    const testOpenAI = new TestOpenAI();
+    const result = await testOpenAI.chat();
+    const spans = memoryExporter.getFinishedSpans();
+    const jokeCreationTaskWrapperSpan = spans.find(
+      (span) => span.name === "joke_creation_task_wrapper.task",
+    );
+    const jokeCreationSpan = spans.find(
+      (span) => span.name === "joke_creation.task",
+    );
+    const openAiChatSpans = spans.find((span) => span.name === "openai.chat");
+
+    assert.ok(result);
+    assert.ok(jokeCreationTaskWrapperSpan);
+    assert.ok(jokeCreationSpan);
+    assert.ok(openAiChatSpans);
+    assert.strictEqual(
+      jokeCreationTaskWrapperSpan.attributes[
+        `${SpanAttributes.TRACELOOP_ENTITY_NAME}`
+      ],
+      "joke_creation_task_wrapper",
+    );
+    assert.strictEqual(
+      jokeCreationSpan.attributes[`${SpanAttributes.TRACELOOP_ENTITY_NAME}`],
+      "joke_creation_task_wrapper.joke_creation",
+    );
+    assert.strictEqual(
+      openAiChatSpans.attributes[`${SpanAttributes.TRACELOOP_ENTITY_NAME}`],
+      "joke_creation_task_wrapper.joke_creation",
+    );
+    assert.strictEqual(
+      jokeCreationSpan.attributes[`${SpanAttributes.TRACELOOP_ENTITY_OUTPUT}`],
+      JSON.stringify(result),
+    );
+  });
 });
