From 5429ff312dea7bd720e57db6cd8d09353fe4eb1e Mon Sep 17 00:00:00 2001 From: Occupying-Mars Date: Sat, 16 Nov 2024 02:23:00 +0530 Subject: [PATCH] workspace bug fixes --- swarms/cli/onboarding_process.py | 33 ++++++++++++++------------- swarms/telemetry/bootup.py | 13 ++++++----- swarms/utils/disable_logging.py | 8 +++---- swarms/utils/loguru_logger.py | 37 ++++++++++++++++++------------- swarms/utils/workspace_manager.py | 25 +++++++++++++++++++++ 5 files changed, 77 insertions(+), 39 deletions(-) create mode 100644 swarms/utils/workspace_manager.py diff --git a/swarms/cli/onboarding_process.py b/swarms/cli/onboarding_process.py index 99018b86..f7c2a395 100644 --- a/swarms/cli/onboarding_process.py +++ b/swarms/cli/onboarding_process.py @@ -9,6 +9,7 @@ from swarms.telemetry.capture_sys_data import ( capture_system_data, log_agent_data, ) +from swarms.utils.workspace_manager import WorkspaceManager class OnboardingProcess: @@ -155,10 +156,11 @@ class OnboardingProcess: def collect_user_info(self) -> None: """ - Initiates the onboarding process by collecting the user's full name, first name, email, - Swarms API key, and system data. Additionally, it reminds the user to set their WORKSPACE_DIR environment variable. + Initiates the onboarding process by collecting user information and setting up workspace. """ logger.info("Initiating swarms cloud onboarding process...") + + # Collect user information self.ask_input( "Enter your first name (or type 'quit' to exit): ", "first_name", @@ -168,24 +170,25 @@ class OnboardingProcess: "last_name", ) self.ask_input( - "Enter your email (or type 'quit' to exit): ", "email" + "Enter your email (or type 'quit' to exit): ", + "email" ) self.ask_input( "Enter your Swarms API key (or type 'quit' to exit): Get this in your swarms dashboard: https://swarms.world/platform/api-keys ", "swarms_api_key", ) - workspace = self.ask_input( - "Enter your WORKSPACE_DIR: This is where logs, errors, and agent configurations will be stored (or type 'quit' to exit). Remember to set this as an environment variable: https://docs.swarms.world/en/latest/swarms/install/quickstart/ || ", - "workspace_dir", - ) - os.environ["WORKSPACE_DIR"] = workspace - logger.info( - "Important: Please ensure you have set your WORKSPACE_DIR environment variable as per the instructions provided." - ) - logger.info( - "Additionally, remember to add your API keys for your respective models in your .env file." - ) - logger.success("Onboarding process completed successfully!") + + # Set up workspace directory + try: + workspace = WorkspaceManager.get_workspace_dir() + logger.info(f"Using workspace directory: {workspace}") + logger.info( + "Additionally, remember to add your API keys for your respective models in your .env file." + ) + logger.success("Onboarding process completed successfully!") + except Exception as e: + logger.error(f"Failed to set up workspace: {e}") + raise def run(self) -> None: """ diff --git a/swarms/telemetry/bootup.py b/swarms/telemetry/bootup.py index 24d7a7c4..6116e99c 100644 --- a/swarms/telemetry/bootup.py +++ b/swarms/telemetry/bootup.py @@ -5,6 +5,7 @@ from concurrent.futures import ThreadPoolExecutor from swarms.telemetry.auto_upgrade_swarms import auto_update from swarms.utils.disable_logging import disable_logging +from swarms.utils.workspace_manager import WorkspaceManager def bootup(): @@ -12,11 +13,13 @@ def bootup(): logging.disable(logging.CRITICAL) os.environ["WANDB_SILENT"] = "true" - # Auto set workspace directory - workspace_dir = os.path.join(os.getcwd(), "agent_workspace") - if not os.path.exists(workspace_dir): - os.makedirs(workspace_dir) - os.environ["WORKSPACE_DIR"] = workspace_dir + # Set workspace directory using WorkspaceManager + try: + workspace_dir = WorkspaceManager.get_workspace_dir() + os.environ["WORKSPACE_DIR"] = workspace_dir + except Exception as e: + print(f"Error setting up workspace directory: {e}") + return warnings.filterwarnings("ignore", category=DeprecationWarning) diff --git a/swarms/utils/disable_logging.py b/swarms/utils/disable_logging.py index c2cc4ab5..2d2e5323 100644 --- a/swarms/utils/disable_logging.py +++ b/swarms/utils/disable_logging.py @@ -3,6 +3,7 @@ import logging import os import warnings from threading import Thread +from swarms.utils.workspace_manager import WorkspaceManager def disable_langchain(): @@ -24,7 +25,9 @@ def disable_logging(): Disables logging for specific modules and sets up file and stream handlers. Runs in a separate thread to avoid blocking the main thread. """ - os.environ["WORKSPACE_DIR"] = "agent_workspace" + # Get workspace directory using WorkspaceManager instead of direct environment variable + workspace_dir = WorkspaceManager.get_workspace_dir() + os.environ["WORKSPACE_DIR"] = workspace_dir warnings.filterwarnings("ignore", category=UserWarning) @@ -61,9 +64,6 @@ def disable_logging(): # Remove all existing handlers logging.getLogger().handlers = [] - # Get the workspace directory from the environment variables - workspace_dir = os.environ["WORKSPACE_DIR"] - # Check if the workspace directory exists, if not, create it if not os.path.exists(workspace_dir): os.makedirs(workspace_dir) diff --git a/swarms/utils/loguru_logger.py b/swarms/utils/loguru_logger.py index b53ec379..1f7f5d4f 100644 --- a/swarms/utils/loguru_logger.py +++ b/swarms/utils/loguru_logger.py @@ -1,23 +1,30 @@ import os from loguru import logger +from swarms.utils.workspace_manager import WorkspaceManager - -WORKSPACE_DIR = os.getenv("WORKSPACE_DIR") - -logger.add( - os.path.join(WORKSPACE_DIR, "swarms.log"), - level="INFO", - colorize=True, - backtrace=True, - diagnose=True, -) - - -def loguru_logger(file_path: str = "swarms.log"): - return logger.add( - os.path.join(WORKSPACE_DIR, file_path), +try: + WORKSPACE_DIR = WorkspaceManager.get_workspace_dir() + + logger.add( + os.path.join(WORKSPACE_DIR, "swarms.log"), level="INFO", colorize=True, backtrace=True, diagnose=True, ) + +except Exception as e: + logger.error(f"Failed to initialize logger: {e}") + +def loguru_logger(file_path: str = "swarms.log"): + try: + return logger.add( + os.path.join(WORKSPACE_DIR, file_path), + level="INFO", + colorize=True, + backtrace=True, + diagnose=True, + ) + except Exception as e: + logger.error(f"Failed to create logger for {file_path}: {e}") + raise diff --git a/swarms/utils/workspace_manager.py b/swarms/utils/workspace_manager.py new file mode 100644 index 00000000..3bd0b352 --- /dev/null +++ b/swarms/utils/workspace_manager.py @@ -0,0 +1,25 @@ +import os +from pathlib import Path +from loguru import logger + +class WorkspaceManager: + DEFAULT_WORKSPACE = "agent_workspace" + + @classmethod + def get_workspace_dir(cls) -> str: + """Get or create workspace directory with proper fallback""" + workspace = os.getenv("WORKSPACE_DIR", cls.DEFAULT_WORKSPACE) + workspace_path = Path(workspace) + + try: + # Create directory if it doesn't exist + workspace_path.mkdir(parents=True, exist_ok=True) + + # Set environment variable if not already set + if "WORKSPACE_DIR" not in os.environ: + os.environ["WORKSPACE_DIR"] = str(workspace_path) + + return str(workspace_path) + except Exception as e: + logger.error(f"Error creating workspace directory: {e}") + raise