diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/__init__.py b/backend/prompt_studio/prompt_studio_document_manager_v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/admin.py b/backend/prompt_studio/prompt_studio_document_manager_v2/admin.py new file mode 100644 index 000000000..a16311f1b --- /dev/null +++ b/backend/prompt_studio/prompt_studio_document_manager_v2/admin.py @@ -0,0 +1,5 @@ +from django.contrib import admin + +from .models import DocumentManager + +admin.site.register(DocumentManager) diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/apps.py b/backend/prompt_studio/prompt_studio_document_manager_v2/apps.py new file mode 100644 index 000000000..12e67143a --- /dev/null +++ b/backend/prompt_studio/prompt_studio_document_manager_v2/apps.py @@ -0,0 +1,5 @@ +from django.apps import AppConfig + + +class PromptStudioDocumentManagerConfig(AppConfig): + name = "prompt_studio.prompt_studio_document_manager_v2" diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/constants.py b/backend/prompt_studio/prompt_studio_document_manager_v2/constants.py new file mode 100644 index 000000000..2ac55ca83 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_document_manager_v2/constants.py @@ -0,0 +1,4 @@ +class PSDMKeys: + DOCUMENT_NAME = "document_name" + TOOL = "tool" + DOCUMENT_ID = "document_id" diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/models.py b/backend/prompt_studio/prompt_studio_document_manager_v2/models.py new file mode 100644 index 000000000..139d67e31 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_document_manager_v2/models.py @@ -0,0 +1,56 @@ +import uuid + +from account_v2.models import User +from django.db import models +from prompt_studio.prompt_studio_core_v2.models import CustomTool +from utils.models.base_model import BaseModel + + +class DocumentManager(BaseModel): + """Model to store the document details.""" + + document_id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) + + document_name = models.CharField( + db_comment="Field to store the document name", + editable=False, + null=False, + blank=False, + ) + + tool = models.ForeignKey( + CustomTool, + on_delete=models.CASCADE, + related_name="document_managers", + null=False, + blank=False, + ) + + created_by = models.ForeignKey( + User, + on_delete=models.SET_NULL, + related_name="document_managers_created", + null=True, + blank=True, + editable=False, + ) + + modified_by = models.ForeignKey( + User, + on_delete=models.SET_NULL, + related_name="document_managers_modified", + null=True, + blank=True, + editable=False, + ) + + class Meta: + verbose_name = "Document Manager" + verbose_name_plural = "Document Managers" + db_table = "document_manager_v2" + constraints = [ + models.UniqueConstraint( + fields=["document_name", "tool"], + name="unique_document_name_tool_index", + ), + ] diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/prompt_studio_document_helper.py b/backend/prompt_studio/prompt_studio_document_manager_v2/prompt_studio_document_helper.py new file mode 100644 index 000000000..618d0e7e0 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_document_manager_v2/prompt_studio_document_helper.py @@ -0,0 +1,24 @@ +import logging + +from prompt_studio.prompt_studio_core_v2.models import CustomTool + +from .models import DocumentManager + +logger = logging.getLogger(__name__) + + +class PromptStudioDocumentHelper: + @staticmethod + def create(tool_id: str, document_name: str) -> DocumentManager: + tool: CustomTool = CustomTool.objects.get(pk=tool_id) + document: DocumentManager = DocumentManager.objects.create( + tool=tool, document_name=document_name + ) + logger.info("Successfully created the record") + return document + + @staticmethod + def delete(document_id: str) -> None: + document: DocumentManager = DocumentManager.objects.get(pk=document_id) + document.delete() + logger.info("Successfully deleted the record") diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/serializers.py b/backend/prompt_studio/prompt_studio_document_manager_v2/serializers.py new file mode 100644 index 000000000..844500fc8 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_document_manager_v2/serializers.py @@ -0,0 +1,21 @@ +from typing import Any + +from backend.serializers import AuditSerializer + +from .constants import PSDMKeys +from .models import DocumentManager + + +class PromptStudioDocumentManagerSerializer(AuditSerializer): + class Meta: + model = DocumentManager + fields = "__all__" + + def to_representation(self, instance: DocumentManager) -> dict[str, Any]: + rep: dict[str, str] = super().to_representation(instance) + required_fields = [ + PSDMKeys.DOCUMENT_NAME, + PSDMKeys.TOOL, + PSDMKeys.DOCUMENT_ID, + ] + return {key: rep[key] for key in required_fields if key in rep} diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/urls.py b/backend/prompt_studio/prompt_studio_document_manager_v2/urls.py new file mode 100644 index 000000000..f9fb9bcd3 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_document_manager_v2/urls.py @@ -0,0 +1,24 @@ +from django.urls import path +from rest_framework.urlpatterns import format_suffix_patterns + +from .views import PromptStudioDocumentManagerView + +prompt_studio_documents_list = PromptStudioDocumentManagerView.as_view( + {"get": "list", "post": "create"} +) + +prompt_studio_documents_detail = PromptStudioDocumentManagerView.as_view( + { + "get": "retrieve", + } +) + +urlpatterns = format_suffix_patterns( + [ + path( + "prompt-document/", + prompt_studio_documents_list, + name="prompt-studio-documents-list", + ), + ] +) diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/views.py b/backend/prompt_studio/prompt_studio_document_manager_v2/views.py new file mode 100644 index 000000000..6867aecb5 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_document_manager_v2/views.py @@ -0,0 +1,31 @@ +from typing import Optional + +from django.db.models import QuerySet +from prompt_studio.prompt_studio_document_manager_v2.serializers import ( + PromptStudioDocumentManagerSerializer, +) +from prompt_studio.prompt_studio_output_manager_v2.constants import ( + PromptStudioOutputManagerKeys, +) +from rest_framework import viewsets +from rest_framework.versioning import URLPathVersioning +from utils.filtering import FilterHelper + +from .models import DocumentManager + + +class PromptStudioDocumentManagerView(viewsets.ModelViewSet): + versioning_class = URLPathVersioning + queryset = DocumentManager.objects.all() + serializer_class = PromptStudioDocumentManagerSerializer + + def get_queryset(self) -> Optional[QuerySet]: + filter_args = FilterHelper.build_filter_args( + self.request, + PromptStudioOutputManagerKeys.TOOL_ID, + ) + queryset = None + if filter_args: + queryset = DocumentManager.objects.filter(**filter_args) + + return queryset diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/__init__.py b/backend/prompt_studio/prompt_studio_index_manager_v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/admin.py b/backend/prompt_studio/prompt_studio_index_manager_v2/admin.py new file mode 100644 index 000000000..188d50ffe --- /dev/null +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/admin.py @@ -0,0 +1,5 @@ +from django.contrib import admin + +from .models import IndexManager + +admin.site.register(IndexManager) diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/apps.py b/backend/prompt_studio/prompt_studio_index_manager_v2/apps.py new file mode 100644 index 000000000..85d25e759 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/apps.py @@ -0,0 +1,5 @@ +from django.apps import AppConfig + + +class PromptStudioIndexManagerConfig(AppConfig): + name = "prompt_studio.prompt_studio_index_manager_v2" diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/constants.py b/backend/prompt_studio/prompt_studio_index_manager_v2/constants.py new file mode 100644 index 000000000..6cf3f5e5b --- /dev/null +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/constants.py @@ -0,0 +1,3 @@ +class IndexManagerKeys: + PROFILE_MANAGER = "profile_manager" + DOCUMENT_MANAGER = "document_manager" diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/models.py b/backend/prompt_studio/prompt_studio_index_manager_v2/models.py new file mode 100644 index 000000000..c160f8561 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/models.py @@ -0,0 +1,132 @@ +import json +import logging +import uuid + +from account_v2.models import User +from django.db import models +from django.db.models.signals import pre_delete +from django.dispatch import receiver +from prompt_studio.prompt_profile_manager_v2.models import ProfileManager +from prompt_studio.prompt_studio_core_v2.prompt_ide_base_tool import PromptIdeBaseTool +from prompt_studio.prompt_studio_document_manager_v2.models import DocumentManager +from unstract.sdk.constants import LogLevel +from unstract.sdk.vector_db import VectorDB +from utils.models.base_model import BaseModel +from utils.user_context import UserContext + +logger = logging.getLogger(__name__) + + +class IndexManager(BaseModel): + """Model to store the index details.""" + + index_manager_id = models.UUIDField( + primary_key=True, default=uuid.uuid4, editable=False + ) + + document_manager = models.ForeignKey( + DocumentManager, + on_delete=models.CASCADE, + related_name="index_managers", + editable=False, + null=False, + blank=False, + ) + + profile_manager = models.ForeignKey( + ProfileManager, + on_delete=models.SET_NULL, + related_name="index_managers", + editable=False, + null=True, + blank=True, + ) + + raw_index_id = models.CharField( + db_comment="Field to store the raw index id", + editable=False, + null=True, + blank=True, + ) + + summarize_index_id = models.CharField( + db_comment="Field to store the summarize index id", + editable=False, + null=True, + blank=True, + ) + + index_ids_history = models.JSONField( + db_comment="List of index ids", + default=list, + null=False, + blank=False, + ) + + created_by = models.ForeignKey( + User, + on_delete=models.SET_NULL, + related_name="index_managers_created", + null=True, + blank=True, + editable=False, + ) + + modified_by = models.ForeignKey( + User, + on_delete=models.SET_NULL, + related_name="index_managers_modified", + null=True, + blank=True, + editable=False, + ) + + class Meta: + verbose_name = "Index Manager" + verbose_name_plural = "Index Managers" + db_table = "index_manager_v2" + constraints = [ + models.UniqueConstraint( + fields=["document_manager", "profile_manager"], + name="unique_document_manager_profile_manager_index", + ), + ] + + +def delete_from_vector_db(index_ids_history, vector_db_instance_id): + organization_identifier = UserContext.get_organization_identifier() + util = PromptIdeBaseTool(log_level=LogLevel.INFO, org_id=organization_identifier) + vector_db = VectorDB( + tool=util, + adapter_instance_id=vector_db_instance_id, + ) + for index_id in index_ids_history: + logger.debug(f"Deleting from VectorDB - index id: {index_id}") + try: + vector_db.delete(ref_doc_id=index_id) + except Exception as e: + # Log error and continue with the next index id + logger.error(f"Error deleting index: {index_id} - {e}") + + +# Function will be executed every time an instance of IndexManager is deleted. +@receiver(pre_delete, sender=IndexManager) +def perform_vector_db_cleanup(sender, instance, **kwargs): + """Signal to perform vector db cleanup.""" + logger.debug( + "Performing vector db cleanup for Document tool id: " + f"{instance.document_manager.tool_id}" + ) + try: + # Get the index_ids_history to clean up from the vector db + index_ids_history = json.loads(instance.index_ids_history) + vector_db_instance_id = str(instance.profile_manager.vector_store.id) + delete_from_vector_db(index_ids_history, vector_db_instance_id) + except Exception as e: + logger.warning( + "Error during vector DB cleanup for deleted document " + "in prompt studio tool %s: %s", + instance.document_manager.tool_id, + e, + exc_info=True, # For additional stack trace + ) diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py b/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py new file mode 100644 index 000000000..084b2c458 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py @@ -0,0 +1,66 @@ +import json +import logging + +from django.db import transaction +from prompt_studio.prompt_profile_manager_v2.models import ProfileManager +from prompt_studio.prompt_studio_core_v2.exceptions import IndexingAPIError +from prompt_studio.prompt_studio_document_manager_v2.models import DocumentManager + +from .models import IndexManager + +logger = logging.getLogger(__name__) + + +class PromptStudioIndexHelper: + @staticmethod + def handle_index_manager( + document_id: str, + is_summary: bool, + profile_manager: ProfileManager, + doc_id: str, + ) -> IndexManager: + try: + + with transaction.atomic(): + + document: DocumentManager = DocumentManager.objects.get(pk=document_id) + + index_id = "raw_index_id" + if is_summary: + index_id = "summarize_index_id" + + args: dict[str, str] = dict() + args["document_manager"] = document + args["profile_manager"] = profile_manager + + # Create or get the existing record for this document and + # profile combo + index_manager, success = IndexManager.objects.get_or_create(**args) + + if success: + logger.info( + f"Index manager doc_id: {doc_id} for " + f"profile {profile_manager.profile_id} created" + ) + else: + logger.info( + f"Index manager doc_id: {doc_id} for " + f"profile {profile_manager.profile_id} updated" + ) + + index_ids = index_manager.index_ids_history + index_ids_list = json.loads(index_ids) if index_ids else [] + if doc_id not in index_ids: + index_ids_list.append(doc_id) + + args[index_id] = doc_id + args["index_ids_history"] = json.dumps(index_ids_list) + + # Update the record with the index id + result: IndexManager = IndexManager.objects.filter( + index_manager_id=index_manager.index_manager_id + ).update(**args) + return result + except Exception as e: + transaction.rollback() + raise IndexingAPIError("Error updating indexing status") from e diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/serializers.py b/backend/prompt_studio/prompt_studio_index_manager_v2/serializers.py new file mode 100644 index 000000000..ae33a186e --- /dev/null +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/serializers.py @@ -0,0 +1,9 @@ +from backend.serializers import AuditSerializer + +from .models import IndexManager + + +class IndexManagerSerializer(AuditSerializer): + class Meta: + model = IndexManager + fields = "__all__" diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/urls.py b/backend/prompt_studio/prompt_studio_index_manager_v2/urls.py new file mode 100644 index 000000000..b082bc1d8 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/urls.py @@ -0,0 +1,22 @@ +from django.urls import path +from rest_framework.urlpatterns import format_suffix_patterns + +from .views import IndexManagerView + +prompt_studio_index_list = IndexManagerView.as_view({"get": "list", "post": "create"}) + +prompt_studio_index_detail = IndexManagerView.as_view( + { + "get": "retrieve", + } +) + +urlpatterns = format_suffix_patterns( + [ + path( + "document-index/", + prompt_studio_index_list, + name="prompt-studio-documents-list", + ), + ] +) diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/views.py b/backend/prompt_studio/prompt_studio_index_manager_v2/views.py new file mode 100644 index 000000000..bfcb5d7e7 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/views.py @@ -0,0 +1,29 @@ +from typing import Optional + +from django.db.models import QuerySet +from prompt_studio.prompt_studio_index_manager_v2.constants import IndexManagerKeys +from prompt_studio.prompt_studio_index_manager_v2.serializers import ( + IndexManagerSerializer, +) +from rest_framework import viewsets +from rest_framework.versioning import URLPathVersioning +from utils.filtering import FilterHelper + +from .models import IndexManager + + +class IndexManagerView(viewsets.ModelViewSet): + versioning_class = URLPathVersioning + queryset = IndexManager.objects.all() + serializer_class = IndexManagerSerializer + + def get_queryset(self) -> Optional[QuerySet]: + filter_args = FilterHelper.build_filter_args( + self.request, + IndexManagerKeys.PROFILE_MANAGER, + IndexManagerKeys.DOCUMENT_MANAGER, + ) + queryset = None + if filter_args: + queryset = IndexManager.objects.filter(**filter_args) + return queryset diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/__init__.py b/backend/prompt_studio/prompt_studio_output_manager_v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/admin.py b/backend/prompt_studio/prompt_studio_output_manager_v2/admin.py new file mode 100644 index 000000000..00fd98264 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/admin.py @@ -0,0 +1,5 @@ +from django.contrib import admin + +from .models import PromptStudioOutputManager + +admin.site.register(PromptStudioOutputManager) diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/apps.py b/backend/prompt_studio/prompt_studio_output_manager_v2/apps.py new file mode 100644 index 000000000..d181ff25f --- /dev/null +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/apps.py @@ -0,0 +1,5 @@ +from django.apps import AppConfig + + +class PromptStudioOutputManager(AppConfig): + name = "prompt_studio.prompt_studio_output_manager_v2" diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/constants.py b/backend/prompt_studio/prompt_studio_output_manager_v2/constants.py new file mode 100644 index 000000000..a9c046aae --- /dev/null +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/constants.py @@ -0,0 +1,7 @@ +class PromptStudioOutputManagerKeys: + TOOL_ID = "tool_id" + PROMPT_ID = "prompt_id" + PROFILE_MANAGER = "profile_manager" + DOCUMENT_MANAGER = "document_manager" + IS_SINGLE_PASS_EXTRACT = "is_single_pass_extract" + NOTES = "NOTES" diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/exceptions.py b/backend/prompt_studio/prompt_studio_output_manager_v2/exceptions.py new file mode 100644 index 000000000..f11530914 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/exceptions.py @@ -0,0 +1,6 @@ +from rest_framework.exceptions import APIException + + +class InternalError(APIException): + status_code = 400 + default_detail = "Internal service error." diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/models.py b/backend/prompt_studio/prompt_studio_output_manager_v2/models.py new file mode 100644 index 000000000..afc133f17 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/models.py @@ -0,0 +1,91 @@ +import uuid + +from account_v2.models import User +from django.db import models +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_document_manager_v2.models import DocumentManager +from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt +from utils.models.base_model import BaseModel + + +class PromptStudioOutputManager(BaseModel): + """Data model to handle output persisitance for Project. + + By default the tools will be added to private tool hub. + """ + + prompt_output_id = models.UUIDField( + primary_key=True, default=uuid.uuid4, editable=False + ) + output = models.CharField( + db_comment="Field to store output", editable=True, null=True, blank=True + ) + context = models.TextField( + db_comment="Field to store chunks used", editable=True, null=True, blank=True + ) + eval_metrics = models.JSONField( + db_column="eval_metrics", + null=False, + blank=False, + default=list, + db_comment="Field to store the evaluation metrics", + ) + is_single_pass_extract = models.BooleanField( + default=False, + db_comment="Is the single pass extraction mode active", + ) + prompt_id = models.ForeignKey( + ToolStudioPrompt, + on_delete=models.CASCADE, + related_name="prompt_studio_outputs", + ) + document_manager = models.ForeignKey( + DocumentManager, + on_delete=models.CASCADE, + related_name="prompt_studio_outputs", + ) + profile_manager = models.ForeignKey( + ProfileManager, + on_delete=models.CASCADE, + related_name="prompt_studio_outputs", + ) + tool_id = models.ForeignKey( + CustomTool, + on_delete=models.CASCADE, + related_name="prompt_studio_outputs", + ) + run_id = models.UUIDField(default=uuid.uuid4, editable=False) + created_by = models.ForeignKey( + User, + on_delete=models.SET_NULL, + related_name="prompt_studio_outputs_created", + null=True, + blank=True, + editable=False, + ) + modified_by = models.ForeignKey( + User, + on_delete=models.SET_NULL, + related_name="prompt_studio_outputs_modified", + null=True, + blank=True, + editable=False, + ) + + class Meta: + verbose_name = "Prompt Studio Output Manager" + verbose_name_plural = "Prompt Studio Output Managers" + db_table = "prompt_studio_output_manager_v2" + constraints = [ + models.UniqueConstraint( + fields=[ + "prompt_id", + "document_manager", + "profile_manager", + "tool_id", + "is_single_pass_extract", + ], + name="unique_prompt_output_index", + ), + ] diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py new file mode 100644 index 000000000..163dd781a --- /dev/null +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py @@ -0,0 +1,148 @@ +import json +import logging +from typing import Any, Optional + +from prompt_studio.prompt_profile_manager_v2.models import ProfileManager +from prompt_studio.prompt_studio_core_v2.exceptions import ( + AnswerFetchError, + DefaultProfileError, +) +from prompt_studio.prompt_studio_core_v2.models import CustomTool +from prompt_studio.prompt_studio_document_manager_v2.models import DocumentManager +from prompt_studio.prompt_studio_output_manager_v2.constants import ( + PromptStudioOutputManagerKeys as PSOMKeys, +) +from prompt_studio.prompt_studio_output_manager_v2.models import ( + PromptStudioOutputManager, +) +from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt + +logger = logging.getLogger(__name__) + + +class OutputManagerHelper: + @staticmethod + def handle_prompt_output_update( + run_id: str, + prompts: list[ToolStudioPrompt], + outputs: Any, + context: Any, + document_id: str, + is_single_pass_extract: bool, + profile_manager_id: Optional[str] = None, + ) -> None: + """Handles updating prompt outputs in the database. + + Args: + run_id (str): ID of the run. + prompts (list[ToolStudioPrompt]): List of prompts to update. + outputs (Any): Outputs corresponding to the prompts. + document_id (str): ID of the document. + profile_manager_id (Optional[str]): UUID of the profile manager. + is_single_pass_extract (bool): + Flag indicating if single pass extract is active. + """ + + def update_or_create_prompt_output( + prompt: ToolStudioPrompt, + profile_manager: ProfileManager, + output: str, + eval_metrics: list[Any], + tool: CustomTool, + context: str, + ): + try: + _, success = PromptStudioOutputManager.objects.get_or_create( + document_manager=document_manager, + tool_id=tool, + profile_manager=profile_manager, + prompt_id=prompt, + is_single_pass_extract=is_single_pass_extract, + defaults={ + "output": output, + "eval_metrics": eval_metrics, + "context": context, + }, + ) + + if success: + logger.info( + f"Created record for prompt_id: {prompt.prompt_id} and " + f"profile {profile_manager.profile_id}" + ) + else: + logger.info( + f"Updated record for prompt_id: {prompt.prompt_id} and " + f"profile {profile_manager.profile_id}" + ) + + args: dict[str, str] = { + "run_id": run_id, + "output": output, + "eval_metrics": eval_metrics, + "context": context, + } + PromptStudioOutputManager.objects.filter( + document_manager=document_manager, + tool_id=tool, + profile_manager=profile_manager, + prompt_id=prompt, + is_single_pass_extract=is_single_pass_extract, + ).update(**args) + + except Exception as e: + raise AnswerFetchError(f"Error updating prompt output {e}") from e + + if not prompts: + return # Return early if prompts list is empty + + tool = prompts[0].tool_id + default_profile = OutputManagerHelper.get_default_profile( + profile_manager_id, tool + ) + document_manager = DocumentManager.objects.get(pk=document_id) + + for prompt in prompts: + if prompt.prompt_type == PSOMKeys.NOTES or not prompt.active: + continue + + if not is_single_pass_extract: + context = json.dumps(context.get(prompt.prompt_key)) + + output = json.dumps(outputs.get(prompt.prompt_key)) + profile_manager = default_profile + eval_metrics = outputs.get(f"{prompt.prompt_key}__evaluation", []) + + update_or_create_prompt_output( + prompt=prompt, + profile_manager=profile_manager, + output=output, + eval_metrics=eval_metrics, + tool=tool, + context=context, + ) + + @staticmethod + def get_default_profile( + profile_manager_id: Optional[str], tool: CustomTool + ) -> ProfileManager: + if profile_manager_id: + return OutputManagerHelper.fetch_profile_manager(profile_manager_id) + else: + return OutputManagerHelper.fetch_default_llm_profile(tool) + + @staticmethod + def fetch_profile_manager(profile_manager_id: str) -> ProfileManager: + try: + return ProfileManager.objects.get(profile_id=profile_manager_id) + except ProfileManager.DoesNotExist: + raise DefaultProfileError( + f"ProfileManager with ID {profile_manager_id} does not exist." + ) + + @staticmethod + def fetch_default_llm_profile(tool: CustomTool) -> ProfileManager: + try: + return ProfileManager.get_default_llm_profile(tool=tool) + except DefaultProfileError: + raise DefaultProfileError("Default ProfileManager does not exist.") diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/serializers.py b/backend/prompt_studio/prompt_studio_output_manager_v2/serializers.py new file mode 100644 index 000000000..f3b769c1b --- /dev/null +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/serializers.py @@ -0,0 +1,28 @@ +import logging + +from usage_v2.helper import UsageHelper + +from backend.serializers import AuditSerializer + +from .models import PromptStudioOutputManager + +logger = logging.getLogger(__name__) + + +class PromptStudioOutputSerializer(AuditSerializer): + class Meta: + model = PromptStudioOutputManager + fields = "__all__" + + def to_representation(self, instance): + data = super().to_representation(instance) + try: + token_usage = UsageHelper.get_aggregated_token_count(instance.run_id) + except Exception as e: + logger.error( + "Error occured while fetching token usage for run_id" + f"{instance.run_id}: {e}" + ) + token_usage = {} + data["token_usage"] = token_usage + return data diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/urls.py b/backend/prompt_studio/prompt_studio_output_manager_v2/urls.py new file mode 100644 index 000000000..45af0ccee --- /dev/null +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/urls.py @@ -0,0 +1,12 @@ +from django.urls import path +from rest_framework.urlpatterns import format_suffix_patterns + +from .views import PromptStudioOutputView + +prompt_doc_list = PromptStudioOutputView.as_view({"get": "list"}) + +urlpatterns = format_suffix_patterns( + [ + path("prompt-output/", prompt_doc_list, name="prompt-doc-list"), + ] +) diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py new file mode 100644 index 000000000..e4dd99b9f --- /dev/null +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py @@ -0,0 +1,51 @@ +import logging +from typing import Optional + +from django.db.models import QuerySet +from prompt_studio.prompt_studio_output_manager_v2.constants import ( + PromptStudioOutputManagerKeys, +) +from prompt_studio.prompt_studio_output_manager_v2.serializers import ( + PromptStudioOutputSerializer, +) +from rest_framework import viewsets +from rest_framework.versioning import URLPathVersioning +from utils.common_utils import CommonUtils +from utils.filtering import FilterHelper + +from .models import PromptStudioOutputManager + +logger = logging.getLogger(__name__) + + +class PromptStudioOutputView(viewsets.ModelViewSet): + versioning_class = URLPathVersioning + queryset = PromptStudioOutputManager.objects.all() + serializer_class = PromptStudioOutputSerializer + + def get_queryset(self) -> Optional[QuerySet]: + filter_args = FilterHelper.build_filter_args( + self.request, + PromptStudioOutputManagerKeys.TOOL_ID, + PromptStudioOutputManagerKeys.PROMPT_ID, + PromptStudioOutputManagerKeys.PROFILE_MANAGER, + PromptStudioOutputManagerKeys.DOCUMENT_MANAGER, + PromptStudioOutputManagerKeys.IS_SINGLE_PASS_EXTRACT, + ) + + # Get the query parameter for "is_single_pass_extract" + is_single_pass_extract_param = self.request.GET.get( + PromptStudioOutputManagerKeys.IS_SINGLE_PASS_EXTRACT, "false" + ) + + # Convert the string representation to a boolean value + is_single_pass_extract = CommonUtils.str_to_bool(is_single_pass_extract_param) + + filter_args[PromptStudioOutputManagerKeys.IS_SINGLE_PASS_EXTRACT] = ( + is_single_pass_extract + ) + + if filter_args: + queryset = PromptStudioOutputManager.objects.filter(**filter_args) + + return queryset