Your Name
2 months ago
parent
47a359ec34
commit
0ce9ad929a
@ -0,0 +1,176 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from swarms.utils.loguru_logger import initialize_logger
|
||||
|
||||
|
||||
logger = initialize_logger("workspace-manager")
|
||||
|
||||
|
||||
class WorkspaceManager:
|
||||
"""
|
||||
Manages the workspace directory and settings for the application.
|
||||
This class is responsible for setting up the workspace directory, logging configuration,
|
||||
and retrieving environment variables for telemetry and API key.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workspace_dir: Optional[str] = "agent_workspace",
|
||||
use_telemetry: Optional[bool] = True,
|
||||
api_key: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Initializes the WorkspaceManager with optional parameters for workspace directory,
|
||||
telemetry usage, and API key.
|
||||
|
||||
Args:
|
||||
workspace_dir (Optional[str]): The path to the workspace directory.
|
||||
use_telemetry (Optional[bool]): A flag indicating whether to use telemetry.
|
||||
api_key (Optional[str]): The API key for the application.
|
||||
"""
|
||||
self.workspace_dir = workspace_dir
|
||||
self.use_telemetry = use_telemetry
|
||||
self.api_key = api_key
|
||||
|
||||
def _create_env_file(self, env_file_path: Path) -> None:
|
||||
"""
|
||||
Create a new .env file with default WORKSPACE_DIR.
|
||||
|
||||
Args:
|
||||
env_file_path (Path): The path to the .env file.
|
||||
"""
|
||||
with env_file_path.open("w") as file:
|
||||
file.write("WORKSPACE_DIR=agent_workspace\n")
|
||||
logger.info(
|
||||
"Created a new .env file with default WORKSPACE_DIR."
|
||||
)
|
||||
|
||||
def _append_to_env_file(self, env_file_path: Path) -> None:
|
||||
"""
|
||||
Append WORKSPACE_DIR to .env if it doesn't exist.
|
||||
|
||||
Args:
|
||||
env_file_path (Path): The path to the .env file.
|
||||
"""
|
||||
with env_file_path.open("r+") as file:
|
||||
content = file.read()
|
||||
if "WORKSPACE_DIR" not in content:
|
||||
file.seek(0, os.SEEK_END)
|
||||
file.write("WORKSPACE_DIR=agent_workspace\n")
|
||||
logger.info("Appended WORKSPACE_DIR to .env file.")
|
||||
|
||||
def _get_workspace_dir(
|
||||
self, workspace_dir: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Get the workspace directory from environment variable or default.
|
||||
|
||||
Args:
|
||||
workspace_dir (Optional[str]): The path to the workspace directory.
|
||||
|
||||
Returns:
|
||||
str: The path to the workspace directory.
|
||||
"""
|
||||
return workspace_dir or os.getenv(
|
||||
"WORKSPACE_DIR", "agent_workspace"
|
||||
)
|
||||
|
||||
def _get_telemetry_status(
|
||||
self, use_telemetry: Optional[bool] = None
|
||||
) -> bool:
|
||||
"""
|
||||
Get telemetry status from environment variable or default.
|
||||
|
||||
Args:
|
||||
use_telemetry (Optional[bool]): A flag indicating whether to use telemetry.
|
||||
|
||||
Returns:
|
||||
bool: The status of telemetry usage.
|
||||
"""
|
||||
return (
|
||||
use_telemetry
|
||||
if use_telemetry is not None
|
||||
else os.getenv("USE_TELEMETRY", "true").lower() == "true"
|
||||
)
|
||||
|
||||
def _get_api_key(
|
||||
self, api_key: Optional[str] = None
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Get API key from environment variable or default.
|
||||
|
||||
Args:
|
||||
api_key (Optional[str]): The API key for the application.
|
||||
|
||||
Returns:
|
||||
Optional[str]: The API key or None if not set.
|
||||
"""
|
||||
return api_key or os.getenv("SWARMS_API_KEY")
|
||||
|
||||
def _init_workspace(self) -> None:
|
||||
"""
|
||||
Initialize the workspace directory if it doesn't exist.
|
||||
"""
|
||||
if not self.workspace_path.exists():
|
||||
self.workspace_path.mkdir(parents=True, exist_ok=True)
|
||||
logger.info("Workspace directory initialized.")
|
||||
|
||||
@property
|
||||
def get_workspace_path(self) -> Path:
|
||||
"""
|
||||
Get the workspace path.
|
||||
|
||||
Returns:
|
||||
Path: The path to the workspace directory.
|
||||
"""
|
||||
return self.workspace_path
|
||||
|
||||
@property
|
||||
def get_telemetry_status(self) -> bool:
|
||||
"""
|
||||
Get telemetry status.
|
||||
|
||||
Returns:
|
||||
bool: The status of telemetry usage.
|
||||
"""
|
||||
return self.use_telemetry
|
||||
|
||||
@property
|
||||
def get_api_key(self) -> Optional[str]:
|
||||
"""
|
||||
Get API key.
|
||||
|
||||
Returns:
|
||||
Optional[str]: The API key or None if not set.
|
||||
"""
|
||||
return self.api_key
|
||||
|
||||
def run(self) -> None:
|
||||
try:
|
||||
# Check if .env file exists and create it if it doesn't
|
||||
env_file_path = Path(".env")
|
||||
if not env_file_path.exists():
|
||||
self._create_env_file(env_file_path)
|
||||
else:
|
||||
# Append WORKSPACE_DIR to .env if it doesn't exist
|
||||
self._append_to_env_file(env_file_path)
|
||||
|
||||
# Set workspace directory
|
||||
self.workspace_dir = self._get_workspace_dir(
|
||||
self.workspace_dir
|
||||
)
|
||||
self.workspace_path = Path(self.workspace_dir)
|
||||
|
||||
# Set telemetry preference
|
||||
self.use_telemetry = self._get_telemetry_status(
|
||||
self.use_telemetry
|
||||
)
|
||||
|
||||
# Set API key
|
||||
self.api_key = self._get_api_key(self.api_key)
|
||||
|
||||
# Initialize workspace
|
||||
self._init_workspace()
|
||||
except Exception as e:
|
||||
logger.error(f"Error initializing WorkspaceManager: {e}")
|
||||
Loading…
Reference in new issue