dify

dify/scripts/stress-test/common/__init__.py

@@ -0,0 +1,6 @@
+"""Common utilities for Dify benchmark suite."""
+
+from .config_helper import config_helper
+from .logger_helper import Logger, ProgressLogger
+
+__all__ = ["Logger", "ProgressLogger", "config_helper"]

dify/scripts/stress-test/common/config_helper.py

@@ -0,0 +1,240 @@
+#!/usr/bin/env python3
+
+import json
+from pathlib import Path
+from typing import Any
+
+
+class ConfigHelper:
+    """Helper class for reading and writing configuration files."""
+
+    def __init__(self, base_dir: Path | None = None):
+        """Initialize ConfigHelper with base directory.
+
+        Args:
+            base_dir: Base directory for config files. If None, uses setup/config
+        """
+        if base_dir is None:
+            # Default to config directory in setup folder
+            base_dir = Path(__file__).parent.parent / "setup" / "config"
+        self.base_dir = base_dir
+        self.state_file = "stress_test_state.json"
+
+    def ensure_config_dir(self) -> None:
+        """Ensure the config directory exists."""
+        self.base_dir.mkdir(exist_ok=True, parents=True)
+
+    def get_config_path(self, filename: str) -> Path:
+        """Get the full path for a config file.
+
+        Args:
+            filename: Name of the config file (e.g., 'admin_config.json')
+
+        Returns:
+            Full path to the config file
+        """
+        if not filename.endswith(".json"):
+            filename += ".json"
+        return self.base_dir / filename
+
+    def read_config(self, filename: str) -> dict[str, Any] | None:
+        """Read a configuration file.
+
+        DEPRECATED: Use read_state() or get_state_section() for new code.
+        This method provides backward compatibility.
+
+        Args:
+            filename: Name of the config file to read
+
+        Returns:
+            Dictionary containing config data, or None if file doesn't exist
+        """
+        # Provide backward compatibility for old config names
+        if filename in ["admin_config", "token_config", "app_config", "api_key_config"]:
+            section_map = {
+                "admin_config": "admin",
+                "token_config": "auth",
+                "app_config": "app",
+                "api_key_config": "api_key",
+            }
+            return self.get_state_section(section_map[filename])
+
+        config_path = self.get_config_path(filename)
+
+        if not config_path.exists():
+            return None
+
+        try:
+            with open(config_path) as f:
+                return json.load(f)
+        except (OSError, json.JSONDecodeError) as e:
+            print(f"❌ Error reading {filename}: {e}")
+            return None
+
+    def write_config(self, filename: str, data: dict[str, Any]) -> bool:
+        """Write data to a configuration file.
+
+        DEPRECATED: Use write_state() or update_state_section() for new code.
+        This method provides backward compatibility.
+
+        Args:
+            filename: Name of the config file to write
+            data: Dictionary containing data to save
+
+        Returns:
+            True if successful, False otherwise
+        """
+        # Provide backward compatibility for old config names
+        if filename in ["admin_config", "token_config", "app_config", "api_key_config"]:
+            section_map = {
+                "admin_config": "admin",
+                "token_config": "auth",
+                "app_config": "app",
+                "api_key_config": "api_key",
+            }
+            return self.update_state_section(section_map[filename], data)
+
+        self.ensure_config_dir()
+        config_path = self.get_config_path(filename)
+
+        try:
+            with open(config_path, "w") as f:
+                json.dump(data, f, indent=2)
+            return True
+        except OSError as e:
+            print(f"❌ Error writing {filename}: {e}")
+            return False
+
+    def config_exists(self, filename: str) -> bool:
+        """Check if a config file exists.
+
+        Args:
+            filename: Name of the config file to check
+
+        Returns:
+            True if file exists, False otherwise
+        """
+        return self.get_config_path(filename).exists()
+
+    def delete_config(self, filename: str) -> bool:
+        """Delete a configuration file.
+
+        Args:
+            filename: Name of the config file to delete
+
+        Returns:
+            True if successful, False otherwise
+        """
+        config_path = self.get_config_path(filename)
+
+        if not config_path.exists():
+            return True  # Already doesn't exist
+
+        try:
+            config_path.unlink()
+            return True
+        except OSError as e:
+            print(f"❌ Error deleting {filename}: {e}")
+            return False
+
+    def read_state(self) -> dict[str, Any] | None:
+        """Read the entire stress test state.
+
+        Returns:
+            Dictionary containing all state data, or None if file doesn't exist
+        """
+        state_path = self.get_config_path(self.state_file)
+        if not state_path.exists():
+            return None
+
+        try:
+            with open(state_path) as f:
+                return json.load(f)
+        except (OSError, json.JSONDecodeError) as e:
+            print(f"❌ Error reading {self.state_file}: {e}")
+            return None
+
+    def write_state(self, data: dict[str, Any]) -> bool:
+        """Write the entire stress test state.
+
+        Args:
+            data: Dictionary containing all state data to save
+
+        Returns:
+            True if successful, False otherwise
+        """
+        self.ensure_config_dir()
+        state_path = self.get_config_path(self.state_file)
+
+        try:
+            with open(state_path, "w") as f:
+                json.dump(data, f, indent=2)
+            return True
+        except OSError as e:
+            print(f"❌ Error writing {self.state_file}: {e}")
+            return False
+
+    def update_state_section(self, section: str, data: dict[str, Any]) -> bool:
+        """Update a specific section of the stress test state.
+
+        Args:
+            section: Name of the section to update (e.g., 'admin', 'auth', 'app', 'api_key')
+            data: Dictionary containing section data to save
+
+        Returns:
+            True if successful, False otherwise
+        """
+        state = self.read_state() or {}
+        state[section] = data
+        return self.write_state(state)
+
+    def get_state_section(self, section: str) -> dict[str, Any] | None:
+        """Get a specific section from the stress test state.
+
+        Args:
+            section: Name of the section to get (e.g., 'admin', 'auth', 'app', 'api_key')
+
+        Returns:
+            Dictionary containing section data, or None if not found
+        """
+        state = self.read_state()
+        if state:
+            return state.get(section)
+        return None
+
+    def get_token(self) -> str | None:
+        """Get the access token from auth section.
+
+        Returns:
+            Access token string or None if not found
+        """
+        auth = self.get_state_section("auth")
+        if auth:
+            return auth.get("access_token")
+        return None
+
+    def get_app_id(self) -> str | None:
+        """Get the app ID from app section.
+
+        Returns:
+            App ID string or None if not found
+        """
+        app = self.get_state_section("app")
+        if app:
+            return app.get("app_id")
+        return None
+
+    def get_api_key(self) -> str | None:
+        """Get the API key token from api_key section.
+
+        Returns:
+            API key token string or None if not found
+        """
+        api_key = self.get_state_section("api_key")
+        if api_key:
+            return api_key.get("token")
+        return None
+
+
+# Create a default instance for convenience
+config_helper = ConfigHelper()

