diff --git a/packages/sample-app/src/sample_decorators.ts b/packages/sample-app/src/sample_decorators.ts index 5cdc437a..b252b032 100644 --- a/packages/sample-app/src/sample_decorators.ts +++ b/packages/sample-app/src/sample_decorators.ts @@ -9,13 +9,15 @@ traceloop.initialize({ const openai = new OpenAI(); class SampleOpenAI { + constructor(private model = "gpt-3.5-turbo") {} + @traceloop.workflow({ name: "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", + model: this.model, }); return chatCompletion.choices[0].message.content; diff --git a/packages/traceloop-sdk/src/lib/tracing/decorators.ts b/packages/traceloop-sdk/src/lib/tracing/decorators.ts index 4ed19cb0..14334968 100644 --- a/packages/traceloop-sdk/src/lib/tracing/decorators.ts +++ b/packages/traceloop-sdk/src/lib/tracing/decorators.ts @@ -187,30 +187,18 @@ function entity( propertyKey: string, descriptor: PropertyDescriptor, ) { - const originalMethod: () => any = descriptor.value; + const originalMethod = descriptor.value; const entityName = config.name ?? originalMethod.name; - if (originalMethod.constructor.name === "AsyncFunction") { - descriptor.value = async function (...args: any[]) { - return await withEntity( - type, - { ...config, name: entityName }, - originalMethod, - target, - ...args, - ); - }; - } else { - descriptor.value = function (...args: any[]) { - return withEntity( - type, - { ...config, name: entityName }, - originalMethod, - target, - ...args, - ); - }; - } + descriptor.value = function (...args: any[]) { + return withEntity( + type, + { ...config, name: entityName }, + originalMethod, + this, + ...args, + ); + }; }; } diff --git a/packages/traceloop-sdk/test/decorators.test.ts b/packages/traceloop-sdk/test/decorators.test.ts index da5f47f9..be6a19e7 100644 --- a/packages/traceloop-sdk/test/decorators.test.ts +++ b/packages/traceloop-sdk/test/decorators.test.ts @@ -134,6 +134,8 @@ describe("Test SDK Decorators", () => { it("should create spans for workflows using decoration syntax", async () => { class TestOpenAI { + constructor(private model = "gpt-3.5-turbo") {} + @traceloop.workflow({ name: "sample_chat" }) async chat(things: Map) { const generations: Map = new Map(); @@ -142,7 +144,7 @@ describe("Test SDK Decorators", () => { messages: [ { role: "user", content: `Tell me a ${key} about ${value}` }, ], - model: "gpt-3.5-turbo", + model: this.model, }); if (chatCompletion.choices[0].message.content) {