aicoder-max-Pillar_For_ComfyUI/client/base_client.py

"""
Base client class for the ComfyUI Extension Service.
This class provides the foundation for service-specific clients 
with common functionality for API communication.
"""
import json
import logging
import socket
import uuid
from enum import Enum
from typing import Dict, Any

import requests

from .exceptions import APIError, RateLimitError, ServiceUnavailableError, ValidationError

logger = logging.getLogger(__name__)

class HttpMethod(str, Enum):
    """HTTP request methods."""
    GET = "GET"
    POST = "POST"

class Endpoint(str, Enum):
    """Common API endpoints."""
    HEALTH_CHECK = "health/direct"
    CLEAR_CACHE = "admin/clear-cache"
    MEMORY_CLEANUP = "admin/cleanup-memory"


class AuthenticationError:
    pass


class BaseClient:
    """
    Base client class for the ComfyUI Extension Service.
    This class provides the foundation for service-specific clients 
    with common functionality for API communication.
    """
    DEFAULT_TIMEOUT = 60
    DEFAULT_USERNAME = "anonymous"
    DEFAULT_IP = "127.0.0.1"
    DEFAULT_HOSTNAME = "localhost"
    CONTENT_TYPE_JSON = "application/json"

    # Map HTTP status codes to exception classes
    ERROR_STATUS_MAP = {
        401: AuthenticationError,
        422: ValidationError,
        429: RateLimitError,
    }

    def __init__(
            self,
            username: str = None,
            timeout: int = DEFAULT_TIMEOUT,
            headers: Dict[str, str] = None,
    ):
        """
        Initialize the base client.
        
        Args:
            username: The username to use for API calls
            timeout: Request timeout in seconds
            headers: Additional headers to include in requests
        """
        self.username = username or self.DEFAULT_USERNAME
        self.timeout = timeout
        self.headers = headers or {}
        self.headers.update({"Content-Type": self.CONTENT_TYPE_JSON})

        # Get client IP address
        self._setup_client_info()

    def _ensure_url_prefix(self,url: str) -> str:
        if not url:
            return url

        # Remove leading/trailing whitespace
        url = url.strip()

        # If URL doesn't start with http:// or https://, add http://
        if not url.startswith("http://") and not url.startswith("https://"):
            url = "http://" + url

        # Convert https:// to http:// for consistency
        if url.startswith("https://"):
            url = "http://" + url[8:]

        return url

    def _setup_client_info(self) -> None:
        """Set up client hostname and IP address information."""
        try:
            self.hostname = socket.gethostname()
            self.ip_address = socket.gethostbyname(self.hostname)
        except Exception as e:
            logger.warning(f"Failed to get client IP address: {e}")
            self.ip_address = self.DEFAULT_IP
            self.hostname = self.DEFAULT_HOSTNAME

    @staticmethod
    def _build_url(base_url: str, endpoint: str) -> str:
        """
        Build a complete URL from .e base URL and endpoint.

        Args:
            base_url: The base URL of the API server
            endpoint: The API endpoint

        Returns:
            The complete URL
        """
        return f"{base_url.rstrip('/')}/{endpoint.lstrip('/')}"

    def _prepare_request_data(self, request_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Prepare request data by adding required fields.
        
        Args:
            request_data: The request data
            
        Returns:
            The prepared request data
        """
        # Create a new dictionary to avoid modifying the original
        prepared_data = request_data.copy()

        # Add required fields if not already present
        if "req_id" not in prepared_data:
            prepared_data["req_id"] = str(uuid.uuid4())
        if "user_name" not in prepared_data:
            prepared_data["user_name"] = self.username
        if "ip_address" not in prepared_data:
            prepared_data["ip_address"] = self.ip_address

        return prepared_data

    def _handle_error_status(self, response: requests.Response) -> None:
        """
        Handle error status codes from the API.
        
        Args:
            response: The response object
            
        Raises:
            APIError: If the API returns an error
            ValidationError: If the response cannot be parsed
        """
        try:
            data = response.json()
        except json.JSONDecodeError:
            raise ValidationError(f"Invalid JSON response: {response.text}")

        # Check if the response is a BaseResponse
        if isinstance(data, dict) and "success" in data:
            if not data["success"]:
                raise APIError(
                    status_code=response.status_code,
                    message=data.get("msg", "Unknown error"),
                    response=data,
                )

        # Check HTTP status
        if response.status_code >= 400:
            error_message = data.get("detail", "Unknown error")

            # Use the error status map to get the appropriate exception class
            exception_class = self.ERROR_STATUS_MAP.get(
                response.status_code,
                ServiceUnavailableError if response.status_code >= 500 else APIError
            )

            if exception_class == APIError:
                raise exception_class(
                    status_code=response.status_code,
                    message=error_message,
                    response=data,
                )
            else:
                raise exception_class(f"{exception_class.__name__}: {error_message}")

    def _request(
            self,
            base_url: str,
            method: HttpMethod,
            endpoint: str,
            data: Dict[str, Any] = None,
            params: Dict[str, Any] = None,
            files: Dict[str, Any] = None,
            headers: Dict[str, str] = None,
    ) -> Dict[str, Any]:
        logger.debug(f"base_url:{base_url}")
        url = self._ensure_url_prefix(base_url)
        logger.debug(f"_ensure_url_prefix:{url}")
        url = self._build_url(url, endpoint)
        logger.debug(f"_build_url:{url}")
        request_headers = self.headers.copy()

        if headers:
            request_headers.update(headers)

        # Prepare request data if it's a dict and not a file upload
        if data and isinstance(data, dict) and not files:
            data = self._prepare_request_data(data)

        try:
            kwargs = {
                "method": method.value,
                "url": url,
                "params": params,
                "headers": request_headers,
                "timeout": self.timeout,
            }

            if files:
                if "Content-Type" in request_headers:
                    del request_headers["Content-Type"]
                kwargs["data"] = data
                kwargs["files"] = files
            else:
                kwargs["json"] = data if data else None

            response = requests.request(**kwargs)

            # Log request details before sending
            logger.info(f"Sending {method.value} request to {url}")
            logger.debug(f"Request headers: {request_headers}")
            if data and not files:  # Only log data for non-file requests
                logger.debug(f"Request data: {data}")
            if files:
                logger.debug(f"Files to upload: {list(files.keys())}")
            # Log response details before returning
            logger.info(f"Received response with status {response.status_code}")
            logger.debug(f"Response headers: {response.headers}")
            logger.debug(f"Response content: {response.text[:1000]}...")  # Limit long responses
            # Handle error status codes
            self._handle_error_status(response)
            return response.json()

        except requests.exceptions.RequestException as e:
            raise ConnectionError(f"Connection error: {str(e)}")