Feat/prompt key required field (#885)

* added required field in prompt key * required fields constant and added required field in promptstudio helper * added frontend required field prompt key functionality * added prompt service required key adding in metadata * added required key in prompt studio registry and removed some unwanted codes * changed required field bool to option with any/all * made changes in frontend to support option required * prompt service changes to support required options field * required field based on enforce type * changed required choices to class * frontend condition fail issue fix in highlight data * required field text changes as per the feedback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * addressed pr comments * added info for required fields --------- Signed-off-by: vishnuszipstack <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Zipstack · Dec 23, 2024 · bd7ec42 · bd7ec42
1 parent ee41dff
commit bd7ec42
Show file tree

Hide file tree

Showing 15 changed files with 162 additions and 7 deletions.
diff --git a/backend/prompt_studio/prompt_studio_core_v2/constants.py b/backend/prompt_studio/prompt_studio_core_v2/constants.py
@@ -96,6 +96,7 @@ class ToolStudioPromptKeys:
     RECORD = "record"
     FILE_PATH = "file_path"
     ENABLE_HIGHLIGHT = "enable_highlight"
+    REQUIRED = "required"
     EXECUTION_SOURCE = "execution_source"
 
 

diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py
@@ -819,6 +819,7 @@ def _fetch_response(
 
         output[TSPKeys.PROMPT] = prompt.prompt
         output[TSPKeys.ACTIVE] = prompt.active
+        output[TSPKeys.REQUIRED] = prompt.required
         output[TSPKeys.CHUNK_SIZE] = profile_manager.chunk_size
         output[TSPKeys.VECTOR_DB] = vector_db
         output[TSPKeys.EMBEDDING] = embedding_model

diff --git a/backend/prompt_studio/prompt_studio_registry_v2/constants.py b/backend/prompt_studio/prompt_studio_registry_v2/constants.py
@@ -98,6 +98,7 @@ class JsonSchemaKey:
     SUMMARIZE_AS_SOURCE = "summarize_as_source"
     ENABLE_HIGHLIGHT = "enable_highlight"
     PLATFORM_POSTAMBLE = "platform_postamble"
+    REQUIRED = "required"
 
 
 class SpecKey:

diff --git a/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py b/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py
@@ -322,6 +322,7 @@ def frame_export_json(
 
             output[JsonSchemaKey.PROMPT] = prompt.prompt
             output[JsonSchemaKey.ACTIVE] = prompt.active
+            output[JsonSchemaKey.REQUIRED] = prompt.required
             output[JsonSchemaKey.CHUNK_SIZE] = prompt.profile_manager.chunk_size
             output[JsonSchemaKey.VECTOR_DB] = vector_db
             output[JsonSchemaKey.EMBEDDING] = embedding_model

diff --git a/backend/prompt_studio/prompt_studio_v2/migrations/0003_toolstudioprompt_required.py b/backend/prompt_studio/prompt_studio_v2/migrations/0003_toolstudioprompt_required.py
@@ -0,0 +1,18 @@
+# Generated by Django 4.2.1 on 2024-12-10 10:07
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("prompt_studio_v2", "0002_alter_toolstudioprompt_enforce_type"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="toolstudioprompt",
+            name="required",
+            field=models.BooleanField(default=False),
+        ),
+    ]
diff --git a/backend/prompt_studio/prompt_studio_v2/migrations/0004_alter_toolstudioprompt_required.py b/backend/prompt_studio/prompt_studio_v2/migrations/0004_alter_toolstudioprompt_required.py
@@ -0,0 +1,24 @@
+# Generated by Django 4.2.1 on 2024-12-12 08:28
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("prompt_studio_v2", "0003_toolstudioprompt_required"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="toolstudioprompt",
+            name="required",
+            field=models.CharField(
+                blank=True,
+                choices=[("all", "All values required"), ("any", "Any value required")],
+                default=None,
+                max_length=3,
+                null=True,
+            ),
+        ),
+    ]
diff --git a/backend/prompt_studio/prompt_studio_v2/migrations/0005_alter_toolstudioprompt_required.py b/backend/prompt_studio/prompt_studio_v2/migrations/0005_alter_toolstudioprompt_required.py
@@ -0,0 +1,24 @@
+# Generated by Django 4.2.1 on 2024-12-20 06:35
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("prompt_studio_v2", "0004_alter_toolstudioprompt_required"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="toolstudioprompt",
+            name="required",
+            field=models.CharField(
+                blank=True,
+                choices=[("all", "All values required"), ("any", "Any value required")],
+                db_comment="Field to store weather the values all values or any         values required. This is used for HQR, based on the value approve or finish         review",
+                default=None,
+                null=True,
+            ),
+        ),
+    ]
diff --git a/backend/prompt_studio/prompt_studio_v2/models.py b/backend/prompt_studio/prompt_studio_v2/models.py
@@ -35,7 +35,15 @@ class PromptType(models.TextChoices):
     class Mode(models.TextChoices):
         DEFAULT = "Default", "Default choice for output"
 
