airia 0.1.13__py3-none-any.whl → 0.1.14__py3-none-any.whl

airia/client/_request_handler/__init__.py

@@ -0,0 +1,4 @@
+from .async_request_handler import AsyncRequestHandler
+from .sync_request_handler import RequestHandler
+
+__all__ = ["RequestHandler", "AsyncRequestHandler"]

airia/client/_request_handler/async_request_handler.py

@@ -0,0 +1,272 @@
+import asyncio
+import weakref
+from typing import Any, AsyncIterator, Dict, Optional
+
+import aiohttp
+import loguru
+
+from ...exceptions import AiriaAPIError
+from ...types._request_data import RequestData
+from ...utils.sse_parser import async_parse_sse_stream_chunked
+from .base_request_handler import BaseRequestHandler
+
+
+class AsyncRequestHandler(BaseRequestHandler):
+    def __init__(
+        self,
+        logger: "loguru.Logger",
+        timeout: float,
+        base_url: str,
+        api_key: Optional[str] = None,
+        bearer_token: Optional[str] = None,
+        log_requests: bool = False,
+    ):
+        self.session = aiohttp.ClientSession()
+
+        self._finalizer = weakref.finalize(self, self._cleanup_session, self.session)
+
+        super().__init__(
+            logger=logger,
+            timeout=timeout,
+            api_key=api_key,
+            base_url=base_url,
+            bearer_token=bearer_token,
+            log_requests=log_requests,
+        )
+
+    @staticmethod
+    def _cleanup_session(session: aiohttp.ClientSession):
+        """Static method to clean up session - called by finalizer"""
+        if session and not session.closed:
+            # Create a new event loop if none exists
+            try:
+                loop = asyncio.get_event_loop()
+                if loop.is_closed():
+                    raise RuntimeError("Event loop is closed")
+            except RuntimeError:
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+
+            # Close the session
+            if not loop.is_running():
+                loop.run_until_complete(session.close())
+            else:
+                # If loop is running, schedule the close operation
+                asyncio.create_task(session.close())
+
+    async def close(self):
+        """
+        Closes the aiohttp session to free up system resources.
+
+        This method should be called when the RequestHandler is no longer needed to ensure
+        proper cleanup of the underlying session and its resources.
+        """
+        if self.session and not self.session.closed:
+            await self.session.close()
+
+    def _handle_exception(
+        self, e: aiohttp.ClientResponseError, url: str, correlation_id: str
+    ):
+        # Log the error response if enabled
+        if self.log_requests:
+            self.logger.error(
+                f"API Error: {e.status} {e.message}\n"
+                f"URL: {url}\n"
+                f"Correlation ID: {correlation_id}"
+            )
+
+        # Extract error details from response
+        error_message = e.message
+
+        # Make sure sensitive auth information is not included in error messages
+        sanitized_message = error_message
+        if self.api_key and self.api_key in sanitized_message:
+            sanitized_message = sanitized_message.replace(self.api_key, "[REDACTED]")
+        if self.bearer_token and self.bearer_token in sanitized_message:
+            sanitized_message = sanitized_message.replace(
+                self.bearer_token, "[REDACTED]"
+            )
+
+        # Raise custom exception with status code and sanitized message
+        raise AiriaAPIError(status_code=e.status, message=sanitized_message) from e
+
+    async def make_request(
+        self, method: str, request_data: RequestData, return_json: bool = True
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Makes an asynchronous HTTP request to the Airia API.
+
+        Args:
+            method (str): The HTTP method (e.g., 'GET', 'POST')
+            request_data: A dictionary containing the following request information:
+                - url: The endpoint URL for the request
+                - headers: HTTP headers to include in the request
+                - payload: The JSON payload/body for the request
+                - params: Optional query parameters to append to the URL
+                - files: Optional file data to be uploaded in the request body
+                - correlation_id: Unique identifier for request tracing
+            return_json (bool): Whether to return the response as JSON. Default is True.
+
+        Returns:
+            resp ([Dict[str, Any]): The JSON response from the API as a dictionary.
+
+        Raises:
+            AiriaAPIError: If the API returns an error response, with details about the error
+            aiohttp.ClientResponseError: For HTTP-related errors
+
+        Note:
+            This is an internal method used by other client methods to make API requests.
+            It handles logging, error handling, and API key redaction in error messages.
+        """
+        try:
+            # Make the request
+            async with self.session.request(
+                method=method,
+                url=request_data.url,
+                json=request_data.payload,
+                params=request_data.params,
+                headers=request_data.headers,
+                timeout=self.timeout,
+            ) as response:
+                # Log the response if enabled
+                if self.log_requests:
+                    self.logger.info(
+                        f"API Response: {response.status} {response.reason}\n"
+                        f"URL: {request_data.url}\n"
+                        f"Correlation ID: {request_data.correlation_id}"
+                    )
+
+                # Check for HTTP errors
+                response.raise_for_status()
+
+                # Return the response as a dictionary
+                if return_json:
+                    return await response.json()
+
+        except aiohttp.ClientResponseError as e:
+            self._handle_exception(e, request_data.url, request_data.correlation_id)
+
+    async def make_request_stream(
+        self, method: str, request_data: RequestData
+    ) -> AsyncIterator[str]:
+        """
+        Makes an asynchronous HTTP request to the Airia API.
+
+        Args:
+            method (str): The HTTP method (e.g., 'GET', 'POST')
+            request_data: A dictionary containing the following request information:
+                - url: The endpoint URL for the request
+                - headers: HTTP headers to include in the request
+                - payload: The JSON payload/body for the request
+                - params: Optional query parameters to append to the URL
+                - files: Optional file data to be uploaded in the request body
+                - correlation_id: Unique identifier for request tracing
+
+        Yields:
+            resp AsyncIterator[str]]: yields chunks of the response as they are received.
+
+        Raises:
+            AiriaAPIError: If the API returns an error response, with details about the error
+            aiohttp.ClientResponseError: For HTTP-related errors
+
+        Note:
+            This is an internal method used by other client methods to make API requests.
+            It handles logging, error handling, and API key redaction in error messages.
+        """
+        try:
+            # Make the request
+            async with self.session.request(
+                method=method,
+                url=request_data.url,
+                json=request_data.payload,
+                params=request_data.params,
+                headers=request_data.headers,
+                timeout=self.timeout,
+                chunked=True,
+            ) as response:
+                # Log the response if enabled
+                if self.log_requests:
+                    self.logger.info(
+                        f"API Response: {response.status} {response.reason}\n"
+                        f"URL: {request_data.url}\n"
+                        f"Correlation ID: {request_data.correlation_id}"
+                    )
+
+                # Check for HTTP errors
+                response.raise_for_status()
+
+                # Yields the response content as a stream if streaming
+                async for message in async_parse_sse_stream_chunked(
+                    response.content.iter_any()
+                ):
+                    yield message
+
+        except aiohttp.ClientResponseError as e:
+            self._handle_exception(e, request_data.url, request_data.correlation_id)
+
+    async def make_request_multipart(
+        self, method: str, request_data: RequestData, return_json: bool = True
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Makes an asynchronous HTTP request with multipart form data to the Airia API.
+
+        Args:
+            method (str): The HTTP method (e.g., 'POST')
+            request_data: A dictionary containing the following request information:
+                - url: The endpoint URL for the request
+                - headers: HTTP headers to include in the request
+                - payload: The form data payload including file content
+                - params: Optional query parameters to append to the URL
+                - files: Optional file data to be uploaded in the request body
+                - correlation_id: Unique identifier for request tracing
+            return_json (bool): Whether to return the response as JSON. Default is True.
+
+        Returns:
+            resp (Optional[Dict[str, Any]]): The JSON response from the API as a dictionary.
+
+        Raises:
+            AiriaAPIError: If the API returns an error response, with details about the error
+            aiohttp.ClientResponseError: For HTTP-related errors
+
+        Note:
+            This is an internal method used by file upload methods to make multipart requests.
+            It handles multipart form data encoding, logging, and error handling.
+        """
+        try:
+            # Prepare multipart form data
+            data = aiohttp.FormData()
+
+            # Add form fields
+            for key, value in request_data.payload.items():
+                data.add_field(key, str(value))
+
+            # Add files
+            for key, value in request_data.files.items():
+                data.add_field(key, value[1], filename=value[0], content_type=value[2])
+
+            # Make the request
+            async with self.session.request(
+                method=method,
+                url=request_data.url,
+                data=data,
+                params=request_data.params,
+                headers=request_data.headers,
+                timeout=self.timeout,
+            ) as response:
+                # Log the response if enabled
+                if self.log_requests:
+                    self.logger.info(
+                        f"API Response: {response.status} {response.reason}\n"
+                        f"URL: {request_data.url}\n"
+                        f"Correlation ID: {request_data.correlation_id}"
+                    )
+
+                # Check for HTTP errors
+                response.raise_for_status()
+
+                # Return the response as a dictionary
+                if return_json:
+                    return await response.json()
+
+        except aiohttp.ClientResponseError as e:
+            self._handle_exception(e, request_data.url, request_data.correlation_id)

airia/client/_request_handler/base_request_handler.py

@@ -0,0 +1,108 @@
+import json
+from typing import Any, Dict, Optional
+
+import loguru
+
+from ...logs import set_correlation_id
+from ...types._request_data import RequestData
+
+
+class BaseRequestHandler:
+    def __init__(
+        self,
+        logger: "loguru.Logger",
+        timeout: float,
+        base_url: str,
+        api_key: Optional[str] = None,
+        bearer_token: Optional[str] = None,
+        log_requests: bool = False,
+    ):
+        self.logger = logger
+        self.timeout = timeout
+        self.base_url = base_url
+        self.api_key = api_key
+        self.bearer_token = bearer_token
+        self.log_requests = log_requests
+
+    def close(self):
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def prepare_request(
+        self,
+        url: str,
+        payload: Optional[Dict[str, Any]] = None,
+        params: Optional[Dict[str, Any]] = None,
+        files: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+    ):
+        """
+        Prepare request data including headers, authentication, and logging.
+
+        This method sets up all the necessary components for an API request including
+        correlation ID, authentication headers, and sanitized logging.
+
+        Args:
+            url: The target URL for the request
+            payload: Optional JSON payload for the request body
+            params: Optional query parameters for the request
+            files: Optional files to be uploaded in the request body
+            correlation_id: Optional correlation ID for request tracing
+
+        Returns:
+            RequestData: A data structure containing all prepared request components
+        """
+        # Set correlation ID if provided or generate a new one
+        correlation_id = set_correlation_id(correlation_id)
+
+        # Set up base headers
+        headers = {"X-Correlation-ID": correlation_id}
+
+        # Add authentication header based on the method used
+        if self.api_key:
+            headers["X-API-KEY"] = self.api_key
+        elif self.bearer_token:
+            headers["Authorization"] = f"Bearer {self.bearer_token}"
+
+        # Log the request if enabled
+        if self.log_requests:
+            # Create a sanitized copy of headers and params for logging
+            log_headers = headers.copy()
+            log_params = params.copy() if params is not None else {}
+            log_files = (
+                {k: (v[0], "[BINARY DATA]", v[2]) for k, v in files.items()}
+                if files is not None
+                else {}
+            )
+
+            # Filter out sensitive headers
+            if "X-API-KEY" in log_headers:
+                log_headers["X-API-KEY"] = "[REDACTED]"
+            if "Authorization" in log_headers:
+                log_headers["Authorization"] = "[REDACTED]"
+
+            # Process payload for logging
+            log_payload = payload.copy() if payload is not None else {}
+            if "images" in log_payload and log_payload["images"] is not None:
+                log_payload["images"] = f"{len(log_payload['images'])} images"
+            if "files" in log_payload and log_payload["files"] is not None:
+                log_payload["files"] = f"{len(log_payload['files'])} files"
+            log_payload = json.dumps(log_payload)
+
+            self.logger.info(
+                f"API Request: POST {url}\n"
+                f"Headers: {json.dumps(log_headers)}\n"
+                f"Payload: {log_payload}\n"
+                f"Files: {log_files}\n"
+                f"Params: {json.dumps(log_params)}\n"
+            )
+
+        return RequestData(
+            **{
+                "url": url,
+                "payload": payload,
+                "headers": headers,
+                "params": params,
+                "files": files,
+                "correlation_id": correlation_id,
+            }
+        )

airia/client/_request_handler/sync_request_handler.py

@@ -0,0 +1,255 @@
+import weakref
+from typing import Any, Dict, Optional
+
+import loguru
+import requests
+
+from ...exceptions import AiriaAPIError
+from ...types._request_data import RequestData
+from ...utils.sse_parser import parse_sse_stream_chunked
+from .base_request_handler import BaseRequestHandler
+
+
+class RequestHandler(BaseRequestHandler):
+    def __init__(
+        self,
+        logger: "loguru.Logger",
+        timeout: float,
+        base_url: str,
+        api_key: Optional[str] = None,
+        bearer_token: Optional[str] = None,
+        log_requests: bool = False,
+    ):
+        # Initialize session for synchronous requests
+        self.session = requests.Session()
+        self._finalizer = weakref.finalize(self, self._cleanup_session, self.session)
+
+        super().__init__(
+            logger=logger,
+            timeout=timeout,
+            api_key=api_key,
+            base_url=base_url,
+            bearer_token=bearer_token,
+            log_requests=log_requests,
+        )
+
+    @staticmethod
+    def _cleanup_session(session: requests.Session):
+        """Static method to clean up session - called by finalizer"""
+        if session:
+            session.close()
+
+    def close(self):
+        """
+        Closes the requests session to free up system resources.
+
+        This method should be called when the RequestHandler is no longer needed to ensure
+        proper cleanup of the underlying session and its resources.
+        """
+        self.session.close()
+
+    def _handle_exception(self, e: requests.HTTPError, url: str, correlation_id: str):
+        # Log the error response if enabled
+        if self.log_requests:
+            self.logger.error(
+                f"API Error: {e.response.status_code} {e.response.reason}\n"
+                f"URL: {url}\n"
+                f"Correlation ID: {correlation_id}"
+            )
+
+        # Extract error details from response if possible
+        error_message = "API request failed"
+        try:
+            error_data = e.response.json()
+            if isinstance(error_data, dict) and "message" in error_data:
+                error_message = error_data["message"]
+            elif isinstance(error_data, dict) and "error" in error_data:
+                error_message = error_data["error"]
+        except (ValueError, KeyError):
+            # If JSON parsing fails or expected keys are missing
+            error_message = f"API request failed: {str(e)}"
+
+        # Make sure sensitive auth information is not included in error messages
+        sanitized_message = error_message
+        if self.api_key and self.api_key in sanitized_message:
+            sanitized_message = sanitized_message.replace(self.api_key, "[REDACTED]")
+        if self.bearer_token and self.bearer_token in sanitized_message:
+            sanitized_message = sanitized_message.replace(
+                self.bearer_token, "[REDACTED]"
+            )
+
+        # Raise custom exception with status code and sanitized message
+        raise AiriaAPIError(
+            status_code=e.response.status_code, message=sanitized_message
+        ) from e
+
+    def make_request(
+        self, method: str, request_data: RequestData, return_json: bool = True
+    ) -> Dict[str, Any]:
+        """
+        Makes a synchronous HTTP request to the Airia API.
+
+        Args:
+            method (str): The HTTP method (e.g., 'GET', 'POST')
+            request_data: A dictionary containing the following request information:
+                - url: The endpoint URL for the request
+                - headers: HTTP headers to include in the request
+                - payload: The JSON payload/body for the request
+                - params: Optional query parameters to append to the URL
+                - files: Optional file data to be uploaded in the request body
+                - correlation_id: Unique identifier for request tracing
+            return_json (bool): Whether to return the response as JSON. Default is True.
+
+        Returns:
+            resp (Dict[str, Any]): The JSON response from the API as a dictionary.
+
+        Raises:
+            AiriaAPIError: If the API returns an error response, with details about the error
+            requests.HTTPError: For HTTP-related errors
+
+        Note:
+            This is an internal method used by other client methods to make API requests.
+            It handles logging, error handling, and API key redaction in error messages.
+        """
+        try:
+            # Make the request
+            response = self.session.request(
+                method=method,
+                url=request_data.url,
+                json=request_data.payload,
+                params=request_data.params,
+                headers=request_data.headers,
+                timeout=self.timeout,
+            )
+
+            # Log the response if enabled
+            if self.log_requests:
+                self.logger.info(
+                    f"API Response: {response.status_code} {response.reason}\n"
+                    f"URL: {request_data.url}\n"
+                    f"Correlation ID: {request_data.correlation_id}\n"
+                )
+
+            # Check for HTTP errors
+            response.raise_for_status()
+
+            # Returns the JSON response
+            if return_json:
+                return response.json()
+
+        except requests.HTTPError as e:
+            self._handle_exception(e, request_data.url, request_data.correlation_id)
+
+    def make_request_stream(self, method: str, request_data: RequestData):
+        """
+        Makes a synchronous HTTP request to the Airia API.
+
+        Args:
+            method (str): The HTTP method (e.g., 'GET', 'POST')
+            request_data: A dictionary containing the following request information:
+                - url: The endpoint URL for the request
+                - headers: HTTP headers to include in the request
+                - payload: The JSON payload/body for the request
+                - params: Optional query parameters to append to the URL
+                - files: Optional file data to be uploaded in the request body
+                - correlation_id: Unique identifier for request tracing
+            stream (bool): If True, the response will be streamed instead of downloaded all at once
+
+        Yields:
+            resp (Iterator[str]): Yields chunks of the response as they are received.
+
+        Raises:
+            AiriaAPIError: If the API returns an error response, with details about the error
+            requests.HTTPError: For HTTP-related errors
+
+        Note:
+            This is an internal method used by other client methods to make API requests.
+            It handles logging, error handling, and API key redaction in error messages.
+        """
+        try:
+            # Make the request
+            response = self.session.request(
+                method=method,
+                url=request_data.url,
+                params=request_data.params,
+                json=request_data.payload,
+                headers=request_data.headers,
+                timeout=self.timeout,
+                stream=True,
+            )
+
+            # Log the response if enabled
+            if self.log_requests:
+                self.logger.info(
+                    f"API Response: {response.status_code} {response.reason}\n"
+                    f"URL: {request_data.url}\n"
+                    f"Correlation ID: {request_data.correlation_id}\n"
+                )
+
+            # Check for HTTP errors
+            response.raise_for_status()
+
+            # Yields the response content as a stream
+            for message in parse_sse_stream_chunked(response.iter_content()):
+                yield message
+
+        except requests.HTTPError as e:
+            self._handle_exception(e, request_data.url, request_data.correlation_id)
+
+    def make_request_multipart(
+        self, method: str, request_data: RequestData, return_json: bool = True
+    ):
+        """
+        Makes a synchronous HTTP request with multipart form data to the Airia API.
+
+        Args:
+            method (str): The HTTP method (e.g., 'POST')
+            request_data: A dictionary containing the following request information:
+                - url: The endpoint URL for the request
+                - headers: HTTP headers to include in the request
+                - payload: The form data payload including file content
+                - params: Optional query parameters to append to the URL
+                - files: Optional file data to be uploaded in the request body
+                - correlation_id: Unique identifier for request tracing
+            return_json (bool): Whether to return the response as JSON. Default is True.
+
+        Returns:
+            resp (Dict[str, Any]): The JSON response from the API as a dictionary.
+
+        Raises:
+            AiriaAPIError: If the API returns an error response, with details about the error
+            requests.HTTPError: For HTTP-related errors
+
+        Note:
+            This is an internal method used by file upload methods to make multipart requests.
+            It handles multipart form data encoding, logging, and error handling.
+        """
+        try:
+            # Make the request
+            response = self.session.request(
+                method=method,
+                url=request_data.url,
+                files=request_data.files,
+                data=request_data.payload,
+                params=request_data.params,
+                headers=request_data.headers,
+                timeout=self.timeout,
+            )
+
+            # Log the response if enabled
+            if self.log_requests:
+                self.logger.info(
+                    f"API Response: {response.status_code} {response.reason}\n"
+                    f"URL: {request_data.url}\n"
+                    f"Correlation ID: {request_data.correlation_id}\n"
+                )
+
+            # Check for HTTP errors
+            response.raise_for_status()
+
+            # Returns the JSON response
+            if return_json:
+                return response.json()
+
+        except requests.HTTPError as e:
+            self._handle_exception(e, request_data.url, request_data.correlation_id)