traceloop
diff --git a/‎packages/sample-app/src/sample_decorators.ts
Lines changed: 3 additions & 1 deletion b/‎packages/sample-app/src/sample_decorators.ts
Lines changed: 3 additions & 1 deletion
diff --git a/‎packages/traceloop-sdk/src/lib/tracing/decorators.ts
Lines changed: 10 additions & 22 deletions b/‎packages/traceloop-sdk/src/lib/tracing/decorators.ts
Lines changed: 10 additions & 22 deletions
diff --git a/‎packages/traceloop-sdk/test/decorators.test.ts
Lines changed: 3 additions & 1 deletion b/‎packages/traceloop-sdk/test/decorators.test.ts
Lines changed: 3 additions & 1 deletion
@@ -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;
 
@@ -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,
+      );
+    };
   };
 }
 
 
@@ -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<string, string>) {
         const generations: Map<string, string> = 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 (cha
38D6
tCompletion.choices[0].message.content) {