Multitenancy V2 porting of Table settings enhancements (#751)

* V2 porting of Table settings enhancements * V2 porting of Table settings enhancements --------- Co-authored-by: Neha <[email protected]>
Zipstack · Oct 1, 2024 · e36941f · e36941f
1 parent 1507085
commit e36941f
Show file tree

Hide file tree

Showing 6 changed files with 33 additions and 8 deletions.
diff --git a/backend/prompt_studio/prompt_studio_core_v2/constants.py b/backend/prompt_studio/prompt_studio_core_v2/constants.py
@@ -93,6 +93,7 @@ class ToolStudioPromptKeys:
     PLATFORM_POSTAMBLE = "platform_postamble"
     SUMMARIZE_AS_SOURCE = "summarize_as_source"
     VARIABLE_MAP = "variable_map"
+    RECORD = "record"
 
 
 class FileViewTypes:

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
@@ -62,7 +62,7 @@
 
 logger = logging.getLogger(__name__)
 
-modifier_loader = load_modifier_plugins()
+modifier_plugins = load_modifier_plugins()
 
 
 class PromptStudioHelper:
@@ -262,6 +262,11 @@ def get_select_fields() -> dict[str, Any]:
         choices = f.read()
         f.close()
         response: dict[str, Any] = json.loads(choices)
+        for modifier_plugin in modifier_plugins:
+            cls = modifier_plugin[ModifierConfig.METADATA][
+                ModifierConfig.METADATA_SERVICE_CLASS
+            ]
+            response = cls.update_select_choices(default_choices=response)
         return response
 
     @staticmethod
@@ -451,7 +456,10 @@ def _execute_single_prompt(
     ):
         prompt_instance = PromptStudioHelper._fetch_prompt_from_id(id)
 
-        if prompt_instance.enforce_type == TSPKeys.TABLE and not modifier_loader:
+        if (
+            prompt_instance.enforce_type == TSPKeys.TABLE
+            or prompt_instance.enforce_type == TSPKeys.RECORD
+        ) and not modifier_plugins:
             raise OperationNotSupported()
 
         prompt_name = prompt_instance.prompt_key
@@ -547,6 +555,7 @@ def _execute_prompts_in_single_pass(
             if prompt.prompt_type != TSPKeys.NOTES
             and prompt.active
             and prompt.enforce_type != TSPKeys.TABLE
+            and prompt.enforce_type != TSPKeys.RECORD
         ]
         if not prompts:
             logger.error(f"[{tool_id or 'NA'}] No prompts found for id: {id}")
@@ -841,13 +850,16 @@ def fetch_table_settings_if_enabled(
         output: dict[str, Any],
     ) -> dict[str, Any]:
 
-        if prompt.enforce_type == TSPKeys.TABLE:
+        if (
+            prompt.enforce_type == TSPKeys.TABLE
+            or prompt.enforce_type == TSPKeys.RECORD
+        ):
             extract_doc_path: str = (
                 PromptStudioHelper._get_extract_or_summary_document_path(
                     org_id, user_id, tool_id, doc_name, TSPKeys.EXTRACT
                 )
             )
-            for modifier_plugin in modifier_loader:
+            for modifier_plugin in modifier_plugins:
                 cls = modifier_plugin[ModifierConfig.METADATA][
                     ModifierConfig.METADATA_SERVICE_CLASS
                 ]
@@ -857,6 +869,7 @@ def fetch_table_settings_if_enabled(
                     prompt_id=str(prompt.prompt_id),
                     prompt=prompt.prompt,
                     input_file=extract_doc_path,
+                    clean_pages=True,
                 )
 
         return output

diff --git a/backend/prompt_studio/prompt_studio_core_v2/static/select_choices.json b/backend/prompt_studio/prompt_studio_core_v2/static/select_choices.json
@@ -14,7 +14,8 @@
         "date":"date",
         "boolean":"boolean",
         "json":"json",
-        "table":"table"
+        "table":"table",
+        "record":"record"
     },
     "output_processing":{
         "DEFAULT":"Default"

diff --git a/backend/prompt_studio/prompt_studio_registry_v2/constants.py b/backend/prompt_studio/prompt_studio_registry_v2/constants.py
@@ -10,6 +10,7 @@ class PromptStudioRegistryKeys:
     FILE_NAME = "file_name"
     UNDEFINED = "undefined"
     TABLE = "table"
+    RECORD = "record"
 
 
 class PromptStudioRegistryErrors:

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
@@ -27,7 +27,7 @@
 from .serializers import PromptStudioRegistrySerializer
 
 logger = logging.getLogger(__name__)
-modifier_loader = load_modifier_plugins()
+modifier_plugins = load_modifier_plugins()
 
 
 class PromptStudioRegistryHelper:
@@ -337,8 +337,11 @@ def frame_export_json(
             output[JsonSchemaKey.REINDEX] = prompt.profile_manager.reindex
             output[JsonSchemaKey.EMBEDDING_SUFFIX] = embedding_suffix
 
-            if prompt.enforce_type == PromptStudioRegistryKeys.TABLE:
-                for modifier_plugin in modifier_loader:
+            if (
+                prompt.enforce_type == PromptStudioRegistryKeys.TABLE
+                or prompt.enforce_type == PromptStudioRegistryKeys.RECORD
+            ):
+                for modifier_plugin in modifier_plugins:
                     cls = modifier_plugin[ModifierConfig.METADATA][
                         ModifierConfig.METADATA_SERVICE_CLASS
                     ]

diff --git a/backend/prompt_studio/prompt_studio_v2/models.py b/backend/prompt_studio/prompt_studio_v2/models.py
@@ -21,6 +21,12 @@ class EnforceType(models.TextChoices):
         BOOLEAN = "boolean", "Response sent as boolean"
         JSON = "json", "Response sent as json"
         TABLE = "table", "Response sent as table"
+        RECORD = "record", (
+            "Response sent for records. "
+            "Entries of records are list of "
+            "logical and organized individual "
+            "entities with distint values"
+        )
 
     class PromptType(models.TextChoices):
         PROMPT = "PROMPT", "Response sent as Text"