From a4aaf8a165d79192bcd1527c9bcd4dbec3217828 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Mon, 21 Apr 2025 23:37:48 -0700 Subject: [PATCH] Remove polling operations. --- comfy_api_nodes/apis/client.py | 126 --------------------------------- 1 file changed, 126 deletions(-) diff --git a/comfy_api_nodes/apis/client.py b/comfy_api_nodes/apis/client.py index 4b585fef2..05706710e 100644 --- a/comfy_api_nodes/apis/client.py +++ b/comfy_api_nodes/apis/client.py @@ -106,12 +106,8 @@ import json import requests from urllib.parse import urljoin -# Import models from your generated stubs - T = TypeVar("T", bound=BaseModel) R = TypeVar("R", bound=BaseModel) -P = TypeVar("P", bound=BaseModel) # For poll response - class EmptyRequest(BaseModel): """Base class for empty request bodies. @@ -383,125 +379,3 @@ class SynchronousOperation(Generic[T, R]): self.response = self.endpoint.response_model.model_validate(resp) logging.debug(f"[DEBUG] Parsed Response: {self.response}") return self.response - - -class TaskStatus(str, Enum): - """Enum for task status values""" - - COMPLETED = "completed" - FAILED = "failed" - PENDING = "pending" - - -class PollingOperation(Generic[T, R]): - """ - Represents an asynchronous API operation that requires polling for completion. - """ - - def __init__( - self, - poll_endpoint: ApiEndpoint[EmptyRequest, R], - completed_statuses: list, - failed_statuses: list, - status_extractor: Callable[[R], str], - request: Optional[T] = None, - api_base: str = "https://stagingapi.comfy.org", - auth_token: Optional[str] = None, - poll_interval: float = 1.0, - ): - self.poll_endpoint = poll_endpoint - self.request = request - self.api_base = api_base - self.auth_token = auth_token - self.poll_interval = poll_interval - - # Polling configuration - self.status_extractor = status_extractor or ( - lambda x: getattr(x, "status", None) - ) - self.completed_statuses = completed_statuses - self.failed_statuses = failed_statuses - - # For storing response data - self.final_response = None - self.error = None - - def execute(self, client: Optional[ApiClient] = None) -> R: - """Execute the polling operation using the provided client. If failed, raise an exception.""" - try: - if client is None: - client = ApiClient( - base_url=self.api_base, - api_key=self.auth_token, - ) - return self._poll_until_complete(client) - except Exception as e: - raise Exception(f"Error during polling: {str(e)}") - - def _check_task_status(self, response: R) -> TaskStatus: - """Check task status using the status extractor function""" - try: - status = self.status_extractor(response) - if status in self.completed_statuses: - return TaskStatus.COMPLETED - elif status in self.failed_statuses: - return TaskStatus.FAILED - return TaskStatus.PENDING - except Exception as e: - logging.debug(f"Error extracting status: {e}") - return TaskStatus.PENDING - - def _poll_until_complete(self, client: ApiClient) -> R: - """Poll until the task is complete""" - poll_count = 0 - while True: - try: - poll_count += 1 - logging.debug(f"[DEBUG] Polling attempt #{poll_count}") - - request_dict = ( - self.request.model_dump(exclude_none=True) - if self.request is not None - else None - ) - - if poll_count == 1: - logging.debug( - f"[DEBUG] Poll Request: {self.poll_endpoint.method.value} {self.poll_endpoint.path}" - ) - logging.debug( - f"[DEBUG] Poll Request Data: {json.dumps(request_dict, indent=2) if request_dict else 'None'}" - ) - - # Query task status - resp = client.request( - method=self.poll_endpoint.method.value, - path=self.poll_endpoint.path, - params=self.poll_endpoint.query_params, - json=request_dict, - ) - - # Parse response - response_obj = self.poll_endpoint.response_model.model_validate(resp) - - # Check if task is complete - status = self._check_task_status(response_obj) - logging.debug(f"[DEBUG] Task Status: {status}") - - if status == TaskStatus.COMPLETED: - logging.debug("[DEBUG] Task completed successfully") - self.final_response = response_obj - return self.final_response - elif status == TaskStatus.FAILED: - logging.debug(f"[DEBUG] Task failed: {json.dumps(resp)}") - raise Exception(f"Task failed: {json.dumps(resp)}") - else: - logging.debug("[DEBUG] Task still pending, continuing to poll...") - - # Wait before polling again - logging.debug(f"[DEBUG] Waiting {self.poll_interval} seconds before next poll") - time.sleep(self.poll_interval) - - except Exception as e: - logging.debug(f"[DEBUG] Polling error: {str(e)}") - raise Exception(f"Error while polling: {str(e)}")