feat/multi llm enhancement

backend/adapter_processor/serializers.py

@@ -120,6 +120,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"] = adapter_metadata["model"]
 
         if instance.is_friction_less:
             rep["created_by_email"] = "Unstract"

backend/prompt_studio/prompt_profile_manager/constants.py

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

backend/prompt_studio/prompt_studio_core/constants.py

@@ -85,6 +85,7 @@ class ToolStudioPromptKeys:
     NOTES = "NOTES"
     OUTPUT = "output"
     SEQUENCE_NUMBER = "sequence_number"
+    PROFILE_MANAGER_ID = "profile_manager"
 
 
 class FileViewTypes:

backend/prompt_studio/prompt_studio_core/exceptions.py

@@ -58,3 +58,10 @@ class PermissionError(APIException):
 class EmptyPromptError(APIException):
     status_code = 422
     default_detail = "Prompt(s) cannot be empty"
+
+
+class MaxProfilesReachedError(APIException):
+    status_code = 403
+    default_detail = (
+        "Maximum number of profiles (max 4) per prompt studio project has been reached."
+    )

backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import time
 import uuid
 from pathlib import Path
 from typing import Any, Optional
@@ -27,6 +28,12 @@
 )
 from prompt_studio.prompt_studio_core.models import CustomTool
 from prompt_studio.prompt_studio_core.prompt_ide_base_tool import PromptIdeBaseTool
