From eb6211f5cc9a2930c6bff1d9ffe3cb995eaf6669 Mon Sep 17 00:00:00 2001
From: Nir Gazit <nirga@users.noreply.github.com>
Date: Thu, 4 Jul 2024 12:42:05 +0200
Subject: [PATCH] fix(sdk): versions on workflows & tasks (#353)

---
 packages/ai-semantic-conventions/src/SemanticAttributes.ts | 1 +
 packages/traceloop-sdk/src/lib/tracing/decorators.ts       | 6 ++++++
 packages/traceloop-sdk/test/decorators.test.ts             | 6 +++++-
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/packages/ai-semantic-conventions/src/SemanticAttributes.ts b/packages/ai-semantic-conventions/src/SemanticAttributes.ts
index bbaad202..828cbd76 100644
--- 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
index a0675db1..c6bf2879 100644
--- 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
index 7db6d038..6ac40c08 100644
--- 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({