feat: workflow annotations

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

@@ -42,17 +42,17 @@ export const SemanticAttributes = {
   TRACELOOP_CORRELATION_ID: "traceloop.correlation.id",
 };
 
-export const LLMRequestTypeValues = {
-  COMPLETION: "completion",
-  CHAT: "chat",
-  RERANK: "rerank",
-  UNKNOWN: "unknown",
-};
+export enum LLMRequestTypeValues {
+  COMPLETION = "completion",
+  CHAT = "chat",
+  RERANK = "rerank",
+  UNKNOWN = "unknown",
+}
 
-export const TraceloopSpanKindValues = {
-  WORKFLOW: "workflow",
-  TASK: "task",
-  AGENT: "agent",
-  TOOL: "tool",
-  UNKNOWN: "unknown",
-};
+export enum TraceloopSpanKindValues {
+  WORKFLOW = "workflow",
+  TASK = "task",
+  AGENT = "agent",
+  TOOL = "tool",
+  UNKNOWN = "unknown",
+}

packages/sample-app/src/sample_openai.ts

@@ -9,23 +9,30 @@ traceloop.initialize({
 });
 const openai = new OpenAI();
 
-async function chat() {
-  const chatCompletion = await openai.chat.completions.create({
-    messages: [{ role: "user", content: "Tell me a joke about OpenTelemetry" }],
-    model: "gpt-3.5-turbo",
-  });
+class SampleOpenAI {
+  @traceloop.workflow("sample_chat")
+  async chat() {
+    const chatCompletion = await openai.chat.completions.create({
+      messages: [
+        { role: "user", content: "Tell me a joke about OpenTelemetry" },
+      ],
+      model: "gpt-3.5-turbo",
+    });
 
-  console.log(chatCompletion.choices[0].message.content);
-}
+    console.log(chatCompletion.choices[0].message.content);
+  }
 
-async function completion() {
-  const completion = await openai.completions.create({
-    prompt: "Tell me a joke about TypeScript",
-    model: "gpt-3.5-turbo-instruct",
-  });
+  @traceloop.workflow("sample_completion")
+  async completion() {
+    const completion = await openai.completions.create({
+      prompt: "Tell me a joke about TypeScript",
+      model: "gpt-3.5-turbo-instruct",
+    });
 
-  console.log(completion.choices[0].text);
+    console.log(completion.choices[0].text);
+  }
 }
 
-chat();
-completion();
+const sampleOpenAI = new SampleOpenAI();
+sampleOpenAI.chat();
+sampleOpenAI.completion();

packages/sample-app/tsconfig.json

@@ -2,7 +2,8 @@
   "extends": "../../tsconfig.base.json",
   "compilerOptions": {
     "outDir": "dist",
-    "rootDir": "."
+    "rootDir": ".",
+    "experimentalDecorators": true
   },
   "files": [],
   "include": ["src/**/*.ts", "test/**/*.ts"],

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

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

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

@@ -0,0 +1,74 @@
+import { Span, context } from "@opentelemetry/api";
+import { getTracer, WORKFLOW_NAME_KEY } from "./tracing";
+import {
+  SemanticAttributes,
+  TraceloopSpanKindValues,
+} from "@traceloop/ai-semantic-conventions";
+
+function entity(type: TraceloopSpanKindValues, name?: string) {
+  return function (
+    target: any,
+    propertyKey: string,
+    descriptor: PropertyDescriptor,
+  ) {
+    const originalMethod: Function = descriptor.value;
+    const entityName = name ?? originalMethod.name;
+
+    if (originalMethod.constructor.name === "AsyncFunction") {
+      descriptor.value = async function (...args: any[]) {
+        const workflowContext =
+          type === TraceloopSpanKindValues.WORKFLOW
+            ? context.active().setValue(WORKFLOW_NAME_KEY, entityName)
+            : context.active();
+        await getTracer().startActiveSpan(
+          `${entityName}.${type}`,
+          {},
+          workflowContext,
+          async (span: Span) => {
+            span.setAttribute(
+              SemanticAttributes.TRACELOOP_WORKFLOW_NAME,
+              entityName,
+            );
+            span.setAttribute(SemanticAttributes.TRACELOOP_SPAN_KIND, type);
+            span.setAttribute(
+              SemanticAttributes.TRACELOOP_ENTITY_NAME,
+              entityName,
+            );
+            const res = await originalMethod.apply(this, args);
+            span.end();
+            return res;
+          },
+        );
+      };
+    } else {
+      descriptor.value = function (...args: any[]) {
+        getTracer().startActiveSpan(`${entityName}.${type}`, (span: Span) => {
+          span.setAttribute(SemanticAttributes.TRACELOOP_SPAN_KIND, type);
+          span.setAttribute(
+            SemanticAttributes.TRACELOOP_ENTITY_NAME,
+            entityName,
+          );
+          const res = originalMethod.apply(this, args);
+          span.end();
+          return res;
+        });
+      };
+    }
+  };
+}
+
+export function workflow(name?: string) {
+  return entity(TraceloopSpanKindValues.WORKFLOW, name);
+}
+
+export function task(name?: string) {
+  return entity(TraceloopSpanKindValues.TASK, name);
+}
+
+export function agent(name?: string) {
+  return entity(TraceloopSpanKindValues.AGENT, name);
+}
+
+export function tool(name?: string) {
+  return entity(TraceloopSpanKindValues.TOOL, name);
+}

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

@@ -2,12 +2,16 @@ import { NodeSDK } from "@opentelemetry/sdk-node";
 import {
   SimpleSpanProcessor,
   BatchSpanProcessor,
+  //   ConsoleSpanExporter,
 } from "@opentelemetry/sdk-trace-node";
+import { Span, Context, context } from "@opentelemetry/api";
 import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-proto";
 import { Resource } from "@opentelemetry/resources";
 import { SemanticResourceAttributes } from "@opentelemetry/semantic-conventions";
 import { InitializeOptions } from "../interfaces";
 import { OpenAIInstrumentation } from "@traceloop/instrumentation-openai";
+import { SemanticAttributes } from "@traceloop/ai-semantic-conventions";
+import { WORKFLOW_NAME_KEY } from "./tracing";
 
 let _sdk: NodeSDK;
 let instrumentations: any[] = [];
@@ -28,14 +32,26 @@ export const startTracing = (options: InitializeOptions) => {
     url: `${options.baseUrl}/v1/traces`,
     headers: { Authorization: `Bearer ${options.apiKey}` },
   });