+from prompt_studio.prompt_studio_core.redis_utils import (
+    get_indexed_document_id,
+    is_document_indexing,
+    mark_document_indexed,
+    set_document_indexing,
+)
 from prompt_studio.prompt_studio_document_manager.models import DocumentManager
 from prompt_studio.prompt_studio_index_manager.prompt_studio_index_helper import (  # noqa: E501
     PromptStudioIndexHelper,
@@ -364,6 +371,7 @@ def prompt_responder(
         document_id: str,
         id: Optional[str] = None,
         run_id: str = None,
+        profile_manager_id: Optional[str] = None,
     ) -> Any:
         """Execute chain/single run of the prompts. Makes a call to prompt
         service and returns the dict of response.
@@ -374,6 +382,7 @@ def prompt_responder(
             user_id (str): User's ID
             document_id (str): UUID of the document uploaded
             id (Optional[str]): ID of the prompt
+            profile_manager_id (Optional[str]): UUID of the profile manager
 
         Raises:
             AnswerFetchError: Error from prompt-service
@@ -442,6 +451,7 @@ def prompt_responder(
                     org_id=org_id,
                     document_id=document_id,
                     run_id=run_id,
+                    profile_manager_id=profile_manager_id,
                 )
 
                 OutputManagerHelper.handle_prompt_output_update(
@@ -450,6 +460,8 @@ def prompt_responder(
                     outputs=response["output"],
                     document_id=document_id,
                     is_single_pass_extract=False,
+                    profile_manager_id=profile_manager_id,
+                    tool=tool,
                 )
             # TODO: Review if this catch-all is required
             except Exception as e:
@@ -562,6 +574,7 @@ def _fetch_response(
         org_id: str,
         document_id: str,
         run_id: str,
+        profile_manager_id: Optional[str] = None,
     ) -> Any:
         """Utility function to invoke prompt service. Used internally.
 
@@ -572,6 +585,8 @@ def _fetch_response(
             prompt (ToolStudioPrompt): ToolStudioPrompt instance to fetch response
             org_id (str): UUID of the organization
             document_id (str): UUID of the document
+            profile_manager_id (Optional[str]): UUID of the profile manager
+
 
         Raises:
             DefaultProfileError: If no default profile is selected
@@ -580,6 +595,20 @@ def _fetch_response(
         Returns:
             Any: Output from LLM
         """
+
+        # Fetch the ProfileManager instance using the profile_manager_id if provided
+        if profile_manager_id:
+            try:
+                profile_manager = ProfileManager.objects.get(
+                    profile_id=profile_manager_id
+                )
+            except ProfileManager.DoesNotExist:
+                raise DefaultProfileError(
+                    f"ProfileManager with ID {profile_manager_id} does not exist."
+                )
+        else:
+            profile_manager = prompt.profile_manager
+
         monitor_llm_instance: Optional[AdapterInstance] = tool.monitor_llm
         monitor_llm: Optional[str] = None
         challenge_llm_instance: Optional[AdapterInstance] = tool.challenge_llm
@@ -600,21 +629,20 @@ def _fetch_response(
             challenge_llm = str(default_profile.llm.id)
 
         # Need to check the user who created profile manager
-        PromptStudioHelper.validate_adapter_status(prompt.profile_manager)
+        PromptStudioHelper.validate_adapter_status(profile_manager)
         # Need to check the user who created profile manager
         # has access to adapters
-        PromptStudioHelper.validate_profile_manager_owner_access(prompt.profile_manager)
+        PromptStudioHelper.validate_profile_manager_owner_access(profile_manager)
         # Not checking reindex here as there might be
         # change in Profile Manager
-        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)
-        prompt_profile_manager: ProfileManager = prompt.profile_manager
-        if not prompt_profile_manager:
+        vector_db = str(profile_manager.vector_store.id)
+        embedding_model = str(profile_manager.embedding_model.id)
+        llm = str(profile_manager.llm.id)
+        x2text = str(profile_manager.x2text.id)
+        if not profile_manager:
             raise DefaultProfileError()
         PromptStudioHelper.dynamic_indexer(
-            profile_manager=prompt_profile_manager,
+            profile_manager=profile_manager,
             file_path=doc_path,
             tool_id=str(tool.tool_id),
             org_id=org_id,
@@ -639,16 +667,16 @@ def _fetch_response(
 
         output[TSPKeys.PROMPT] = prompt.prompt
         output[TSPKeys.ACTIVE] = prompt.active
-        output[TSPKeys.CHUNK_SIZE] = prompt.profile_manager.chunk_size
+        output[TSPKeys.CHUNK_SIZE] = profile_manager.chunk_size
         output[TSPKeys.VECTOR_DB] = vector_db
         output[TSPKeys.EMBEDDING] = embedding_model
-        output[TSPKeys.CHUNK_OVERLAP] = prompt.profile_manager.chunk_overlap
+        output[TSPKeys.CHUNK_OVERLAP] = profile_manager.chunk_overlap
         output[TSPKeys.LLM] = llm
         output[TSPKeys.TYPE] = prompt.enforce_type
         output[TSPKeys.NAME] = prompt.prompt_key
-        output[TSPKeys.RETRIEVAL_STRATEGY] = prompt.profile_manager.retrieval_strategy
-        output[TSPKeys.SIMILARITY_TOP_K] = prompt.profile_manager.similarity_top_k
-        output[TSPKeys.SECTION] = prompt.profile_manager.section
+        output[TSPKeys.RETRIEVAL_STRATEGY] = profile_manager.retrieval_strategy
+        output[TSPKeys.SIMILARITY_TOP_K] = profile_manager.similarity_top_k
+        output[TSPKeys.SECTION] = profile_manager.section
         output[TSPKeys.X2TEXT_ADAPTER] = x2text
         # Eval settings for the prompt
         output[TSPKeys.EVAL_SETTINGS] = {}
@@ -750,9 +778,43 @@ def dynamic_indexer(
             profile_manager.chunk_size = 0
 
         try:
+
             usage_kwargs = {"run_id": run_id}
             util = PromptIdeBaseTool(log_level=LogLevel.INFO, org_id=org_id)
             tool_index = Index(tool=util)
+            doc_id_key = tool_index.generate_file_id(
+                tool_id=tool_id,
+                vector_db=vector_db,
+                embedding=embedding_model,
+                x2text=x2text_adapter,
+                chunk_size=str(profile_manager.chunk_size),
+                chunk_overlap=str(profile_manager.chunk_overlap),
+                file_path=file_path,
+                file_hash=None,
+            )
+            indexed_doc_id = get_indexed_document_id(doc_id_key)
+            if indexed_doc_id:
+                return indexed_doc_id
+
+            # Polling if document is already being indexed
+            if is_document_indexing(doc_id_key):
+                max_wait_time = 1800  # 30 minutes
+                wait_time = 0
+                polling_interval = 5  # Poll every 5 seconds
+                while is_document_indexing(doc_id_key):
+                    if wait_time >= max_wait_time:
+                        raise IndexingAPIError(
+                            "Indexing timed out. Please try again later."
+                        )
+                    time.sleep(polling_interval)
+                    wait_time += polling_interval
+
+                # After waiting, check if the document is indexed
+                indexed_doc_id = get_indexed_document_id(doc_id_key)
+                if indexed_doc_id:
+                    return indexed_doc_id
+            # Set the document as being indexed
+            set_document_indexing(doc_id_key)
             doc_id: str = tool_index.index(
                 tool_id=tool_id,
                 embedding_instance_id=embedding_model,
@@ -772,6 +834,7 @@ def dynamic_indexer(
                 profile_manager=profile_manager,
                 doc_id=doc_id,
             )
+            mark_document_indexed(doc_id_key, doc_id)
             return doc_id
         except (IndexingError, IndexingAPIError, SdkError) as e:
             doc_name = os.path.split(file_path)[1]

backend/prompt_studio/prompt_studio_core/redis_utils.py

@@ -0,0 +1,24 @@
+from utils.cache_service import CacheService
+
+
+def set_document_indexing(doc_id_key, ttl=1800):
+    CacheService.set_key(f"document_indexing:{doc_id_key}", "started", expire=ttl)
+
+
+def is_document_indexing(doc_id_key):
+    return CacheService.get_key(f"document_indexing:{doc_id_key}") == b"started"
+
+
+def mark_document_indexed(doc_id_key, doc_id):
+    CacheService.set_key(f"document_indexing:{doc_id_key}", doc_id, expire=3600)
+
+
+def get_indexed_document_id(doc_id_key):
+    result = CacheService.get_key(f"document_indexing:{doc_id_key}")
+    if result and result != b"started":
+        return result
+    return None
+
+
+def remove_document_indexing(doc_id_key):
+    CacheService.delete_a_key(f"document_indexing:{doc_id_key}")

backend/prompt_studio/prompt_studio_core/views.py

@@ -10,7 +10,10 @@
 from file_management.file_management_helper import FileManagerHelper
 from permissions.permission import IsOwner, IsOwnerOrSharedUser
 from prompt_studio.processor_loader import ProcessorConfig, load_plugins
-from prompt_studio.prompt_profile_manager.constants import ProfileManagerErrors
+from prompt_studio.prompt_profile_manager.constants import (
+    ProfileManagerErrors,
+    ProfileManagerKeys,
+)
 from prompt_studio.prompt_profile_manager.models import ProfileManager
 from prompt_studio.prompt_profile_manager.serializers import ProfileManagerSerializer
 from prompt_studio.prompt_studio.constants import ToolStudioPromptErrors
@@ -23,13 +26,16 @@
 )
 from prompt_studio.prompt_studio_core.exceptions import (
     IndexingAPIError,
+    MaxProfilesReachedError,
     ToolDeleteError,
 )
 from prompt_studio.prompt_studio_core.prompt_studio_helper import PromptStudioHelper
+from prompt_studio.prompt_studio_core.redis_utils import remove_document_indexing
 from prompt_studio.prompt_studio_document_manager.models import DocumentManager
 from prompt_studio.prompt_studio_document_manager.prompt_studio_document_helper import (  # noqa: E501
     PromptStudioDocumentHelper,
 )
+from prompt_studio.prompt_studio_index_manager.models import IndexManager
 from prompt_studio.prompt_studio_registry.prompt_studio_registry_helper import (
     PromptStudioRegistryHelper,
 )
@@ -264,6 +270,9 @@ def fetch_response(self, request: HttpRequest, pk: Any = None) -> Response:
         document_id: str = request.data.get(ToolStudioPromptKeys.DOCUMENT_ID)
         id: str = request.data.get(ToolStudioPromptKeys.ID)
         run_id: str = request.data.get(ToolStudioPromptKeys.RUN_ID)
+        profile_manager: str = request.data.get(
+            ToolStudioPromptKeys.PROFILE_MANAGER_ID, None
+        )
         if not run_id:
             # Generate a run_id
             run_id = CommonUtils.generate_uuid()
@@ -275,6 +284,7 @@ def fetch_response(self, request: HttpRequest, pk: Any = None) -> Response:
             user_id=custom_tool.created_by.user_id,
             document_id=document_id,
             run_id=run_id,
+            profile_manager_id=profile_manager,
         )
         return Response(response, status=status.HTTP_200_OK)
 
@@ -339,6 +349,17 @@ def create_profile_manager(self, request: HttpRequest, pk: Any = None) -> Respon
         serializer = ProfileManagerSerializer(data=request.data, context=context)
 
         serializer.is_valid(raise_exception=True)
+
+        # Check for the maximum number of profiles constraint
+        prompt_studio_tool = serializer.validated_data[
+            ProfileManagerKeys.PROMPT_STUDIO_TOOL
+        ]
+        profile_count = ProfileManager.objects.filter(
+            prompt_studio_tool=prompt_studio_tool
+        ).count()
+
+        if profile_count >= 4:
+            raise MaxProfilesReachedError()
         try:
             self.perform_create(serializer)
         except IntegrityError:
@@ -457,6 +478,11 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response:
         path = file_path
         file_system = LocalStorageFS(settings={"path": path})
         try:
+            # Delete indexed flags in redis
+            index_managers = IndexManager.objects.filter(document_manager=document_id)
+            for index_manager in index_managers:
+                raw_index_id = index_manager.raw_index_id
+                remove_document_indexing(raw_index_id)
             # Delete the document record
             document.delete()
             # Delete the files
@@ -466,6 +492,8 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response:
             FileManagerHelper.delete_related_files(
                 file_system, path, file_name, directories
             )
+            # Delete indexed flags in redis
+
             return Response(
                 {"data": "File deleted succesfully."},
                 status=status.HTTP_200_OK,

backend/prompt_studio/prompt_studio_index_manager/urls.py

@@ -10,6 +10,9 @@
         "get": "retrieve",
     }
 )
+prompt_studio_index_data = IndexManagerView.as_view(
+    {"get": "get_indexed_data_for_profile"}
+)
 
 urlpatterns = format_suffix_patterns(
     [
@@ -18,5 +21,10 @@
             prompt_studio_index_list,
             name="prompt-studio-documents-list",
         ),
+        path(
+            "indexed-result/",
+            prompt_studio_index_data,
+            name="prompt-studio-indexed-list",
+        ),
     ]
 )