minor fixes

MagnivOrg · Oct 11, 2023 · d7e90c7 · d7e90c7
1 parent f1f476f
commit d7e90c7
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 13 deletions.
diff --git 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
@@ -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}`