Upload 5 files

aworld/utils/async_func.py

@@ -0,0 +1,66 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+
+import asyncio
+from functools import wraps
+from typing import Callable, Optional, Union, Any, Dict
+
+
+class Functionable:
+    def __init__(self, function: Callable[..., Any], *args: Any, **kwargs: Dict[str, Any]) -> None:
+        self.function = function
+        self.args = args
+        self.kwargs = kwargs
+        self.done: bool = False
+        self.error: bool = False
+        self.result: Optional[Any] = None
+        self.exception: Optional[Exception] = None
+
+    def __call__(self) -> None:
+        try:
+            self.result = self.function(*self.args, **self.kwargs)
+        except Exception as e:
+            self.error = True
+            self.exception = e
+        self.done = True
+
+    def call(self):
+        self.__call__()
+
+
+def async_decorator(*func, delay: Optional[Union[int, float]] = 0.5) -> Callable:
+    def wrapper(function: Callable[..., Any]) -> Callable[..., Any]:
+        @wraps(function)
+        async def inner_wrapper(*args: Any, **kwargs: Any) -> Any:
+            sleep_time = 0 if delay is None else delay
+            task = Functionable(function, *args, **kwargs)
+            # TODO: Use thread pool to process task
+            task.call()
+            if task.error:
+                raise task.exception
+            await asyncio.sleep(sleep_time)
+            return task.result
+
+        return inner_wrapper
+
+    if not func:
+        return wrapper
+    else:
+        if asyncio.iscoroutinefunction(func[0]):
+            # coroutine function, return itself
+            return func[0]
+        return wrapper(func[0])
+
+def async_func(function: Callable[..., Any]) -> Callable[..., Any]:
+    @wraps(function)
+    async def inner_wrapper(*args: Any, **kwargs: Any) -> Any:
+        task = Functionable(function, *args, **kwargs)
+        task.call()
+        if task.error:
+            raise task.exception
+        return task.result
+
+    if asyncio.iscoroutinefunction(function):
+        # coroutine function, return itself
+        return function
+    return inner_wrapper

aworld/utils/common.py

