[FEATURE] Added enable highlight flag (#754)

* Added enable highlight flag Signed-off-by: Deepak <[email protected]> * Added highlight-manager in FE Signed-off-by: Deepak <[email protected]> * AMinor fixes Signed-off-by: Deepak <[email protected]> * Updated structure tool version Signed-off-by: Deepak <[email protected]> * Removed console.log() * Minor fix Signed-off-by: Deepak <[email protected]> --------- Signed-off-by: Deepak <[email protected]> Signed-off-by: Deepak K <[email protected]> Co-authored-by: Gayathri <[email protected]>
Zipstack · Nov 8, 2024 · ee8309f · ee8309f
1 parent 8037dbe
commit ee8309f
Show file tree

Hide file tree

Showing 9 changed files with 48 additions and 12 deletions.
diff --git a/backend/prompt_studio/prompt_studio_core/constants.py b/backend/prompt_studio/prompt_studio_core/constants.py
@@ -95,6 +95,7 @@ class ToolStudioPromptKeys:
     SUMMARIZE_AS_SOURCE = "summarize_as_source"
     VARIABLE_MAP = "variable_map"
     RECORD = "record"
+    ENABLE_HIGHLIGHT = "enable_highlight"
 
 
 class FileViewTypes:

diff --git a/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py
@@ -804,6 +804,7 @@ def _fetch_response(
         tool_settings[TSPKeys.PREAMBLE] = tool.preamble
         tool_settings[TSPKeys.POSTAMBLE] = tool.postamble
         tool_settings[TSPKeys.GRAMMAR] = grammar_list
+        tool_settings[TSPKeys.ENABLE_HIGHLIGHT] = tool.enable_highlight
         tool_settings[TSPKeys.PLATFORM_POSTAMBLE] = getattr(
             settings, TSPKeys.PLATFORM_POSTAMBLE.upper(), ""
         )
@@ -1062,6 +1063,7 @@ def _fetch_single_pass_response(
         tool_settings[TSPKeys.CHUNK_SIZE] = default_profile.chunk_size
         tool_settings[TSPKeys.CHUNK_OVERLAP] = default_profile.chunk_overlap
         tool_settings[TSPKeys.ENABLE_CHALLENGE] = tool.enable_challenge
+        tool_settings[TSPKeys.ENABLE_HIGHLIGHT] = tool.enable_highlight
         tool_settings[TSPKeys.CHALLENGE_LLM] = challenge_llm
 
         for prompt in prompts:

diff --git a/backend/prompt_studio/prompt_studio_registry/prompt_studio_registry_helper.py b/backend/prompt_studio/prompt_studio_registry/prompt_studio_registry_helper.py
@@ -71,6 +71,12 @@ def frame_spec(tool: CustomTool) -> Spec:
                 "default": False,
                 "description": "Enables SinglePass Extraction",
             },
+            "enable_highlight": {
+                "type": "boolean",
+                "title": "Enable highlight",
+                "default": False,
+                "description": "Enables highlight",
+            },
         }
 
         spec = Spec(

diff --git a/backend/sample.env b/backend/sample.env
@@ -81,9 +81,9 @@ PROMPT_STUDIO_FILE_PATH=/app/prompt-studio-data
 
 # Structure Tool Image (Runs prompt studio exported tools)
 # https://hub.docker.com/r/unstract/tool-structure
-STRUCTURE_TOOL_IMAGE_URL="docker:unstract/tool-structure:0.0.46"
+STRUCTURE_TOOL_IMAGE_URL="docker:unstract/tool-structure:0.0.47"
 STRUCTURE_TOOL_IMAGE_NAME="unstract/tool-structure"
-STRUCTURE_TOOL_IMAGE_TAG="0.0.46"
+STRUCTURE_TOOL_IMAGE_TAG="0.0.47"
 
 # Feature Flags
 EVALUATION_SERVER_IP=unstract-flipt

diff --git a/frontend/src/components/agency/ds-settings-card/DsSettingsCard.jsx b/frontend/src/components/agency/ds-settings-card/DsSettingsCard.jsx
@@ -14,6 +14,7 @@ import {
   Space,
   Tooltip,
   Typography,
+  Modal,
 } from "antd";
 import PropTypes from "prop-types";
 import { useEffect, useState } from "react";
@@ -320,6 +321,16 @@ function DsSettingsCard({ type, endpointDetails, message }) {
         } else {
           updatedData["destination"] = data;
         }
+        if (
+          type === "output" &&
+          updatedData?.destination?.connection_type === "MANUALREVIEW"
+        ) {
+          Modal.warning({
+            title: "Warning",
+            content:
+              "Please ensure that the tool in use is has highlight enabled in the tool settings.",
+          });
+        }
         updateWorkflow(updatedData);
         if (showSuccess) {
           setAlertDetails({

diff --git a/frontend/src/components/custom-tools/settings-modal/SettingsModal.jsx b/frontend/src/components/custom-tools/settings-modal/SettingsModal.jsx
@@ -20,11 +20,14 @@ import "./SettingsModal.css";
 let SummarizeManager = null;
 const EvaluationManager = null;
 let ChallengeManager = null;
+let HighlightManager = null;
 try {
   SummarizeManager =
     require("../../../plugins/summarize-manager/SummarizeManager").SummarizeManager;
   ChallengeManager =
     require("../../../plugins/challenge-manager/ChallengeManager").ChallengeManager;
+  HighlightManager =
+    require("../../../plugins/highlight-manager/HighlightManager").HighlightManager;
 } catch {
   // Component will remain null if it is not present.
 }
@@ -95,10 +98,22 @@ function SettingsModal({ open, setOpen, handleUpdateTool }) {
         getMenuItem("LLMChallenge", 4, <FileTextOutlined />)
       );
       listOfComponents[4] = (
-        <ChallengeManager handleUpdateTool={handleUpdateTool} />
+        <ChallengeManager
+          handleUpdateTool={handleUpdateTool}
+          type="challenge"
+        />
+      );
+      position++;
+    }
+    if (HighlightManager) {
+      items.push(getMenuItem("Highlight Manager", 8, <FileTextOutlined />));
+      listOfComponents[8] = (
+        <HighlightManager
+          handleUpdateTool={handleUpdateTool}
+          type="highlight"
+        />
       );
     }
-
     setMenuItems(items);
     setComponents(listOfComponents);
   }, [llmItems]);

diff --git a/prompt-service/src/unstract/prompt_service/constants.py b/prompt-service/src/unstract/prompt_service/constants.py
@@ -68,6 +68,7 @@ class PromptServiceContants:
     VARIABLE_MAP = "variable_map"
     RECORD = "record"
     TEXT = "text"
+    ENABLE_HIGHLIGHT = "enable_highlight"
 
 
 class RunLevel(Enum):

diff --git a/prompt-service/src/unstract/prompt_service/helper.py b/prompt-service/src/unstract/prompt_service/helper.py
@@ -217,7 +217,9 @@ def construct_and_run_prompt(
     metadata: dict[str, Any],
 ) -> str:
     platform_postamble = tool_settings.get(PSKeys.PLATFORM_POSTAMBLE, "")
-    if tool_settings.get(PSKeys.SUMMARIZE_AS_SOURCE):
+    summarize_as_source = tool_settings.get(PSKeys.SUMMARIZE_AS_SOURCE)
+    enable_highlight = tool_settings.get(PSKeys.ENABLE_HIGHLIGHT, False)
+    if not enable_highlight or summarize_as_source:
         platform_postamble = ""
     prompt = construct_prompt(
         preamble=tool_settings.get(PSKeys.PREAMBLE, ""),
@@ -233,6 +235,7 @@ def construct_and_run_prompt(
         metadata=metadata,
         prompt_key=output[PSKeys.NAME],
         prompt_type=output.get(PSKeys.TYPE, PSKeys.TEXT),
+        enable_highlight=enable_highlight,
     )
 
 
@@ -272,14 +275,15 @@ def run_completion(
     metadata: Optional[dict[str, str]] = None,
     prompt_key: Optional[str] = None,
     prompt_type: Optional[str] = PSKeys.TEXT,
+    enable_highlight: bool = False,
 ) -> str:
     logger: Logger = current_app.logger
     try:
         extract_epilogue_plugin: dict[str, Any] = plugins.get(
             PSKeys.EXTRACT_EPILOGUE, {}
         )
         extract_epilogue = None
-        if extract_epilogue_plugin:
+        if extract_epilogue_plugin and enable_highlight:
             extract_epilogue = extract_epilogue_plugin["entrypoint_cls"].run
         completion = llm.complete(
             prompt=prompt,

diff --git a/tools/structure/src/main.py b/tools/structure/src/main.py
@@ -41,6 +41,7 @@ def run(
             SettingsKeys.SINGLE_PASS_EXTRACTION_MODE, False
         )
         challenge_llm: str = settings.get(SettingsKeys.CHALLENGE_LLM_ADAPTER_ID, "")
+        enable_highlight: bool = settings.get(SettingsKeys.ENABLE_HIGHLIGHT, False)
         responder: PromptTool = PromptTool(
             tool=self,
             prompt_port=self.get_env_or_die(SettingsKeys.PROMPT_PORT),
@@ -67,13 +68,13 @@ def run(
         tool_id = tool_metadata[SettingsKeys.TOOL_ID]
         tool_settings = tool_metadata[SettingsKeys.TOOL_SETTINGS]
         outputs = tool_metadata[SettingsKeys.OUTPUTS]
-        enable_highlight: bool = tool_settings.get(SettingsKeys.ENABLE_HIGHLIGHT, False)
         tool_settings[SettingsKeys.CHALLENGE_LLM] = challenge_llm
         tool_settings[SettingsKeys.ENABLE_CHALLENGE] = enable_challenge
         tool_settings[SettingsKeys.ENABLE_SINGLE_PASS_EXTRACTION] = (
             single_pass_extraction_mode
         )
         tool_settings[SettingsKeys.SUMMARIZE_AS_SOURCE] = summarize_as_source
+        tool_settings[SettingsKeys.ENABLE_HIGHLIGHT] = enable_highlight
         prompt_service_resp = None
         _, file_name = os.path.split(input_file)
         if summarize_as_source:
@@ -122,7 +123,6 @@ def run(
                 output_file_path=tool_data_dir / SettingsKeys.EXTRACT,
                 reindex=True,
                 usage_kwargs=usage_kwargs,
-                enable_highlight=enable_highlight,
                 process_text=process_text,
             )
             if summarize_as_source:
@@ -134,7 +134,6 @@ def run(
                     outputs=outputs,
                     index=index,
                     usage_kwargs=usage_kwargs,
-                    enable_highlight=enable_highlight,
                 )
                 payload[SettingsKeys.FILE_HASH] = summarize_file_hash
             self.stream_log("Fetching response for single pass extraction")
@@ -160,7 +159,6 @@ def run(
                             output_file_path=tool_data_dir / SettingsKeys.EXTRACT,
                             reindex=reindex,
                             usage_kwargs=usage_kwargs,
-                            enable_highlight=enable_highlight,
                             process_text=process_text,
                         )
 
@@ -262,7 +260,6 @@ def _summarize_and_index(
         outputs: dict[str, Any],
         index: Index,
         usage_kwargs: dict[Any, Any] = {},
-        enable_highlight: bool = False,
     ) -> str:
         """Summarizes the context of the file and indexes the summarized
         content.
@@ -336,7 +333,6 @@ def _summarize_and_index(
             chunk_size=0,
             chunk_overlap=0,
             usage_kwargs=usage_kwargs,
-            enable_highlight=enable_highlight,
         )
         return summarize_file_hash