diff --git a/backend/prompt_studio/prompt_studio_registry_v2/__init__.py b/backend/prompt_studio/prompt_studio_registry_v2/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/prompt_studio/prompt_studio_registry_v2/admin.py b/backend/prompt_studio/prompt_studio_registry_v2/admin.py new file mode 100644 index 0000000000..9f6bdeb831 --- /dev/null +++ b/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) diff --git a/backend/prompt_studio/prompt_studio_registry_v2/apps.py b/backend/prompt_studio/prompt_studio_registry_v2/apps.py new file mode 100644 index 0000000000..e70823fb72 --- /dev/null +++ b/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" diff --git a/backend/prompt_studio/prompt_studio_registry_v2/constants.py b/backend/prompt_studio/prompt_studio_registry_v2/constants.py new file mode 100644 index 0000000000..6dbde1c096 --- /dev/null +++ b/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" diff --git a/backend/prompt_studio/prompt_studio_registry_v2/exceptions.py b/backend/prompt_studio/prompt_studio_registry_v2/exceptions.py new file mode 100644 index 0000000000..12a1829918 --- /dev/null +++ b/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." + ) diff --git a/backend/prompt_studio/prompt_studio_registry_v2/fields.py b/backend/prompt_studio/prompt_studio_registry_v2/fields.py new file mode 100644 index 0000000000..fcf8330b79 --- /dev/null +++ b/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 diff --git a/backend/prompt_studio/prompt_studio_registry_v2/models.py b/backend/prompt_studio/prompt_studio_registry_v2/models.py new file mode 100644 index 0000000000..ec768a854a --- /dev/null +++ b/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" 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 new file mode 100644 index 0000000000..e9d267e4d0 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py @@ -0,0 +1,382 @@ +import logging +from typing import Any, Optional + +from account_v2.models import User +from adapter_processor_v2.models import AdapterInstance +from django.conf import settings +from django.db import IntegrityError +from prompt_studio.prompt_profile_manager_v2.models import ProfileManager +from prompt_studio.prompt_studio_core_v2.models import CustomTool +from prompt_studio.prompt_studio_core_v2.prompt_studio_helper import PromptStudioHelper +from prompt_studio.prompt_studio_output_manager_v2.models import ( + PromptStudioOutputManager, +) +from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt +from unstract.tool_registry.dto import Properties, Spec, Tool + +from .constants import JsonSchemaKey +from .exceptions import ( + EmptyToolExportError, + InternalError, + InValidCustomToolError, + ToolSaveError, +) +from .models import PromptStudioRegistry +from .serializers import PromptStudioRegistrySerializer + +logger = logging.getLogger(__name__) + + +class PromptStudioRegistryHelper: + """Class to register custom tools to tool studio registry. + + By default the exported tools will be private and will be executed + with the help of a proto tool. + """ + + @staticmethod + def frame_spec(tool: CustomTool) -> Spec: + """Method to return spec of the Custom tool. + + Args: + tool (CustomTool): Saved tool data + + Returns: + dict: spec dict + """ + properties = { + "challenge_llm": { + "type": "string", + "title": "Challenge LLM", + "adapterType": "LLM", + "description": "LLM to use for challenge", + "adapterIdKey": "challenge_llm_adapter_id", + }, + "enable_challenge": { + "type": "boolean", + "title": "Enable challenge", + "default": False, + "description": "Enables Challenge", + }, + "summarize_as_source": { + "type": "boolean", + "title": "Summarize and use summary as source", + "default": False, + "description": "Enables summary and use summarized content as source", + }, + "single_pass_extraction_mode": { + "type": "boolean", + "title": "Enable Single pass extraction", + "default": False, + "description": "Enables single pass extraction", + }, + } + + spec = Spec( + title=str(tool.tool_id), + description=tool.description, + required=[JsonSchemaKey.CHALLENGE_LLM], + properties=properties, + ) + return spec + + @staticmethod + def frame_properties(tool: CustomTool) -> Properties: + """Method to return properties of the tool. + + Args: + tool (CustomTool): Saved custom tool data. + + Returns: + dict: Properties dict + """ + # TODO: Update for new architecture + tool_props = Properties( + display_name=tool.tool_name, + function_name=str(tool.tool_id), + description=tool.description, + ) + return tool_props + + @staticmethod + def get_tool_by_prompt_registry_id( + prompt_registry_id: str, + ) -> Optional[Tool]: + """Gets the `Tool` associated with a prompt registry ID if it exists. + + Args: + prompt_registry_id (str): Prompt registry ID to fetch for + + Returns: + Optional[Tool]: The `Tool` exported from Prompt Studio + """ + try: + prompt_registry_tool = PromptStudioRegistry.objects.get( + pk=prompt_registry_id + ) + # Suppress all exceptions to allow processing + except Exception as e: + logger.warning( + "Error while fetching for prompt registry " + f"ID {prompt_registry_id}: {e} " + ) + return None + # The below properties are introduced after 0.20.0 + # So defaulting to 0.20.0 if the properties are not found + image_url = prompt_registry_tool.tool_metadata.get( + JsonSchemaKey.IMAGE_URL, "docker:unstract/tool-structure:0.0.20" + ) + image_name = prompt_registry_tool.tool_metadata.get( + JsonSchemaKey.IMAGE_NAME, "unstract/tool-structure" + ) + image_tag = prompt_registry_tool.tool_metadata.get( + JsonSchemaKey.IMAGE_TAG, "0.0.20" + ) + return Tool( + tool_uid=prompt_registry_tool.prompt_registry_id, + properties=Properties.from_dict(prompt_registry_tool.tool_property), + spec=Spec.from_dict(prompt_registry_tool.tool_spec), + icon=prompt_registry_tool.icon, + image_url=image_url, + image_name=image_name, + image_tag=image_tag, + ) + + @staticmethod + def update_or_create_psr_tool( + custom_tool: CustomTool, shared_with_org: bool, user_ids: set[int] + ) -> PromptStudioRegistry: + """Updates or creates the PromptStudioRegistry record. + + This appears as a separate tool in the workflow and is mapped + 1:1 with the `CustomTool`. + + Args: + tool_id (str): ID of the custom tool. + + Raises: + ToolSaveError + InternalError + + Returns: + obj: PromptStudioRegistry instance that was updated or created + """ + try: + properties: Properties = PromptStudioRegistryHelper.frame_properties( + tool=custom_tool + ) + spec: Spec = PromptStudioRegistryHelper.frame_spec(tool=custom_tool) + prompts: list[ToolStudioPrompt] = PromptStudioHelper.fetch_prompt_from_tool( + tool_id=custom_tool.tool_id + ) + metadata = PromptStudioRegistryHelper.frame_export_json( + tool=custom_tool, prompts=prompts + ) + + obj: PromptStudioRegistry + created: bool + obj, created = PromptStudioRegistry.objects.update_or_create( + custom_tool=custom_tool, + created_by=custom_tool.created_by, + modified_by=custom_tool.modified_by, + defaults={ + "name": custom_tool.tool_name, + "tool_property": properties.to_dict(), + "tool_spec": spec.to_dict(), + "tool_metadata": metadata, + "icon": custom_tool.icon, + "description": custom_tool.description, + }, + ) + if created: + logger.info(f"PSR {obj.prompt_registry_id} was created") + else: + logger.info(f"PSR {obj.prompt_registry_id} was updated") + + obj.shared_to_org = shared_with_org + if not shared_with_org: + obj.shared_users.clear() + obj.shared_users.add(*user_ids) + # add prompt studio users + # for shared_user in custom_tool.shared_users: + obj.shared_users.add( + *custom_tool.shared_users.all().values_list("id", flat=True) + ) + # add prompt studio owner + obj.shared_users.add(custom_tool.created_by) + else: + obj.shared_users.clear() + obj.save() + return obj + except IntegrityError as error: + logger.error( + "Integrity Error - Error occurred while " + f"exporting custom tool : {error}" + ) + raise ToolSaveError + + @staticmethod + def frame_export_json( + tool: CustomTool, prompts: list[ToolStudioPrompt] + ) -> dict[str, Any]: + export_metadata = {} + + prompt_grammer = tool.prompt_grammer + grammar_list = [] + grammer_dict = {} + outputs: list[dict[str, Any]] = [] + output: dict[str, Any] = {} + invalidated_prompts: list[str] = [] + invalidated_outputs: list[str] = [] + + if not prompts: + raise EmptyToolExportError() + + if prompt_grammer: + for word, synonyms in prompt_grammer.items(): + synonyms = prompt_grammer[word] + grammer_dict[JsonSchemaKey.WORD] = word + grammer_dict[JsonSchemaKey.SYNONYMS] = synonyms + grammar_list.append(grammer_dict) + grammer_dict = {} + + export_metadata[JsonSchemaKey.NAME] = tool.tool_name + export_metadata[JsonSchemaKey.DESCRIPTION] = tool.description + export_metadata[JsonSchemaKey.AUTHOR] = tool.author + export_metadata[JsonSchemaKey.TOOL_ID] = str(tool.tool_id) + export_metadata[JsonSchemaKey.IMAGE_URL] = settings.STRUCTURE_TOOL_IMAGE_URL + export_metadata[JsonSchemaKey.IMAGE_NAME] = settings.STRUCTURE_TOOL_IMAGE_NAME + export_metadata[JsonSchemaKey.IMAGE_TAG] = settings.STRUCTURE_TOOL_IMAGE_TAG + + default_llm_profile = ProfileManager.get_default_llm_profile(tool) + challenge_llm_instance: Optional[AdapterInstance] = tool.challenge_llm + challenge_llm: Optional[str] = None + # Using default profile manager llm if challenge_llm is None + if challenge_llm_instance: + challenge_llm = str(challenge_llm_instance.id) + else: + challenge_llm = str(default_llm_profile.llm.id) + + embedding_suffix = "" + adapter_id = "" + vector_db = str(default_llm_profile.vector_store.id) + embedding_model = str(default_llm_profile.embedding_model.id) + llm = str(default_llm_profile.llm.id) + x2text = str(default_llm_profile.x2text.id) + + # Tool settings + tool_settings = {} + tool_settings[JsonSchemaKey.SUMMARIZE_PROMPT] = tool.summarize_prompt + tool_settings[JsonSchemaKey.SUMMARIZE_AS_SOURCE] = tool.summarize_as_source + tool_settings[JsonSchemaKey.PREAMBLE] = tool.preamble + tool_settings[JsonSchemaKey.POSTAMBLE] = tool.postamble + tool_settings[JsonSchemaKey.GRAMMAR] = grammar_list + tool_settings[JsonSchemaKey.LLM] = llm + tool_settings[JsonSchemaKey.X2TEXT_ADAPTER] = x2text + tool_settings[JsonSchemaKey.VECTOR_DB] = vector_db + tool_settings[JsonSchemaKey.EMBEDDING] = embedding_model + tool_settings[JsonSchemaKey.CHUNK_SIZE] = default_llm_profile.chunk_size + tool_settings[JsonSchemaKey.CHUNK_OVERLAP] = default_llm_profile.chunk_overlap + tool_settings[JsonSchemaKey.ENABLE_CHALLENGE] = tool.enable_challenge + tool_settings[JsonSchemaKey.CHALLENGE_LLM] = challenge_llm + tool_settings[JsonSchemaKey.ENABLE_SINGLE_PASS_EXTRACTION] = ( + tool.single_pass_extraction_mode + ) + + for prompt in prompts: + + if not prompt.prompt: + invalidated_prompts.append(prompt.prompt_key) + continue + + prompt_output = PromptStudioOutputManager.objects.filter( + tool_id=tool.tool_id, + prompt_id=prompt.prompt_id, + profile_manager=prompt.profile_manager, + ).all() + + if not prompt_output: + invalidated_outputs.append(prompt.prompt_key) + continue + + if prompt.prompt_type == JsonSchemaKey.NOTES: + continue + if not prompt.profile_manager: + prompt.profile_manager = default_llm_profile + + vector_db = str(prompt.profile_manager.vector_store.id) + embedding_model = str(prompt.profile_manager.embedding_model.id) + llm = str(prompt.profile_manager.llm.id) + x2text = str(prompt.profile_manager.x2text.id) + adapter_id = str(prompt.profile_manager.embedding_model.adapter_id) + embedding_suffix = adapter_id.split("|")[0] + + output[JsonSchemaKey.PROMPT] = prompt.prompt + output[JsonSchemaKey.ACTIVE] = prompt.active + output[JsonSchemaKey.CHUNK_SIZE] = prompt.profile_manager.chunk_size + output[JsonSchemaKey.VECTOR_DB] = vector_db + output[JsonSchemaKey.EMBEDDING] = embedding_model + output[JsonSchemaKey.X2TEXT_ADAPTER] = x2text + output[JsonSchemaKey.CHUNK_OVERLAP] = prompt.profile_manager.chunk_overlap + output[JsonSchemaKey.LLM] = llm + output[JsonSchemaKey.PREAMBLE] = tool.preamble + output[JsonSchemaKey.POSTAMBLE] = tool.postamble + output[JsonSchemaKey.GRAMMAR] = grammar_list + output[JsonSchemaKey.TYPE] = prompt.enforce_type + output[JsonSchemaKey.NAME] = prompt.prompt_key + output[JsonSchemaKey.RETRIEVAL_STRATEGY] = ( + prompt.profile_manager.retrieval_strategy + ) + output[JsonSchemaKey.SIMILARITY_TOP_K] = ( + prompt.profile_manager.similarity_top_k + ) + output[JsonSchemaKey.SECTION] = prompt.profile_manager.section + output[JsonSchemaKey.REINDEX] = prompt.profile_manager.reindex + output[JsonSchemaKey.EMBEDDING_SUFFIX] = embedding_suffix + outputs.append(output) + output = {} + vector_db = "" + embedding_suffix = "" + adapter_id = "" + llm = "" + embedding_model = "" + + if invalidated_prompts: + raise InValidCustomToolError( + f"Cannot export tool. Prompt(s): {', '.join(invalidated_prompts)} " + "are empty. Please enter a valid prompt." + ) + if invalidated_outputs: + raise InValidCustomToolError( + f"Cannot export tool. Prompt(s): {', '.join(invalidated_outputs)} " + "were not run. Please run them before exporting." + ) + export_metadata[JsonSchemaKey.TOOL_SETTINGS] = tool_settings + export_metadata[JsonSchemaKey.OUTPUTS] = outputs + return export_metadata + + @staticmethod + def fetch_json_for_registry(user: User) -> list[dict[str, Any]]: + try: + # filter the Prompt studio registry based on the users and org flag + prompt_studio_tools = PromptStudioRegistry.objects.list_tools(user) + pi_serializer = PromptStudioRegistrySerializer( + instance=prompt_studio_tools, many=True + ) + except Exception as error: + logger.error(f"Error occured while fetching tool for tool_id: {error}") + raise InternalError() + tool_metadata: dict[str, Any] = {} + tool_list = [] + for prompts in pi_serializer.data: + tool_metadata[JsonSchemaKey.NAME] = prompts.get(JsonSchemaKey.NAME) + tool_metadata[JsonSchemaKey.DESCRIPTION] = prompts.get( + JsonSchemaKey.DESCRIPTION + ) + tool_metadata[JsonSchemaKey.ICON] = prompts.get(JsonSchemaKey.ICON) + tool_metadata[JsonSchemaKey.FUNCTION_NAME] = prompts.get( + JsonSchemaKey.PROMPT_REGISTRY_ID + ) + tool_list.append(tool_metadata) + tool_metadata = {} + return tool_list diff --git a/backend/prompt_studio/prompt_studio_registry_v2/serializers.py b/backend/prompt_studio/prompt_studio_registry_v2/serializers.py new file mode 100644 index 0000000000..e995bb2663 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_registry_v2/serializers.py @@ -0,0 +1,38 @@ +from typing import Any + +from account_v2.serializer import UserSerializer +from rest_framework import serializers + +from backend.serializers import AuditSerializer + +from .models import PromptStudioRegistry + + +class PromptStudioRegistrySerializer(AuditSerializer): + class Meta: + model = PromptStudioRegistry + fields = "__all__" + + +class PromptStudioRegistryInfoSerializer(AuditSerializer): + shared_users = UserSerializer(many=True) + prompt_studio_users = serializers.SerializerMethodField() + + class Meta: + model = PromptStudioRegistry + fields = ( + "name", + "shared_users", + "shared_to_org", + "prompt_studio_users", + ) + + def get_prompt_studio_users(self, obj: PromptStudioRegistry) -> Any: + + prompt_studio_users = obj.custom_tool.shared_users + return UserSerializer(prompt_studio_users, many=True).data + + +class ExportToolRequestSerializer(serializers.Serializer): + is_shared_with_org = serializers.BooleanField(default=False) + user_id = serializers.ListField(child=serializers.IntegerField(), required=False) diff --git a/backend/prompt_studio/prompt_studio_registry_v2/urls.py b/backend/prompt_studio/prompt_studio_registry_v2/urls.py new file mode 100644 index 0000000000..9cc93e05ec --- /dev/null +++ b/backend/prompt_studio/prompt_studio_registry_v2/urls.py @@ -0,0 +1,3 @@ +from rest_framework.urlpatterns import format_suffix_patterns + +urlpatterns = format_suffix_patterns([]) diff --git a/backend/prompt_studio/prompt_studio_registry_v2/views.py b/backend/prompt_studio/prompt_studio_registry_v2/views.py new file mode 100644 index 0000000000..6b8a5d7e7f --- /dev/null +++ b/backend/prompt_studio/prompt_studio_registry_v2/views.py @@ -0,0 +1,35 @@ +import logging +from typing import Optional + +from django.db.models import QuerySet +from prompt_studio.prompt_studio_registry_v2.constants import PromptStudioRegistryKeys +from prompt_studio.prompt_studio_registry_v2.serializers import ( + PromptStudioRegistrySerializer, +) +from rest_framework import viewsets +from rest_framework.versioning import URLPathVersioning +from utils.filtering import FilterHelper + +from .models import PromptStudioRegistry + +logger = logging.getLogger(__name__) + + +class PromptStudioRegistryView(viewsets.ModelViewSet): + """Driver class to handle export and registering of custom tools to private + tool hub.""" + + versioning_class = URLPathVersioning + queryset = PromptStudioRegistry.objects.all() + serializer_class = PromptStudioRegistrySerializer + + def get_queryset(self) -> Optional[QuerySet]: + filterArgs = FilterHelper.build_filter_args( + self.request, + PromptStudioRegistryKeys.PROMPT_REGISTRY_ID, + ) + queryset = None + if filterArgs: + queryset = PromptStudioRegistry.objects.filter(**filterArgs) + + return queryset