Feat/agentic prompt studio v2

backend/platform_settings_v2/models.py

@@ -0,0 +1,73 @@
+import uuid
+
+from adapter_processor_v2.models import AdapterInstance
+from django.db import models
+from utils.models.base_model import BaseModel
+from utils.models.organization_mixin import (
+    DefaultOrganizationManagerMixin,
+    DefaultOrganizationMixin,
+)
+
+
+class PlatformSettingsModelManager(DefaultOrganizationManagerMixin, models.Manager):
+    """Manager for PlatformSettings model."""
+
+    pass
+
+
+class PlatformSettings(DefaultOrganizationMixin, BaseModel):
+    """Platform-level settings for an organization.
+
+    This model stores organization-wide settings including the system LLM
+    adapter that will be used for platform operations like vibe extractor
+    prompt generation.
+    """
+
+    id = models.UUIDField(
+        primary_key=True,
+        default=uuid.uuid4,
+        editable=False,
+        db_comment="Unique identifier for the platform settings",
+    )
+
+    # System LLM for platform operations (e.g., vibe extractor, prompt generation)
+    system_llm_adapter = models.ForeignKey(
+        AdapterInstance,
+        on_delete=models.SET_NULL,
+        null=True,
+        blank=True,
+        related_name="platform_system_llm",
+        db_comment="System LLM adapter used for platform-level AI operations like prompt generation",
+    )
+
+    objects = PlatformSettingsModelManager()
+
+    class Meta:
+        verbose_name = "Platform Setting"
+        verbose_name_plural = "Platform Settings"
+        db_table = "platform_settings"
+        constraints = [
+            models.UniqueConstraint(
+                fields=["organization"],
+                name="unique_organization_platform_settings",
+            ),
+        ]
+
+    def __str__(self) -> str:
+        return f"PlatformSettings({self.organization})"
+
+    @classmethod
+    def get_for_organization(cls, organization):
+        """Get or create platform settings for an organization.
+
+        Args:
+            organization: Organization instance
+
+        Returns:
+            PlatformSettings instance
+        """
+        settings, created = cls.objects.get_or_create(
+            organization=organization,
+            defaults={},
+        )
+        return settings

backend/platform_settings_v2/serializers.py

@@ -1,7 +1,16 @@
 from account_v2.models import PlatformKey
+from adapter_processor_v2.models import AdapterInstance
 from rest_framework import serializers
+from rest_framework.exceptions import ValidationError
 
 from backend.serializers import AuditSerializer
+from platform_settings_v2.models import PlatformSettings
+from unstract.flags.feature_flag import check_feature_flag_status
+
+if check_feature_flag_status("sdk1"):
+    from unstract.sdk1.adapters.enums import AdapterTypes
+else:
+    from unstract.sdk.adapters.enums import AdapterTypes
 
 
 class PlatformKeySerializer(AuditSerializer):
@@ -22,3 +31,50 @@ class PlatformKeyIDSerializer(serializers.Serializer):
     key_name = serializers.CharField()
     key = serializers.CharField()
     is_active = serializers.BooleanField()
+
+
+class PlatformSettingsSerializer(AuditSerializer):
+    """Serializer for PlatformSettings model."""
+
+    system_llm_adapter = serializers.PrimaryKeyRelatedField(
+        queryset=AdapterInstance.objects.all(),
+        required=False,
+        allow_null=True,
+    )
+
+    class Meta:
+        model = PlatformSettings
+        fields = [
+            "id",
+            "organization",
+            "system_llm_adapter",
+            "created_at",
+            "modified_at",
+        ]
+        read_only_fields = ["id", "organization", "created_at", "modified_at"]
+
+    def validate_system_llm_adapter(self, value):
+        """Validate that the adapter type is LLM and is accessible to the user."""
+        if value is None:
+            return value
+
+        # Check if user has access to this adapter
+        request = self.context.get("request")
+        if request and hasattr(request, "user"):
+            try:
+                adapter = AdapterInstance.objects.for_user(request.user).get(id=value.id)
+                # Validate that the adapter type is LLM
+                if adapter.adapter_type != AdapterTypes.LLM.value:
+                    raise ValidationError("Only LLM adapters are allowed for system LLM")
+
+                # Validate that adapter is usable and active
+                if not adapter.is_usable:
+                    raise ValidationError("Selected LLM adapter is not usable")
+
+                if not adapter.is_active:
+                    raise ValidationError("Selected LLM adapter is not active")
+
+            except AdapterInstance.DoesNotExist:
+                raise ValidationError("Selected LLM adapter not found or not accessible")
+
+        return value

