feat: association properties; another API for workflows

packages/ai-semantic-conventions/src/SemanticAttributes.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-export const SemanticAttributes = {
+export const SpanAttributes = {
   LLM_VENDOR: "llm.vendor",
   LLM_REQUEST_TYPE: "llm.request.type",
   LLM_REQUEST_MODEL: "llm.request.model",
@@ -39,7 +39,7 @@ export const SemanticAttributes = {
   TRACELOOP_SPAN_KIND: "traceloop.span.kind",
   TRACELOOP_WORKFLOW_NAME: "traceloop.workflow.name",
   TRACELOOP_ENTITY_NAME: "traceloop.entity.name",
-  TRACELOOP_CORRELATION_ID: "traceloop.correlation.id",
+  TRACELOOP_ASSOCIATION_PROPERTIES: "traceloop.association.properties",
 };
 
 export enum LLMRequestTypeValues {

packages/instrumentation-openai/src/instrumentation.ts

@@ -28,7 +28,7 @@ import {
   InstrumentationNodeModuleDefinition,
   safeExecuteInTheMiddle,
 } from "@opentelemetry/instrumentation";
-import { SemanticAttributes } from "@traceloop/ai-semantic-conventions";
+import { SpanAttributes } from "@traceloop/ai-semantic-conventions";
 import { OpenAIInstrumentationConfig } from "./types";
 import {
   ChatCompletion,
@@ -133,42 +133,40 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
         params: CompletionCreateParamsNonStreaming;
       }): Span {
     const attributes: Attributes = {
-      [SemanticAttributes.LLM_VENDOR]: "OpenAI",
-      [SemanticAttributes.LLM_REQUEST_TYPE]: type,
+      [SpanAttributes.LLM_VENDOR]: "OpenAI",
+      [SpanAttributes.LLM_REQUEST_TYPE]: type,
     };
 
-    attributes[SemanticAttributes.LLM_REQUEST_MODEL] = params.model;
+    attributes[SpanAttributes.LLM_REQUEST_MODEL] = params.model;
     if (params.max_tokens) {
-      attributes[SemanticAttributes.LLM_REQUEST_MAX_TOKENS] = params.max_tokens;
+      attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] = params.max_tokens;
     }
     if (params.temperature) {
-      attributes[SemanticAttributes.LLM_TEMPERATURE] = params.temperature;
+      attributes[SpanAttributes.LLM_TEMPERATURE] = params.temperature;
     }
     if (params.top_p) {
-      attributes[SemanticAttributes.LLM_TOP_P] = params.top_p;
+      attributes[SpanAttributes.LLM_TOP_P] = params.top_p;
     }
     if (params.frequency_penalty) {
-      attributes[SemanticAttributes.LLM_FREQUENCY_PENALTY] =
+      attributes[SpanAttributes.LLM_FREQUENCY_PENALTY] =
         params.frequency_penalty;
     }
     if (params.presence_penalty) {
-      attributes[SemanticAttributes.LLM_PRESENCE_PENALTY] =
-        params.presence_penalty;
+      attributes[SpanAttributes.LLM_PRESENCE_PENALTY] = params.presence_penalty;
     }
 
     if (shouldSendPrompts()) {
       if (type === "chat") {
         params.messages.forEach((message, index) => {
-          attributes[`${SemanticAttributes.LLM_PROMPTS}.${index}.role`] =
+          attributes[`${SpanAttributes.LLM_PROMPTS}.${index}.role`] =
             message.role;
-          attributes[`${SemanticAttributes.LLM_PROMPTS}.${index}.content`] =
+          attributes[`${SpanAttributes.LLM_PROMPTS}.${index}.content`] =
             message.content || "";
         });
       } else {
         if (typeof params.prompt === "string") {
-          attributes[`${SemanticAttributes.LLM_PROMPTS}.0.role`] = "user";
-          attributes[`${SemanticAttributes.LLM_PROMPTS}.0.content`] =
-            params.prompt;
+          attributes[`${SpanAttributes.LLM_PROMPTS}.0.role`] = "user";
+          attributes[`${SpanAttributes.LLM_PROMPTS}.0.content`] = params.prompt;
         }
       }
     }
@@ -217,18 +215,18 @@ function endSpan({
 }:
   | { span: Span; type: "chat"; result: ChatCompletion }
   | { span: Span; type: "completion"; result: Completion }) {
-  span.setAttribute(SemanticAttributes.LLM_RESPONSE_MODEL, result.model);
+  span.setAttribute(SpanAttributes.LLM_RESPONSE_MODEL, result.model);
   if (result.usage) {
     span.setAttribute(
-      SemanticAttributes.LLM_USAGE_TOTAL_TOKENS,
+      SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
       result.usage?.total_tokens,
     );
     span.setAttribute(
-      SemanticAttributes.LLM_USAGE_COMPLETION_TOKENS,
+      SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
       result.usage?.completion_tokens,
     );
     span.setAttribute(
-      SemanticAttributes.LLM_USAGE_PROMPT_TOKENS,
+      SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
       result.usage?.prompt_tokens,
     );
   }
@@ -237,41 +235,41 @@ function endSpan({
     if (type === "chat") {
       result.choices.forEach((choice, index) => {
         span.setAttribute(
-          `${SemanticAttributes.LLM_COMPLETIONS}.${index}.finish_reason`,
+          `${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`,
           choice.finish_reason,
         );
         span.setAttribute(
-          `${SemanticAttributes.LLM_COMPLETIONS}.${index}.role`,
+          `${SpanAttributes.LLM_COMPLETIONS}.${index}.role`,
           choice.message.role,
         );
         span.setAttribute(
-          `${SemanticAttributes.LLM_COMPLETIONS}.${index}.content`,
+          `${SpanAttributes.LLM_COMPLETIONS}.${index}.content`,
           choice.message.content ?? "",
         );
 
         if (choice.message.function_call) {
           span.setAttribute(
-            `${SemanticAttributes.LLM_COMPLETIONS}.${index}.function_call.name`,
+            `${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.name`,
             choice.message.function_call.name,
           );
           span.setAttribute(
-            `${SemanticAttributes.LLM_COMPLETIONS}.${index}.function_call.arguments`,
+            `${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.arguments`,
             choice.message.function_call.arguments,
           );
         }
       });
     } else {
       result.choices.forEach((choice, index) => {
         span.setAttribute(
-          `${SemanticAttributes.LLM_COMPLETIONS}.${index}.finish_reason`,
+          `${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`,
           choice.finish_reason,
         );
         span.setAttribute(
-          `${SemanticAttributes.LLM_COMPLETIONS}.${index}.role`,
+          `${SpanAttributes.LLM_COMPLETIONS}.${index}.role`,
           "assistant",
         );
         span.setAttribute(
-          `${SemanticAttributes.LLM_COMPLETIONS}.${index}.content`,
+          `${SpanAttributes.LLM_COMPLETIONS}.${index}.content`,
           choice.text,
         );
       });

packages/sample-app/package.json

@@ -1,10 +1,11 @@
 {
   "name": "sample-app",
-  "version": "0.0.13-alpha.0",
+  "version": "0.0.1",
   "description": "Sample app for using Traceloop SDK",
   "scripts": {
     "build": "tsc --build tsconfig.json",
-    "run:openai": "npm run build && node dist/src/sample_openai.js",
+    "run:decorators": "npm run build && node dist/src/sample_decorators.js",
+    "run:with": "npm run build && node dist/src/sample_with.js",
     "lint": "eslint . --ext .ts",
     "lint:fix": "eslint . --ext .ts --fix"
   },

packages/sample-app/src/sample_openai.ts → packages/sample-app/src/sample_decorators.ts

@@ -18,7 +18,7 @@ class SampleOpenAI {
       model: "gpt-3.5-turbo",
     });
 
-    console.log(chatCompletion.choices[0].message.content);
+    return chatCompletion.choices[0].message.content;
   }
 
   @traceloop.workflow("sample_completion")
@@ -28,10 +28,15 @@ class SampleOpenAI {
       model: "gpt-3.5-turbo-instruct",
     });
 
-    console.log(completion.choices[0].text);
+    return completion.choices[0].text;
   }
 }
 
-const sampleOpenAI = new SampleOpenAI();
-sampleOpenAI.chat();
-sampleOpenAI.completion();
+traceloop.withAssociationProperties({ userId: "12345" }, async () => {
+  const sampleOpenAI = new SampleOpenAI();
+  const chat = await sampleOpenAI.chat();
+  console.log(chat);
+
+  const completion = await sampleOpenAI.completion();
+  console.log(completion);
+});

packages/sample-app/src/sample_with.ts

@@ -0,0 +1,41 @@
+import * as traceloop from "@traceloop/node-server-sdk";
+import OpenAI from "openai";
+
+traceloop.initialize({
+  appName: "sample_openai",
+  apiKey: process.env.TRACELOOP_API_KEY,
+  disableBatch: true,
+});
+const openai = new OpenAI();
+
+async function chat() {
+  return traceloop.withWorkflow("sample_chat", async () => {
+    const chatCompletion = await openai.chat.completions.create({
+      messages: [
+        { role: "user", content: "Tell me a joke about OpenTelemetry" },
+      ],
+      model: "gpt-3.5-turbo",
+    });
+
+    return chatCompletion.choices[0].message.content;
+  });
+}
+
+async function completion() {
+  return traceloop.withWorkflow("sample_completion", async () => {
+    const completion = await openai.completions.create({
+      prompt: "Tell me a joke about TypeScript",
+      model: "gpt-3.5-turbo-instruct",
+    });
+
+    return completion.choices[0].text;
+  });
+}
+
+traceloop.withAssociationProperties({ userId: "12345" }, async () => {
+  const chatResponse = await chat();
+  console.log(chatResponse);
+
+  const completionResponse = await completion();
+  console.log(completionResponse);
+});

packages/traceloop-sdk/src/lib/node-server-sdk.ts

@@ -5,5 +5,6 @@ export { InitializeOptions } from "./interfaces";
 export { initialize } from "./configuration";
 export { forceFlush } from "./tracing";
 export * from "./tracing/decorators";
+export * from "./tracing/association";
 
 initInstrumentations();

packages/traceloop-sdk/src/lib/tracing/association.ts

@@ -0,0 +1,17 @@
+import { context } from "@opentelemetry/api";
+import { ASSOCATION_PROPERTIES_KEY } from "./tracing";
+
+export function withAssociationProperties<
+  A extends unknown[],
+  F extends (...args: A) => ReturnType<F>,
+>(
+  properties: { [name: string]: string },
+  fn: F,
+  thisArg?: ThisParameterType<F>,
+  ...args: A
+) {
+  const newContext = context
+    .active()
+    .setValue(ASSOCATION_PROPERTIES_KEY, properties);
+  return context.with(newContext, fn, thisArg, ...args);
+}