From 3386976d0bb8595d654a24b66bbcb1dada6a8bb0 Mon Sep 17 00:00:00 2001 From: harini-venkataraman Date: Tue, 17 Dec 2024 15:50:03 +0530 Subject: [PATCH 1/8] UN-1920 Fix:Dynamic passing of File storage init --- .../prompt_studio_core_v2/constants.py | 13 ++++++ .../prompt_studio_helper.py | 7 ++- .../src/unstract/prompt_service/constants.py | 16 +++++++ .../unstract/prompt_service/env_manager.py | 18 +++++++- .../src/unstract/prompt_service/helper.py | 43 +++++++++++++++++-- .../src/unstract/prompt_service/main.py | 3 ++ .../prompt_service_file_helper.py | 31 +++++++++++++ tools/structure/src/constants.py | 2 + tools/structure/src/main.py | 1 + 9 files changed, 129 insertions(+), 5 deletions(-) create mode 100644 prompt-service/src/unstract/prompt_service/prompt_service_file_helper.py diff --git a/backend/prompt_studio/prompt_studio_core_v2/constants.py b/backend/prompt_studio/prompt_studio_core_v2/constants.py index cb335b90f..2c6a80ac6 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/constants.py +++ b/backend/prompt_studio/prompt_studio_core_v2/constants.py @@ -96,6 +96,7 @@ class ToolStudioPromptKeys: RECORD = "record" FILE_PATH = "file_path" ENABLE_HIGHLIGHT = "enable_highlight" + EXECUTION_SOURCE = "execution_source" class FileViewTypes: @@ -132,3 +133,15 @@ class DefaultPrompts: "Do not include any explanation in the reply. " "Only include the extracted information in the reply." ) + + +class ExecutionSource(Enum): + """Enum to indicate the source of invocation. + Any new sources can be added to this enum. + This is to indicate the prompt service. + + Args: + Enum (_type_): ide/tool + """ + + IDE = "ide" diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 6bcacf340..7985173a1 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -19,7 +19,11 @@ from prompt_studio.prompt_profile_manager_v2.profile_manager_helper import ( ProfileManagerHelper, ) -from prompt_studio.prompt_studio_core_v2.constants import IndexingStatus, LogLevels +from prompt_studio.prompt_studio_core_v2.constants import ( + ExecutionSource, + IndexingStatus, + LogLevels, +) from prompt_studio.prompt_studio_core_v2.constants import ( ToolStudioPromptKeys as TSPKeys, ) @@ -1176,6 +1180,7 @@ def _fetch_single_pass_response( TSPKeys.FILE_HASH: file_hash, TSPKeys.FILE_NAME: doc_name, Common.LOG_EVENTS_ID: StateStore.get(Common.LOG_EVENTS_ID), + TSPKeys.EXECUTION_SOURCE: ExecutionSource.IDE.value, } util = PromptIdeBaseTool(log_level=LogLevel.INFO, org_id=org_id) diff --git a/prompt-service/src/unstract/prompt_service/constants.py b/prompt-service/src/unstract/prompt_service/constants.py index 96807e4d8..0d4dc4521 100644 --- a/prompt-service/src/unstract/prompt_service/constants.py +++ b/prompt-service/src/unstract/prompt_service/constants.py @@ -72,6 +72,7 @@ class PromptServiceContants: FILE_PATH = "file_path" HIGHLIGHT_DATA = "highlight_data" CONFIDENCE_DATA = "confidence_data" + EXECUTION_SOURCE = "execution_source" class RunLevel(Enum): @@ -100,3 +101,18 @@ class DBTableV2: PROMPT_STUDIO_REGISTRY = "prompt_studio_registry" PLATFORM_KEY = "platform_key" TOKEN_USAGE = "usage" + + +class FileStorageKeys: + FILE_STORAGE_PROVIDER = "FILE_STORAGE_PROVIDER" + FILE_STORAGE_CREDENTIALS = "FILE_STORAGE_CREDENTIALS" + + +class FileStorageType(Enum): + PERMANENT = "permanent" + TEMPORARY = "temporary" + + +class ExecutionSource(Enum): + IDE = "ide" + TOOL = "tool" diff --git a/prompt-service/src/unstract/prompt_service/env_manager.py b/prompt-service/src/unstract/prompt_service/env_manager.py index c51f5cf80..745b9083d 100644 --- a/prompt-service/src/unstract/prompt_service/env_manager.py +++ b/prompt-service/src/unstract/prompt_service/env_manager.py @@ -1,5 +1,8 @@ +import json import os -from typing import Optional +from typing import Any, Optional + +from unstract.prompt_service.constants import FileStorageKeys class EnvLoader: @@ -9,3 +12,16 @@ def get_env_or_die(env_key: str, default: Optional[str] = None) -> str: if env_value is None or env_value == "": raise ValueError(f"Env variable {env_key} is required") return env_value + + @staticmethod + def load_provider_credentials() -> dict[str, Any]: + cred_env_data: str = EnvLoader.get_env_or_die( + env_key=FileStorageKeys.FILE_STORAGE_CREDENTIALS + ) + credentials = json.loads(cred_env_data) + provider_data: dict[str, Any] = {} + provider_data[FileStorageKeys.FILE_STORAGE_PROVIDER] = credentials["provider"] + provider_data[FileStorageKeys.FILE_STORAGE_CREDENTIALS] = credentials[ + "credentials" + ] + return provider_data diff --git a/prompt-service/src/unstract/prompt_service/helper.py b/prompt-service/src/unstract/prompt_service/helper.py index 62a24f2d1..b64522542 100644 --- a/prompt-service/src/unstract/prompt_service/helper.py +++ b/prompt-service/src/unstract/prompt_service/helper.py @@ -7,7 +7,12 @@ from dotenv import load_dotenv from flask import Flask, current_app from unstract.prompt_service.config import db -from unstract.prompt_service.constants import DBTableV2 +from unstract.prompt_service.constants import ( + DBTableV2, + ExecutionSource, + FeatureFlag, + FileStorageType, +) from unstract.prompt_service.constants import PromptServiceContants as PSKeys from unstract.prompt_service.db_utils import DBUtils from unstract.prompt_service.env_manager import EnvLoader @@ -16,6 +21,14 @@ from unstract.sdk.exceptions import SdkError from unstract.sdk.llm import LLM +from unstract.flags.src.unstract.flags.feature_flag import check_feature_flag_status + +if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE): + from unstract.prompt_service.prompt_service_file_helper import ( + PromptServiceFileHelper, + ) + from unstract.sdk.file_storage import FileStorage + load_dotenv() # Global variable to store plugins @@ -278,6 +291,7 @@ def run_completion( prompt_type: Optional[str] = PSKeys.TEXT, enable_highlight: bool = False, file_path: str = "", + execution_source: Optional[str] = None, ) -> str: logger: Logger = current_app.logger try: @@ -286,8 +300,17 @@ def run_completion( ) highlight_data = None if highlight_data_plugin and enable_highlight: + fs_instance: FileStorage + if execution_source == ExecutionSource.IDE.value: + fs_instance = PromptServiceFileHelper.initialize_file_storage( + type=FileStorageType.PERMANENT + ) + if execution_source == ExecutionSource.TOOL.value: + fs_instance = PromptServiceFileHelper.initialize_file_storage( + type=FileStorageType.TEMPORARY + ) highlight_data = highlight_data_plugin["entrypoint_cls"]( - logger=current_app.logger, file_path=file_path + logger=current_app.logger, file_path=file_path, fs_instance=fs_instance ).run completion = llm.complete( prompt=prompt, @@ -325,6 +348,7 @@ def extract_table( structured_output: dict[str, Any], llm: LLM, enforce_type: str, + execution_source: str, ) -> dict[str, Any]: table_settings = output[PSKeys.TABLE_SETTINGS] table_extractor: dict[str, Any] = plugins.get("table-extractor", {}) @@ -333,9 +357,22 @@ def extract_table( "Unable to extract table details. " "Please contact admin to resolve this issue." ) + if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE): + fs_instance: FileStorage + if execution_source == ExecutionSource.IDE.value: + fs_instance = PromptServiceFileHelper.initialize_file_storage( + type=FileStorageType.PERMANENT + ) + if execution_source == ExecutionSource.TOOL.value: + fs_instance = PromptServiceFileHelper.initialize_file_storage( + type=FileStorageType.TEMPORARY + ) try: answer = table_extractor["entrypoint_cls"].extract_large_table( - llm=llm, table_settings=table_settings, enforce_type=enforce_type + llm=llm, + table_settings=table_settings, + enforce_type=enforce_type, + fs_instance=fs_instance, ) structured_output[output[PSKeys.NAME]] = answer # We do not support summary and eval for table. diff --git a/prompt-service/src/unstract/prompt_service/main.py b/prompt-service/src/unstract/prompt_service/main.py index 531617cf8..5d9840641 100644 --- a/prompt-service/src/unstract/prompt_service/main.py +++ b/prompt-service/src/unstract/prompt_service/main.py @@ -110,6 +110,8 @@ def prompt_processor() -> Any: PSKeys.CONTEXT: {}, } variable_names: list[str] = [] + # Identifier for source of invocation + execution_source = payload.get(PSKeys.EXECUTION_SOURCE, "") publish_log( log_events_id, {"tool_id": tool_id, "run_id": run_id, "doc_name": doc_name}, @@ -225,6 +227,7 @@ def prompt_processor() -> Any: structured_output=structured_output, llm=llm, enforce_type=output[PSKeys.TYPE], + execution_source=execution_source, ) metadata = query_usage_metadata(token=platform_key, metadata=metadata) response = { diff --git a/prompt-service/src/unstract/prompt_service/prompt_service_file_helper.py b/prompt-service/src/unstract/prompt_service/prompt_service_file_helper.py new file mode 100644 index 000000000..b194c5170 --- /dev/null +++ b/prompt-service/src/unstract/prompt_service/prompt_service_file_helper.py @@ -0,0 +1,31 @@ +from typing import Union + +from unstract.prompt_service.constants import FileStorageKeys, FileStorageType +from unstract.prompt_service.env_manager import EnvLoader +from unstract.sdk.file_storage import ( + FileStorageProvider, + PermanentFileStorage, + SharedTemporaryFileStorage, +) + + +class PromptServiceFileHelper: + @staticmethod + def initialize_file_storage( + type: FileStorageType, + ) -> Union[PermanentFileStorage, SharedTemporaryFileStorage]: + provider_data = EnvLoader.load_provider_credentials() + provider = provider_data[FileStorageKeys.FILE_STORAGE_PROVIDER] + provider_value = FileStorageProvider(provider) + credentials = provider_data[FileStorageKeys.FILE_STORAGE_CREDENTIALS] + if type.value == FileStorageType.PERMANENT.value: + file_storage = PermanentFileStorage(provider=provider_value, **credentials) + elif type.value == FileStorageType.TEMPORARY.value: + file_storage = SharedTemporaryFileStorage( + provider=provider_value, **credentials + ) + else: + file_storage = PermanentFileStorage( + provider=FileStorageProvider.LOCAL, **credentials + ) + return file_storage diff --git a/tools/structure/src/constants.py b/tools/structure/src/constants.py index 8cf7c8653..cf8999905 100644 --- a/tools/structure/src/constants.py +++ b/tools/structure/src/constants.py @@ -75,3 +75,5 @@ class SettingsKeys: CONFIDENCE_DATA = "confidence_data" EXECUTION_RUN_DATA_FOLDER = "EXECUTION_RUN_DATA_FOLDER" FILE_PATH = "file_path" + EXECUTION_SOURCE = "execution_source" + TOOL = "tool" diff --git a/tools/structure/src/main.py b/tools/structure/src/main.py index a1736c554..776af010d 100644 --- a/tools/structure/src/main.py +++ b/tools/structure/src/main.py @@ -115,6 +115,7 @@ def run( SettingsKeys.FILE_HASH: file_hash, SettingsKeys.FILE_NAME: file_name, SettingsKeys.FILE_PATH: extracted_input_file, + SettingsKeys.EXECUTION_SOURCE: SettingsKeys.TOOL, } # TODO: Need to split extraction and indexing # to avoid unwanted indexing From a83cb4dd696bd5d76163a9e256dee5310a94b47b Mon Sep 17 00:00:00 2001 From: harini-venkataraman Date: Wed, 18 Dec 2024 10:59:10 +0530 Subject: [PATCH 2/8] Initilizing fs_instances --- prompt-service/src/unstract/prompt_service/helper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/prompt-service/src/unstract/prompt_service/helper.py b/prompt-service/src/unstract/prompt_service/helper.py index b64522542..6eb38604b 100644 --- a/prompt-service/src/unstract/prompt_service/helper.py +++ b/prompt-service/src/unstract/prompt_service/helper.py @@ -22,12 +22,12 @@ from unstract.sdk.llm import LLM from unstract.flags.src.unstract.flags.feature_flag import check_feature_flag_status +from unstract.sdk.file_storage import FileStorage, FileStorageProvider if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE): from unstract.prompt_service.prompt_service_file_helper import ( PromptServiceFileHelper, ) - from unstract.sdk.file_storage import FileStorage load_dotenv() @@ -300,7 +300,7 @@ def run_completion( ) highlight_data = None if highlight_data_plugin and enable_highlight: - fs_instance: FileStorage + fs_instance: FileStorage = FileStorage(FileStorageProvider.LOCAL) if execution_source == ExecutionSource.IDE.value: fs_instance = PromptServiceFileHelper.initialize_file_storage( type=FileStorageType.PERMANENT @@ -358,7 +358,7 @@ def extract_table( "Please contact admin to resolve this issue." ) if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE): - fs_instance: FileStorage + fs_instance: FileStorage = FileStorage(FileStorageProvider.LOCAL) if execution_source == ExecutionSource.IDE.value: fs_instance = PromptServiceFileHelper.initialize_file_storage( type=FileStorageType.PERMANENT From 954465186b20b8a6daf2dd6f3c29a847d5d239b4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 18 Dec 2024 05:29:56 +0000 Subject: [PATCH 3/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- prompt-service/src/unstract/prompt_service/helper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/prompt-service/src/unstract/prompt_service/helper.py b/prompt-service/src/unstract/prompt_service/helper.py index 6eb38604b..73dd30dd3 100644 --- a/prompt-service/src/unstract/prompt_service/helper.py +++ b/prompt-service/src/unstract/prompt_service/helper.py @@ -19,10 +19,10 @@ from unstract.prompt_service.exceptions import APIError, RateLimitError from unstract.sdk.exceptions import RateLimitError as SdkRateLimitError from unstract.sdk.exceptions import SdkError +from unstract.sdk.file_storage import FileStorage, FileStorageProvider from unstract.sdk.llm import LLM from unstract.flags.src.unstract.flags.feature_flag import check_feature_flag_status -from unstract.sdk.file_storage import FileStorage, FileStorageProvider if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE): from unstract.prompt_service.prompt_service_file_helper import ( From ba8def609b16ad3945f61feba6d43c6803112808 Mon Sep 17 00:00:00 2001 From: harini-venkataraman Date: Wed, 18 Dec 2024 14:06:36 +0530 Subject: [PATCH 4/8] Adding feature flag to init --- .../src/unstract/prompt_service/constants.py | 1 + .../src/unstract/prompt_service/helper.py | 67 ++++++++++++------- 2 files changed, 43 insertions(+), 25 deletions(-) diff --git a/prompt-service/src/unstract/prompt_service/constants.py b/prompt-service/src/unstract/prompt_service/constants.py index 0d4dc4521..8f122a809 100644 --- a/prompt-service/src/unstract/prompt_service/constants.py +++ b/prompt-service/src/unstract/prompt_service/constants.py @@ -106,6 +106,7 @@ class DBTableV2: class FileStorageKeys: FILE_STORAGE_PROVIDER = "FILE_STORAGE_PROVIDER" FILE_STORAGE_CREDENTIALS = "FILE_STORAGE_CREDENTIALS" + REMOTE_STORAGE = "REMOTE_STORAGE" class FileStorageType(Enum): diff --git a/prompt-service/src/unstract/prompt_service/helper.py b/prompt-service/src/unstract/prompt_service/helper.py index 73dd30dd3..994d62094 100644 --- a/prompt-service/src/unstract/prompt_service/helper.py +++ b/prompt-service/src/unstract/prompt_service/helper.py @@ -11,6 +11,7 @@ DBTableV2, ExecutionSource, FeatureFlag, + FileStorageKeys, FileStorageType, ) from unstract.prompt_service.constants import PromptServiceContants as PSKeys @@ -25,9 +26,9 @@ from unstract.flags.src.unstract.flags.feature_flag import check_feature_flag_status if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE): - from unstract.prompt_service.prompt_service_file_helper import ( - PromptServiceFileHelper, - ) + from unstract.sdk.file_storage import FileStorage, FileStorageProvider + from unstract.sdk.file_storage.env_helper import EnvHelper + from unstract.sdk.file_storage.constants import StorageType load_dotenv() @@ -300,18 +301,25 @@ def run_completion( ) highlight_data = None if highlight_data_plugin and enable_highlight: - fs_instance: FileStorage = FileStorage(FileStorageProvider.LOCAL) - if execution_source == ExecutionSource.IDE.value: - fs_instance = PromptServiceFileHelper.initialize_file_storage( - type=FileStorageType.PERMANENT - ) - if execution_source == ExecutionSource.TOOL.value: - fs_instance = PromptServiceFileHelper.initialize_file_storage( - type=FileStorageType.TEMPORARY - ) - highlight_data = highlight_data_plugin["entrypoint_cls"]( - logger=current_app.logger, file_path=file_path, fs_instance=fs_instance - ).run + if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE): + fs_instance: FileStorage = FileStorage(FileStorageProvider.LOCAL) + if execution_source == ExecutionSource.IDE.value: + fs_instance = EnvHelper.get_storage( + storage_type=StorageType.PERMANENT, + env_name=FileStorageKeys.REMOTE_STORAGE + ) + if execution_source == ExecutionSource.TOOL.value: + fs_instance = EnvHelper.get_storage( + storage_type=StorageType.TEMPORARY, + env_name=FileStorageKeys.REMOTE_STORAGE, + ) + highlight_data = highlight_data_plugin["entrypoint_cls"]( + logger=current_app.logger, file_path=file_path, fs_instance=fs_instance + ).run + else: + highlight_data = highlight_data_plugin["entrypoint_cls"]( + logger=current_app.logger, file_path=file_path + ).run completion = llm.complete( prompt=prompt, process_text=highlight_data, @@ -360,20 +368,29 @@ def extract_table( if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE): fs_instance: FileStorage = FileStorage(FileStorageProvider.LOCAL) if execution_source == ExecutionSource.IDE.value: - fs_instance = PromptServiceFileHelper.initialize_file_storage( - type=FileStorageType.PERMANENT + fs_instance = EnvHelper.get_storage( + storage_type=StorageType.PERMANENT, + env_name=FileStorageKeys.REMOTE_STORAGE, ) if execution_source == ExecutionSource.TOOL.value: - fs_instance = PromptServiceFileHelper.initialize_file_storage( - type=FileStorageType.TEMPORARY + fs_instance = EnvHelper.get_storage( + storage_type=StorageType.TEMPORARY, + env_name=FileStorageKeys.REMOTE_STORAGE, ) try: - answer = table_extractor["entrypoint_cls"].extract_large_table( - llm=llm, - table_settings=table_settings, - enforce_type=enforce_type, - fs_instance=fs_instance, - ) + if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE): + answer = table_extractor["entrypoint_cls"].extract_large_table( + llm=llm, + table_settings=table_settings, + enforce_type=enforce_type, + fs_instance=fs_instance, + ) + else: + answer = table_extractor["entrypoint_cls"].extract_large_table( + llm=llm, + table_settings=table_settings, + enforce_type=enforce_type, + ) structured_output[output[PSKeys.NAME]] = answer # We do not support summary and eval for table. # Hence returning the result From 03174efc0198297d7a712ed0fb0bf85292440de6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 18 Dec 2024 08:37:16 +0000 Subject: [PATCH 5/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- prompt-service/src/unstract/prompt_service/helper.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/prompt-service/src/unstract/prompt_service/helper.py b/prompt-service/src/unstract/prompt_service/helper.py index 994d62094..70030fcf3 100644 --- a/prompt-service/src/unstract/prompt_service/helper.py +++ b/prompt-service/src/unstract/prompt_service/helper.py @@ -12,7 +12,6 @@ ExecutionSource, FeatureFlag, FileStorageKeys, - FileStorageType, ) from unstract.prompt_service.constants import PromptServiceContants as PSKeys from unstract.prompt_service.db_utils import DBUtils @@ -27,8 +26,8 @@ if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE): from unstract.sdk.file_storage import FileStorage, FileStorageProvider - from unstract.sdk.file_storage.env_helper import EnvHelper from unstract.sdk.file_storage.constants import StorageType + from unstract.sdk.file_storage.env_helper import EnvHelper load_dotenv() @@ -305,8 +304,8 @@ def run_completion( fs_instance: FileStorage = FileStorage(FileStorageProvider.LOCAL) if execution_source == ExecutionSource.IDE.value: fs_instance = EnvHelper.get_storage( - storage_type=StorageType.PERMANENT, - env_name=FileStorageKeys.REMOTE_STORAGE + storage_type=StorageType.PERMANENT, + env_name=FileStorageKeys.REMOTE_STORAGE, ) if execution_source == ExecutionSource.TOOL.value: fs_instance = EnvHelper.get_storage( @@ -314,7 +313,9 @@ def run_completion( env_name=FileStorageKeys.REMOTE_STORAGE, ) highlight_data = highlight_data_plugin["entrypoint_cls"]( - logger=current_app.logger, file_path=file_path, fs_instance=fs_instance + logger=current_app.logger, + file_path=file_path, + fs_instance=fs_instance, ).run else: highlight_data = highlight_data_plugin["entrypoint_cls"]( From db6c1dda5c298b292f2e4c155ffaa488600a38b8 Mon Sep 17 00:00:00 2001 From: harini-venkataraman Date: Wed, 18 Dec 2024 14:19:53 +0530 Subject: [PATCH 6/8] Removing unused helpers --- .../unstract/prompt_service/env_manager.py | 18 +---------- .../prompt_service_file_helper.py | 31 ------------------- 2 files changed, 1 insertion(+), 48 deletions(-) delete mode 100644 prompt-service/src/unstract/prompt_service/prompt_service_file_helper.py diff --git a/prompt-service/src/unstract/prompt_service/env_manager.py b/prompt-service/src/unstract/prompt_service/env_manager.py index 745b9083d..c51f5cf80 100644 --- a/prompt-service/src/unstract/prompt_service/env_manager.py +++ b/prompt-service/src/unstract/prompt_service/env_manager.py @@ -1,8 +1,5 @@ -import json import os -from typing import Any, Optional - -from unstract.prompt_service.constants import FileStorageKeys +from typing import Optional class EnvLoader: @@ -12,16 +9,3 @@ def get_env_or_die(env_key: str, default: Optional[str] = None) -> str: if env_value is None or env_value == "": raise ValueError(f"Env variable {env_key} is required") return env_value - - @staticmethod - def load_provider_credentials() -> dict[str, Any]: - cred_env_data: str = EnvLoader.get_env_or_die( - env_key=FileStorageKeys.FILE_STORAGE_CREDENTIALS - ) - credentials = json.loads(cred_env_data) - provider_data: dict[str, Any] = {} - provider_data[FileStorageKeys.FILE_STORAGE_PROVIDER] = credentials["provider"] - provider_data[FileStorageKeys.FILE_STORAGE_CREDENTIALS] = credentials[ - "credentials" - ] - return provider_data diff --git a/prompt-service/src/unstract/prompt_service/prompt_service_file_helper.py b/prompt-service/src/unstract/prompt_service/prompt_service_file_helper.py deleted file mode 100644 index b194c5170..000000000 --- a/prompt-service/src/unstract/prompt_service/prompt_service_file_helper.py +++ /dev/null @@ -1,31 +0,0 @@ -from typing import Union - -from unstract.prompt_service.constants import FileStorageKeys, FileStorageType -from unstract.prompt_service.env_manager import EnvLoader -from unstract.sdk.file_storage import ( - FileStorageProvider, - PermanentFileStorage, - SharedTemporaryFileStorage, -) - - -class PromptServiceFileHelper: - @staticmethod - def initialize_file_storage( - type: FileStorageType, - ) -> Union[PermanentFileStorage, SharedTemporaryFileStorage]: - provider_data = EnvLoader.load_provider_credentials() - provider = provider_data[FileStorageKeys.FILE_STORAGE_PROVIDER] - provider_value = FileStorageProvider(provider) - credentials = provider_data[FileStorageKeys.FILE_STORAGE_CREDENTIALS] - if type.value == FileStorageType.PERMANENT.value: - file_storage = PermanentFileStorage(provider=provider_value, **credentials) - elif type.value == FileStorageType.TEMPORARY.value: - file_storage = SharedTemporaryFileStorage( - provider=provider_value, **credentials - ) - else: - file_storage = PermanentFileStorage( - provider=FileStorageProvider.LOCAL, **credentials - ) - return file_storage From 8f8f6e77655c663dd746e7baa1d8f6a224692cfe Mon Sep 17 00:00:00 2001 From: harini-venkataraman Date: Wed, 18 Dec 2024 15:02:12 +0530 Subject: [PATCH 7/8] Adding env constants --- prompt-service/src/unstract/prompt_service/constants.py | 3 ++- prompt-service/src/unstract/prompt_service/helper.py | 9 ++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/prompt-service/src/unstract/prompt_service/constants.py b/prompt-service/src/unstract/prompt_service/constants.py index 8f122a809..f2ac67439 100644 --- a/prompt-service/src/unstract/prompt_service/constants.py +++ b/prompt-service/src/unstract/prompt_service/constants.py @@ -106,7 +106,8 @@ class DBTableV2: class FileStorageKeys: FILE_STORAGE_PROVIDER = "FILE_STORAGE_PROVIDER" FILE_STORAGE_CREDENTIALS = "FILE_STORAGE_CREDENTIALS" - REMOTE_STORAGE = "REMOTE_STORAGE" + PERMANENT_REMOTE_STORAGE = "PERMANENT_REMOTE_STORAGE" + TEMPORARY_REMOTE_STORAGE = "TEMPORARY_REMOTE_STORAGE" class FileStorageType(Enum): diff --git a/prompt-service/src/unstract/prompt_service/helper.py b/prompt-service/src/unstract/prompt_service/helper.py index 70030fcf3..126a41c93 100644 --- a/prompt-service/src/unstract/prompt_service/helper.py +++ b/prompt-service/src/unstract/prompt_service/helper.py @@ -19,7 +19,6 @@ from unstract.prompt_service.exceptions import APIError, RateLimitError from unstract.sdk.exceptions import RateLimitError as SdkRateLimitError from unstract.sdk.exceptions import SdkError -from unstract.sdk.file_storage import FileStorage, FileStorageProvider from unstract.sdk.llm import LLM from unstract.flags.src.unstract.flags.feature_flag import check_feature_flag_status @@ -305,12 +304,12 @@ def run_completion( if execution_source == ExecutionSource.IDE.value: fs_instance = EnvHelper.get_storage( storage_type=StorageType.PERMANENT, - env_name=FileStorageKeys.REMOTE_STORAGE, + env_name=FileStorageKeys.PERMANENT_REMOTE_STORAGE, ) if execution_source == ExecutionSource.TOOL.value: fs_instance = EnvHelper.get_storage( storage_type=StorageType.TEMPORARY, - env_name=FileStorageKeys.REMOTE_STORAGE, + env_name=FileStorageKeys.TEMPORARY_REMOTE_STORAGE, ) highlight_data = highlight_data_plugin["entrypoint_cls"]( logger=current_app.logger, @@ -371,12 +370,12 @@ def extract_table( if execution_source == ExecutionSource.IDE.value: fs_instance = EnvHelper.get_storage( storage_type=StorageType.PERMANENT, - env_name=FileStorageKeys.REMOTE_STORAGE, + env_name=FileStorageKeys.PERMANENT_REMOTE_STORAGE, ) if execution_source == ExecutionSource.TOOL.value: fs_instance = EnvHelper.get_storage( storage_type=StorageType.TEMPORARY, - env_name=FileStorageKeys.REMOTE_STORAGE, + env_name=FileStorageKeys.TEMPORARY_REMOTE_STORAGE, ) try: if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE): From 1baa6afea6c964622373eead48ecd67be05b323f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 19 Dec 2024 07:56:58 +0000 Subject: [PATCH 8/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- prompt-service/src/unstract/prompt_service/constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/prompt-service/src/unstract/prompt_service/constants.py b/prompt-service/src/unstract/prompt_service/constants.py index eb32ac253..4f30dd11d 100644 --- a/prompt-service/src/unstract/prompt_service/constants.py +++ b/prompt-service/src/unstract/prompt_service/constants.py @@ -74,7 +74,7 @@ class PromptServiceContants: CONFIDENCE_DATA = "confidence_data" EXECUTION_SOURCE = "execution_source" METRICS = "metrics" - + class RunLevel(Enum): """Different stages of prompt execution.