airia 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

airia/client/__init__.py

@@ -0,0 +1,4 @@
+from .async_client import AiriaAsyncClient
+from .sync_client import AiriaClient
+
+__all__ = ["AiriaClient", "AiriaAsyncClient"]

airia/client/async_client.py

@@ -0,0 +1,448 @@
+from typing import Any, AsyncIterator, Dict, List, Literal, Optional, overload
+from urllib.parse import urljoin
+
+import aiohttp
+import loguru
+
+from ..exceptions import AiriaAPIError
+from ..types import (
+    ApiVersion,
+    PipelineExecutionDebugResponse,
+    PipelineExecutionResponse,
+    PipelineExecutionV1StreamedResponse,
+    PipelineExecutionV2AsyncStreamedResponse,
+    RequestData,
+)
+from .base_client import AiriaBaseClient
+
+
+class AiriaAsyncClient(AiriaBaseClient):
+    """Asynchronous client for interacting with the Airia API."""
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        timeout: float = 30.0,
+        log_requests: bool = False,
+        custom_logger: Optional["loguru.Logger"] = None,
+    ):
+        """
+        Initialize the asynchronous Airia API client.
+
+        Args:
+            api_key: API key for authentication. If not provided, will attempt to use AIRIA_API_KEY environment variable.
+            timeout: Request timeout in seconds.
+            log_requests: Whether to log API requests and responses. Default is False.
+            custom_logger: Optional custom logger object to use for logging. If not provided, will use a default logger when `log_requests` is True.
+        """
+        super().__init__(api_key, timeout, log_requests, custom_logger)
+
+        # Session will be initialized in __aenter__
+        self.session = None
+        self.headers = {"Content-Type": "application/json"}
+    
+    @classmethod
+    def with_openai_gateway(
+        cls,
+        api_key: Optional[str] = None,
+        timeout: float = 30.0,
+        log_requests: bool = False,
+        custom_logger: Optional["loguru.Logger"] = None,
+        **kwargs,
+    ):
+        """
+        Initialize the asynchronous Airia API client with AsyncOpenAI gateway capabilities.
+
+        Args:
+            api_key: API key for authentication. If not provided, will attempt to use AIRIA_API_KEY environment variable.
+            timeout: Request timeout in seconds.
+            log_requests: Whether to log API requests and responses. Default is False.
+            custom_logger: Optional custom logger object to use for logging. If not provided, will use a default logger when `log_requests` is True.
+            **kwargs: Additional keyword arguments to pass to the AsyncOpenAI client initialization.
+        """
+        from openai import AsyncOpenAI
+
+        api_key = cls._get_api_key(api_key)
+        cls.openai = AsyncOpenAI(
+            api_key=api_key,
+            base_url="https://gateway.airia.ai/openai/v1",
+            **kwargs,
+        )
+
+        return cls(api_key, timeout, log_requests, custom_logger)
+
+    @classmethod
+    def with_anthropic_gateway(
+        cls,
+        api_key: Optional[str] = None,
+        timeout: float = 30.0,
+        log_requests: bool = False,
+        custom_logger: Optional["loguru.Logger"] = None,
+        **kwargs,
+    ):
+        """
+        Initialize the asynchronous Airia API client with AsyncAnthropic gateway capabilities.
+
+        Args:
+            api_key: API key for authentication. If not provided, will attempt to use AIRIA_API_KEY environment variable.
+            timeout: Request timeout in seconds.
+            log_requests: Whether to log API requests and responses. Default is False.
+            custom_logger: Optional custom logger object to use for logging. If not provided, will use a default logger when `log_requests` is True.
+            **kwargs: Additional keyword arguments to pass to the AsyncAnthropic client initialization.
+        """
+        from anthropic import AsyncAnthropic
+
+        api_key = cls._get_api_key(api_key)
+        cls.anthropic = AsyncAnthropic(
+            api_key=api_key,
+            base_url="https://gateway.airia.ai/anthropic",
+            **kwargs,
+        )
+
+        return cls(api_key, timeout, log_requests, custom_logger)
+
+    async def __aenter__(self):
+        """Async context manager entry point."""
+        self.session = aiohttp.ClientSession(headers=self.headers)
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit point."""
+        if self.session:
+            await self.session.close()
+            self.session = None
+
+    def _check_session(self):
+        """Check if the client session is initialized."""
+        if not self.session:
+            raise RuntimeError(
+                "Client session not initialized. Use async with AiriaAsyncClient() as client: ..."
+            )
+
+    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 API key is not included in error messages
+        sanitized_message = (
+            error_message.replace(self.api_key, "[REDACTED]")
+            if self.api_key in error_message
+            else error_message
+        )
+
+        # 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
+    ) -> 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
+                - correlation_id: Unique identifier for request tracing
+
+        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,
+                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
+                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
+                - 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,
+                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 chunk in response.content.iter_any():
+                    yield chunk.decode("utf-8")
+
+        except aiohttp.ClientResponseError as e:
+            self._handle_exception(e, request_data.url, request_data.correlation_id)
+
+    @overload
+    async def execute_pipeline(
+        self,
+        pipeline_id: str,
+        user_input: str,
+        debug: Literal[False] = False,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        async_output: Literal[False] = False,
+        include_tools_response: bool = False,
+        images: Optional[List[str]] = None,
+        files: Optional[List[str]] = None,
+        data_source_folders: Optional[Dict[str, Any]] = None,
+        data_source_files: Optional[Dict[str, Any]] = None,
+        in_memory_messages: Optional[List[Dict[str, str]]] = None,
+        current_date_time: Optional[str] = None,
+        save_history: bool = True,
+        additional_info: Optional[List[Any]] = None,
+        prompt_variables: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+        api_version: str = ApiVersion.V2.value,
+    ) -> PipelineExecutionResponse: ...
+
+    @overload
+    async def execute_pipeline(
+        self,
+        pipeline_id: str,
+        user_input: str,
+        debug: Literal[True] = True,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        async_output: Literal[False] = False,
+        include_tools_response: bool = False,
+        images: Optional[List[str]] = None,
+        files: Optional[List[str]] = None,
+        data_source_folders: Optional[Dict[str, Any]] = None,
+        data_source_files: Optional[Dict[str, Any]] = None,
+        in_memory_messages: Optional[List[Dict[str, str]]] = None,
+        current_date_time: Optional[str] = None,
+        save_history: bool = True,
+        additional_info: Optional[List[Any]] = None,
+        prompt_variables: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+        api_version: str = ApiVersion.V2.value,
+    ) -> PipelineExecutionDebugResponse: ...
+
+    @overload
+    async def execute_pipeline(
+        self,
+        pipeline_id: str,
+        user_input: str,
+        debug: bool = False,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        async_output: Literal[True] = True,
+        include_tools_response: bool = False,
+        images: Optional[List[str]] = None,
+        files: Optional[List[str]] = None,
+        data_source_folders: Optional[Dict[str, Any]] = None,
+        data_source_files: Optional[Dict[str, Any]] = None,
+        in_memory_messages: Optional[List[Dict[str, str]]] = None,
+        current_date_time: Optional[str] = None,
+        save_history: bool = True,
+        additional_info: Optional[List[Any]] = None,
+        prompt_variables: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+        api_version: Literal["v2"] = ApiVersion.V2.value,
+    ) -> PipelineExecutionV2AsyncStreamedResponse: ...
+
+    @overload
+    async def execute_pipeline(
+        self,
+        pipeline_id: str,
+        user_input: str,
+        debug: bool = False,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        async_output: Literal[True] = True,
+        include_tools_response: bool = False,
+        images: Optional[List[str]] = None,
+        files: Optional[List[str]] = None,
+        data_source_folders: Optional[Dict[str, Any]] = None,
+        data_source_files: Optional[Dict[str, Any]] = None,
+        in_memory_messages: Optional[List[Dict[str, str]]] = None,
+        current_date_time: Optional[str] = None,
+        save_history: bool = True,
+        additional_info: Optional[List[Any]] = None,
+        prompt_variables: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+        api_version: Literal["v1"] = ApiVersion.V1.value,
+    ) -> PipelineExecutionV1StreamedResponse: ...
+
+    async def execute_pipeline(
+        self,
+        pipeline_id: str,
+        user_input: str,
+        debug: bool = False,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        async_output: bool = False,
+        include_tools_response: bool = False,
+        images: Optional[List[str]] = None,
+        files: Optional[List[str]] = None,
+        data_source_folders: Optional[Dict[str, Any]] = None,
+        data_source_files: Optional[Dict[str, Any]] = None,
+        in_memory_messages: Optional[List[Dict[str, str]]] = None,
+        current_date_time: Optional[str] = None,
+        save_history: bool = True,
+        additional_info: Optional[List[Any]] = None,
+        prompt_variables: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+        api_version: str = ApiVersion.V2.value,
+    ) -> Dict[str, Any]:
+        """
+        Execute a pipeline with the provided input asynchronously.
+
+        Args:
+            pipeline_id: The ID of the pipeline to execute.
+            user_input: input text to process.
+            debug: Whether debug mode execution is enabled. Default is False.
+            user_id: Optional ID of the user making the request (guid).
+            conversation_id: Optional conversation ID (guid).
+            async_output: Whether to stream the response. Default is False.
+            include_tools_response: Whether to return the initial LLM tool result. Default is False.
+            images: Optional list of images formatted as base64 strings.
+            files: Optional list of files formatted as base64 strings.
+            data_source_folders: Optional data source folders information.
+            data_source_files: Optional data source files information.
+            in_memory_messages: Optional list of in-memory messages, each with a role and message.
+            current_date_time: Optional current date and time in ISO format.
+            save_history: Whether to save the userInput and output to conversation history. Default is True.
+            additional_info: Optional additional information.
+            prompt_variables: Optional variables to be used in the prompt.
+            correlation_id: Optional correlation ID for request tracing. If not provided,
+                        one will be generated automatically.
+            api_version: API version to use. Default is `v2`
+
+        Returns:
+            The API response as a dictionary.
+
+        Raises:
+            AiriaAPIError: If the API request fails with details about the error.
+            aiohttp.ClientError: For other request-related errors.
+
+        Example:
+            >>> async with AiriaAsyncClient(api_key="your_api_key") as client:
+            ...     response = await client.execute_pipeline(
+            ...         pipeline_id="pipeline_123",
+            ...         user_input="Tell me about quantum computing"
+            ...     )
+            >>> print(response.result)
+        """
+        self._check_session()
+
+        request_data = self._pre_execute_pipeline(
+            pipeline_id=pipeline_id,
+            user_input=user_input,
+            debug=debug,
+            user_id=user_id,
+            conversation_id=conversation_id,
+            async_output=async_output,
+            include_tools_response=include_tools_response,
+            images=images,
+            files=files,
+            data_source_folders=data_source_folders,
+            data_source_files=data_source_files,
+            in_memory_messages=in_memory_messages,
+            current_date_time=current_date_time,
+            save_history=save_history,
+            additional_info=additional_info,
+            prompt_variables=prompt_variables,
+            correlation_id=correlation_id,
+            api_version=api_version,
+        )
+        stream = async_output and api_version == ApiVersion.V2.value
+        if stream:
+            resp = self._make_request_stream(method="POST", request_data=request_data)
+        else:
+            resp = await self._make_request("POST", request_data)
+
+        if not async_output:
+            if not debug:
+                return PipelineExecutionResponse(**resp)
+            return PipelineExecutionDebugResponse(**resp)
+
+        if api_version == ApiVersion.V1.value:
+            url = urljoin(
+                self.base_url, f"{api_version}/StreamSocketConfig/GenerateUrl"
+            )
+            request_data = self._prepare_request(
+                url,
+                {"socketIdentifier": resp},
+                request_data.headers["X-Correlation-ID"],
+            )
+            resp = await self._make_request("POST", request_data)
+
+            return PipelineExecutionV1StreamedResponse(**resp)
+
+        return PipelineExecutionV2AsyncStreamedResponse(stream=resp)

airia/client/base_client.py

@@ -0,0 +1,187 @@
+import json
+import os
+from typing import Any, Dict, List, Optional
+from urllib.parse import urljoin
+
+import loguru
+
+from ..logs import configure_logging, set_correlation_id
+from ..types import ApiVersion, RequestData
+
+
+class AiriaBaseClient:
+    """Base client containing shared functionality for Airia API clients."""
+    openai = None
+    anthropic = None
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        timeout: float = 30.0,
+        log_requests: bool = False,
+        custom_logger: Optional["loguru.Logger"] = None,
+    ):
+        """
+        Initialize the Airia API client base class.
+
+        Args:
+            api_key: API key for authentication. If not provided, will attempt to use AIRIA_API_KEY environment variable.
+            timeout: Request timeout in seconds.
+            log_requests: Whether to log API requests and responses. Default is False.
+            custom_logger: Optional custom logger object to use for logging. If not provided, will use a default logger when `log_requests` is True.
+        """
+        # Resolve API key: parameter takes precedence over environment variable
+        self.api_key = self.__class__._get_api_key(api_key)
+
+        # Store configuration
+        self.base_url = "https://api.airia.ai/"
+        self.timeout = timeout
+        self.log_requests = log_requests
+
+        # Initialize logger
+        self.logger = configure_logging() if custom_logger is None else custom_logger
+
+    @staticmethod
+    def _get_api_key(api_key: Optional[str] = None):
+        """
+        Get the API key from either the provided parameter or environment variable.
+
+        Args:
+            api_key (Optional[str]): The API key provided as a parameter. Defaults to None.
+
+        Returns:
+            str: The resolved API key.
+
+        Raises:
+            ValueError: If no API key is provided through either method.
+        """
+        api_key = api_key or os.environ.get("AIRIA_API_KEY")
+
+        if not api_key:
+            raise ValueError(
+                "API key must be provided either as a parameter or through the AIRIA_API_KEY environment variable."
+            )
+
+        return api_key
+
+    def _prepare_request(
+        self,
+        url: str,
+        payload: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+    ):
+        """
+        Prepare the request parameters for an API call.
+
+        Args:
+            url (str): The endpoint URL for the API request.
+            payload (Optional[Dict[str, Any]]): The request payload/body to be sent.
+            correlation_id (Optional[str]): A unique identifier for tracing the request. If None, one will be generated.
+
+        Returns:
+            dict: A dictionary containing the prepared request parameters with the following keys:
+                - url: The request URL
+                - payload: The request payload
+                - headers: Request headers including API key and correlation ID
+                - correlation_id: The correlation ID used for the request
+
+        Note:
+            This method handles:
+            - Setting/generating correlation IDs
+            - Adding authentication headers
+            - Logging requests (if enabled) with sensitive information redacted
+        """
+        # Set correlation ID if provided or generate a new one
+        correlation_id = set_correlation_id(correlation_id)
+
+        # Add the X-API-KEY header and correlation ID
+        headers = {
+            "X-API-KEY": self.api_key,
+            "X-Correlation-ID": correlation_id,
+            "Content-Type": "application/json",
+        }
+
+        # Log the request if enabled
+        if self.log_requests:
+            # Create a sanitized copy of headers for logging
+            log_headers = headers.copy()
+
+            # Filter out sensitive headers
+            if "X-API-KEY" in log_headers:
+                log_headers["X-API-KEY"] = "[REDACTED]"
+
+            # Process payload for logging
+            log_payload = None
+            if payload is not None:
+                log_payload = payload.copy()
+
+                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}"
+            )
+
+        return RequestData(
+            **{
+                "url": url,
+                "payload": payload,
+                "headers": headers,
+                "correlation_id": correlation_id,
+            }
+        )
+
+    def _pre_execute_pipeline(
+        self,
+        pipeline_id: str,
+        user_input: str,
+        debug: bool = False,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        async_output: bool = False,
+        include_tools_response: bool = False,
+        images: Optional[List[str]] = None,
+        files: Optional[List[str]] = None,
+        data_source_folders: Optional[Dict[str, Any]] = None,
+        data_source_files: Optional[Dict[str, Any]] = None,
+        in_memory_messages: Optional[List[Dict[str, str]]] = None,
+        current_date_time: Optional[str] = None,
+        save_history: bool = True,
+        additional_info: Optional[List[Any]] = None,
+        prompt_variables: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+        api_version: str = ApiVersion.V2.value,
+    ):
+        if api_version not in ApiVersion.as_list():
+            raise ValueError(
+                f"Invalid API version: {api_version}. Valid versions are: {', '.join(ApiVersion.as_list())}"
+            )
+        url = urljoin(self.base_url, f"{api_version}/PipelineExecution/{pipeline_id}")
+
+        payload = {
+            "userInput": user_input,
+            "debug": debug,
+            "userId": user_id,
+            "conversationId": conversation_id,
+            "asyncOutput": async_output,
+            "includeToolsResponse": include_tools_response,
+            "images": images,
+            "files": files,
+            "dataSourceFolders": data_source_folders,
+            "dataSourceFiles": data_source_files,
+            "inMemoryMessages": in_memory_messages,
+            "currentDateTime": current_date_time,
+            "saveHistory": save_history,
+            "additionalInfo": additional_info,
+            "promptVariables": prompt_variables,
+        }
+
+        request_data = self._prepare_request(url, payload, correlation_id)
+
+        return request_data

airia/client/sync_client.py

@@ -0,0 +1,436 @@
+from typing import Any, Dict, List, Literal, Optional, overload
+from urllib.parse import urljoin
+
+import loguru
+import requests
+
+from ..exceptions import AiriaAPIError
+from ..types import (
+    ApiVersion,
+    PipelineExecutionDebugResponse,
+    PipelineExecutionResponse,
+    PipelineExecutionV1StreamedResponse,
+    PipelineExecutionV2StreamedResponse,
+    RequestData,
+)
+from .base_client import AiriaBaseClient
+
+
+class AiriaClient(AiriaBaseClient):
+    """Synchronous client for interacting with the Airia API."""
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        timeout: float = 30.0,
+        log_requests: bool = False,
+        custom_logger: Optional["loguru.Logger"] = None,
+    ):
+        """
+        Initialize the synchronous Airia API client.
+
+        Args:
+            api_key: API key for authentication. If not provided, will attempt to use AIRIA_API_KEY environment variable.
+            timeout: Request timeout in seconds.
+            log_requests: Whether to log API requests and responses. Default is False.
+            custom_logger: Optional custom logger object to use for logging. If not provided, will use a default logger when `log_requests` is True.
+        """
+        super().__init__(api_key, timeout, log_requests, custom_logger)
+
+        # Initialize session for synchronous requests
+        self.session = requests.Session()
+        self.session.headers.update({"Content-Type": "application/json"})
+
+    @classmethod
+    def with_openai_gateway(
+        cls,
+        api_key: Optional[str] = None,
+        timeout: float = 30.0,
+        log_requests: bool = False,
+        custom_logger: Optional["loguru.Logger"] = None,
+        **kwargs,
+    ):
+        """
+        Initialize the synchronous Airia API client with OpenAI gateway capabilities.
+
+        Args:
+            api_key: API key for authentication. If not provided, will attempt to use AIRIA_API_KEY environment variable.
+            timeout: Request timeout in seconds.
+            log_requests: Whether to log API requests and responses. Default is False.
+            custom_logger: Optional custom logger object to use for logging. If not provided, will use a default logger when `log_requests` is True.
+            **kwargs: Additional keyword arguments to pass to the OpenAI client initialization.
+        """
+        from openai import OpenAI
+
+        api_key = cls._get_api_key(api_key)
+        cls.openai = OpenAI(
+            api_key=api_key,
+            base_url="https://gateway.airia.ai/openai/v1",
+            **kwargs,
+        )
+
+        return cls(api_key, timeout, log_requests, custom_logger)
+
+    @classmethod
+    def with_anthropic_gateway(
+        cls,
+        api_key: Optional[str] = None,
+        timeout: float = 30.0,
+        log_requests: bool = False,
+        custom_logger: Optional["loguru.Logger"] = None,
+        **kwargs,
+    ):
+        """
+        Initialize the synchronous Airia API client with Anthropic gateway capabilities.
+
+        Args:
+            api_key: API key for authentication. If not provided, will attempt to use AIRIA_API_KEY environment variable.
+            timeout: Request timeout in seconds.
+            log_requests: Whether to log API requests and responses. Default is False.
+            custom_logger: Optional custom logger object to use for logging. If not provided, will use a default logger when `log_requests` is True.
+            **kwargs: Additional keyword arguments to pass to the Anthropic client initialization.
+        """
+        from anthropic import Anthropic
+
+        api_key = cls._get_api_key(api_key)
+        cls.anthropic = Anthropic(
+            api_key=api_key,
+            base_url="https://gateway.airia.ai/anthropic",
+            **kwargs,
+        )
+
+        return cls(api_key, timeout, log_requests, custom_logger)
+
+    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 API key is not included in error messages
+        sanitized_message = (
+            error_message.replace(self.api_key, "[REDACTED]")
+            if self.api_key in error_message
+            else error_message
+        )
+
+        # 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):
+        """
+        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
+                - correlation_id: Unique identifier for request tracing
+
+        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,
+                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
+            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
+                - 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,
+                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 chunk in response.iter_content():
+                yield chunk.decode("utf-8")
+
+        except requests.HTTPError as e:
+            self._handle_exception(e, request_data.url, request_data.correlation_id)
+
+    @overload
+    def execute_pipeline(
+        self,
+        pipeline_id: str,
+        user_input: str,
+        debug: Literal[False] = False,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        async_output: Literal[False] = False,
+        include_tools_response: bool = False,
+        images: Optional[List[str]] = None,
+        files: Optional[List[str]] = None,
+        data_source_folders: Optional[Dict[str, Any]] = None,
+        data_source_files: Optional[Dict[str, Any]] = None,
+        in_memory_messages: Optional[List[Dict[str, str]]] = None,
+        current_date_time: Optional[str] = None,
+        save_history: bool = True,
+        additional_info: Optional[List[Any]] = None,
+        prompt_variables: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+        api_version: str = ApiVersion.V2.value,
+    ) -> PipelineExecutionResponse: ...
+
+    @overload
+    def execute_pipeline(
+        self,
+        pipeline_id: str,
+        user_input: str,
+        debug: Literal[True] = True,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        async_output: Literal[False] = False,
+        include_tools_response: bool = False,
+        images: Optional[List[str]] = None,
+        files: Optional[List[str]] = None,
+        data_source_folders: Optional[Dict[str, Any]] = None,
+        data_source_files: Optional[Dict[str, Any]] = None,
+        in_memory_messages: Optional[List[Dict[str, str]]] = None,
+        current_date_time: Optional[str] = None,
+        save_history: bool = True,
+        additional_info: Optional[List[Any]] = None,
+        prompt_variables: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+        api_version: str = ApiVersion.V2.value,
+    ) -> PipelineExecutionDebugResponse: ...
+
+    @overload
+    def execute_pipeline(
+        self,
+        pipeline_id: str,
+        user_input: str,
+        debug: bool = False,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        async_output: Literal[True] = True,
+        include_tools_response: bool = False,
+        images: Optional[List[str]] = None,
+        files: Optional[List[str]] = None,
+        data_source_folders: Optional[Dict[str, Any]] = None,
+        data_source_files: Optional[Dict[str, Any]] = None,
+        in_memory_messages: Optional[List[Dict[str, str]]] = None,
+        current_date_time: Optional[str] = None,
+        save_history: bool = True,
+        additional_info: Optional[List[Any]] = None,
+        prompt_variables: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+        api_version: Literal["v2"] = ApiVersion.V2.value,
+    ) -> PipelineExecutionV2StreamedResponse: ...
+
+    @overload
+    def execute_pipeline(
+        self,
+        pipeline_id: str,
+        user_input: str,
+        debug: bool = False,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        async_output: Literal[True] = True,
+        include_tools_response: bool = False,
+        images: Optional[List[str]] = None,
+        files: Optional[List[str]] = None,
+        data_source_folders: Optional[Dict[str, Any]] = None,
+        data_source_files: Optional[Dict[str, Any]] = None,
+        in_memory_messages: Optional[List[Dict[str, str]]] = None,
+        current_date_time: Optional[str] = None,
+        save_history: bool = True,
+        additional_info: Optional[List[Any]] = None,
+        prompt_variables: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+        api_version: Literal["v1"] = ApiVersion.V1.value,
+    ) -> PipelineExecutionV1StreamedResponse: ...
+
+    def execute_pipeline(
+        self,
+        pipeline_id: str,
+        user_input: str,
+        debug: bool = False,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        async_output: bool = False,
+        include_tools_response: bool = False,
+        images: Optional[List[str]] = None,
+        files: Optional[List[str]] = None,
+        data_source_folders: Optional[Dict[str, Any]] = None,
+        data_source_files: Optional[Dict[str, Any]] = None,
+        in_memory_messages: Optional[List[Dict[str, str]]] = None,
+        current_date_time: Optional[str] = None,
+        save_history: bool = True,
+        additional_info: Optional[List[Any]] = None,
+        prompt_variables: Optional[Dict[str, Any]] = None,
+        correlation_id: Optional[str] = None,
+        api_version: str = ApiVersion.V2.value,
+    ):
+        """
+        Execute a pipeline with the provided input.
+
+        Args:
+            pipeline_id: The ID of the pipeline to execute.
+            user_input: input text to process.
+            debug: Whether debug mode execution is enabled. Default is False.
+            user_id: Optional ID of the user making the request (guid).
+            conversation_id: Optional conversation ID (guid).
+            async_output: Whether to stream the response. Default is False.
+            include_tools_response: Whether to return the initial LLM tool result. Default is False.
+            images: Optional list of images formatted as base64 strings.
+            files: Optional list of files formatted as base64 strings.
+            data_source_folders: Optional data source folders information.
+            data_source_files: Optional data source files information.
+            in_memory_messages: Optional list of in-memory messages, each with a role and message.
+            current_date_time: Optional current date and time in ISO format.
+            save_history: Whether to save the userInput and output to conversation history. Default is True.
+            additional_info: Optional additional information.
+            prompt_variables: Optional variables to be used in the prompt.
+            correlation_id: Optional correlation ID for request tracing. If not provided,
+                        one will be generated automatically.
+            api_version: API version to use. Default is `v2`
+
+        Returns:
+            The API response as a dictionary.
+
+        Raises:
+            AiriaAPIError: If the API request fails with details about the error.
+            requests.RequestException: For other request-related errors.
+
+        Example:
+            >>> client = AiriaClient(api_key="your_api_key")
+            >>> response = client.execute_pipeline(
+            ...     pipeline_id="pipeline_123",
+            ...     user_input="Tell me about quantum computing"
+            ... )
+            >>> print(response.result)
+        """
+        request_data = self._pre_execute_pipeline(
+            pipeline_id=pipeline_id,
+            user_input=user_input,
+            debug=debug,
+            user_id=user_id,
+            conversation_id=conversation_id,
+            async_output=async_output,
+            include_tools_response=include_tools_response,
+            images=images,
+            files=files,
+            data_source_folders=data_source_folders,
+            data_source_files=data_source_files,
+            in_memory_messages=in_memory_messages,
+            current_date_time=current_date_time,
+            save_history=save_history,
+            additional_info=additional_info,
+            prompt_variables=prompt_variables,
+            correlation_id=correlation_id,
+            api_version=api_version,
+        )
+        stream = async_output and api_version == ApiVersion.V2.value
+        if stream:
+            resp = self._make_request_stream("POST", request_data)
+        else:
+            resp = self._make_request("POST", request_data)
+
+        if not async_output:
+            if not debug:
+                return PipelineExecutionResponse(**resp)
+            return PipelineExecutionDebugResponse(**resp)
+
+        if api_version == ApiVersion.V1.value:
+            url = urljoin(
+                self.base_url, f"{api_version}/StreamSocketConfig/GenerateUrl"
+            )
+            request_data = self._prepare_request(
+                url,
+                {"socketIdentifier": resp},
+                request_data.headers["X-Correlation-ID"],
+            )
+            resp = self._make_request("POST", request_data)
+
+            return PipelineExecutionV1StreamedResponse(**resp)
+
+        return PipelineExecutionV2StreamedResponse(stream=resp)