backend/platform_settings_v2/urls.py

@@ -1,7 +1,7 @@
 from django.urls import path
 from rest_framework.urlpatterns import format_suffix_patterns
 
-from .views import PlatformKeyViewSet
+from .views import PlatformKeyViewSet, PlatformSettingsViewSet
 
 platform_key_list = PlatformKeyViewSet.as_view(
     {"post": "create", "put": "refresh", "get": "list"}
@@ -10,6 +10,12 @@
     {"put": "toggle_platform_key", "delete": "destroy"}
 )
 
+platform_settings_view = PlatformSettingsViewSet.as_view(
+    {"get": "list", "put": "update", "patch": "update"}
+)
+
+platform_settings_system_llm = PlatformSettingsViewSet.as_view({"get": "system_llm"})
+
 urlpatterns = format_suffix_patterns(
     [
         path(
@@ -22,5 +28,15 @@
             platform_key_update,
             name="update_platform_key",
         ),
+        path(
+            "settings/",
+            platform_settings_view,
+            name="platform_settings",
+        ),
+        path(
+            "settings/system-llm/",
+            platform_settings_system_llm,
+            name="platform_settings_system_llm",
+        ),
     ]
 )

backend/platform_settings_v2/views.py

@@ -5,17 +5,20 @@
 
 from account_v2.models import Organization, PlatformKey
 from rest_framework import status, viewsets
+from rest_framework.decorators import action
 from rest_framework.request import Request
 from rest_framework.response import Response
 from utils.user_context import UserContext
 
 from platform_settings_v2.constants import PlatformServiceConstants
+from platform_settings_v2.models import PlatformSettings
 from platform_settings_v2.platform_auth_helper import PlatformAuthHelper
 from platform_settings_v2.platform_auth_service import PlatformAuthenticationService
 from platform_settings_v2.serializers import (
     PlatformKeyGenerateSerializer,
     PlatformKeyIDSerializer,
     PlatformKeySerializer,
+    PlatformSettingsSerializer,
 )
 
 logger = logging.getLogger(__name__)
@@ -123,3 +126,79 @@ def create(self, request: Request) -> Response:
             status=status.HTTP_201_CREATED,
             data=serialized_data,
         )
+
+
+class PlatformSettingsViewSet(viewsets.ModelViewSet):
+    """ViewSet for managing platform settings."""
+
+    serializer_class = PlatformSettingsSerializer
+
+    def get_queryset(self):
+        """Get platform settings for the user's organization."""
+        organization = UserContext.get_organization()
+        return PlatformSettings.objects.filter(organization=organization)
+
+    def get_object(self):
+        """Get or create platform settings for the user's organization."""
+        organization = UserContext.get_organization()
+        settings, created = PlatformSettings.objects.get_or_create(
+            organization=organization
+        )
+        return settings
+
+    def list(
+        self, request: Request, *args: tuple[Any], **kwargs: dict[str, Any]
+    ) -> Response:
+        """List platform settings for the organization."""
+        settings = self.get_object()
+        serializer = self.get_serializer(settings)
+        return Response(serializer.data)
+
+    def retrieve(
+        self, request: Request, *args: tuple[Any], **kwargs: dict[str, Any]
+    ) -> Response:
+        """Retrieve platform settings for the organization."""
+        settings = self.get_object()
+        serializer = self.get_serializer(settings)
+        return Response(serializer.data)
+
+    def update(
+        self, request: Request, *args: tuple[Any], **kwargs: dict[str, Any]
+    ) -> Response:
+        """Update platform settings."""
+        settings = self.get_object()
+        serializer = self.get_serializer(
+            settings, data=request.data, partial=True, context={"request": request}
+        )
+        serializer.is_valid(raise_exception=True)
+        serializer.save()
+        return Response(serializer.data)
+
+    @action(detail=False, methods=["get"])
+    def system_llm(self, request: Request) -> Response:
+        """Get the configured system LLM adapter for the organization.
+
+        Returns:
+            Response with system LLM adapter details or null if not configured
+        """
+        settings = self.get_object()
+        if settings.system_llm_adapter:
+            from adapter_processor_v2.serializers import AdapterInstanceSerializer
+
+            adapter_serializer = AdapterInstanceSerializer(settings.system_llm_adapter)
+            return Response(
+                {
+                    "system_llm_adapter": adapter_serializer.data,
+                    "is_configured": True,
+                },
+                status=status.HTTP_200_OK,
+            )
+        else:
+            return Response(
+                {
+                    "system_llm_adapter": None,
+                    "is_configured": False,
+                    "message": "No system LLM adapter configured for this organization",
+                },
+                status=status.HTTP_200_OK,
+            )

