diff --git a/packages/ai-semantic-conventions/src/SemanticAttributes.ts b/packages/ai-semantic-conventions/src/SemanticAttributes.ts index 41750747..bfb083c8 100644 --- a/packages/ai-semantic-conventions/src/SemanticAttributes.ts +++ b/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", +} diff --git a/packages/sample-app/src/sample_openai.ts b/packages/sample-app/src/sample_openai.ts index b35d0ad9..4c1d9310 100644 --- a/packages/sample-app/src/sample_openai.ts +++ b/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(); diff --git a/packages/sample-app/tsconfig.json b/packages/sample-app/tsconfig.json index 3c3ab535..26c3d572 100644 --- a/packages/sample-app/tsconfig.json +++ b/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"], diff --git a/packages/traceloop-sdk/src/lib/node-server-sdk.ts b/packages/traceloop-sdk/src/lib/node-server-sdk.ts index 6f67cb81..1d518637 100644 --- a/packages/traceloop-sdk/src/lib/node-server-sdk.ts +++ b/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(); diff --git a/packages/traceloop-sdk/src/lib/tracing/decorators.ts b/packages/traceloop-sdk/src/lib/tracing/decorators.ts new file mode 100644 index 00000000..3030ad4e --- /dev/null +++ b/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); +} diff --git a/packages/traceloop-sdk/src/lib/tracing/index.ts b/packages/traceloop-sdk/src/lib/tracing/index.ts index ed46c9f0..4831cbfd 100644 --- a/packages/traceloop-sdk/src/lib/tracing/index.ts +++ b/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, }); diff --git a/packages/traceloop-sdk/src/lib/tracing/tracing.ts b/packages/traceloop-sdk/src/lib/tracing/tracing.ts new file mode 100644 index 00000000..f44ffd55 --- /dev/null +++ b/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); +};