merged main

.pre-commit-config.yaml

@@ -57,7 +57,7 @@ repos:
               unstract/flags/src/unstract/flags/evaluation_.*\.py|
           )$
   - repo: https://github.com/pycqa/flake8
-    rev: 7.0.0
+    rev: 7.1.0
     hooks:
       - id: flake8
         args: [--max-line-length=88]
@@ -96,7 +96,7 @@ repos:
         types:
           - python
   - repo: https://github.com/gitleaks/gitleaks
-    rev: v8.18.3
+    rev: v8.18.4
     hooks:
       - id: gitleaks
   # - repo: https://github.com/hadolint/hadolint

backend/backend/urls.py

@@ -30,6 +30,7 @@
     path("", include("tool_instance.urls")),
     path("", include("pipeline.urls")),
     path("", include("apps.urls")),
+    path("", include("feature_flag.urls")),
     path("workflow/", include("workflow_manager.urls")),
     path("platform/", include("platform_settings.urls")),
     path("api/", include("api.urls")),

backend/feature_flag/urls.py

@@ -3,9 +3,20 @@
 This module defines the URL patterns for the feature_flags app.
 """
 
-import feature_flag.views as views
 from django.urls import path
+from feature_flag.views import FeatureFlagViewSet
+from rest_framework.urlpatterns import format_suffix_patterns
 
-urlpatterns = [
-    path("evaluate/", views.evaluate_feature_flag, name="evaluate_feature_flag"),
-]
+feature_flags_list = FeatureFlagViewSet.as_view(
+    {
+        "post": "evaluate",
+        "get": "list",
+    }
+)
+
+urlpatterns = format_suffix_patterns(
+    [
+        path("evaluate/", feature_flags_list, name="evaluate_feature_flag"),
+        path("flags/", feature_flags_list, name="list_feature_flags"),
+    ]
+)

backend/feature_flag/views.py

@@ -6,52 +6,43 @@
 
 import logging
 
-from rest_framework import status
-from rest_framework.decorators import api_view
-from rest_framework.request import Request
+from rest_framework import status, viewsets
 from rest_framework.response import Response
-
-from unstract.flags.client import EvaluationClient
+from utils.request.feature_flag import check_feature_flag_status, list_all_flags
 
 logger = logging.getLogger(__name__)
 
 
-@api_view(["POST"])
-def evaluate_feature_flag(request: Request) -> Response:
-    """Function to evaluate the feature flag.
-
-    To-Do: Refactor to a class based view, use serializers (DRF).
+class FeatureFlagViewSet(viewsets.ViewSet):
+    """A simple ViewSet for evaluating feature flag."""
 
-    Args:
-        request: request object
+    def evaluate(self, request):
+        try:
+            flag_key = request.data.get("flag_key")
 
-    Returns:
-        evaluate response
-    """
-    try:
-        namespace_key = request.data.get("namespace_key")
-        flag_key = request.data.get("flag_key")
-        entity_id = request.data.get("entity_id")
-        context = request.data.get("context")
+            if not flag_key:
+                return Response(
+                    {"message": "Request parameters are missing."},
+                    status=status.HTTP_400_BAD_REQUEST,
+                )
 
-        if not namespace_key or not flag_key or not entity_id:
+            flag_enabled = check_feature_flag_status(flag_key)
+            return Response({"flag_status": flag_enabled}, status=status.HTTP_200_OK)
+        except Exception as e:
+            logger.error("No response from server: %s", e)
             return Response(
-                {"message": "Request paramteres are missing."},
-                status=status.HTTP_400_BAD_REQUEST,
+                {"message": "No response from server"},
+                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
             )
 
-        evaluation_client = EvaluationClient()
-        response = evaluation_client.boolean_evaluate_feature_flag(
-            namespace_key=namespace_key,
-            flag_key=flag_key,
-            entity_id=entity_id,
-            context=context,
-        )
-
-        return Response({"enabled": response}, status=status.HTTP_200_OK)
-    except Exception as e:
-        logger.error("No response from server: %s", e)
-        return Response(
-            {"message": "No response from server"},
-            status=status.HTTP_500_INTERNAL_SERVER_ERROR,
-        )
+    def list(self, request):
+        try:
+            namespace_key = request.query_params.get("namespace", "default")
+            feature_flags = list_all_flags(namespace_key)
+            return Response({"feature_flags": feature_flags}, status=status.HTTP_200_OK)
+        except Exception as e:
+            logger.error("No response from server: %s", e)
+            return Response(
+                {"message": "No response from server"},
+                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            )

backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py

@@ -302,7 +302,7 @@ def index_document(
         """
         tool: CustomTool = CustomTool.objects.get(pk=tool_id)
         if is_summary:
-            profile_manager = ProfileManager.objects.get(
+            profile_manager: ProfileManager = ProfileManager.objects.get(
                 prompt_studio_tool=tool, is_summarize_llm=True
             )
             default_profile = profile_manager

backend/prompt_studio/prompt_studio_core/urls.py

@@ -1,4 +1,6 @@
+from django.db import transaction
 from django.urls import path
+from django.utils.decorators import method_decorator
 from rest_framework.urlpatterns import format_suffix_patterns
 
 from .views import PromptStudioCoreView
@@ -77,7 +79,9 @@
         ),
         path(
             "prompt-studio/index-document/<uuid:pk>",
-            prompt_studio_prompt_index,
+            method_decorator(transaction.non_atomic_requests)(
+                prompt_studio_prompt_index
+            ),
             name="prompt-studio-prompt-index",
         ),
         path(

backend/prompt_studio/prompt_studio_core/views.py

@@ -189,7 +189,7 @@ def make_profile_default(self, request: HttpRequest, pk: Any = None) -> Response
             data={"default_profile": profile_manager.profile_id},
         )
 
