diff --git a/backend/utils/file_storage/helpers/prompt_studio_file_helper.py b/backend/utils/file_storage/helpers/prompt_studio_file_helper.py index 8b6811069..5bd01c73d 100644 --- a/backend/utils/file_storage/helpers/prompt_studio_file_helper.py +++ b/backend/utils/file_storage/helpers/prompt_studio_file_helper.py @@ -2,7 +2,7 @@ import logging import os from pathlib import Path -from typing import Any, Union +from typing import Any from file_management.file_management_helper import FileManagerHelper from unstract.sdk.file_storage import FileStorage @@ -77,7 +77,7 @@ def upload_for_ide( @staticmethod def fetch_file_contents( org_id: str, user_id: str, tool_id: str, file_name: str - ) -> Union[bytes, str]: + ) -> dict[str, Any]: """Method to fetch file contents from the remote location. The path is constructed in runtime based on the args""" fs_instance = EnvHelper.get_storage( @@ -125,7 +125,7 @@ def fetch_file_contents( encoding="utf-8", ) encoded_string = base64.b64encode(bytes(text_content_bytes)) - return encoded_string + return {"data": encoded_string, "mime_type": file_content_type} elif file_content_type == "text/plain": text_content_string: str = fs_instance.read( @@ -134,7 +134,7 @@ def fetch_file_contents( legacy_storage_path=legacy_file_path, encoding="utf-8", ) - return text_content_string + return {"data": text_content_string, "mime_type": file_content_type} else: raise ValueError(f"Unsupported file type: {file_content_type}") diff --git a/unstract/workflow-execution/src/unstract/workflow_execution/constants.py b/unstract/workflow-execution/src/unstract/workflow_execution/constants.py index c00145620..e570a3f00 100644 --- a/unstract/workflow-execution/src/unstract/workflow_execution/constants.py +++ b/unstract/workflow-execution/src/unstract/workflow_execution/constants.py @@ -23,6 +23,10 @@ class ToolRuntimeVariable: EXECUTION_BY_TOOL = "EXECUTION_BY_TOOL" WORKFLOW_EXECUTION_DIR_PREFIX = "WORKFLOW_EXECUTION_DIR_PREFIX" API_EXECUTION_DIR_PREFIX = "API_EXECUTION_DIR_PREFIX" + REDIS_HOST = "REDIS_HOST" + REDIS_PORT = "REDIS_PORT" + REDIS_USER = "REDIS_USER" + REDIS_PASSWORD = "REDIS_PASSWORD" class WorkflowFileType: diff --git a/unstract/workflow-execution/src/unstract/workflow_execution/tools_utils.py b/unstract/workflow-execution/src/unstract/workflow_execution/tools_utils.py index ba7287f19..60ffccb85 100644 --- a/unstract/workflow-execution/src/unstract/workflow_execution/tools_utils.py +++ b/unstract/workflow-execution/src/unstract/workflow_execution/tools_utils.py @@ -52,6 +52,12 @@ def __init__( self.llmw_max_polls = ToolsUtils.get_env( ToolRV.ADAPTER_LLMW_MAX_POLLS, raise_exception=False ) + self.redis_host = ToolsUtils.get_env(ToolRV.REDIS_HOST, raise_exception=True) + self.redis_port = ToolsUtils.get_env(ToolRV.REDIS_PORT, raise_exception=True) + self.redis_user = ToolsUtils.get_env(ToolRV.REDIS_USER, raise_exception=True) + self.redis_password = ToolsUtils.get_env( + ToolRV.REDIS_PASSWORD, raise_exception=True + ) def set_messaging_channel(self, messaging_channel: str) -> None: self.messaging_channel = messaging_channel @@ -219,6 +225,10 @@ def get_tool_environment_variables(self) -> dict[str, Any]: ToolRV.X2TEXT_HOST: self.x2text_host, ToolRV.X2TEXT_PORT: self.x2text_port, ToolRV.EXECUTION_BY_TOOL: True, + ToolRV.REDIS_HOST: self.redis_host, + ToolRV.REDIS_PORT: self.redis_port, + ToolRV.REDIS_USER: self.redis_user, + ToolRV.REDIS_PASSWORD: self.redis_password, } # For async LLM Whisperer extraction if self.llmw_poll_interval: @@ -243,6 +253,6 @@ def get_env(env_key: str, raise_exception: bool = False) -> Optional[str]: Optional[str]: Value for the env variable """ env_value = os.environ.get(env_key) - if (env_value is None or env_value == "") and raise_exception: + if (env_value is None) and raise_exception: raise MissingEnvVariable(f"Env variable {env_key} is required") return env_value