feat/multi llm enhancement

backend/adapter_processor/serializers.py

@@ -124,6 +124,10 @@ def to_representation(self, instance: AdapterInstance) -> dict[str, str]:
         rep[common.ICON] = AdapterProcessor.get_adapter_data_with_key(
             instance.adapter_id, common.ICON
         )
+        adapter_metadata = instance.get_adapter_meta_data()
+        model = adapter_metadata.get("model")
+        if model:
+            rep["model"] = model
 
         if instance.is_friction_less:
             rep["created_by_email"] = "Unstract"

backend/backend/settings/base.py

@@ -167,6 +167,7 @@ def get_required_setting(
     "CELERY_BROKER_URL", f"redis://{REDIS_HOST}:{REDIS_PORT}"
 )
 
+INDEXING_FLAG_TTL = int(get_required_setting("INDEXING_FLAG_TTL"))
 # Flag to Enable django admin
 ADMIN_ENABLED = False
 

backend/prompt_studio/prompt_profile_manager/constants.py

@@ -7,6 +7,8 @@ class ProfileManagerKeys:
     VECTOR_STORE = "vector_store"
     EMBEDDING_MODEL = "embedding_model"
     X2TEXT = "x2text"
+    PROMPT_STUDIO_TOOL = "prompt_studio_tool"
+    MAX_PROFILE_COUNT = 4
 
 
 class ProfileManagerErrors:

backend/prompt_studio/prompt_profile_manager/profile_manager_helper.py

@@ -0,0 +1,11 @@
+from prompt_studio.prompt_profile_manager.models import ProfileManager
+
+
+class ProfileManagerHelper:
+
+    @classmethod
+    def get_profile_manager(cls, profile_manager_id: str) -> ProfileManager:
+        try:
+            return ProfileManager.objects.get(profile_id=profile_manager_id)
+        except ProfileManager.DoesNotExist:
+            raise ValueError("ProfileManager does not exist.")

backend/prompt_studio/prompt_profile_manager/serializers.py

@@ -2,6 +2,7 @@
 
 from adapter_processor.adapter_processor import AdapterProcessor
 from prompt_studio.prompt_profile_manager.constants import ProfileManagerKeys
+from prompt_studio.prompt_studio_core.exceptions import MaxProfilesReachedError
 
 from backend.serializers import AuditSerializer
 
@@ -38,3 +39,15 @@ def to_representation(self, instance):  # type: ignore
                 AdapterProcessor.get_adapter_instance_by_id(x2text)
             )
         return rep
+
+    def validate(self, data):
+        prompt_studio_tool = data.get(ProfileManagerKeys.PROMPT_STUDIO_TOOL)
+
+        profile_count = ProfileManager.objects.filter(
+            prompt_studio_tool=prompt_studio_tool
+        ).count()
+
+        if profile_count >= ProfileManagerKeys.MAX_PROFILE_COUNT:
+            raise MaxProfilesReachedError()
+
+        return data

backend/prompt_studio/prompt_studio_core/constants.py

@@ -85,6 +85,9 @@ class ToolStudioPromptKeys:
     NOTES = "NOTES"
     OUTPUT = "output"
     SEQUENCE_NUMBER = "sequence_number"
+    PROFILE_MANAGER_ID = "profile_manager"
+    CONTEXT = "context"
+    METADATA = "metadata"
 
 
 class FileViewTypes:
@@ -108,6 +111,13 @@ class LogLevel(Enum):
     FATAL = "FATAL"
 
 
+class IndexingStatus(Enum):
+    PENDING_STATUS = "pending"
+    COMPLETED_STATUS = "completed"
+    STARTED_STATUS = "started"
+    DOCUMENT_BEING_INDEXED = "Document is being indexed"
+
+
 class DefaultPrompts:
     PREAMBLE = (
         "Your ability to extract and summarize this context accurately "

backend/prompt_studio/prompt_studio_core/document_indexing_service.py

@@ -0,0 +1,53 @@
+from typing import Optional
+
+from django.conf import settings
+from prompt_studio.prompt_studio_core.constants import IndexingStatus
+from utils.cache_service import CacheService
+
+
+class DocumentIndexingService:
+    CACHE_PREFIX = "document_indexing:"
+
+    @classmethod
+    def set_document_indexing(cls, org_id: str, user_id: str, doc_id_key: str) -> None:
+        CacheService.set_key(
+            cls._cache_key(org_id, user_id, doc_id_key),
+            IndexingStatus.STARTED_STATUS.value,
+            expire=settings.INDEXING_FLAG_TTL,
+        )
+
+    @classmethod
+    def is_document_indexing(cls, org_id: str, user_id: str, doc_id_key: str) -> bool:
+        return (
+            CacheService.get_key(cls._cache_key(org_id, user_id, doc_id_key))
+            == IndexingStatus.STARTED_STATUS.value
+        )
+
+    @classmethod
+    def mark_document_indexed(
+        cls, org_id: str, user_id: str, doc_id_key: str, doc_id: str
+    ) -> None:
+        CacheService.set_key(
+            cls._cache_key(org_id, user_id, doc_id_key),
+            doc_id,
+            expire=settings.INDEXING_FLAG_TTL,
+        )
+
+    @classmethod
+    def get_indexed_document_id(
+        cls, org_id: str, user_id: str, doc_id_key: str
+    ) -> Optional[str]:
+        result = CacheService.get_key(cls._cache_key(org_id, user_id, doc_id_key))
+        if result and result != IndexingStatus.STARTED_STATUS.value:
+            return result
+        return None
+
+    @classmethod
+    def remove_document_indexing(
+        cls, org_id: str, user_id: str, doc_id_key: str
+    ) -> None:
+        CacheService.delete_a_key(cls._cache_key(org_id, user_id, doc_id_key))
+
+    @classmethod
+    def _cache_key(cls, org_id: str, user_id: str, doc_id_key: str) -> str:
+        return f"{cls.CACHE_PREFIX}{org_id}:{user_id}:{doc_id_key}"

backend/prompt_studio/prompt_studio_core/exceptions.py

@@ -1,3 +1,4 @@
+from prompt_studio.prompt_profile_manager.constants import ProfileManagerKeys
 from prompt_studio.prompt_studio_core.constants import ToolStudioErrors
 from rest_framework.exceptions import APIException
 
@@ -58,3 +59,11 @@ class PermissionError(APIException):
 class EmptyPromptError(APIException):
     status_code = 422
     default_detail = "Prompt(s) cannot be empty"
+
+
+class MaxProfilesReachedError(APIException):
+    status_code = 403
+    default_detail = (
+        f"Maximum number of profiles (max {ProfileManagerKeys.MAX_PROFILE_COUNT})"
+        " per prompt studio project has been reached."
+    )