fix(sdk): versions on workflows & tasks (#353)

traceloop · Jul 4, 2024 · eb6211f · eb6211f
1 parent 75c68ce
commit eb6211f
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 1 deletion.
diff --git a/packages/ai-semantic-conventions/src/SemanticAttributes.ts b/packages/ai-semantic-conventions/src/SemanticAttributes.ts
@@ -51,6 +51,7 @@ export const SpanAttributes = {
   TRACELOOP_SPAN_KIND: "traceloop.span.kind",
   TRACELOOP_WORKFLOW_NAME: "traceloop.workflow.name",
   TRACELOOP_ENTITY_NAME: "traceloop.entity.name",
+  TRACELOOP_ENTITY_VERSION: "traceloop.entity.version",
   TRACELOOP_ASSOCIATION_PROPERTIES: "traceloop.association.properties",
   TRACELOOP_ENTITY_INPUT: "traceloop.entity.input",
   TRACELOOP_ENTITY_OUTPUT: "traceloop.entity.output",

diff --git a/packages/traceloop-sdk/src/lib/tracing/decorators.ts b/packages/traceloop-sdk/src/lib/tracing/decorators.ts
@@ -14,6 +14,7 @@ import { Telemetry } from "../telemetry/telemetry";
 
 export type DecoratorConfig = {
   name: string;
+  version?: number;
   associationProperties?: { [name: string]: string };
   traceContent?: boolean;
   inputParameters?: unknown[];
@@ -26,6 +27,7 @@ function withEntity<
   type: TraceloopSpanKindValues,
   {
     name,
+    version,
     associationProperties,
     traceContent: overrideTraceContent,
     inputParameters,
@@ -69,6 +71,10 @@ function withEntity<
         span.setAttribute(SpanAttributes.TRACELOOP_SPAN_KIND, type);
         span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name);
 
+        if (version) {
+          span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_VERSION, version);
+        }
+
         if (shouldSendTraces()) {
           try {
             const input = inputParameters ?? args;

diff --git a/packages/traceloop-sdk/test/decorators.test.ts b/packages/traceloop-sdk/test/decorators.test.ts
@@ -147,7 +147,7 @@ describe("Test SDK Decorators", () => {
 
   it("should create spans for workflows using decoration syntax", async () => {
     class TestOpenAI {
-      @traceloop.workflow({ name: "sample_chat" })
+      @traceloop.workflow({ name: "sample_chat", version: 2 })
       async chat(things: Map<string, string>) {
         const generations: Map<string, string> = new Map();
         for await (const [key, value] of things) {
@@ -195,6 +195,10 @@ describe("Test SDK Decorators", () => {
       workflowSpan.attributes[`${SpanAttributes.TRACELOOP_ENTITY_NAME}`],
       "sample_chat",
     );
+    assert.strictEqual(
+      workflowSpan.attributes[`${SpanAttributes.TRACELOOP_ENTITY_VERSION}`],
+      2,
+    );
     assert.strictEqual(
       workflowSpan.attributes[`${SpanAttributes.TRACELOOP_ENTITY_INPUT}`],
       JSON.stringify({