feat: association property in workflows (#36)

traceloop · Jan 17, 2024 · 1ed3a6a · 1ed3a6a
1 parent 018b242
commit 1ed3a6a
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 30 deletions.
diff --git a/packages/sample-app/src/sample_with.ts b/packages/sample-app/src/sample_with.ts
@@ -9,7 +9,7 @@ traceloop.initialize({
 const openai = new OpenAI();
 
 async function chat() {
-  return await traceloop.withWorkflow("sample_chat", async () => {
+  return await traceloop.withWorkflow("sample_chat", {}, async () => {
     const chatCompletion = await openai.chat.completions.create({
       messages: [
         { role: "user", content: "Tell me a joke about OpenTelemetry" },
@@ -22,7 +22,7 @@ async function chat() {
 }
 
 async function completion() {
-  return await traceloop.withWorkflow("sample_completion", async () => {
+  return await traceloop.withWorkflow("sample_completion", {}, async () => {
     const completion = await openai.completions.create({
       prompt: "Tell me a joke about TypeScript",
       model: "gpt-3.5-turbo-instruct",

diff --git a/packages/traceloop-sdk/src/lib/tracing/association.ts b/packages/traceloop-sdk/src/lib/tracing/association.ts
@@ -10,6 +10,10 @@ export function withAssociationProperties<
   thisArg?: ThisParameterType<F>,
   ...args: A
 ) {
+  if (Object.keys(properties).length === 0) {
+    return fn.apply(thisArg, args);
+  }
+
   const newContext = context
     .active()
     .setValue(ASSOCATION_PROPERTIES_KEY, properties);

diff --git a/packages/traceloop-sdk/src/lib/tracing/decorators.ts b/packages/traceloop-sdk/src/lib/tracing/decorators.ts
@@ -4,13 +4,15 @@ import {
   SpanAttributes,
   TraceloopSpanKindValues,
 } from "@traceloop/ai-semantic-conventions";
+import { withAssociationProperties } from "./association";
 
 function withEntity<
   A extends unknown[],
   F extends (...args: A) => ReturnType<F>,
 >(
   type: TraceloopSpanKindValues,
   name: string,
+  associationProperties: { [name: string]: string },
   fn: F,
   thisArg?: ThisParameterType<F>,
   ...args: A
@@ -22,46 +24,57 @@ function withEntity<
       : context.active();
 
   if (fn.constructor.name === "AsyncFunction") {
-    return getTracer().startActiveSpan(
+    return withAssociationProperties(associationProperties, () =>
+      getTracer().startActiveSpan(
+        `${name}.${type}`,
+        {},
+        workflowContext,
+        async (span: Span) => {
+          if (
+            type === TraceloopSpanKindValues.WORKFLOW ||
+            type === TraceloopSpanKindValues.AGENT
+          ) {
+            span.setAttribute(SpanAttributes.TRACELOOP_WORKFLOW_NAME, name);
+          }
+          span.setAttribute(SpanAttributes.TRACELOOP_SPAN_KIND, type);
+          span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name);
+          const res = await fn.apply(thisArg, args);
+          span.end();
+          return res;
+        },
+      ),
+    );
+  }
+  return withAssociationProperties(associationProperties, () =>
+    getTracer().startActiveSpan(
       `${name}.${type}`,
       {},
       workflowContext,
-      async (span: Span) => {
-        if (
-          type === TraceloopSpanKindValues.WORKFLOW ||
-          type === TraceloopSpanKindValues.AGENT
-        ) {
-          span.setAttribute(SpanAttributes.TRACELOOP_WORKFLOW_NAME, name);
-        }
+      (span) => {
         span.setAttribute(SpanAttributes.TRACELOOP_SPAN_KIND, type);
         span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name);
-        const res = await fn.apply(thisArg, args);
+        const res = fn.apply(thisArg, args);
         span.end();
         return res;
       },
-    );
-  }
-  return getTracer().startActiveSpan(
-    `${name}.${type}`,
-    {},
-    workflowContext,
-    (span) => {
-      span.setAttribute(SpanAttributes.TRACELOOP_SPAN_KIND, type);
-      span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name);
-      const res = fn.apply(thisArg, args);
-      span.end();
-      return res;
-    },
+    ),
   );
 }
 
 export function withWorkflow<
   A extends unknown[],
   F extends (...args: A) => ReturnType<F>,
->(name: string, fn: F, thisArg?: ThisParameterType<F>, ...args: A) {
+>(
+  name: string,
+  associationProperties: { [name: string]: string },
+  fn: F,
+  thisArg?: ThisParameterType<F>,
+  ...args: A
+) {
   return withEntity(
     TraceloopSpanKindValues.WORKFLOW,
     name,
+    associationProperties,
     fn,
     thisArg,
     ...args,
@@ -72,21 +85,48 @@ export function withTask<
   A extends unknown[],
   F extends (...args: A) => ReturnType<F>,
 >(name: string, fn: F, thisArg?: ThisParameterType<F>, ...args: A) {
-  return withEntity(TraceloopSpanKindValues.TASK, name, fn, thisArg, ...args);
+  return withEntity(
+    TraceloopSpanKindValues.TASK,
+    name,
+    {},
+    fn,
+    thisArg,
+    ...args,
+  );
 }
 
 export function withAgent<
   A extends unknown[],
   F extends (...args: A) => ReturnType<F>,
->(name: string, fn: F, thisArg?: ThisParameterType<F>, ...args: A) {
-  return withEntity(TraceloopSpanKindValues.AGENT, name, fn, thisArg, ...args);
+>(
+  name: string,
+  associationProperties: { [name: string]: string },
+  fn: F,
+  thisArg?: ThisParameterType<F>,
+  ...args: A
+) {
+  return withEntity(
+    TraceloopSpanKindValues.AGENT,
+    name,
+    associationProperties,
+    fn,
+    thisArg,
+    ...args,
+  );
 }
 
 export function withTool<
   A extends unknown[],
   F extends (...args: A) => ReturnType<F>,
 >(name: string, fn: F, thisArg?: ThisParameterType<F>, ...args: A) {
-  return withEntity(TraceloopSpanKindValues.TOOL, name, fn, thisArg, ...args);
+  return withEntity(
+    TraceloopSpanKindValues.TOOL,
+    name,
+    {},
+    fn,
+    thisArg,
+    ...args,
+  );
 }
 
 function entity(type: TraceloopSpanKindValues, name?: string) {
@@ -103,14 +143,22 @@ function entity(type: TraceloopSpanKindValues, name?: string) {
         return await withEntity(
           type,
           entityName,
+          {},
           originalMethod,
           target,
           ...args,
         );
       };
     } else {
       descriptor.value = function (...args: any[]) {
-        return withEntity(type, entityName, originalMethod, target, ...args);
+        return withEntity(
+          type,
+          entityName,
+          {},
+          originalMethod,
+          target,
+          ...args,
+        );
       };
     }
   };