[Feat] Highlight support in prompt studio (#887)

jaags-dev · ritwik-g · web-flow · commit eff8f546d9e8 · 2025-01-06T11:27:04.000+05:30
* added highlight feature in prompt studio

* fixed hightligh UI

* code refactor/optimisation

* added highlight data to store

---------

Co-authored-by: Ritwik G &lt;100672805+ritwik-g@users.noreply.github.com&gt;
diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/migrations/0002_promptstudiooutputmanager_highlight_data.py b/backend/prompt_studio/prompt_studio_output_manager_v2/migrations/0002_promptstudiooutputmanager_highlight_data.py
@@ -0,0 +1,20 @@
+# Generated by Django 4.2.1 on 2024-12-05 10:21
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("prompt_studio_output_manager_v2", "0001_initial"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="promptstudiooutputmanager",
+            name="highlight_data",
+            field=models.JSONField(
+                blank=True, db_comment="Field to store highlight data", null=True
+            ),
+        ),
+    ]
diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/models.py b/backend/prompt_studio/prompt_studio_output_manager_v2/models.py
@@ -27,6 +27,9 @@ class PromptStudioOutputManager(BaseModel):
     challenge_data = models.JSONField(
         db_comment="Field to store challenge data", editable=True, null=True, blank=True
     )
+    highlight_data = models.JSONField(
+        db_comment="Field to store highlight data", editable=True, null=True, blank=True
+    )
     eval_metrics = models.JSONField(
         db_column="eval_metrics",
         null=False,
diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py
@@ -60,6 +60,7 @@ def update_or_create_prompt_output(
             tool: CustomTool,
             context: str,
             challenge_data: Optional[dict[str, Any]],
+            highlight_data: Optional[dict[str, Any]],
         ) -> PromptStudioOutputManager:
             """Handles creating or updating a single prompt output and returns
             the instance."""
@@ -76,6 +77,7 @@ def update_or_create_prompt_output(
                             "eval_metrics": eval_metrics,
                             "context": context,
                             "challenge_data": challenge_data,
+                            "highlight_data": highlight_data,
                         },
                     )
                 )