@@ -0,0 +1,352 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import asyncio
+import inspect
+import os
+import pkgutil
+import re
+import socket
+import sys
+import threading
+import time
+
+from functools import wraps
+from types import FunctionType
+from typing import Callable, Any, Tuple, List, Iterator, Dict, Union
+
+from aworld.logs.util import logger
+
+
+def convert_to_snake(name: str) -> str:
+    """Class name convert to snake."""
+    if '_' not in name:
+        name = re.sub(r'([a-z])([A-Z])', r'\1_\2', name)
+    return name.lower()
+
+
+def snake_to_camel(snake):
+    words = snake.split('_')
+    return ''.join([w.capitalize() for w in words])
+
+
+def is_abstract_method(cls, method_name):
+    method = getattr(cls, method_name)
+    return (hasattr(method, '__isabstractmethod__') and method.__isabstractmethod__) or (
+            isinstance(method, FunctionType) and hasattr(
+        method, '__abstractmethods__') and method in method.__abstractmethods__)
+
+
+def override_in_subclass(name: str, sub_cls: object, base_cls: object) -> bool:
+    """Judge whether a subclass overrides a specified method.
+
+    Args:
+        name: The method name of sub class and base class
+        sub_cls: Specify subclasses of the base class.
+        base_cls: The parent class of the subclass.
+
+    Returns:
+        Overwrite as true in subclasses, vice versa.
+    """
+    if not issubclass(sub_cls, base_cls):
+        logger.warning(f"{sub_cls} is not sub class of {base_cls}")
+        return False
+
+    if sub_cls == base_cls and hasattr(sub_cls, name) and not is_abstract_method(sub_cls, name):
+        return True
+
+    this_method = getattr(sub_cls, name)
+    base_method = getattr(base_cls, name)
+    return this_method is not base_method
+
+
+def convert_to_subclass(obj, subclass):
+    obj.__class__ = subclass
+    return obj
+
+
+def _walk_to_root(path: str) -> Iterator[str]:
+    """Yield directories starting from the given directory up to the root."""
+    if not os.path.exists(path):
+        yield ''
+
+    if os.path.isfile(path):
+        path = os.path.dirname(path)
+
+    last_dir = None
+    current_dir = os.path.abspath(path)
+    while last_dir != current_dir:
+        yield current_dir
+        parent_dir = os.path.abspath(os.path.join(current_dir, os.path.pardir))
+        last_dir, current_dir = current_dir, parent_dir
+
+
+def find_file(filename: str) -> str:
+    """Find file from the folders for the given file.
+
+    NOTE: Current running path priority, followed by the execution file path, and finally the aworld package path.
+
+    Args:
+        filename: The file name that you want to search.
+    """
+
+    def run_dir():
+        try:
+            main = __import__('__main__', None, None, fromlist=['__file__'])
+            return os.path.dirname(main.__file__)
+        except ModuleNotFoundError:
+            return os.getcwd()
+
+    path = os.getcwd()
+    if os.path.exists(os.path.join(path, filename)):
+        path = os.getcwd()
+    elif os.path.exists(os.path.join(run_dir(), filename)):
+        path = run_dir()
+    else:
+        frame = inspect.currentframe()
+        current_file = __file__
+
+        while frame.f_code.co_filename == current_file or not os.path.exists(
+                frame.f_code.co_filename
+        ):
+            assert frame.f_back is not None
+            frame = frame.f_back
+        frame_filename = frame.f_code.co_filename
+        path = os.path.dirname(os.path.abspath(frame_filename))
+
+    for dirname in _walk_to_root(path):
+        if not dirname:
+            continue
+        check_path = os.path.join(dirname, filename)
+        if os.path.isfile(check_path):
+            return check_path
+
+    return ''
+
+
+def search_in_module(module: object, base_classes: List[type]) -> List[Tuple[str, type]]:
+    """Find all classes that inherit from a specific base class in the module."""
+    results = []
+    for name, obj in inspect.getmembers(module, inspect.isclass):
+        for base_class in base_classes:
+            if issubclass(obj, base_class) and obj is not base_class:
+                results.append((name, obj))
+    return results
+
+
+def _scan_package(package_name: str, base_classes: List[type], results: List[Tuple[str, type]] = []):
+    try:
+        package = sys.modules[package_name]
+    except:
+        return
+
+    try:
+        for sub_package, name, is_pkg in pkgutil.walk_packages(package.__path__):
+            try:
+                __import__(f"{package_name}.{name}")
+            except:
+                continue
+
+            if is_pkg:
+                _scan_package(package_name + "." + name, base_classes, results)
+            try:
+                module = __import__(f"{package_name}.{name}", fromlist=[name])
+                results.extend(search_in_module(module, base_classes))
+            except:
+                continue
+    except:
+        pass
+
+
+def scan_packages(package: str, base_classes: List[type]) -> List[Tuple[str, type]]:
+    results = []
+    _scan_package(package, base_classes, results)
+    return results
+
+
+class ReturnThread(threading.Thread):
+    def __init__(self, func, *args, **kwargs):
+        threading.Thread.__init__(self)
+        self.func = func
+        self.args = args
+        self.kwargs = kwargs
+        self.result = None
+        self.daemon = True
+
+    def run(self):
+        self.result = asyncio.run(self.func(*self.args, **self.kwargs))
+
+
+def asyncio_loop():
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+    return loop
+
+
+def sync_exec(async_func: Callable[..., Any], *args, **kwargs):
+    """Async function to sync execution."""
+    if not asyncio.iscoroutinefunction(async_func):
+        return async_func(*args, **kwargs)
+
+    loop = asyncio_loop()
+    if loop and loop.is_running():
+        thread = ReturnThread(async_func, *args, **kwargs)
+        thread.start()
+        thread.join()
+        result = thread.result
+    else:
+        loop = asyncio.get_event_loop()
+        result = loop.run_until_complete(async_func(*args, **kwargs))
+    return result
+
+
+def nest_dict_counter(usage: Dict[str, Union[int, Dict[str, int]]],
+                      other: Dict[str, Union[int, Dict[str, int]]],
+                      ignore_zero: bool = True):
+    """Add counts from two dicts or nest dicts."""
+    result = {}
+    for elem, count in usage.items():
+        # nest dict
+        if isinstance(count, Dict):
+            res = nest_dict_counter(usage[elem], other.get(elem, {}))
+            result[elem] = res
+            continue
+
+        newcount = count + other.get(elem, 0)
+        if not ignore_zero or newcount > 0:
+            result[elem] = newcount
+
+    for elem, count in other.items():
+        if elem not in usage and not ignore_zero:
+            result[elem] = count
+    return result
+
+
+def get_class(module_class: str):
+    import importlib
+
+    assert module_class
+    module_class = module_class.strip()
+    idx = module_class.rfind('.')
+    if idx != -1:
+        module = importlib.import_module(module_class[0:idx])
+        class_names = module_class[idx + 1:].split(":")
+        cls_obj = getattr(module, class_names[0])
+        for inner_class_name in class_names[1:]:
+            cls_obj = getattr(cls_obj, inner_class_name)
+        return cls_obj
+    else:
+        raise Exception("{} can not find!".format(module_class))
+
+
+def new_instance(module_class: str, *args, **kwargs):
+    """Create module class instance based on module name."""
+    return get_class(module_class)(*args, **kwargs)
+
+
+def retryable(tries: int = 3, delay: int = 1):
+    def inner_retry(f):
+        @wraps(f)
+        def f_retry(*args, **kwargs):
+            mtries, mdelay = tries, delay
+            while mtries > 0:
+                try:
+                    return f(*args, **kwargs)
+                except Exception as e:
+                    msg = f"{str(e)}, Retrying in {mdelay} seconds..."
+                    logger.warning(msg)
+                    time.sleep(mdelay)
+                    mtries -= 1
+            return f(*args, **kwargs)
+
+        return f_retry
+
+    return inner_retry
+
+
+def get_local_ip():
+    try:
+        # build UDP socket
+        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        # connect to an external address (no need to connect)
+        s.connect(("8.8.8.8", 80))
+        # get local IP
+        local_ip = s.getsockname()[0]
+        s.close()
+        return local_ip
+    except Exception:
+        return "127.0.0.1"
+
+
+def replace_env_variables(config) -> Any:
+    """Replace environment variables in configuration.
+
+    Environment variables should be in the format ${ENV_VAR_NAME}.
+
+    Args:
+        config: Configuration to process (dict, list, or other value)
+
+    Returns:
+        Processed configuration with environment variables replaced
+    """
+    if isinstance(config, dict):
+        for key, value in config.items():
+            if isinstance(value, str):
+                if "${" in value and "}" in value:
+                    pattern = r'\${([^}]+)}'
+                    matches = re.findall(pattern, value)
+                    result = value
+                    for env_var_name in matches:
+                        env_var_value = os.getenv(env_var_name, f"${{{env_var_name}}}")
+                        result = result.replace(f"${{{env_var_name}}}", env_var_value)
+                    config[key] = result
+                    logger.info(f"Replaced {value} with {config[key]}")
+            if isinstance(value, dict) or isinstance(value, list):
+                replace_env_variables(value)
+    elif isinstance(config, list):
+        for index, item in enumerate(config):
+            if isinstance(item, str):
+                if "${" in item and "}" in item:
+                    pattern = r'\${([^}]+)}'
+                    matches = re.findall(pattern, item)
+                    result = item
+                    for env_var_name in matches:
+                        env_var_value = os.getenv(env_var_name, f"${{{env_var_name}}}")
+                        result = result.replace(f"${{{env_var_name}}}", env_var_value)
+                    config[index] = result
+                    logger.info(f"Replaced {item} with {config[index]}")
+            if isinstance(item, dict) or isinstance(item, list):
+                replace_env_variables(item)
+    return config
+
+
+def get_local_hostname():
+    """
+    Get the local hostname.
+    First try `socket.gethostname()`, if it fails or returns an invalid value,
+    then try reverse DNS lookup using local IP.
+    """
+    try:
+        hostname = socket.gethostname()
+        # Simple validation - if hostname contains '.', consider it a valid FQDN (Fully Qualified Domain Name)
+        if hostname and '.' in hostname:
+            return hostname
+
+        # If hostname is not qualified, try reverse lookup via IP
+        local_ip = get_local_ip()
+        if local_ip:
+            try:
+                # Get hostname from IP
+                hostname, _, _ = socket.gethostbyaddr(local_ip)
+                return hostname
+            except (socket.herror, socket.gaierror):
+                # Reverse lookup failed, return original hostname or IP
+                pass
+
+        # If all methods fail, return original gethostname() result or IP
+        return hostname if hostname else local_ip
+
+    except Exception:
+        # Final fallback strategy
+        return "localhost"