airia/types/__init__.py

@@ -0,0 +1,19 @@
+from .api_version import ApiVersion
+from .pipeline_execution import (
+    PipelineExecutionDebugResponse,
+    PipelineExecutionResponse,
+    PipelineExecutionV1StreamedResponse,
+    PipelineExecutionV2AsyncStreamedResponse,
+    PipelineExecutionV2StreamedResponse,
+)
+from .request_data import RequestData
+
+__all__ = [
+    "ApiVersion",
+    "PipelineExecutionDebugResponse",
+    "PipelineExecutionResponse",
+    "PipelineExecutionV1StreamedResponse",
+    "PipelineExecutionV2AsyncStreamedResponse",
+    "PipelineExecutionV2StreamedResponse",
+    "RequestData",
+]

airia/types/api_version.py

@@ -0,0 +1,13 @@
+from enum import Enum
+
+
+class ApiVersion(Enum):
+    """Enum for API Versions."""
+
+    V1 = "v1"
+    V2 = "v2"
+
+    @classmethod
+    def as_list(cls):
+        """Return a list of all API Version values."""
+        return [version.value for version in cls]

airia/types/pipeline_execution.py

@@ -0,0 +1,31 @@
+from typing import Any, AsyncIterator, Dict, Iterator
+
+from pydantic import BaseModel, ConfigDict
+
+
+class PipelineExecutionResponse(BaseModel):
+    result: str
+    report: None
+    isBackupPipeline: bool
+
+
+class PipelineExecutionDebugResponse(BaseModel):
+    result: str
+    report: Dict[str, Any]
+    isBackupPipeline: bool
+
+
+class PipelineExecutionV1StreamedResponse(BaseModel):
+    webSocketUrl: str
+
+
+class PipelineExecutionV2StreamedResponse(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    stream: Iterator[str]
+
+
+class PipelineExecutionV2AsyncStreamedResponse(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    stream: AsyncIterator[str]

airia/types/request_data.py

@@ -0,0 +1,10 @@
+from typing import Any, Dict
+
+from pydantic import BaseModel
+
+
+class RequestData(BaseModel):
+    url: str
+    payload: Dict[str, Any]
+    headers: Dict[str, Any]
+    correlation_id: str

{airia-0.1.3.dist-info → airia-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airia
-Version: 0.1.3
+Version: 0.1.4
 Summary: Python SDK for Airia API
 Author-email: Airia LLC <support@airia.com>
 License: MIT

airia-0.1.4.dist-info/RECORD

@@ -0,0 +1,16 @@
+airia/__init__.py,sha256=T39gO8E5T5zxlw-JP78ruxOu7-LeKOJCJzz6t40kdQo,231
+airia/exceptions.py,sha256=4Z55n-cRJrtTa5-pZBIK2oZD4-Z99aUtKx_kfTFYY5o,1146
+airia/logs.py,sha256=17YZ4IuzOF0m5bgofj9-QYlJ2BYR2kRZbBVQfFSLFEk,5441
+airia/client/__init__.py,sha256=6gSQ9bl7j79q1HPE0o5py3IRdkwWWuU_7J4h05Dd2o8,127
+airia/client/async_client.py,sha256=Xd1GAGxQvXVh0-hH8vE4HOeaNY_Fawh5iRxeS26CInY,18173
+airia/client/base_client.py,sha256=3oUcCWWq-DmPbiJke_HHPXm9UzDvD-MzrnoiKnBjcdk,6810
+airia/client/sync_client.py,sha256=p_auqkfyP-9YuCoppl43SyMdKek30kA6M05QI0LYunE,17689
+airia/types/__init__.py,sha256=OOFtJ0VIbO98px89u7cq645iL7CYDOel43g85IAjRRw,562
+airia/types/api_version.py,sha256=Uzom6O2ZG92HN_Z2h-lTydmO2XYX9RVs4Yi4DJmXytE,255
+airia/types/pipeline_execution.py,sha256=obp8KOz-ShNxEciF1YQCSpHXPoJUyEa_9uPv-VHf6YA,699
+airia/types/request_data.py,sha256=RbVPWPRIYuT7FaolJKn19IVzuQKbLM4VhXM5r7UXTUY,186
+airia-0.1.4.dist-info/licenses/LICENSE,sha256=R3ClUMMKPRItIcZ0svzyj2taZZnFYw568YDNzN9KQ1Q,1066
+airia-0.1.4.dist-info/METADATA,sha256=VOWxozh4xZXAHUXH6_ZyCUTIU95LXxly3Yfs-ayi0sk,9966
+airia-0.1.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+airia-0.1.4.dist-info/top_level.txt,sha256=qUQEKfs_hdOYTwjKj1JZbRhS5YeXDNaKQaVTrzabS6w,6
+airia-0.1.4.dist-info/RECORD,,

airia-0.1.3.dist-info/RECORD

@@ -1,8 +0,0 @@
-airia/__init__.py,sha256=T39gO8E5T5zxlw-JP78ruxOu7-LeKOJCJzz6t40kdQo,231
-airia/exceptions.py,sha256=4Z55n-cRJrtTa5-pZBIK2oZD4-Z99aUtKx_kfTFYY5o,1146
-airia/logs.py,sha256=17YZ4IuzOF0m5bgofj9-QYlJ2BYR2kRZbBVQfFSLFEk,5441
-airia-0.1.3.dist-info/licenses/LICENSE,sha256=R3ClUMMKPRItIcZ0svzyj2taZZnFYw568YDNzN9KQ1Q,1066
-airia-0.1.3.dist-info/METADATA,sha256=9LEQFb6HGNqoTfhn6aX8elSCa5Eaogk43cj2j_pcx4o,9966
-airia-0.1.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-airia-0.1.3.dist-info/top_level.txt,sha256=qUQEKfs_hdOYTwjKj1JZbRhS5YeXDNaKQaVTrzabS6w,6
-airia-0.1.3.dist-info/RECORD,,