Merge pull request #1 from Zipstack/fix/prompt-studio-index-and-run

jaseemjaskp · web-flow · commit c64479ca665a · 2024-02-26T12:55:28.000+05:30
fix: Prompt studio index and run issues
diff --git a/backend/prompt_studio/prompt_studio_core/prompt_ide_base_tool.py b/backend/prompt_studio/prompt_studio_core/prompt_ide_base_tool.py
@@ -1,6 +1,8 @@
 import os
 
-from account.models import PlatformKey
+from platform_settings.platform_auth_service import (
+    PlatformAuthenticationService,
+)
 from prompt_studio.prompt_studio_core.constants import LogLevel, ToolStudioKeys
 from unstract.sdk.tool.stream import StreamMixin
 
@@ -32,7 +34,11 @@ def get_env_or_die(self, env_key: str) -> str:
         """
         # HACK: Adding platform key for multitenancy
         if env_key == ToolStudioKeys.PLATFORM_SERVICE_API_KEY:
-            platform_key = PlatformKey.objects.get(organization=self.org_id)
+            platform_key = (
+                PlatformAuthenticationService.get_active_platform_key(
+                    self.org_id
+                )
+            )
             key: str = str(platform_key.key)
             return key
         else:
diff --git a/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py
@@ -4,8 +4,6 @@
 from pathlib import Path
 from typing import Any, Optional
 
-from account.models import Organization
-from account.organization import OrganizationService
 from django.conf import settings
 from file_management.file_management_helper import FileManagerHelper
 from prompt_studio.prompt_profile_manager.models import ProfileManager
@@ -218,14 +216,8 @@ def prompt_responder(
                 ),
             )
             logger.info(f"Invoking prompt service for prompt id {id}")
-            org: Optional[
-                Organization
-            ] = OrganizationService.get_organization_by_org_id(org_id=org_id)
             response = PromptStudioHelper._fetch_response(
-                path=file_path,
-                tool=tool,
-                prompts=prompts,
-                org_id=org.id,  # type:ignore
+                path=file_path, tool=tool, prompts=prompts, org_id=org_id
             )
             stream_log.publish(
                 tool.tool_id,
diff --git a/prompt-service/src/unstract/prompt_service/prompt_ide_base_tool.py b/prompt-service/src/unstract/prompt_service/prompt_ide_base_tool.py
@@ -30,6 +30,7 @@ def get_env_or_die(self, env_key: str) -> str:
         Returns:
             str: Value of the env
         """
+        # HACK: Adding platform key for multitenancy
         if env_key == PromptServiceContants.PLATFORM_SERVICE_API_KEY:
             if not self.platform_key:
                 current_app.logger.error(f"{env_key} is required")