aworld/utils/import_package.py

@@ -0,0 +1,256 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import os.path
+import time
+import sys
+import importlib
+import subprocess
+from importlib import metadata
+from aworld.logs.util import logger
+
+
+class ModuleAlias:
+    def __init__(self, module):
+        self.module = module
+
+    def __getattr__(self, name):
+        return getattr(self.module, name)
+
+
+def is_package_installed(package_name: str, version: str = "") -> bool:
+    """
+    Check if package is already installed and matches version if specified.
+
+    Args:
+        package_name: Name of the package to check
+        version: Required version of the package
+
+    Returns:
+        bool: True if package is installed (and version matches if specified), False otherwise
+    """
+    try:
+        dist = metadata.distribution(package_name)
+
+        if version and dist.version != version:
+            logger.info(f"Package {package_name} is installed but version {dist.version} "
+                        f"does not match required version {version}")
+            return False
+
+        logger.info(f"Package {package_name} is already installed (version: {dist.version})")
+        return True
+
+    except metadata.PackageNotFoundError:
+        logger.info(f"Package {package_name} is not installed")
+        return False
+    except Exception as e:
+        logger.warning(f"Error checking if {package_name} is installed: {str(e)}")
+        return False
+
+
+def import_packages(packages: list[str]) -> dict:
+    """
+    Import and install multiple packages
+
+    Args:
+        packages: List of packages to import
+
+    Returns:
+        dict: Dictionary mapping package names to imported modules
+    """
+    modules = {}
+    for package in packages:
+        package_ = import_package(package)
+        if package_:
+            modules[package] = package_
+    return modules
+
+
+def import_package(
+        package_name: str,
+        alias: str = '',
+        install_name: str = '',
+        version: str = '',
+        installer: str = 'pip',
+        timeout: int = 300,
+        retry_count: int = 3,
+        retry_delay: int = 5
+) -> object:
+    """
+    Import and install package if not available.
+
+    Args:
+        package_name: Name of the package to import
+        alias: Alias to use for the imported module
+        install_name: Name of the package to install (if different from import name)
+        version: Required version of the package
+        installer: Package installer to use ('pip' or 'conda')
+        timeout: Installation timeout in seconds
+        retry_count: Number of installation retries if install fails
+        retry_delay: Delay between retries in seconds
+
+    Returns:
+        Imported module
+
+    Raises:
+        ValueError: If input parameters are invalid
+        ImportError: If package cannot be imported or installed
+        TimeoutError: If installation exceeds timeout
+    """
+    # Validate input parameters
+    if not package_name:
+        raise ValueError("Package name cannot be empty")
+
+    if installer not in ['pip', 'conda']:
+        raise ValueError(f"Unsupported installer: {installer}")
+
+    # Use package_name as install_name if not provided
+    real_install_name = install_name if install_name else package_name
+
+    # First, check if we need to install the package
+    need_install = False
+
+    # Try to import the module first
+    try:
+        logger.debug(f"Attempting to import {package_name}")
+        module = importlib.import_module(package_name)
+        logger.debug(f"Successfully imported {package_name}")
+
+        # If we successfully imported the module, check version if specified
+        if version:
+            try:
+                # For packages with different import and install names,
+                # we need to check the install name for version info
+                installed_version = metadata.version(real_install_name)
+                if installed_version != version:
+                    logger.warning(
+                        f"Package {real_install_name} version mismatch. "
+                        f"Required: {version}, Installed: {installed_version}"
+                    )
+                    need_install = True
+            except metadata.PackageNotFoundError:
+                logger.warning(f"Could not determine version for {real_install_name}")
+
+        # If no need to reinstall for version mismatch, return the module
+        if not need_install:
+            return ModuleAlias(module) if alias else module
+
+    except ImportError as import_err:
+        logger.info(f"Could not import {package_name}: {str(import_err)}")
+        # Check if the package is installed
+        if not is_package_installed(real_install_name, version):
+            need_install = True
+        else:
+            # If package is installed but import failed, there might be an issue with dependencies
+            # or the package itself. Still, let's try to reinstall it.
+            logger.warning(f"Package {real_install_name} is installed but import of {package_name} failed. "
+                           f"Will attempt reinstallation.")
+            need_install = True
+
+    # Install the package if needed
+    if need_install:
+        logger.info(f"Installation needed for {real_install_name}")
+
+        # Attempt installation with retries
+        for attempt in range(retry_count):
+            try:
+                cmd = _get_install_command(installer, real_install_name, version)
+                logger.info(f"Installing {real_install_name} with command: {' '.join(cmd)}")
+                _execute_install_command(cmd, timeout)
+
+                # Break out of retry loop if installation succeeds
+                break
+
+            except (ImportError, TimeoutError, subprocess.SubprocessError) as e:
+                if attempt < retry_count - 1:
+                    logger.warning(
+                        f"Installation attempt {attempt + 1} failed: {str(e)}. Retrying in {retry_delay} seconds...")
+                    time.sleep(retry_delay)
+                else:
+                    logger.error(f"All installation attempts failed for {real_install_name}")
+                    raise ImportError(f"Failed to install {real_install_name} after {retry_count} attempts: {str(e)}")
+
+    # Try importing after installation
+    try:
+        logger.debug(f"Attempting to import {package_name} after installation")
+        module = importlib.import_module(package_name)
+        logger.debug(f"Successfully imported {package_name}")
+        return ModuleAlias(module) if alias else module
+    except ImportError as e:
+        error_msg = f"Failed to import {package_name} even after installation of {real_install_name}: {str(e)}"
+        logger.error(error_msg)
+
+
+def _get_install_command(installer: str, package_name: str, version: str = "") -> list:
+    """
+    Generate installation command based on specified installer.
+
+    Args:
+        installer: Package installer to use ('pip' or 'conda')
+        package_name: Name of the package to install
+        version: Required version of the package
+
+    Returns:
+        list: Command as a list of strings
+
+    Raises:
+        ValueError: If unsupported installer is specified
+    """
+    if installer == 'pip':
+        # Use sys.executable to ensure the right Python interpreter is used
+        pytho3 = os.path.basename(sys.executable)
+        cmd = [sys.executable, '-m', 'pip', 'install', '--upgrade']
+        if version:
+            cmd.append(f'{package_name}=={version}')
+        else:
+            cmd.append(package_name)
+    elif installer == 'conda':
+        cmd = ['conda', 'install', '-y', package_name]
+        if version:
+            cmd.extend([f'={version}'])
+    else:
+        raise ValueError(f"Unsupported installer: {installer}")
+
+    return cmd
+
+
+def _execute_install_command(cmd: list, timeout: int) -> None:
+    """
+    Execute package installation command.
+
+    Args:
+        cmd: Installation command as list of strings
+        timeout: Installation timeout in seconds
+
+    Raises:
+        TimeoutError: If installation exceeds timeout
+        ImportError: If installation fails
+    """
+    logger.info(f"Executing: {' '.join(cmd)}")
+
+    process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE
+    )
+
+    try:
+        stdout, stderr = process.communicate(timeout=timeout)
+
+        # Log installation output for debugging
+        if stdout:
+            logger.debug(f"Installation stdout: {stdout.decode()}")
+        if stderr:
+            logger.debug(f"Installation stderr: {stderr.decode()}")
+
+    except subprocess.TimeoutExpired:
+        process.kill()
+        error_msg = f"Package installation timed out after {timeout} seconds"
+        logger.error(error_msg)
+        raise TimeoutError(error_msg)
+
+    if process.returncode != 0:
+        error_msg = f"Installation failed with code {process.returncode}: {stderr.decode()}"
+        logger.error(error_msg)
+        raise ImportError(error_msg)
+
+    logger.info("Installation completed successfully")