+  //   const traceExporter = new ConsoleSpanExporter();
+  const spanProcessor = options.disableBatch
+    ? new SimpleSpanProcessor(traceExporter)
+    : new BatchSpanProcessor(traceExporter);
+
+  spanProcessor.onStart = (span: Span, parentContext: Context) => {
+    const workflowName = context.active().getValue(WORKFLOW_NAME_KEY);
+    if (workflowName) {
+      span.setAttribute(
+        SemanticAttributes.TRACELOOP_WORKFLOW_NAME,
+        workflowName as string,
+      );
+    }
+  };
 
   _sdk = new NodeSDK({
     resource: new Resource({
       [SemanticResourceAttributes.SERVICE_NAME]: options.appName,
     }),
-    spanProcessor: options.disableBatch
-      ? new SimpleSpanProcessor(traceExporter)
-      : new BatchSpanProcessor(traceExporter),
+    spanProcessor,
     traceExporter,
     instrumentations,
   });

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

@@ -0,0 +1,8 @@
+import { trace, createContextKey } from "@opentelemetry/api";
+
+const TRACER_NAME = "traceloop.tracer";
+export const WORKFLOW_NAME_KEY = createContextKey("workflow_name");
+
+export const getTracer = () => {
+  return trace.getTracer(TRACER_NAME);
+};