[FEAT] Support for table extraction in Prompt studio (#564)

* Support for table extraction in Prompt studio

* Refactors-Table extraction

* Sonar issues

* Removing plugins

* Adding missed urls

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Enrich error message

Co-authored-by: Chandrasekharan M <[email protected]>
Signed-off-by: harini-venkataraman <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Warn cloud exceptions

Co-authored-by: Chandrasekharan M <[email protected]>
Signed-off-by: harini-venkataraman <[email protected]>

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

Co-authored-by: Chandrasekharan M <[email protected]>
Signed-off-by: harini-venkataraman <[email protected]>

* Skip table for single pass

* Adding method alias

* Refactors

* Method refactors

* Constants

* FEAT: Table Extractor (FE) (#590)

* FE changes for table extraction

* UI improvements related to table extraction

* Fixed Eslint issues

* Passed the output prop

* Fixed component rendering and sonar issue

* Added optional chaining

* Output in payload

* Sonar fixes

* Sonar fixes

* Bump structure tool version

* Remove metadata includions

* Add TODO for plugging logic to tools

---------

Signed-off-by: harini-venkataraman <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Chandrasekharan M <[email protected]>
Co-authored-by: Tahier Hussain <[email protected]>

Author: 4 people

backend/backend/urls.py

@@ -118,10 +118,17 @@
 # Clone urls
 try:
 
-    import pluggable_apps.clone.urls  # noqa # pylint: disable=unused-import
-
     urlpatterns += [
         path("", include("pluggable_apps.clone.urls")),
     ]
 except ImportError:
     pass
+
+try:
+    import pluggable_apps.apps.table_settings  # noqa # pylint: disable=unused-import
+
+    urlpatterns += [
+        path("", include("pluggable_apps.apps.table_settings.urls")),
+    ]
+except ImportError:
+    pass

backend/prompt_studio/modifier_loader.py

@@ -0,0 +1,72 @@
+import logging
+import os
+from importlib import import_module
+from typing import Any
+
+from django.apps import apps
+
+logger = logging.getLogger(__name__)
+
+
+class ModifierConfig:
+    """Loader config for extraction plugins."""
+
+    PLUGINS_APP = "plugins"
+    PLUGIN_DIR = "modifier"
+    MODULE = "module"
+    METADATA = "metadata"
+    METADATA_NAME = "name"
+    METADATA_SERVICE_CLASS = "service_class"
+    METADATA_IS_ACTIVE = "is_active"
+
+
+def load_plugins() -> list[Any]:
+    """Iterate through the extraction plugins and register them."""
+    plugins_app = apps.get_app_config(ModifierConfig.PLUGINS_APP)
+    package_path = plugins_app.module.__package__
+    modifier_dir = os.path.join(plugins_app.path, ModifierConfig.PLUGIN_DIR)
+    modifier_package_path = f"{package_path}.{ModifierConfig.PLUGIN_DIR}"
+    modifier_plugins: list[Any] = []
+
+    if not os.path.exists(modifier_dir):
+        return modifier_plugins
+
+    for item in os.listdir(modifier_dir):
+        # Loads a plugin if it is in a directory.
+        if os.path.isdir(os.path.join(modifier_dir, item)):
+            modifier_module_name = item
+        # Loads a plugin if it is a shared library.
+        # Module name is extracted from shared library name.
+        elif item.endswith(".so"):
+            modifier_module_name = item.split(".")[0]
+        else:
+            continue
+        try:
+            full_module_path = f"{modifier_package_path}.{modifier_module_name}"
+            module = import_module(full_module_path)
+            metadata = getattr(module, ModifierConfig.METADATA, {})
+
+            if metadata.get(ModifierConfig.METADATA_IS_ACTIVE, False):
+                modifier_plugins.append(
+                    {
+                        ModifierConfig.MODULE: module,
+                        ModifierConfig.METADATA: module.metadata,
+                    }
+                )
+                logger.info(
+                    "Loaded modifier plugin: %s, is_active: %s",
+                    module.metadata[ModifierConfig.METADATA_NAME],
+                    module.metadata[ModifierConfig.METADATA_IS_ACTIVE],
+                )
+            else:
+                logger.info(
+                    "modifier plugin %s is not active.",
+                    modifier_module_name,
+                )
+        except ModuleNotFoundError:
+            logger.warning("No prompt modifier plugins loaded")
+
+    if len(modifier_plugins) == 0:
+        logger.info("No modifier plugins found.")
+
+    return modifier_plugins

backend/prompt_studio/prompt_studio/migrations/0007_alter_toolstudioprompt_enforce_type.py

@@ -0,0 +1,31 @@
+# Generated by Django 4.2.1 on 2024-08-07 14:20
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("prompt_studio", "0006_alter_toolstudioprompt_prompt_key_and_more"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="toolstudioprompt",
+            name="enforce_type",
+            field=models.TextField(
+                blank=True,
+                choices=[
+                    ("Text", "Response sent as Text"),
+                    ("number", "Response sent as number"),
+                    ("email", "Response sent as email"),
+                    ("date", "Response sent as date"),
+                    ("boolean", "Response sent as boolean"),
+                    ("json", "Response sent as json"),
+                    ("table", "Response sent as table"),
+                ],
+                db_comment="Field to store the type in             which the response to be returned.",
+                default="Text",
+            ),
+        ),
+    ]

backend/prompt_studio/prompt_studio/models.py

@@ -20,6 +20,7 @@ class EnforceType(models.TextChoices):
         DATE = "date", "Response sent as date"
         BOOLEAN = "boolean", "Response sent as boolean"
         JSON = "json", "Response sent as json"
+        TABLE = "table", "Response sent as table"
 
     class PromptType(models.TextChoices):
         PROMPT = "PROMPT", "Response sent as Text"

backend/prompt_studio/prompt_studio_core/constants.py

@@ -88,6 +88,9 @@ class ToolStudioPromptKeys:
     CONTEXT = "context"
     METADATA = "metadata"
     INCLUDE_METADATA = "include_metadata"
+    TXT_EXTENTION = ".txt"
+    TABLE = "table"
+    EXTRACT = "extract"
     PLATFORM_POSTAMBLE = "platform_postamble"
     SUMMARIZE_AS_SOURCE = "summarize_as_source"
 

backend/prompt_studio/prompt_studio_core/exceptions.py

@@ -67,3 +67,13 @@ class MaxProfilesReachedError(APIException):
         f"Maximum number of profiles (max {ProfileManagerKeys.MAX_PROFILE_COUNT})"
         " per prompt studio project has been reached."
     )