aworld/utils/json_encoder.py

@@ -0,0 +1,12 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import json
+
+import numpy as np
+
+
+class NumpyEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return super().default(obj)

aworld/utils/oss.py

@@ -0,0 +1,628 @@
+# coding: utf-8
+"""
+oss.py
+Utility class for OSS (Object Storage Service) operations.
+Provides simple methods for data operations: upload, read, delete, update.
+"""
+import os
+import json
+import tempfile
+from typing import Optional, Dict, List, Any, Tuple, Union, BinaryIO, TextIO, IO, AnyStr
+
+from aworld.utils import import_package
+from aworld.logs.util import logger
+
+
+class OSSClient:
+    """
+    A utility class for OSS (Object Storage Service) operations.
+    Provides methods for data operations: upload, read, delete, update.
+    """
+
+    def __init__(self,
+                 access_key_id: Optional[str] = None,
+                 access_key_secret: Optional[str] = None,
+                 endpoint: Optional[str] = None,
+                 bucket_name: Optional[str] = None,
+                 enable_export: Optional[bool] = None):
+        """
+        Initialize OSSClient with credentials.
+
+        Args:
+            access_key_id: OSS access key ID. If None, will try to get from environment variable OSS_ACCESS_KEY_ID
+            access_key_secret: OSS access key secret. If None, will try to get from environment variable OSS_ACCESS_KEY_SECRET
+            endpoint: OSS endpoint. If None, will try to get from environment variable OSS_ENDPOINT
+            bucket_name: OSS bucket name. If None, will try to get from environment variable OSS_BUCKET_NAME
+            enable_export: Whether to enable OSS export. If None, will try to get from environment variable EXPORT_REPLAY_TRACE_TO_OSS
+        """
+        self.access_key_id = access_key_id or os.getenv('OSS_ACCESS_KEY_ID')
+        self.access_key_secret = access_key_secret or os.getenv('OSS_ACCESS_KEY_SECRET')
+        self.endpoint = endpoint or os.getenv('OSS_ENDPOINT')
+        self.bucket_name = bucket_name or os.getenv('OSS_BUCKET_NAME')
+        self.enable_export = enable_export if enable_export is not None else os.getenv("EXPORT_REPLAY_TRACE_TO_OSS",
+                                                                                       "false").lower() == "true"
+        self.bucket = None
+        self._initialized = False
+
+    def initialize(self) -> bool:
+        """
+        Initialize the OSS client with the provided or environment credentials.
+
+        Returns:
+            bool: True if initialization is successful, False otherwise
+        """
+        if self._initialized:
+            return True
+
+        if not self.enable_export:
+            logger.info("OSS export is disabled. Set EXPORT_REPLAY_TRACE_TO_OSS=true to enable.")
+            return False
+
+        if not all([self.access_key_id, self.access_key_secret, self.endpoint, self.bucket_name]):
+            logger.warn(
+                "Missing required OSS credentials. Please provide all required parameters or set environment variables.")
+            return False
+
+        try:
+            import_package("oss2")
+            import oss2
+            auth = oss2.Auth(self.access_key_id, self.access_key_secret)
+            self.bucket = oss2.Bucket(auth, self.endpoint, self.bucket_name)
+            self._initialized = True
+            return True
+        except ImportError:
+            logger.warn("Failed to import oss2 module. Please install it with 'pip install oss2'.")
+            return False
+        except Exception as e:
+            logger.warn(f"Failed to initialize OSS client. Error: {str(e)}")
+            return False
+
+    # ---- Basic Data Operation Methods ----
+
+    def upload_data(self, data: Union[IO[AnyStr], str, bytes, dict], oss_key: str) -> Optional[str]:
+        """
+        Upload data to OSS. Supports various types of data:
+        - In-memory file objects (IO[AnyStr])
+        - Strings (str)
+        - Bytes (bytes)
+        - Dictionaries (dict), will be automatically converted to JSON
+        - File paths (str)
+
+        Args:
+            data: Data to upload, can be a file object or other supported types
+            oss_key: The key (path) in OSS where the data will be stored
+
+        Returns:
+            str: The OSS key if successful, None otherwise
+        """
+        if not self.initialize():
+            logger.warn("OSS client not initialized or export is disabled")
+            return None
+
+        try:
+            # Handle file objects
+            if hasattr(data, 'read'):
+                content = data.read()
+                if isinstance(content, str):
+                    content = content.encode('utf-8')
+                self.bucket.put_object(oss_key, content)
+                logger.info(f"Successfully uploaded memory file to OSS: {oss_key}")
+                return oss_key
+
+            # Handle dictionaries
+            if isinstance(data, dict):
+                content = json.dumps(data, ensure_ascii=False).encode('utf-8')
+                self.bucket.put_object(oss_key, content)
+                return oss_key
+
+            # Handle strings
+            if isinstance(data, str):
+                # Check if it's a file path
+                if os.path.isfile(data):
+                    self.bucket.put_object_from_file(oss_key, data)
+                    logger.info(f"Successfully uploaded file {data} to OSS: {oss_key}")
+                    return oss_key
+                # Otherwise treat as string content
+                content = data.encode('utf-8')
+                self.bucket.put_object(oss_key, content)
+                return oss_key
+
+            # Handle bytes
+            self.bucket.put_object(oss_key, data)
+            logger.info(f"Successfully uploaded data to OSS: {oss_key}")
+            return oss_key
+        except Exception as e:
+            logger.warn(f"Failed to upload data to OSS: {str(e)}")
+            return None
+
+    def read_data(self, oss_key: str, as_json: bool = False) -> Union[bytes, dict, str, None]:
+        """
+        Read data from OSS.
+
+        Args:
+            oss_key: The key (path) in OSS of the data to read
+            as_json: If True, parse the data as JSON and return a dict
+
+        Returns:
+            The data as bytes, dict (if as_json=True), or None if failed
+        """
+        if not self.initialize():
+            logger.warn("OSS client not initialized or export is disabled")
+            return None
+
+        try:
+            # Read data
+            result = self.bucket.get_object(oss_key)
+            data = result.read()
+
+            # Convert to string or JSON if requested
+            if as_json:
+                return json.loads(data)
+
+            return data
+        except Exception as e:
+            logger.warn(f"Failed to read data from OSS: {str(e)}")
+            return None
+
+    def read_text(self, oss_key: str) -> Optional[str]:
+        """
+        Read text data from OSS.
+
+        Args:
+            oss_key: The key (path) in OSS of the text to read
+
+        Returns:
+            str: The text data, or None if failed
+        """
+        data = self.read_data(oss_key)
+        if data is not None:
+            try:
+                return data.decode('utf-8')
+            except Exception as e:
+                logger.warn(f"Failed to decode data as UTF-8: {str(e)}")
+        return None
+
+    def delete_data(self, oss_key: str) -> bool:
+        """
+        Delete data from OSS.
+
+        Args:
+            oss_key: The key (path) in OSS of the data to delete
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        if not self.initialize():
+            logger.warn("OSS client not initialized or export is disabled")
+            return False
+
+        try:
+            self.bucket.delete_object(oss_key)
+            logger.info(f"Successfully deleted data from OSS: {oss_key}")
+            return True
+        except Exception as e:
+            logger.warn(f"Failed to delete data from OSS: {str(e)}")
+            return False
+
+    def update_data(self, oss_key: str, data: Union[IO[AnyStr], str, bytes, dict]) -> Optional[str]:
+        """
+        Update data in OSS (delete and upload).
+
+        Args:
+            oss_key: The key (path) in OSS of the data to update
+            data: New data to upload, can be a file object or other supported types
+
+        Returns:
+            str: The OSS key if successful, None otherwise
+        """
+        # For OSS, update is the same as upload (it overwrites)
+        return self.upload_data(data, oss_key)
+
+    def update_json(self, oss_key: str, update_dict: dict) -> Optional[str]:
+        """
+        Update JSON data in OSS by merging with existing data.
+
+        Args:
+            oss_key: The key (path) in OSS of the JSON data to update
+            update_dict: Dictionary with fields to update
+
+        Returns:
+            str: The OSS key if successful, None otherwise
+        """
+        if not self.initialize():
+            logger.warn("OSS client not initialized or export is disabled")
+            return None
+
+        try:
+            # Read existing data
+            existing_data = self.read_data(oss_key, as_json=True)
+            if existing_data is None:
+                existing_data = {}
+
+            # Update data
+            if isinstance(existing_data, dict):
+                existing_data.update(update_dict)
+            else:
+                logger.warn(f"Existing data is not a dictionary: {oss_key}")
+                return None
+
+            # Upload updated data
+            return self.upload_data(existing_data, oss_key)
+        except Exception as e:
+            logger.warn(f"Failed to update JSON data in OSS: {str(e)}")
+            return None
+
+    # ---- File Operation Methods ----
+
+    def upload_file(self, local_file: str, oss_key: Optional[str] = None) -> Optional[str]:
+        """
+        Upload a local file to OSS.
+
+        Args:
+            local_file: Path to the local file
+            oss_key: The key (path) in OSS where the file will be stored.
+                     If None, will use the basename of the local file
+
+        Returns:
+            str: The OSS key if successful, None otherwise
+        """
+        if not self.initialize():
+            logger.warn("OSS client not initialized or export is disabled")
+            return None
+
+        try:
+            if not os.path.exists(local_file):
+                logger.warn(f"Local file {local_file} does not exist")
+                return None
+
+            if oss_key is None:
+                oss_key = f"uploads/{os.path.basename(local_file)}"
+
+            self.bucket.put_object_from_file(oss_key, local_file)
+            logger.info(f"Successfully uploaded {local_file} to OSS: {oss_key}")
+            return oss_key
+        except Exception as e:
+            logger.warn(f"Failed to upload {local_file} to OSS: {str(e)}")
+            return None
+
+    def download_file(self, oss_key: str, local_file: str) -> bool:
+        """
+        Download a file from OSS to local.
+
+        Args:
+            oss_key: The key (path) in OSS of the file to download
+            local_file: Path where the downloaded file will be saved
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        if not self.initialize():
+            logger.warn("OSS client not initialized or export is disabled")
+            return False
+
+        try:
+            # Ensure the directory exists
+            os.makedirs(os.path.dirname(os.path.abspath(local_file)), exist_ok=True)
+
+            # Download the file
+            self.bucket.get_object_to_file(oss_key, local_file)
+            logger.info(f"Successfully downloaded {oss_key} to {local_file}")
+            return True
+        except Exception as e:
+            logger.warn(f"Failed to download {oss_key} from OSS: {str(e)}")
+            return False
+
+    def list_objects(self, prefix: str = "", delimiter: str = "") -> List[Dict[str, Any]]:
+        """
+        List objects in the OSS bucket with the given prefix.
+
+        Args:
+            prefix: Prefix to filter objects
+            delimiter: Delimiter for hierarchical listing
+
+        Returns:
+            List of objects with their properties
+        """
+        if not self.initialize():
+            logger.warn("OSS client not initialized or export is disabled")
+            return []
+
+        try:
+            result = []
+            for obj in self.bucket.list_objects(prefix=prefix, delimiter=delimiter).object_list:
+                result.append({
+                    'key': obj.key,
+                    'size': obj.size,
+                    'last_modified': obj.last_modified
+                })
+            return result
+        except Exception as e:
+            logger.warn(f"Failed to list objects with prefix {prefix}: {str(e)}")
+            return []
+
+    # ---- Advanced Operation Methods ----
+
+    def exists(self, oss_key: str) -> bool:
+        """
+        Check if an object exists in OSS.
+
+        Args:
+            oss_key: The key (path) in OSS to check
+
+        Returns:
+            bool: True if the object exists, False otherwise
+        """
+        if not self.initialize():
+            logger.warn("OSS client not initialized or export is disabled")
+            return False
+
+        try:
+            # Use head_object to check if the object exists
+            self.bucket.head_object(oss_key)
+            return True
+        except:
+            return False
+
+    def copy_object(self, source_key: str, target_key: str) -> bool:
+        """
+        Copy an object within the same bucket.
+
+        Args:
+            source_key: The source object key
+            target_key: The target object key
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        if not self.initialize():
+            logger.warn("OSS client not initialized or export is disabled")
+            return False
+
+        try:
+            self.bucket.copy_object(self.bucket_name, source_key, target_key)
+            logger.info(f"Successfully copied {source_key} to {target_key}")
+            return True
+        except Exception as e:
+            logger.warn(f"Failed to copy {source_key} to {target_key}: {str(e)}")
+            return False
+
+    def get_object_url(self, oss_key: str, expires: int = 3600) -> Optional[str]:
+        """
+        Generate a temporary URL for accessing an object.
+
+        Args:
+            oss_key: The key (path) in OSS of the object
+            expires: URL expiration time in seconds
+
+        Returns:
+            str: The signed URL, or None if failed
+        """
+        if not self.initialize():
+            logger.warn("OSS client not initialized or export is disabled")
+            return None
+
+        try:
+            url = self.bucket.sign_url('GET', oss_key, expires)
+            return url
+        except Exception as e:
+            logger.warn(f"Failed to generate URL for {oss_key}: {str(e)}")
+            return None
+
+    def upload_directory(self, local_dir: str, oss_prefix: str = "") -> Tuple[bool, List[str]]:
+        """
+        Upload an entire directory to OSS.
+
+        Args:
+            local_dir: Path to the local directory
+            oss_prefix: Prefix to prepend to all uploaded files
+
+        Returns:
+            Tuple of (success, list of uploaded files)
+        """
+        if not self.initialize():
+            logger.warn("OSS client not initialized or export is disabled")
+            return False, []
+
+        if not os.path.isdir(local_dir):
+            logger.warn(f"Local directory {local_dir} does not exist or is not a directory")
+            return False, []
+
+        uploaded_files = []
+        errors = []
+
+        for root, _, files in os.walk(local_dir):
+            for file in files:
+                local_file = os.path.join(root, file)
+                rel_path = os.path.relpath(local_file, local_dir)
+                oss_key = os.path.join(oss_prefix, rel_path).replace("\\", "/")
+
+                result = self.upload_file(local_file, oss_key)
+                if result:
+                    uploaded_files.append(result)
+                else:
+                    errors.append(local_file)
+
+        if errors:
+            logger.warn(f"Failed to upload {len(errors)} files")
+            return False, uploaded_files
+        return True, uploaded_files
+
+
+def get_oss_client(access_key_id: Optional[str] = None,
+                   access_key_secret: Optional[str] = None,
+                   endpoint: Optional[str] = None,
+                   bucket_name: Optional[str] = None,
+                   enable_export: Optional[bool] = None) -> OSSClient:
+    """
+    Factory function to create and initialize an OSSClient.
+
+    Args:
+        access_key_id: OSS access key ID
+        access_key_secret: OSS access key secret
+        endpoint: OSS endpoint
+        bucket_name: OSS bucket name
+        enable_export: Whether to enable OSS export
+
+    Returns:
+        OSSClient: An initialized OSSClient instance
+    """
+    client = OSSClient(
+        access_key_id=access_key_id,
+        access_key_secret=access_key_secret,
+        endpoint=endpoint,
+        bucket_name=bucket_name,
+        enable_export=enable_export
+    )
+    client.initialize()
+    return client
+
+
+# ---- Test Cases ----
+if __name__ == "__main__":
+    """
+    OSS tool class test cases
+    Note: Before running the tests, you need to set the following environment variables,
+    or provide the parameters directly in the test code:
+    - OSS_ACCESS_KEY_ID
+    - OSS_ACCESS_KEY_SECRET
+    - OSS_ENDPOINT
+    - OSS_BUCKET_NAME
+    - EXPORT_REPLAY_TRACE_TO_OSS=true
+    """
+    import io
+    import time
+
+    # Test configuration
+    TEST_PREFIX = f"test/oss_utils_123"  # Use timestamp to avoid conflicts
+
+    # Initialize client
+    # Method 1: Using environment variables
+    # oss_client = get_oss_client(enable_export=True)
+
+    # Method 2: Provide parameters directly
+    oss_client = get_oss_client(
+        access_key_id="",  # Replace with your actual access key ID
+        access_key_secret="",  # Replace with your actual access key secret
+        endpoint="",  # Replace with your actual OSS endpoint
+        bucket_name="",  # Replace with your actual bucket name
+        enable_export=True
+    )
+
+    text_key = f"{TEST_PREFIX}/text.txt"
+    result = oss_client.upload_data("malai This is a test text", text_key)
+    print(f"Upload string data: {'Success: ' + result if result else 'Failed'}")
+    print("\nTest 6: Read text data")
+    content = oss_client.read_text(text_key)
+    print(f"Read text data: {content}")
+
+
+    # Test 1: Upload string data
+    print("\nTest 1: Upload string data")
+    text_key = f"{TEST_PREFIX}/text.txt"
+    result = oss_client.upload_data("This is a test text", text_key)
+    print(f"Upload string data: {'Success: ' + result if result else 'Failed'}")
+
+    # Test 2: Upload dictionary data (automatically converted to JSON)
+    print("\nTest 2: Upload dictionary data")
+    json_key = f"{TEST_PREFIX}/data.json"
+    data = {
+        "name": "Test data",
+        "values": [1, 2, 3],
+        "nested": {
+            "key": "value"
+        }
+    }
+    result = oss_client.upload_data(data, json_key)
+    print(f"Upload dictionary data: {'Success: ' + result if result else 'Failed'}")
+
+    # Test 3: Upload in-memory binary file object
+    print("\nTest 3: Upload in-memory binary file object")
+    binary_key = f"{TEST_PREFIX}/binary.dat"
+    binary_data = io.BytesIO(b"\x00\x01\x02\x03\x04")
+    result = oss_client.upload_data(binary_data, binary_key)
+    print(f"Upload binary file object: {'Success: ' + result if result else 'Failed'}")
+
+    # Test 4: Upload in-memory text file object
+    print("\nTest 4: Upload in-memory text file object")
+    text_file_key = f"{TEST_PREFIX}/text_file.txt"
+    text_file = io.StringIO("This is the content of an in-memory text file")
+    result = oss_client.upload_data(text_file, text_file_key)
+    print(f"Upload text file object: {'Success: ' + result if result else 'Failed'}")
+
+    # Test 5: Create and upload temporary file
+    print("\nTest 5: Create and upload temporary file")
+    with tempfile.NamedTemporaryFile(delete=False) as tmp:
+        tmp.write(b"This is the content of a temporary file")
+        tmp_path = tmp.name
+
+    file_key = f"{TEST_PREFIX}/temp_file.txt"
+    result = oss_client.upload_file(tmp_path, file_key)
+    print(f"Upload temporary file: {'Success: ' + result if result else 'Failed'}")
+    os.unlink(tmp_path)  # Delete temporary file
+
+    # Test 6: Read text data
+    print("\nTest 6: Read text data")
+    content = oss_client.read_text(text_key)
+    print(f"Read text data: {content}")
+
+    # Test 7: Read JSON data
+    print("\nTest 7: Read JSON data")
+    json_content = oss_client.read_data(json_key, as_json=True)
+    print(f"Read JSON data: {json_content}")
+
+    # Test 8: Update JSON data (merge method)
+    print("\nTest 8: Update JSON data")
+    update_data = {"updated": True, "timestamp": time.time()}
+    result = oss_client.update_json(json_key, update_data)
+    print(f"Update JSON data: {'Success: ' + result if result else 'Failed'}")
+
+    # View updated JSON data
+    updated_json = oss_client.read_data(json_key, as_json=True)
+    print(f"Updated JSON data: {updated_json}")
+
+    # Test 9: Overwrite existing data
+    print("\nTest 9: Overwrite existing data")
+    result = oss_client.upload_data("This is the overwritten text", text_key)
+    print(f"Overwrite existing data: {'Success: ' + result if result else 'Failed'}")
+
+    # View overwritten data
+    new_content = oss_client.read_text(text_key)
+    print(f"Overwritten text data: {new_content}")
+
+    # Test 10: List objects
+    print("\nTest 10: List objects")
+    objects = oss_client.list_objects(prefix=TEST_PREFIX)
+    print(f"Found {len(objects)} objects:")
+    for obj in objects:
+        print(f"  - {obj['key']} (Size: {obj['size']} bytes, Modified: {obj['last_modified']})")
+
+    # Test 11: Generate temporary URL
+    print("\nTest 11: Generate temporary URL")
+    url = oss_client.get_object_url(text_key, expires=300)  # 5 minutes expiration
+    print(f"Temporary URL: {url}")
+
+    # Test 12: Copy object
+    print("\nTest 12: Copy object")
+    copy_key = f"{TEST_PREFIX}/copy_of_text.txt"
+    result = oss_client.copy_object(text_key, copy_key)
+    print(f"Copy object: {'Success: ' + copy_key if result else 'Failed'}")
+
+    # Test 13: Check if object exists
+    print("\nTest 13: Check if object exists")
+    exists = oss_client.exists(text_key)
+    print(f"Object {text_key} exists: {exists}")
+
+    non_existent_key = f"{TEST_PREFIX}/non_existent.txt"
+    exists = oss_client.exists(non_existent_key)
+    print(f"Object {non_existent_key} exists: {exists}")
+
+    # Test 14: Delete objects
+    print("\nTest 14: Delete objects")
+    for obj in objects:
+        success = oss_client.delete_data(obj['key'])
+        print(f"Delete object {obj['key']}: {'Success' if success else 'Failed'}")
+
+    # Cleanup: Delete copied object (may not be included in the previous list)
+    oss_client.delete_data(copy_key)
+
+    print("\nTests completed!")