-    prompt_id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
+    class RequiredType(models.TextChoices):
+        ALL = "all", "All values required"
+        ANY = "any", "Any value required"
+
+    prompt_id = models.UUIDField(
+        primary_key=True,
+        default=uuid.uuid4,
+        editable=False,
+    )
     prompt_key = models.TextField(
         blank=False,
         db_comment="Field to store the prompt key",
@@ -84,6 +92,15 @@ class Mode(models.TextChoices):
         db_comment="Field to store the prompt key",
         unique=False,
     )
+    required = models.CharField(
+        choices=RequiredType.choices,
+        null=True,  # Allows the field to store NULL in the database
+        blank=True,  # Allows the field to be optional in forms
+        default=None,  # Sets the default value to None
+        db_comment="Field to store weather the values all values or any \
+        values required. This is used for HQR, based on the value approve or finish \
+        review",
+    )
     is_assert = models.BooleanField(default=False)
     active = models.BooleanField(default=True, null=False, blank=False)
     output_metadata = models.JSONField(

diff --git a/frontend/src/components/custom-tools/pdf-viewer/PdfViewer.jsx b/frontend/src/components/custom-tools/pdf-viewer/PdfViewer.jsx
@@ -21,7 +21,7 @@ function PdfViewer({ fileUrl, highlightData }) {
   const { jumpToPage } = pageNavigationPluginInstance;
   const parentRef = useRef(null);
   function removeZerosAndDeleteIfAllZero(highlightData) {
-    return highlightData.filter((innerArray) =>
+    return highlightData?.filter((innerArray) =>
       innerArray.some((value) => value !== 0)
     ); // Keep arrays that contain at least one non-zero value
   }
@@ -36,8 +36,8 @@ function PdfViewer({ fileUrl, highlightData }) {
 
   // Jump to page when highlightData changes
   useEffect(() => {
+    highlightData = removeZerosAndDeleteIfAllZero(highlightData); // Removing zeros before checking the highlight data condition
     if (highlightData && highlightData.length > 0) {
-      highlightData = removeZerosAndDeleteIfAllZero(highlightData);
       const pageNumber = highlightData[0][0]; // Assume highlightData[0][0] is the page number
       if (pageNumber !== null && jumpToPage) {
         setTimeout(() => {

diff --git a/frontend/src/components/custom-tools/prompt-card/Header.jsx b/frontend/src/components/custom-tools/prompt-card/Header.jsx
@@ -6,6 +6,7 @@ import {
   PlayCircleFilled,
   PlayCircleOutlined,
   SyncOutlined,
+  InfoCircleOutlined,
 } from "@ant-design/icons";
 import { useEffect, useState } from "react";
 import { Button, Checkbox, Col, Dropdown, Row, Tag, Tooltip } from "antd";
@@ -45,6 +46,7 @@ function Header({
   setExpandCard,
   spsLoading,
   handleSpsLoading,
+  enforceType,
 }) {
   const {
     selectedDoc,
@@ -58,6 +60,7 @@ function Header({
   const [items, setItems] = useState([]);
 
   const [isDisablePrompt, setIsDisablePrompt] = useState(null);
+  const [required, setRequired] = useState(false);
 
   const handleRunBtnClick = (promptRunType, docId = null) => {
     setExpandCard(true);
@@ -73,8 +76,22 @@ function Header({
       }
     );
   };
+  const handleRequiredChange = (value) => {
+    const newValue = value === required ? null : value; // Allow deselection
+    setRequired(newValue);
+    handleChange(
+      newValue,
+      promptDetails?.prompt_id,
+      "required",
+      true,
+      true
+    ).catch(() => {
+      setRequired(promptDetails?.required || null); // Rollback state in case of error
+    });
+  };
   useEffect(() => {
     setIsDisablePrompt(promptDetails?.active);
+    setRequired(promptDetails?.required);
   }, [promptDetails, details]);
 
   useEffect(() => {
@@ -87,6 +104,47 @@ function Header({
         ),
         key: "enable",
       },
+      {
+        label: (
+          <div>
+            {["json", "table", "record"].indexOf(enforceType) === -1 && (
+              <Checkbox
+                checked={required === "all"}
+                onChange={() => handleRequiredChange("all")}
+              >
+                Value Required{" "}
+                <Tooltip title="Marks this as a required field. Saving this record won't be allowed in Human Quality Review should this field be empty.">
+                  <InfoCircleOutlined />
+                </Tooltip>
+              </Checkbox>
+            )}
+            {enforceType === "json" && (
+              <>
+                <Checkbox
+                  checked={required === "all"}
+                  onChange={() => handleRequiredChange("all")}
+                >
+                  All JSON Values Required
+                </Checkbox>
+                <Tooltip title="When set, saving this record won't be allowed in Human Quality Review without all key/values filled in this JSON structure.">
+                  <InfoCircleOutlined />
+                </Tooltip>
+                <Checkbox
+                  checked={required === "any"}
+                  onChange={() => handleRequiredChange("any")}
+                  className="required-checkbox-padding"
+                >
+                  Atleast 1 JSON Value Required
+                </Checkbox>
+                <Tooltip title="When set, saving this record won't be allowed in Human Quality Review without at least one value filled in this JSON structure.">
+                  <InfoCircleOutlined />
+                </Tooltip>
+              </>
+            )}
+          </div>
+        ),
+        key: "required",
+      },
       {
         label: (
           <ConfirmModal
@@ -109,7 +167,7 @@ function Header({
     }
 
     setItems(dropdownItems);
-  }, [isDisablePrompt]);
+  }, [isDisablePrompt, required, enforceType]);
 
   return (
     <Row>
@@ -270,6 +328,7 @@ Header.propTypes = {
   setExpandCard: PropTypes.func.isRequired,
   spsLoading: PropTypes.object,
   handleSpsLoading: PropTypes.func.isRequired,
+  enforceType: PropTypes.text,
 };
 
 export { Header };
diff --git a/frontend/src/components/custom-tools/prompt-card/PromptCard.css b/frontend/src/components/custom-tools/prompt-card/PromptCard.css
@@ -299,3 +299,7 @@
 .info-circle-colored {
   color: #f0ad4e;
 }
+
+.required-checkbox-padding{
+  padding-left: 5px;
+}
diff --git a/frontend/src/components/custom-tools/prompt-card/PromptCardItems.jsx b/frontend/src/components/custom-tools/prompt-card/PromptCardItems.jsx
@@ -180,6 +180,7 @@ function PromptCardItems({
             enabledProfiles={enabledProfiles}
             spsLoading={spsLoading}
             handleSpsLoading={handleSpsLoading}
+            enforceType={enforceType}
           />
         </Space>
       </div>

diff --git a/prompt-service/src/unstract/prompt_service/constants.py b/prompt-service/src/unstract/prompt_service/constants.py
@@ -72,6 +72,8 @@ class PromptServiceContants:
     FILE_PATH = "file_path"
     HIGHLIGHT_DATA = "highlight_data"
     CONFIDENCE_DATA = "confidence_data"
+    REQUIRED_FIELDS = "required_fields"
+    REQUIRED = "required"
     EXECUTION_SOURCE = "execution_source"
     METRICS = "metrics"
 

diff --git a/prompt-service/src/unstract/prompt_service/helper.py b/prompt-service/src/unstract/prompt_service/helper.py
@@ -328,7 +328,6 @@ def run_completion(
         answer: str = completion[PSKeys.RESPONSE].text
         highlight_data = completion.get(PSKeys.HIGHLIGHT_DATA)
         confidence_data = completion.get(PSKeys.CONFIDENCE_DATA)
-
         if metadata is not None and prompt_key:
             if highlight_data:
                 metadata.setdefault(PSKeys.HIGHLIGHT_DATA, {})[
@@ -339,7 +338,6 @@ def run_completion(
                 metadata.setdefault(PSKeys.CONFIDENCE_DATA, {})[
                     prompt_key
                 ] = confidence_data
-
         return answer
     # TODO: Catch and handle specific exception here
     except SdkRateLimitError as e:

diff --git a/prompt-service/src/unstract/prompt_service/main.py b/prompt-service/src/unstract/prompt_service/main.py
@@ -108,6 +108,7 @@ def prompt_processor() -> Any:
         PSKeys.RUN_ID: run_id,
         PSKeys.FILE_NAME: doc_name,
         PSKeys.CONTEXT: {},
+        PSKeys.REQUIRED_FIELDS: {},
     }
     metrics: dict = {}
     variable_names: list[str] = []
@@ -120,10 +121,13 @@ def prompt_processor() -> Any:
         RunLevel.RUN,
         f"Preparing to execute {len(prompts)} prompt(s)",
     )
-
     # TODO: Rename "output" to "prompt"
     for output in prompts:  # type:ignore
         variable_names.append(output[PSKeys.NAME])
+        metadata[PSKeys.REQUIRED_FIELDS][output[PSKeys.NAME]] = output.get(
+            PSKeys.REQUIRED, None
+        )
+
     for output in prompts:  # type:ignore
         prompt_name = output[PSKeys.NAME]
         prompt_text = output[PSKeys.PROMPT]