+
+
+class OperationNotSupported(APIException):
+    status_code = 403
+    default_detail = (
+        "This feature is not supported "
+        "in the open-source version. "
+        "Please check our cloud or enterprise on-premise offering  "
+        "for access to this functionality."
+    )

backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py

@@ -12,6 +12,8 @@
 from django.conf import settings
 from django.db.models.manager import BaseManager
 from file_management.file_management_helper import FileManagerHelper
+from prompt_studio.modifier_loader import ModifierConfig
+from prompt_studio.modifier_loader import load_plugins as load_modifier_plugins
 from prompt_studio.prompt_profile_manager.models import ProfileManager
 from prompt_studio.prompt_profile_manager.profile_manager_helper import (
     ProfileManagerHelper,
@@ -28,6 +30,7 @@
     EmptyPromptError,
     IndexingAPIError,
     NoPromptsFound,
+    OperationNotSupported,
     PermissionError,
 )
 from prompt_studio.prompt_studio_core.models import CustomTool
@@ -53,6 +56,8 @@
 
 logger = logging.getLogger(__name__)
 
+modifier_loader = load_modifier_plugins()
+
 
 class PromptStudioHelper:
     """Helper class for Custom tool operations."""
@@ -434,6 +439,10 @@ def _execute_single_prompt(
         text_processor: Optional[type[Any]] = None,
     ):
         prompt_instance = PromptStudioHelper._fetch_prompt_from_id(id)