@@ -97,6 +99,7 @@ def update_or_create_prompt_output(
                     "eval_metrics": eval_metrics,
                     "context": context,
                     "challenge_data": challenge_data,
+                    "highlight_data": highlight_data,
                 }
                 PromptStudioOutputManager.objects.filter(
                     document_manager=document_manager,
@@ -118,6 +121,7 @@ def update_or_create_prompt_output(
         serialized_data: list[dict[str, Any]] = []
         context = metadata.get("context")
         challenge_data = metadata.get("challenge_data")
+        highlight_data = metadata.get("highlight_data")
 
         if not prompts:
             return serialized_data
@@ -134,6 +138,8 @@ def update_or_create_prompt_output(
 
             if not is_single_pass_extract:
                 context = context.get(prompt.prompt_key)
+                if highlight_data:
+                    highlight_data = highlight_data.get(prompt.prompt_key)
                 if challenge_data:
                     challenge_data = challenge_data.get(prompt.prompt_key)
 
@@ -156,6 +162,7 @@ def update_or_create_prompt_output(
                 tool=tool,
                 context=json.dumps(context),
                 challenge_data=challenge_data,
+                highlight_data=highlight_data,
             )
 
             # Serialize the instance
diff --git a/frontend/src/components/custom-tools/document-manager/DocumentManager.jsx b/frontend/src/components/custom-tools/document-manager/DocumentManager.jsx
@@ -95,11 +95,13 @@ function DocumentManager({ generateIndex, handleUpdateTool, handleDocChange }) {
     isSimplePromptStudio,
     isPublicSource,
     refreshRawView,
+    selectedHighlight,
   } = useCustomToolStore();
   const { sessionDetails } = useSessionStore();
   const axiosPrivate = useAxiosPrivate();
   const { setPostHogCustomEvent } = usePostHogEvents();
   const { id } = useParams();
+  const highlightData = selectedHighlight?.highlight || [];
 
   useEffect(() => {
     if (isSimplePromptStudio) {
@@ -386,7 +388,7 @@ function DocumentManager({ generateIndex, handleUpdateTool, handleDocChange }) {
           setOpenManageDocsModal={setOpenManageDocsModal}
           errMsg={fileErrMsg}
         >
-          <PdfViewer fileUrl={fileUrl} />
+          <PdfViewer fileUrl={fileUrl} highlightData={highlightData} />
         </DocumentViewer>
       )}
       {activeKey === "2" && (
diff --git a/frontend/src/components/custom-tools/pdf-viewer/PdfViewer.jsx b/frontend/src/components/custom-tools/pdf-viewer/PdfViewer.jsx
@@ -1,4 +1,4 @@
-import { useEffect, useRef } from "react";
+import { useEffect, useRef, useMemo } from "react";
 import { Viewer, Worker } from "@react-pdf-viewer/core";
 import { defaultLayoutPlugin } from "@react-pdf-viewer/default-layout";
 import { pageNavigationPlugin } from "@react-pdf-viewer/page-navigation";
@@ -23,29 +23,43 @@ function PdfViewer({ fileUrl, highlightData }) {
   function removeZerosAndDeleteIfAllZero(highlightData) {
     return highlightData?.filter((innerArray) =>
       innerArray.some((value) => value !== 0)
-    ); // Keep arrays that contain at least one non-zero value
-  }
-  let highlightPluginInstance = "";
-  if (RenderHighlights && highlightData) {
-    highlightPluginInstance = highlightPlugin({
-      renderHighlights: (props) => (
-        <RenderHighlights {...props} highlightData={highlightData} />
-      ),
-    });
+    );
   }
 
+  const processHighlightData = highlightData
+    ? removeZerosAndDeleteIfAllZero(highlightData)
+    : [];
+
+  const processedHighlightData =
+    processHighlightData?.length > 0 ? processHighlightData : [[0, 0, 0, 0]];
+
+  const highlightPluginInstance = useMemo(() => {
+    if (
+      RenderHighlights &&
+      Array.isArray(processedHighlightData) &&
+      processedHighlightData?.length > 0
+    ) {
+      return highlightPlugin({
+        renderHighlights: (props) => (
+          <RenderHighlights {...props} highlightData={processedHighlightData} />
+        ),
+      });
+    }
+    return "";
+  }, [RenderHighlights, processedHighlightData]);
+
   // Jump to page when highlightData changes
   useEffect(() => {
     highlightData = removeZerosAndDeleteIfAllZero(highlightData); // Removing zeros before checking the highlight data condition
     if (highlightData && highlightData.length > 0) {
       const pageNumber = highlightData[0][0]; // Assume highlightData[0][0] is the page number
       if (pageNumber !== null && jumpToPage) {
         setTimeout(() => {
-          jumpToPage(pageNumber); // jumpToPage is 0-indexed, so subtract 1
+          jumpToPage(pageNumber); // jumpToPage is 0-indexed, so subtract 1 if necessary
         }, 100); // Add a slight delay to ensure proper page rendering
       }
     }
-  }, [highlightData, jumpToPage]);
+  }, [processedHighlightData, jumpToPage]);
 
   return (
     <div ref={parentRef} className="doc-manager-body">
diff --git a/frontend/src/components/custom-tools/prompt-card/PromptCard.css b/frontend/src/components/custom-tools/prompt-card/PromptCard.css
@@ -300,6 +300,15 @@
   color: #f0ad4e;
 }
 
-.required-checkbox-padding{
-  padding-left: 5px;
+.highlighted-prompt {
+  border-width: 1.2px;
+  border-style: solid;
+  border-color: #4096ff;
+  box-shadow: 4px 4px 12.5px 0px rgba(0, 0, 0, 0.08);
+}
+
+.highlighted-prompt-cell {
+  border-width: 1.2px;
+  border-style: solid;
+  border-color: #ffb400;
 }
diff --git a/frontend/src/components/custom-tools/prompt-card/PromptCard.jsx b/frontend/src/components/custom-tools/prompt-card/PromptCard.jsx
@@ -43,8 +43,13 @@ const PromptCard = memo(
     const [openOutputForDoc, setOpenOutputForDoc] = useState(false);
     const [progressMsg, setProgressMsg] = useState({});
     const [spsLoading, setSpsLoading] = useState({});
-    const { llmProfiles, selectedDoc, details, summarizeIndexStatus } =
-      useCustomToolStore();
+    const {
+      llmProfiles,
+      selectedDoc,
+      details,
+      summarizeIndexStatus,
+      updateCustomTool,
+    } = useCustomToolStore();
     const { messages } = useSocketCustomToolStore();
     const { setAlertDetails } = useAlertStore();
     const { setPostHogCustomEvent } = usePostHogEvents();
@@ -161,6 +166,22 @@ const PromptCard = memo(
       );
     };
 
+    const handleSelectHighlight = (
+      highlightData,
+      highlightedPrompt,
+      highlightedProfile
+    ) => {
+      if (details?.enable_highlight) {
+        updateCustomTool({
+          selectedHighlight: {
+            highlight: highlightData,
+            highlightedPrompt: highlightedPrompt,
+            highlightedProfile: highlightedProfile,
+          },
+        });
+      }
+    };
+
     const handleTypeChange = (value) => {
       handleChange(value, promptDetailsState?.prompt_id, "enforce_type", true);
     };
@@ -261,6 +282,7 @@ const PromptCard = memo(
           promptRunStatus={promptRunStatus}
           coverageCountData={coverageCountData}
           isChallenge={isChallenge}
+          handleSelectHighlight={handleSelectHighlight}
         />
         <OutputForDocModal
           open={openOutputForDoc}
diff --git a/frontend/src/components/custom-tools/prompt-card/PromptCardItems.jsx b/frontend/src/components/custom-tools/prompt-card/PromptCardItems.jsx
@@ -52,6 +52,7 @@ function PromptCardItems({
   promptRunStatus,
   coverageCountData,
   isChallenge,
+  handleSelectHighlight,
 }) {
   const {
     llmProfiles,
@@ -62,6 +63,8 @@ function PromptCardItems({
     isSimplePromptStudio,
     isPublicSource,
     adapters,
+    selectedHighlight,
+    details,
     singlePassExtractMode,
   } = useCustomToolStore();
 
@@ -171,7 +174,13 @@ function PromptCardItems({
   }, [llmProfiles, selectedLlmProfileId, enabledProfiles]);
 
   return (
-    <Card className="prompt-card">
+    <Card
+      className={`prompt-card ${
+        details?.enable_highlight &&
+        selectedHighlight?.highlightedPrompt === promptDetails?.prompt_id &&
+        "highlighted-prompt"
+      }`}
+    >
       <div className="prompt-card-div prompt-card-bg-col1 prompt-card-rad">
         <Space direction="vertical" className="width-100" ref={divRef}>
           <Header
@@ -289,6 +298,7 @@ function PromptCardItems({
               promptOutputs={promptOutputs}
               promptRunStatus={promptRunStatus}
               isChallenge={isChallenge}
+              handleSelectHighlight={handleSelectHighlight}
             />
           </Row>
         </Collapse.Panel>
@@ -326,6 +336,7 @@ PromptCardItems.propTypes = {
   promptRunStatus: PropTypes.object.isRequired,
   coverageCountData: PropTypes.object.isRequired,
   isChallenge: PropTypes.bool.isRequired,
+  handleSelectHighlight: PropTypes.func.isRequired,
 };
 
 export { PromptCardItems };
diff --git a/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx b/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx
diff --git a/frontend/src/components/custom-tools/settings-modal/SettingsModal.jsx b/frontend/src/components/custom-tools/settings-modal/SettingsModal.jsx
diff --git a/frontend/src/hooks/usePromptOutput.js b/frontend/src/hooks/usePromptOutput.js
diff --git a/frontend/src/store/custom-tool-store.js b/frontend/src/store/custom-tool-store.js
