From c04ed2a72fc8f78f29e3aaa34828e6e3d8283841 Mon Sep 17 00:00:00 2001
From: jagadeeswaran-zipstack <jagadeeswaran@zipstack.com>
Date: Mon, 23 Dec 2024 10:21:37 +0530
Subject: [PATCH] [Fix] Singlepass coverage issue (#910)

* fixes for prompt sudio coverage

* fixed prompt studio local variable issue

* updated return type

* added optional chaining

* added missing type

* added comment for error suppression

* handled edge cases

* added fix on run all prompt coverage issue

* added optional chaining
---
 frontend/src/hooks/usePromptOutput.js | 29 ++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/frontend/src/hooks/usePromptOutput.js b/frontend/src/hooks/usePromptOutput.js
index 5517ec968..4c6fba3cf 100644
--- a/frontend/src/hooks/usePromptOutput.js
+++ b/frontend/src/hooks/usePromptOutput.js
@@ -24,8 +24,10 @@ try {
 const usePromptOutput = () => {
   const { sessionDetails } = useSessionStore();
   const { setTokenUsage, updateTokenUsage } = useTokenUsageStore();
-  const { setPromptOutput, updatePromptOutput } = usePromptOutputStore();
-  const { isSimplePromptStudio, isPublicSource } = useCustomToolStore();
+  const { setPromptOutput, updatePromptOutput, promptOutputs } =
+    usePromptOutputStore();
+  const { isSimplePromptStudio, isPublicSource, selectedDoc } =
+    useCustomToolStore();
   const axiosPrivate = useAxiosPrivate();
   const { id } = useParams();
 
@@ -152,7 +154,28 @@ const usePromptOutput = () => {
       setPromptOutput(outputs);
       setTokenUsage(tokenUsageDetails);
     } else {
-      updatePromptOutput(outputs);
+      let updatedPromptOutputs = promptOutputs;
+      Object.keys(outputs).forEach((key) => {
+        const [keyPromptId, keyDoctId, keyLlmProfile, keyIsSinglePass] =
+          key.split("__");
+        // only add output of selected document
+        if (keyDoctId === selectedDoc?.document_id) {
+          const currentOutput = { [key]: outputs[key] };
+          updatedPromptOutputs = { ...promptOutputs, ...currentOutput };
+        }
+        Object.keys(updatedPromptOutputs).forEach((innerKey) => {
+          const [existingPromptId, , existingLlmProfile, existingIsSinglePass] =
+            innerKey.split("__"); // Extract promptId from key
+          if (
+            keyPromptId === existingPromptId &&
+            keyLlmProfile === existingLlmProfile &&
+            keyIsSinglePass === existingIsSinglePass
+          ) {
+            updatedPromptOutputs[innerKey].coverage = outputs[key]?.coverage;
+          }
+        });
+      });
+      updatePromptOutput(updatedPromptOutputs);
       updateTokenUsage(tokenUsageDetails);
     }
   };