+
+        if prompt_instance.enforce_type == TSPKeys.TABLE and not modifier_loader:
+            raise OperationNotSupported()
+
         prompt_name = prompt_instance.prompt_key
         PromptStudioHelper._publish_log(
             {
@@ -523,7 +532,9 @@ def _execute_prompts_in_single_pass(
         prompts = [
             prompt
             for prompt in prompts
-            if prompt.prompt_type != TSPKeys.NOTES and prompt.active
+            if prompt.prompt_type != TSPKeys.NOTES
+            and prompt.active
+            and prompt.enforce_type != TSPKeys.TABLE
         ]
         if not prompts:
             logger.error(f"[{tool_id or 'NA'}] No prompts found for id: {id}")
@@ -583,6 +594,19 @@ def _get_document_path(org_id, user_id, tool_id, doc_name):
         )
         return str(Path(doc_path) / doc_name)
 
+    @staticmethod
+    def _get_extract_or_summary_document_path(
+        org_id, user_id, tool_id, doc_name, doc_type
+    ) -> str:
+        doc_path = FileManagerHelper.handle_sub_directory_for_tenants(
+            org_id=org_id,
+            user_id=user_id,
+            tool_id=tool_id,
+            is_create=False,
+        )
+        extracted_doc_name = Path(doc_name).stem + TSPKeys.TXT_EXTENTION
+        return str(Path(doc_path) / doc_type / extracted_doc_name)
+
     @staticmethod
     def _handle_response(
         response,
@@ -698,7 +722,7 @@ def _fetch_response(
                 "status": IndexingStatus.PENDING_STATUS.value,
                 "message": IndexingStatus.DOCUMENT_BEING_INDEXED.value,
             }
-
+        tool_id = str(tool.tool_id)
         output: dict[str, Any] = {}
         outputs: list[dict[str, Any]] = []
         grammer_dict = {}
@@ -738,6 +762,10 @@ def _fetch_response(
                 attr_val = getattr(prompt, attr)
                 output[TSPKeys.EVAL_SETTINGS][attr] = attr_val
 
+        output = PromptStudioHelper.fetch_table_settings_if_enabled(
+            doc_name, prompt, org_id, user_id, tool_id, output
+        )
+
         outputs.append(output)
 
         tool_settings = {}
@@ -754,8 +782,6 @@ def _fetch_response(
             settings, TSPKeys.PLATFORM_POSTAMBLE.upper(), ""
         )
 
-        tool_id = str(tool.tool_id)
-
         file_hash = ToolUtils.get_hash_from_file(file_path=doc_path)
 
         payload = {
@@ -789,6 +815,36 @@ def _fetch_response(
         output_response = json.loads(answer["structure_output"])
         return output_response
 
+    @staticmethod
+    def fetch_table_settings_if_enabled(
+        doc_name: str,
+        prompt: ToolStudioPrompt,
+        org_id: str,
+        user_id: str,
+        tool_id: str,
+        output: dict[str, Any],
+    ) -> dict[str, Any]:
+
+        if prompt.enforce_type == TSPKeys.TABLE:
+            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:
+                cls = modifier_plugin[ModifierConfig.METADATA][
+                    ModifierConfig.METADATA_SERVICE_CLASS
+                ]
+                output = cls.update(
+                    output=output,
+                    tool_id=tool_id,
+                    prompt_id=str(prompt.prompt_id),
+                    prompt=prompt.prompt,
+                    input_file=extract_doc_path,
+                )
+
+        return output
+
     @staticmethod
     def dynamic_indexer(
         profile_manager: ProfileManager,

backend/prompt_studio/prompt_studio_core/static/select_choices.json

@@ -13,7 +13,8 @@
         "email":"email",
         "date":"date",
         "boolean":"boolean",
-        "json":"json"
+        "json":"json",
+        "table":"table"
     },
     "output_processing":{
         "DEFAULT":"Default"

backend/prompt_studio/prompt_studio_registry/constants.py

@@ -9,6 +9,7 @@ class PromptStudioRegistryKeys:
     PROMPT_REGISTRY_ID = "prompt_registry_id"
     FILE_NAME = "file_name"
     UNDEFINED = "undefined"
+    TABLE = "table"
 
 
 class PromptStudioRegistryErrors:

backend/prompt_studio/prompt_studio_registry/prompt_studio_registry_helper.py

@@ -5,14 +5,16 @@
 from adapter_processor.models import AdapterInstance
 from django.conf import settings
 from django.db import IntegrityError
+from prompt_studio.modifier_loader import ModifierConfig
+from prompt_studio.modifier_loader import load_plugins as load_modifier_plugins
 from prompt_studio.prompt_profile_manager.models import ProfileManager
 from prompt_studio.prompt_studio.models import ToolStudioPrompt
 from prompt_studio.prompt_studio_core.models import CustomTool
 from prompt_studio.prompt_studio_core.prompt_studio_helper import PromptStudioHelper
 from prompt_studio.prompt_studio_output_manager.models import PromptStudioOutputManager
 from unstract.tool_registry.dto import Properties, Spec, Tool
 
-from .constants import JsonSchemaKey
+from .constants import JsonSchemaKey, PromptStudioRegistryKeys
 from .exceptions import (
     EmptyToolExportError,
     InternalError,
@@ -23,6 +25,7 @@
 from .serializers import PromptStudioRegistrySerializer
 
 logger = logging.getLogger(__name__)
+modifier_loader = load_modifier_plugins()
 
 
 class PromptStudioRegistryHelper:
@@ -320,6 +323,19 @@ def frame_export_json(
             output[JsonSchemaKey.SECTION] = prompt.profile_manager.section
             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:
+                    cls = modifier_plugin[ModifierConfig.METADATA][
+                        ModifierConfig.METADATA_SERVICE_CLASS
+                    ]
+                    output = cls.update(
+                        output=output,
+                        tool_id=tool.tool_id,
+                        prompt_id=prompt.prompt_id,
+                        prompt=prompt.prompt,
+                    )
+
             outputs.append(output)
             output = {}
             vector_db = ""