[Fix] Prompt Studio Output update issue

frontend/src/components/custom-tools/prompt-card/PromptCardItems.jsx

@@ -19,7 +19,6 @@ import { Header } from "./Header";
 import { OutputForIndex } from "./OutputForIndex";
 import { PromptOutput } from "./PromptOutput";
 import { TABLE_ENFORCE_TYPE, RECORD_ENFORCE_TYPE } from "./constants";
-import usePromptOutput from "../../../hooks/usePromptOutput";
 
 let TableExtractionSettingsBtn;
 try {
@@ -63,11 +62,9 @@ function PromptCardItems({
     isSimplePromptStudio,
     isPublicSource,
     adapters,
-    defaultLlmProfile,
     singlePassExtractMode,
   } = useCustomToolStore();
 
-  const { generatePromptOutputKey } = usePromptOutput();
   const [isEditingPrompt, setIsEditingPrompt] = useState(false);
   const [isEditingTitle, setIsEditingTitle] = useState(false);
   const [expandCard, setExpandCard] = useState(true);
@@ -81,17 +78,6 @@ function PromptCardItems({
   const divRef = useRef(null);
   const [enforceType, setEnforceType] = useState("");
   const promptId = promptDetails?.prompt_id;
-  const docId = selectedDoc?.document_id;
-  const promptProfile = promptDetails?.profile_manager || defaultLlmProfile;
-  const promptOutputKey = generatePromptOutputKey(
-    promptId,
-    docId,
-    promptProfile,
-    singlePassExtractMode,
-    true
-  );
-  const promptCoverage =
-    promptOutputs[promptOutputKey]?.coverage || coverageCountData;
 
   useEffect(() => {
     if (enforceType !== promptDetails?.enforce_type) {
@@ -123,6 +109,32 @@ function PromptCardItems({
     return result;
   };
 
+  const getUpdatedCoverage = (promptId, singlePass, promptOutputs) => {
+    let updatedCoverage = null;
+    Object.keys(promptOutputs).forEach((key) => {
+      const [keyPromptId, , , keyIsSinglePass] = key.split("__"); // Destructure the key parts
+
+      // Check if the key matches the criteria
+      if (keyPromptId === promptId && keyIsSinglePass === String(singlePass)) {
+        const currentCoverage = promptOutputs[key]?.coverage || [];
+
+        // Update the highestCoverage if the current one is longer
+        if (
+          !updatedCoverage ||
+          currentCoverage.length > updatedCoverage.length
+        ) {
+          updatedCoverage = currentCoverage;
+        }
+      }
+    });
+
+    return updatedCoverage;
+  };
+
+  const promptCoverage =
+    getUpdatedCoverage(promptId, singlePassExtractMode, promptOutputs) ||
+    coverageCountData;
+
   const getAdapterInfo = async (adapterData) => {
     // If simple prompt studio, return early
     if (isSimplePromptStudio) {

frontend/src/components/custom-tools/tool-ide/ToolIde.jsx

@@ -216,13 +216,18 @@ function ToolIde() {
     const body = {
       output: doc?.document_id,
     };
-    handleUpdateTool(body).catch((err) => {
-      const revertSelectedDoc = {
-        selectedDoc: prevSelectedDoc,
-      };
-      updateCustomTool(revertSelectedDoc);
-      setAlertDetails(handleException(err, "Failed to select the document"));
-    });
+    handleUpdateTool(body)
+      .then((res) => {
+        const updatedToolData = res?.data;
+        updateCustomTool({ details: updatedToolData });
+      })
+      .catch((err) => {
+        const revertSelectedDoc = {
+          selectedDoc: prevSelectedDoc,
+        };
+        updateCustomTool(revertSelectedDoc);
+        setAlertDetails(handleException(err, "Failed to select the document"));
+      });
   };
 
   return (

frontend/src/hooks/usePromptOutput.js

@@ -24,8 +24,7 @@ try {
 const usePromptOutput = () => {
   const { sessionDetails } = useSessionStore();
   const { setTokenUsage, updateTokenUsage } = useTokenUsageStore();
-  const { setPromptOutput, updatePromptOutput, promptOutputs } =
-    usePromptOutputStore();
+  const { setPromptOutput, updatePromptOutput } = usePromptOutputStore();
   const { isSimplePromptStudio, isPublicSource, selectedDoc } =
     useCustomToolStore();
   const axiosPrivate = useAxiosPrivate();
@@ -154,32 +153,35 @@ const usePromptOutput = () => {
       setPromptOutput(outputs);
       setTokenUsage(tokenUsageDetails);
     } else {
-      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);
+      const prevOutputs = usePromptOutputStore.getState().promptOutputs;
+      updatePromptOutput(updateCoverage(prevOutputs, outputs));
       updateTokenUsage(tokenUsageDetails);
     }
   };
 
+  const updateCoverage = (promptOutputs, outputs) => {
+    let updatedPromptOutputs = promptOutputs;
+    Object.keys(outputs).forEach((key) => {
+      const [keyPromptId, keyDoctId, , 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, , , existingIsSinglePass] =
+          innerKey.split("__"); // Extract promptId from key
+        if (
+          keyPromptId === existingPromptId &&
+          keyIsSinglePass === existingIsSinglePass
+        ) {
+          updatedPromptOutputs[innerKey].coverage = outputs[key]?.coverage;
+        }
+      });
+    });
+    return updatedPromptOutputs;
+  };
+
   const promptOutputApi = async (
     toolId,
     docId = null,