-    @action(detail=True, methods=["get"])
+    @action(detail=True, methods=["post"])
     def index_document(self, request: HttpRequest, pk: Any = None) -> Response:
         """API Entry point method to index input file.
 

backend/prompt_studio/prompt_studio_index_manager/prompt_studio_index_helper.py

@@ -1,6 +1,7 @@
 import json
 import logging
 
+from django.db import transaction
 from prompt_studio.prompt_profile_manager.models import ProfileManager
 from prompt_studio.prompt_studio_core.exceptions import IndexingAPIError
 from prompt_studio.prompt_studio_document_manager.models import DocumentManager
@@ -18,46 +19,48 @@ def handle_index_manager(
         profile_manager: ProfileManager,
         doc_id: str,
     ) -> IndexManager:
-        document: DocumentManager = DocumentManager.objects.get(pk=document_id)
+        try:
 
-        index_id = "raw_index_id"
-        if is_summary:
-            index_id = "summarize_index_id"
+            with transaction.atomic():
 
-        args: dict[str, str] = dict()
-        args["document_manager"] = document
-        args["profile_manager"] = profile_manager
+                document: DocumentManager = DocumentManager.objects.get(pk=document_id)
 
-        try:
-            # 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)
+                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
-
-        return result

backend/sample.env

@@ -89,9 +89,9 @@ PROMPT_PORT=3003
 PROMPT_STUDIO_FILE_PATH=/app/prompt-studio-data
 
 # Structure Tool
-STRUCTURE_TOOL_IMAGE_URL="docker:unstract/tool-structure:0.0.25"
+STRUCTURE_TOOL_IMAGE_URL="docker:unstract/tool-structure:0.0.26"
 STRUCTURE_TOOL_IMAGE_NAME="unstract/tool-structure"
-STRUCTURE_TOOL_IMAGE_TAG="0.0.25"
+STRUCTURE_TOOL_IMAGE_TAG="0.0.26"
 
 # Feature Flags
 EVALUATION_SERVER_IP=localhost

backend/tenant_account/migrations/0001_initial.py

@@ -14,8 +14,9 @@ class Migration(migrations.Migration):
     ]
 
     operations = [
+        # Updated the name here as the 002, 0002 step is just name change
         migrations.CreateModel(
-            name="User",
+            name="OrganizationMember",
             fields=[
                 (
                     "user_ptr",
@@ -28,6 +29,8 @@ class Migration(migrations.Migration):
                         to=settings.AUTH_USER_MODEL,
                     ),
                 ),
+                # Added column which is used in 0002 here
+                ("role", models.CharField(default="admin")),
             ],
             options={
                 "verbose_name": "user",

backend/tenant_account/migrations/0002_organizationmember_delete_user.py

@@ -1,9 +1,6 @@
 # Generated by Django 4.2.1 on 2023-08-21 11:12
 
-import django.contrib.auth.models
-import django.db.models.deletion
-from django.conf import settings
-from django.db import migrations, models
+from django.db import migrations
 
 
 class Migration(migrations.Migration):
@@ -13,33 +10,36 @@ class Migration(migrations.Migration):
     ]
 
     operations = [
-        migrations.CreateModel(
-            name="OrganizationMember",
-            fields=[
-                (
-                    "user_ptr",
-                    models.OneToOneField(
-                        auto_created=True,
-                        on_delete=django.db.models.deletion.CASCADE,
-                        parent_link=True,
-                        primary_key=True,
-                        serialize=False,
-                        to=settings.AUTH_USER_MODEL,
-                    ),
-                ),
-                ("role", models.CharField(default="admin")),
-            ],
-            options={
-                "verbose_name": "user",
-                "verbose_name_plural": "users",
-                "abstract": False,
-            },
-            bases=("account.user",),
-            managers=[
-                ("objects", django.contrib.auth.models.UserManager()),
-            ],
-        ),
-        migrations.DeleteModel(
-            name="User",
-        ),
+        # # Commenting out here as this is taken care in 0001
+        # migrations.CreateModel(
+        #     name="OrganizationMember",
+        #     fields=[
+        #         (
+        #             "user_ptr",
+        #             models.OneToOneField(
+        #                 auto_created=True,
+        #                 on_delete=django.db.models.deletion.CASCADE,
+        #                 parent_link=True,
+        #                 primary_key=True,
+        #                 serialize=False,
+        #                 to=settings.AUTH_USER_MODEL,
+        #             ),
+        #         ),
+        #         ("role", models.CharField(default="admin")),
+        #     ],
+        #     options={
+        #         "verbose_name": "user",
+        #         "verbose_name_plural": "users",
+        #         "abstract": False,
+        #     },
+        #     bases=("account.user",),
+        #     managers=[
+        #         ("objects", django.contrib.auth.models.UserManager()),
+        #     ],
+        # ),
+        # # https://www.geeksforgeeks.org/what-is-access-exclusive-lock-mode-in-postgreysql/
+        # # commenting drop table to ignore AccesExclusive Lock
+        # migrations.DeleteModel(
+        #     name="User",
+        # ),
     ]

backend/tool_instance/exceptions.py

@@ -38,4 +38,9 @@ class ToolInstantiationError(ToolInstanceBaseException):
 
 class BadRequestException(ToolInstanceBaseException):
     status_code = 400
-    default_detail = "Invalid input"
+    default_detail = "Invalid input."
+
+
+class ToolSettingValidationError(APIException):
+    status_code = 400
+    default_detail = "Error while validating tool's setting."