v2 changes of PromptStudo-Registry

backend/prompt_studio/prompt_studio_registry_v2/admin.py

@@ -0,0 +1,5 @@
+from django.contrib import admin
+
+from .models import PromptStudioRegistry
+
+admin.site.register(PromptStudioRegistry)

backend/prompt_studio/prompt_studio_registry_v2/apps.py

@@ -0,0 +1,5 @@
+from django.apps import AppConfig
+
+
+class PromptStudioRegistry(AppConfig):
+    name = "prompt_studio.prompt_studio_registry_v2"

backend/prompt_studio/prompt_studio_registry_v2/constants.py

@@ -0,0 +1,109 @@
+class PromptStudioRegistryKeys:
+    CREATED_BY = "created_by"
+    TOOL_ID = "tool_id"
+    NUMBER = "Number"
+    FLOAT = "Float"
+    PG_VECTOR = "Postgres pg_vector"
+    ANSWERS = "answers"
+    UNIQUE_FILE_ID = "unique_file_id"
+    PROMPT_REGISTRY_ID = "prompt_registry_id"
+    FILE_NAME = "file_name"
+    UNDEFINED = "undefined"
+
+
+class PromptStudioRegistryErrors:
+    SERIALIZATION_FAILED = "Data Serialization Failed."
+    DUPLICATE_API = "It appears that a duplicate call may have been made."
+    CUSTOM_TOOL_EXISTS = "Custom tool with similiar configuration already exists"
+
+
+class LogLevels:
+    INFO = "INFO"
+    ERROR = "ERROR"
+    DEBUG = "DEBUG"
+    RUN = "RUN"
+
+
+# TODO: Update prompt studio constants to have a single source of truth
+class JsonSchemaKey:
+    TYPE = "type"
+    TITLE = "title"
+    DEFAULT = "default"
+    ENUM = "enum"
+    DESCRIPTION = "description"
+    REQUIRED = "required"
+    STRING = "string"
+    PROCESSOR_TO_USE = "Processor to use"
+    AZURE_OPEN_AI = "Azure OpenAI"
+    PROPERTIES = "properties"
+    DISPLAY_NAME = "display_name"
+    FUNCTION_NAME = "function_name"
+    PARAMETERS = "parameters"
+    VERSIONS = "versions"
+    OUTPUT_TYPE = "output_type"
+    INPUT_TYPE = "input_type"
+    IS_CACHABLE = "is_cacheable"
+    REQUIRES = "requires"
+    DEFAULT_DESCRIPTION_PROCESSOR = "Use Unstract processor \
+        if you do not want to use a cloud provider for privacy reasons"
+    NAME = "name"
+    ACTIVE = "active"
+    PROMPT = "prompt"
+    CHUNK_SIZE = "chunk-size"
+    PROMPTX = "promptx"
+    VECTOR_DB = "vector-db"
+    EMBEDDING = "embedding"
+    X2TEXT_ADAPTER = "x2text_adapter"
+    CHUNK_OVERLAP = "chunk-overlap"
+    LLM = "llm"
+    RETRIEVAL_STRATEGY = "retrieval-strategy"
+    SIMPLE = "simple"
+    TYPE = "type"
+    NUMBER = "number"
+    EMAIL = "email"
+    DATE = "date"
+    BOOLEAN = "boolean"
+    JSON = "json"
+    PREAMBLE = "preamble"
+    SIMILARITY_TOP_K = "similarity-top-k"
+    PROMPT_TOKENS = "prompt_tokens"
+    COMPLETION_TOKENS = "completion_tokens"
+    TOTAL_TOKENS = "total_tokens"
+    RESPONSE = "response"
+    POSTAMBLE = "postamble"
+    GRAMMAR = "grammar"
+    WORD = "word"
+    SYNONYMS = "synonyms"
+    OUTPUTS = "outputs"
+    SECTION = "section"
+    DEFAULT = "default"
+    AUTHOR = "author"
+    ICON = "icon"
+    REINDEX = "reindex"
+    TOOL_ID = "tool_id"
+    EMBEDDING_SUFFIX = "embedding_suffix"
+    FUNCTION_NAME = "function_name"
+    PROMPT_REGISTRY_ID = "prompt_registry_id"
+    NOTES = "NOTES"
+    TOOL_SETTINGS = "tool_settings"
+    ENABLE_CHALLENGE = "enable_challenge"
+    CHALLENGE_LLM = "challenge_llm"
+    ENABLE_SINGLE_PASS_EXTRACTION = "enable_single_pass_extraction"
+    IMAGE_URL = "image_url"
+    IMAGE_NAME = "image_name"
+    IMAGE_TAG = "image_tag"
+    SUMMARIZE_PROMPT = "summarize_prompt"
+    SUMMARIZE_AS_SOURCE = "summarize_as_source"
+    ENABLE_HIGHLIGHT = "enable_highlight"
+
+
+class SpecKey:
+    PROCESSOR = "processor"
+    SPEC = "spec"
+    OUTPUT_FOLDER = "outputFolder"
+    CREATE_OUTPUT_DOCUMENT = "createOutputDocument"
+    USE_CACHE = "useCache"
+    EMBEDDING_TRANSFORMER = "embeddingTransformer"
+    VECTOR_STORE = "vectorstore"
+    OUTPUT_TYPE = "outputType"
+    OUTPUT_PROCESSING = "outputProcessing"

backend/prompt_studio/prompt_studio_registry_v2/exceptions.py

