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 index d1ea56510..111a891d7 100644 --- 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 index 3e7818751..a2a4165a1 100644 --- 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 index 727f4771c..346df91a0 100644 --- 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")