backend/prompt_studio/prompt_studio_vibe_extractor_v2/admin.py

@@ -0,0 +1,30 @@
+from django.contrib import admin
+
+from prompt_studio.prompt_studio_vibe_extractor_v2.models import (
+    VibeExtractorProject,
+)
+
+
+@admin.register(VibeExtractorProject)
+class VibeExtractorProjectAdmin(admin.ModelAdmin):
+    """Admin interface for VibeExtractorProject."""
+
+    list_display = [
+        "project_id",
+        "document_type",
+        "status",
+        "tool_id",
+        "created_at",
+        "modified_at",
+    ]
+    list_filter = ["status", "created_at"]
+    search_fields = ["document_type", "project_id"]
+    readonly_fields = [
+        "project_id",
+        "generation_output_path",
+        "generation_progress",
+        "created_by",
+        "modified_by",
+        "created_at",
+        "modified_at",
+    ]

backend/prompt_studio/prompt_studio_vibe_extractor_v2/apps.py

@@ -0,0 +1,7 @@
+from django.apps import AppConfig
+
+
+class PromptStudioVibeExtractorV2Config(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "prompt_studio.prompt_studio_vibe_extractor_v2"
+    verbose_name = "Prompt Studio Vibe Extractor V2"

backend/prompt_studio/prompt_studio_vibe_extractor_v2/constants.py

@@ -0,0 +1,44 @@
+"""Constants for Vibe Extractor."""
+
+
+class VibeExtractorKeys:
+    """Keys for Vibe Extractor API requests and responses."""
+
+    PROJECT_ID = "project_id"
+    DOCUMENT_TYPE = "document_type"
+    STATUS = "status"
+    GENERATION_OUTPUT_PATH = "generation_output_path"
+    ERROR_MESSAGE = "error_message"
+    GENERATION_PROGRESS = "generation_progress"
+    TOOL_ID = "tool_id"
+
+
+class VibeExtractorFileNames:
+    """File names for generated files."""
+
+    METADATA_YAML = "metadata.yaml"
+    EXTRACTION_YAML = "extraction.yaml"
+    PAGE_EXTRACTION_SYSTEM_MD = "page-extraction-system.md"
+    PAGE_EXTRACTION_USER_MD = "page-extraction-user.md"
+    SCALARS_EXTRACTION_SYSTEM_MD = "extraction-scalars-system.md"
+    SCALARS_EXTRACTION_USER_MD = "extraction-scalars-user.md"
+    TABLES_EXTRACTION_SYSTEM_MD = "extraction-table-system.md"
+    TABLES_EXTRACTION_USER_MD = "extraction-table-user.md"
+
+
+class VibeExtractorPaths:
+    """Path constants for Vibe Extractor."""
+
+    PROMPTS_DIR = "prompts"
+    STAGING_DIR = "staging"
+    REFERENCE_DIR = "reference"
+
+
+class GenerationSteps:
+    """Steps in the generation process."""
+
+    METADATA = "metadata"
+    EXTRACTION_FIELDS = "extraction_fields"
+    PAGE_EXTRACTION_PROMPTS = "page_extraction_prompts"
+    SCALARS_EXTRACTION_PROMPTS = "scalars_extraction_prompts"
+    TABLES_EXTRACTION_PROMPTS = "tables_extraction_prompts"

backend/prompt_studio/prompt_studio_vibe_extractor_v2/exceptions.py

@@ -0,0 +1,31 @@
+"""Exceptions for Vibe Extractor."""
+
+
+class VibeExtractorError(Exception):
+    """Base exception for Vibe Extractor errors."""
+
+    pass
+
+
+class ProjectNotFoundError(VibeExtractorError):
+    """Raised when a project is not found."""
+
+    pass
+
+
+class GenerationError(VibeExtractorError):
+    """Raised when generation fails."""
+
+    pass
+
+
+class FileReadError(VibeExtractorError):
+    """Raised when reading a generated file fails."""
+
+    pass
+
+
+class InvalidDocumentTypeError(VibeExtractorError):
+    """Raised when document type is invalid."""
+
+    pass