dify/scripts/stress-test/common/logger_helper.py

@@ -0,0 +1,218 @@
+#!/usr/bin/env python3
+
+import sys
+import time
+from enum import Enum
+
+
+class LogLevel(Enum):
+    """Log levels with associated colors and symbols."""
+
+    DEBUG = ("🔍", "\033[90m")  # Gray
+    INFO = ("ℹ️ ", "\033[94m")  # Blue
+    SUCCESS = ("✅", "\033[92m")  # Green
+    WARNING = ("⚠️ ", "\033[93m")  # Yellow
+    ERROR = ("❌", "\033[91m")  # Red
+    STEP = ("🚀", "\033[96m")  # Cyan
+    PROGRESS = ("📋", "\033[95m")  # Magenta
+
+
+class Logger:
+    """Logger class for formatted console output."""
+
+    def __init__(self, name: str | None = None, use_colors: bool = True):
+        """Initialize logger.
+
+        Args:
+            name: Optional name for the logger (e.g., script name)
+            use_colors: Whether to use ANSI color codes
+        """
+        self.name = name
+        self.use_colors = use_colors and sys.stdout.isatty()
+        self._reset_color = "\033[0m" if self.use_colors else ""
+
+    def _format_message(self, level: LogLevel, message: str, indent: int = 0) -> str:
+        """Format a log message with level, color, and indentation.
+
+        Args:
+            level: Log level
+            message: Message to log
+            indent: Number of spaces to indent
+
+        Returns:
+            Formatted message string
+        """
+        symbol, color = level.value
+        color = color if self.use_colors else ""
+        reset = self._reset_color
+
+        prefix = " " * indent
+
+        if self.name and level in [LogLevel.STEP, LogLevel.ERROR]:
+            return f"{prefix}{color}{symbol} [{self.name}] {message}{reset}"
+        else:
+            return f"{prefix}{color}{symbol} {message}{reset}"
+
+    def debug(self, message: str, indent: int = 0) -> None:
+        """Log debug message."""
+        print(self._format_message(LogLevel.DEBUG, message, indent))
+
+    def info(self, message: str, indent: int = 0) -> None:
+        """Log info message."""
+        print(self._format_message(LogLevel.INFO, message, indent))
+
+    def success(self, message: str, indent: int = 0) -> None:
+        """Log success message."""
+        print(self._format_message(LogLevel.SUCCESS, message, indent))
+
+    def warning(self, message: str, indent: int = 0) -> None:
+        """Log warning message."""
+        print(self._format_message(LogLevel.WARNING, message, indent))
+
+    def error(self, message: str, indent: int = 0) -> None:
+        """Log error message."""
+        print(self._format_message(LogLevel.ERROR, message, indent), file=sys.stderr)
+
+    def step(self, message: str, indent: int = 0) -> None:
+        """Log a step in a process."""
+        print(self._format_message(LogLevel.STEP, message, indent))
+
+    def progress(self, message: str, indent: int = 0) -> None:
+        """Log progress information."""
+        print(self._format_message(LogLevel.PROGRESS, message, indent))
+
+    def separator(self, char: str = "-", length: int = 60) -> None:
+        """Print a separator line."""
+        print(char * length)
+
+    def header(self, title: str, width: int = 60) -> None:
+        """Print a formatted header."""
+        if self.use_colors:
+            print(f"\n\033[1m{'=' * width}\033[0m")  # Bold
+            print(f"\033[1m{title.center(width)}\033[0m")
+            print(f"\033[1m{'=' * width}\033[0m\n")
+        else:
+            print(f"\n{'=' * width}")
+            print(title.center(width))
+            print(f"{'=' * width}\n")
+
+    def box(self, title: str, width: int = 60) -> None:
+        """Print a title in a box."""
+        border = "═" * (width - 2)
+        if self.use_colors:
+            print(f"\033[1m╔{border}╗\033[0m")
+            print(f"\033[1m║{title.center(width - 2)}║\033[0m")
+            print(f"\033[1m╚{border}╝\033[0m")
+        else:
+            print(f"╔{border}╗")
+            print(f"║{title.center(width - 2)}║")
+            print(f"╚{border}╝")
+
+    def list_item(self, item: str, indent: int = 2) -> None:
+        """Print a list item."""
+        prefix = " " * indent
+        print(f"{prefix}• {item}")
+
+    def key_value(self, key: str, value: str, indent: int = 2) -> None:
+        """Print a key-value pair."""
+        prefix = " " * indent
+        if self.use_colors:
+            print(f"{prefix}\033[1m{key}:\033[0m {value}")
+        else:
+            print(f"{prefix}{key}: {value}")
+
+    def spinner_start(self, message: str) -> None:
+        """Start a spinner (simple implementation)."""
+        sys.stdout.write(f"\r{message}... ")
+        sys.stdout.flush()
+
+    def spinner_stop(self, success: bool = True, message: str | None = None) -> None:
+        """Stop the spinner and show result."""
+        if success:
+            symbol = "✅" if message else "Done"
+            sys.stdout.write(f"\r{symbol} {message or ''}\n")
+        else:
+            symbol = "❌" if message else "Failed"
+            sys.stdout.write(f"\r{symbol} {message or ''}\n")
+        sys.stdout.flush()
+
+
+class ProgressLogger:
+    """Logger for tracking progress through multiple steps."""
+
+    def __init__(self, total_steps: int, logger: Logger | None = None):
+        """Initialize progress logger.
+
+        Args:
+            total_steps: Total number of steps
+            logger: Logger instance to use (creates new if None)
+        """
+        self.total_steps = total_steps
+        self.current_step = 0
+        self.logger = logger or Logger()
+        self.start_time = time.time()
+
+    def next_step(self, description: str) -> None:
+        """Move to next step and log it."""
+        self.current_step += 1
+        elapsed = time.time() - self.start_time
+
+        if self.logger.use_colors:
+            progress_bar = self._create_progress_bar()
+            print(f"\n\033[1m[Step {self.current_step}/{self.total_steps}]\033[0m {progress_bar}")
+            self.logger.step(f"{description} (Elapsed: {elapsed:.1f}s)")
+        else:
+            print(f"\n[Step {self.current_step}/{self.total_steps}]")
+            self.logger.step(f"{description} (Elapsed: {elapsed:.1f}s)")
+
+    def _create_progress_bar(self, width: int = 20) -> str:
+        """Create a simple progress bar."""
+        filled = int(width * self.current_step / self.total_steps)
+        bar = "█" * filled + "░" * (width - filled)
+        percentage = int(100 * self.current_step / self.total_steps)
+        return f"[{bar}] {percentage}%"
+
+    def complete(self) -> None:
+        """Mark progress as complete."""
+        elapsed = time.time() - self.start_time
+        self.logger.success(f"All steps completed! Total time: {elapsed:.1f}s")
+
+
+# Create default logger instance
+logger = Logger()
+
+
+# Convenience functions using default logger
+def debug(message: str, indent: int = 0) -> None:
+    """Log debug message using default logger."""
+    logger.debug(message, indent)
+
+
+def info(message: str, indent: int = 0) -> None:
+    """Log info message using default logger."""
+    logger.info(message, indent)
+
+
+def success(message: str, indent: int = 0) -> None:
+    """Log success message using default logger."""
+    logger.success(message, indent)
+
+
+def warning(message: str, indent: int = 0) -> None:
+    """Log warning message using default logger."""
+    logger.warning(message, indent)
+
+
+def error(message: str, indent: int = 0) -> None:
+    """Log error message using default logger."""
+    logger.error(message, indent)
+
+
+def step(message: str, indent: int = 0) -> None:
+    """Log step using default logger."""
+    logger.step(message, indent)
+
+
+def progress(message: str, indent: int = 0) -> None:
+    """Log progress using default logger."""
+    logger.progress(message, indent)