Feat/prompt key required field

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"
 
 
 class FileViewTypes:

backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py

@@ -815,6 +815,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

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:

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

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),
+        ),
+    ]

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,
+            ),
+        ),
+    ]

backend/prompt_studio/prompt_studio_v2/models.py

@@ -84,6 +84,17 @@ class Mode(models.TextChoices):
         db_comment="Field to store the prompt key",
         unique=False,
     )
+    REQUIRED_CHOICES = [
+        ("all", "All values required"),
+        ("any", "Any value required"),
+    ]
+    required = models.CharField(
+        max_length=3,
+        choices=REQUIRED_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
+    )
     is_assert = models.BooleanField(default=False)
     active = models.BooleanField(default=True, null=False, blank=False)
     output_metadata = models.JSONField(

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

@@ -45,6 +45,7 @@ function Header({
   setExpandCard,
   spsLoading,
   handleSpsLoading,
+  enforceType,
 }) {
   const {
     selectedDoc,
@@ -58,6 +59,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 +75,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 +103,37 @@ function Header({
         ),
         key: "enable",
       },
+      {
+        label: (
+          <div>
+            {["json", "table", "record"].indexOf(enforceType) === -1 && (
+              <Checkbox
+                checked={required === "all"}
+                onChange={() => handleRequiredChange("all")}
+              >
+                Required
+              </Checkbox>
+            )}
+            {enforceType === "json" && (
+              <>
+                <Checkbox
+                  checked={required === "all"}
+                  onChange={() => handleRequiredChange("all")}
+                >
+                  All Required
+                </Checkbox>
+                <Checkbox
+                  checked={required === "any"}
+                  onChange={() => handleRequiredChange("any")}
+                >
+                  Any Required
+                </Checkbox>
+              </>
+            )}
+          </div>
+        ),
+        key: "required",
+      },
       {
         label: (
           <ConfirmModal
@@ -109,7 +156,7 @@ function Header({
     }
 
     setItems(dropdownItems);
-  }, [isDisablePrompt]);
+  }, [isDisablePrompt, required, enforceType]);
 
   return (
     <Row>
@@ -270,6 +317,7 @@ Header.propTypes = {
   setExpandCard: PropTypes.func.isRequired,
   spsLoading: PropTypes.object,
   handleSpsLoading: PropTypes.func.isRequired,
+  enforceType: PropTypes.text,
 };
 
 export { Header };

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

@@ -169,6 +169,7 @@ function PromptCardItems({
             enabledProfiles={enabledProfiles}
             spsLoading={spsLoading}
             handleSpsLoading={handleSpsLoading}
+            enforceType={enforceType}
           />
         </Space>
       </div>

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"
 
 
 class RunLevel(Enum):

prompt-service/src/unstract/prompt_service/helper.py

@@ -297,7 +297,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, {})[
@@ -308,7 +307,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:
@@ -344,3 +342,24 @@ def extract_table(
     except table_extractor["exception_cls"] as e:
         msg = f"Couldn't extract table. {e}"
         raise APIError(message=msg)
+
+
+def add_required_field(
+    required_fields: dict[str, str], key: str, required: str
+) -> dict[str, str]:
+    """
+    Add or update a required field in the required_fields dictionary.
+
+    Args:
+        required_fields (Dict[str, str]): The dictionary
+        containing existing required fields,
+        where keys are field names and values are boolean flags.
+        key (str): The field key to add or update in the dictionary.
+        required (str): A indicating whether the
+        all/any values in field is required.
+
+    Returns:
+        Dict[str, str]: The updated dictionary of required fields.
+    """
+    required_fields[key] = required
+    return required_fields

prompt-service/src/unstract/prompt_service/main.py

@@ -11,6 +11,7 @@
 from unstract.prompt_service.constants import RunLevel
 from unstract.prompt_service.exceptions import APIError, ErrorResponse, NoPayloadError
 from unstract.prompt_service.helper import (
+    add_required_field,
     construct_and_run_prompt,
     extract_table,
     extract_variable,
@@ -108,6 +109,7 @@ def prompt_processor() -> Any:
         PSKeys.RUN_ID: run_id,
         PSKeys.FILE_NAME: doc_name,
         PSKeys.CONTEXT: {},
+        PSKeys.REQUIRED_FIELDS: {},
     }
     variable_names: list[str] = []
     publish_log(
@@ -117,10 +119,14 @@ 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] = add_required_field(
+            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]