@@ -0,0 +1,32 @@
+from rest_framework.exceptions import APIException
+
+
+class InternalError(APIException):
+    status_code = 500
+    default_detail = "Internal service error."
+
+
+class ToolDoesNotExist(APIException):
+    status_code = 500
+    default_detail = "Tool does not exist."
+
+
+class ToolSaveError(APIException):
+    status_code = 500
+    default_detail = "Error while saving the tool."
+
+
+class EmptyToolExportError(APIException):
+    status_code = 500
+    default_detail = (
+        "Empty Prompt Studio project without prompts cannot be exported. "
+        "Try adding a prompt and executing it."
+    )
+
+
+class InValidCustomToolError(APIException):
+    status_code = 500
+    default_detail = (
+        "This prompt studio project cannot be exported. It probably "
+        "has some empty or unexecuted prompts."
+    )

backend/prompt_studio/prompt_studio_registry_v2/fields.py

@@ -0,0 +1,31 @@
+import logging
+
+from django.db import models
+
+logger = logging.getLogger(__name__)
+
+
+class ToolBaseJSONField(models.JSONField):
+    def from_db_value(self, value, expression, connection):  # type: ignore
+        metadata = super().from_db_value(value, expression, connection)
+        return metadata
+
+
+# TODO: Investigate if ToolBaseJSONField can replace the need for ToolPropertyJSONField,
+#       ToolSpecJSONField, ToolVariablesJSONField, and ToolMetadataJSONField classes.
+
+
+class ToolPropertyJSONField(ToolBaseJSONField, models.JSONField):
+    pass
+
+
+class ToolSpecJSONField(ToolBaseJSONField, models.JSONField):
+    pass
+
+
+class ToolVariablesJSONField(ToolBaseJSONField, models.JSONField):
+    pass
+
+
+class ToolMetadataJSONField(ToolBaseJSONField, models.JSONField):
+    pass

backend/prompt_studio/prompt_studio_registry_v2/models.py

@@ -0,0 +1,110 @@
+import logging
+import uuid
+from typing import Any
+
+from account_v2.models import User
+from django.db import models
+from django.db.models import QuerySet
+from prompt_studio.prompt_studio_registry_v2.fields import (
+    ToolMetadataJSONField,
+    ToolPropertyJSONField,
+    ToolSpecJSONField,
+)
+from prompt_studio.prompt_studio_v2.models import CustomTool
+from utils.models.base_model import BaseModel
+from utils.models.organization_mixin import (
+    DefaultOrganizationManagerMixin,
+    DefaultOrganizationMixin,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class PromptStudioRegistryModelManager(DefaultOrganizationManagerMixin, models.Manager):
+    def get_queryset(self) -> QuerySet[Any]:
+        return super().get_queryset()
+
+    def list_tools(self, user: User) -> QuerySet[Any]:
+        return (
+            self.get_queryset()
+            .filter(
+                models.Q(created_by=user)
+                | models.Q(shared_users=user)
+                | models.Q(shared_to_org=True)
+            )
+            .distinct("prompt_registry_id")
+        )
+
+
+class PromptStudioRegistry(DefaultOrganizationMixin, BaseModel):
+    """Data model to export JSON fields needed for registering the Custom tool
+    to the tool registry.
+
+    By default the tools will be added to private tool hub.
+    """
+
+    prompt_registry_id = models.UUIDField(
+        primary_key=True, default=uuid.uuid4, editable=False
+    )
+    name = models.CharField(editable=False, default="")
+    description = models.CharField(editable=False, default="")
+    tool_property = ToolPropertyJSONField(
+        db_column="tool_property",
+        db_comment="PROPERTIES of the tool",
+        null=False,
+        blank=False,
+        default=dict,
+    )
+    tool_spec = ToolSpecJSONField(
+        db_column="tool_spec",
+        db_comment="SPEC of the tool",
+        null=False,
+        blank=False,
+        default=dict,
+    )
+    tool_metadata = ToolMetadataJSONField(
+        db_column="tool_metadata",
+        db_comment="Metadata from Prompt Studio",
+        null=False,
+        blank=False,
+        default=dict,
+    )
+    icon = models.CharField(db_comment="Tool icon in svg format", editable=False)
+    url = models.CharField(editable=False)
+    custom_tool = models.OneToOneField(
+        CustomTool,
+        on_delete=models.CASCADE,
+        related_name="prompt_studio_registries",
+        editable=False,
+        null=True,
+    )
+    created_by = models.ForeignKey(
+        User,
+        on_delete=models.SET_NULL,
+        related_name="prompt_registries_created",
+        null=True,
+        blank=True,
+        editable=False,
+    )
+    modified_by = models.ForeignKey(
+        User,
+        on_delete=models.SET_NULL,
+        related_name="prompt_registries_modified",
+        null=True,
+        blank=True,
+        editable=False,
+    )
+    shared_to_org = models.BooleanField(
+        default=False,
+        db_comment="Is the exported tool shared with entire org",
+    )
+    # Introduced field to establish M2M relation between users and tools.
+    # This will introduce intermediary table which relates both the models.
+    shared_users = models.ManyToManyField(User, related_name="prompt_registries")
+
+    objects = PromptStudioRegistryModelManager()
+
+    class Meta:
+        verbose_name = "Prompt Studio Registry"
+        verbose_name_plural = "Prompt Studio Registries"
+        db_table = "prompt_studio_registry_v2"