From d7e90c779064786d97383a26b7e45cef4fdd2c37 Mon Sep 17 00:00:00 2001
From: AbuBakar <abubakarsohail.dev@gmail.com>
Date: Wed, 11 Oct 2023 18:24:14 +0500
Subject: [PATCH] minor fixes

---
 src/promptlayer/index.ts | 32 ++++++++++++++++++++------------
 src/utils/index.ts       |  3 ++-
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/src/promptlayer/index.ts b/src/promptlayer/index.ts
index c0feec6..4d39dfe 100644
--- a/src/promptlayer/index.ts
+++ b/src/promptlayer/index.ts
@@ -38,20 +38,28 @@ const promptLayerBase = (
       const request_start_time = new Date().toISOString();
       const function_name = Reflect.get(target, "function_name");
       const provider_type = Reflect.get(target, "provider");
+      const return_pl_id = argArray[0]?.return_pl_id;
+      delete argArray[0]?.return_pl_id;
       const response = Reflect.apply(target, thisArg, argArray);
       if (response instanceof Promise) {
-        response.then((request_response) => {
-          const request_end_time = new Date().toISOString();
-          promptLayerApiRequest({
-            api_key: getApiKey(),
-            provider_type,
-            function_name,
-            request_start_time,
-            request_end_time,
-            request_response,
-            kwargs: argArray[0],
-          });
-          return request_response;
+        return new Promise((resolve, reject) => {
+          response
+            .then(async (request_response) => {
+              const response = await promptLayerApiRequest({
+                api_key: getApiKey(),
+                provider_type,
+                function_name,
+                request_start_time,
+                request_end_time: new Date().toISOString(),
+                request_response,
+                kwargs: argArray[0],
+                return_pl_id,
+              });
+              resolve(response);
+            })
+            .catch((error) => {
+              reject(error);
+            });
         });
       }
       return response;
diff --git a/src/utils/index.ts b/src/utils/index.ts
index 1287479..0447ca7 100644
--- a/src/utils/index.ts
+++ b/src/utils/index.ts
@@ -31,8 +31,9 @@ const promptLayerApiRequest = async (body: TrackRequest) => {
       );
     }
     if (data && body.return_pl_id) {
-      return data.request_id;
+      return [body.request_response, data.request_id];
     }
+    return body.request_response;
   } catch (e) {
     console.warn(
       `WARNING: While logging your request PromptLayer had the following error: ${e}`