airia 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

airia/client/async_client.py

@@ -1,21 +1,26 @@
 import asyncio
 import weakref
 from typing import Any, AsyncIterator, Dict, List, Literal, Optional, overload
-from urllib.parse import urljoin
 
 import aiohttp
 import loguru
 
+from ..constants import (
+    DEFAULT_ANTHROPIC_GATEWAY_URL,
+    DEFAULT_BASE_URL,
+    DEFAULT_OPENAI_GATEWAY_URL,
+    DEFAULT_TIMEOUT,
+)
 from ..exceptions import AiriaAPIError
-from ..types import (
-    ApiVersion,
+from ..types._api_version import ApiVersion
+from ..types._request_data import RequestData
+from ..types.api import (
+    CreateConversationResponse,
     GetPipelineConfigResponse,
+    PipelineExecutionAsyncStreamedResponse,
     PipelineExecutionDebugResponse,
     PipelineExecutionResponse,
-    PipelineExecutionV1StreamedResponse,
-    PipelineExecutionV2AsyncStreamedResponse,
     ProjectItem,
-    RequestData,
 )
 from ..utils.sse_parser import async_parse_sse_stream_chunked
 from .base_client import AiriaBaseClient
@@ -26,9 +31,10 @@ class AiriaAsyncClient(AiriaBaseClient):
 
     def __init__(
         self,
-        base_url: str = "https://api.airia.ai/",
+        base_url: str = DEFAULT_BASE_URL,
         api_key: Optional[str] = None,
-        timeout: float = 30.0,
+        bearer_token: Optional[str] = None,
+        timeout: float = DEFAULT_TIMEOUT,
         log_requests: bool = False,
         custom_logger: Optional["loguru.Logger"] = None,
     ):
@@ -38,6 +44,7 @@ class AiriaAsyncClient(AiriaBaseClient):
         Args:
             base_url: Base URL of the Airia API.
             api_key: API key for authentication. If not provided, will attempt to use AIRIA_API_KEY environment variable.
+            bearer_token: Bearer token for authentication. Must be provided explicitly (no environment variable fallback).
             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.
@@ -45,6 +52,7 @@ class AiriaAsyncClient(AiriaBaseClient):
         super().__init__(
             base_url=base_url,
             api_key=api_key,
+            bearer_token=bearer_token,
             timeout=timeout,
             log_requests=log_requests,
             custom_logger=custom_logger,
@@ -80,10 +88,10 @@ class AiriaAsyncClient(AiriaBaseClient):
     @classmethod
     def with_openai_gateway(
         cls,
-        base_url: str = "https://api.airia.ai/",
-        gateway_url: str = "https://gateway.airia.ai/openai/v1",
+        base_url: str = DEFAULT_BASE_URL,
+        gateway_url: str = DEFAULT_OPENAI_GATEWAY_URL,
         api_key: Optional[str] = None,
-        timeout: float = 30.0,
+        timeout: float = DEFAULT_TIMEOUT,
         log_requests: bool = False,
         custom_logger: Optional["loguru.Logger"] = None,
         **kwargs,
@@ -102,22 +110,28 @@ class AiriaAsyncClient(AiriaBaseClient):
         """
         from openai import AsyncOpenAI
 
-        api_key = cls._get_api_key(api_key)
-        cls.openai = AsyncOpenAI(
+        client = cls(
+            base_url=base_url,
             api_key=api_key,
+            timeout=timeout,
+            log_requests=log_requests,
+            custom_logger=custom_logger,
+        )
+        cls.openai = AsyncOpenAI(
+            api_key=client.api_key,
             base_url=gateway_url,
             **kwargs,
         )
 
-        return cls(base_url, api_key, timeout, log_requests, custom_logger)
+        return client
 
     @classmethod
     def with_anthropic_gateway(
         cls,
-        base_url: str = "https://api.airia.ai/",
-        gateway_url: str = "https://gateway.airia.ai/anthropic",
+        base_url: str = DEFAULT_BASE_URL,
+        gateway_url: str = DEFAULT_ANTHROPIC_GATEWAY_URL,
         api_key: Optional[str] = None,
-        timeout: float = 30.0,
+        timeout: float = DEFAULT_TIMEOUT,
         log_requests: bool = False,
         custom_logger: Optional["loguru.Logger"] = None,
         **kwargs,
@@ -136,14 +150,47 @@ class AiriaAsyncClient(AiriaBaseClient):
         """
         from anthropic import AsyncAnthropic
 
-        api_key = cls._get_api_key(api_key)
-        cls.anthropic = AsyncAnthropic(
+        client = cls(
+            base_url=base_url,
             api_key=api_key,
+            timeout=timeout,
+            log_requests=log_requests,
+            custom_logger=custom_logger,
+        )
+        cls.anthropic = AsyncAnthropic(
+            api_key=client.api_key,
             base_url=gateway_url,
             **kwargs,
         )
 
-        return cls(base_url, api_key, timeout, log_requests, custom_logger)
+        return client
+
+    @classmethod
+    def with_bearer_token(
+        cls,
+        bearer_token: str,
+        base_url: str = DEFAULT_BASE_URL,
+        timeout: float = DEFAULT_TIMEOUT,
+        log_requests: bool = False,
+        custom_logger: Optional["loguru.Logger"] = None,
+    ):
+        """
+        Initialize the asynchronous Airia API client with bearer token authentication.
+
+        Args:
+            bearer_token: Bearer token for authentication.
+            base_url: Base URL of the Airia API.
+            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.
+        """
+        return cls(
+            base_url=base_url,
+            bearer_token=bearer_token,
+            timeout=timeout,
+            log_requests=log_requests,
+            custom_logger=custom_logger,
+        )
 
     def _handle_exception(
         self, e: aiohttp.ClientResponseError, url: str, correlation_id: str
@@ -159,12 +206,14 @@ class AiriaAsyncClient(AiriaBaseClient):
         # 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
-        )
+        # 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
@@ -297,7 +346,6 @@ class AiriaAsyncClient(AiriaBaseClient):
         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
@@ -320,7 +368,6 @@ class AiriaAsyncClient(AiriaBaseClient):
         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
@@ -343,31 +390,7 @@ class AiriaAsyncClient(AiriaBaseClient):
         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: ...
+    ) -> PipelineExecutionAsyncStreamedResponse: ...
 
     async def execute_pipeline(
         self,
@@ -388,7 +411,6 @@ class AiriaAsyncClient(AiriaBaseClient):
         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.
@@ -412,7 +434,6 @@ class AiriaAsyncClient(AiriaBaseClient):
             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.
@@ -447,38 +468,23 @@ class AiriaAsyncClient(AiriaBaseClient):
             additional_info=additional_info,
             prompt_variables=prompt_variables,
             correlation_id=correlation_id,
-            api_version=api_version,
+            api_version=ApiVersion.V2.value,
+        )
+        resp = (
+            self._make_request_stream(method="POST", request_data=request_data)
+            if async_output
+            else await self._make_request("POST", request_data=request_data)
         )
-        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=url,
-                payload={"socketIdentifier": resp},
-                correlation_id=request_data.headers["X-Correlation-ID"],
-            )
-            resp = await self._make_request("POST", request_data)
-
-            return PipelineExecutionV1StreamedResponse(**resp)
-
-        return PipelineExecutionV2AsyncStreamedResponse(stream=resp)
+        return PipelineExecutionAsyncStreamedResponse(stream=resp)
 
     async def get_projects(
-        self,
-        correlation_id: Optional[str] = None,
-        api_version: str = ApiVersion.V1.value,
+        self, correlation_id: Optional[str] = None
     ) -> List[ProjectItem]:
         """
         Retrieve a list of all projects accessible to the authenticated user.
@@ -490,8 +496,6 @@ class AiriaAsyncClient(AiriaBaseClient):
         Args:
             correlation_id (str, optional): A unique identifier for request tracing
                 and logging. If not provided, one will be automatically generated.
-            api_version (str, optional): The API version to use for the request.
-                Defaults to "v1". Valid versions are defined in ApiVersion enum.
 
         Returns:
             List[ProjectItem]: A list of ProjectItem objects containing project
@@ -499,7 +503,6 @@ class AiriaAsyncClient(AiriaBaseClient):
                 or found.
 
         Raises:
-            ValueError: If the provided api_version is not valid.
             AiriaAPIError: If the API request fails, including cases where:
                 - Authentication fails (401)
                 - Access is forbidden (403)
@@ -528,7 +531,7 @@ class AiriaAsyncClient(AiriaBaseClient):
             access to.
         """
         request_data = self._pre_get_projects(
-            correlation_id=correlation_id, api_version=api_version
+            correlation_id=correlation_id, api_version=ApiVersion.V1.value
         )
         resp = await self._make_request("GET", request_data)
 
@@ -538,10 +541,7 @@ class AiriaAsyncClient(AiriaBaseClient):
         return [ProjectItem(**item) for item in resp["items"]]
 
     async def get_active_pipelines_ids(
-        self,
-        project_id: Optional[str] = None,
-        correlation_id: Optional[str] = None,
-        api_version: str = ApiVersion.V1.value,
+        self, project_id: Optional[str] = None, correlation_id: Optional[str] = None
     ) -> List[str]:
         """
         Retrieve a list of active pipeline IDs.
@@ -555,15 +555,12 @@ class AiriaAsyncClient(AiriaBaseClient):
                 accessible to the authenticated user.
             correlation_id (str, optional): A unique identifier for request tracing
                 and logging. If not provided, one will be automatically generated.
-            api_version (str, optional): The API version to use for the request.
-                Defaults to "v1". Valid versions are defined in ApiVersion enum.
 
         Returns:
             List[str]: A list of pipeline IDs that are currently active. Returns an
                 empty list if no active pipelines are found.
 
         Raises:
-            ValueError: If the provided api_version is not valid.
             AiriaAPIError: If the API request fails, including cases where:
                 - The project_id doesn't exist (404)
                 - Authentication fails (401)
@@ -594,7 +591,7 @@ class AiriaAsyncClient(AiriaBaseClient):
         request_data = self._pre_get_active_pipelines_ids(
             project_id=project_id,
             correlation_id=correlation_id,
-            api_version=api_version,
+            api_version=ApiVersion.V1.value,
         )
         resp = await self._make_request("GET", request_data)
 
@@ -606,10 +603,7 @@ class AiriaAsyncClient(AiriaBaseClient):
         return pipeline_ids
 
     async def get_pipeline_config(
-        self,
-        pipeline_id: str,
-        correlation_id: Optional[str] = None,
-        api_version: str = ApiVersion.V1.value,
+        self, pipeline_id: str, correlation_id: Optional[str] = None
     ) -> GetPipelineConfigResponse:
         """
         Retrieve configuration details for a specific pipeline.
@@ -620,8 +614,6 @@ class AiriaAsyncClient(AiriaBaseClient):
         Args:
             pipeline_id (str): The unique identifier of the pipeline to retrieve
                 configuration for.
-            api_version (str, optional): The API version to use for the request.
-                Defaults to "v1". Valid versions are defined in ApiVersion enum.
             correlation_id (str, optional): A unique identifier for request tracing
                 and logging. If not provided, one will be automatically generated.
 
@@ -630,7 +622,6 @@ class AiriaAsyncClient(AiriaBaseClient):
                 configuration.
 
         Raises:
-            ValueError: If the provided api_version is not valid.
             AiriaAPIError: If the API request fails, including cases where:
                 - The pipeline_id doesn't exist (404)
                 - Authentication fails (401)
@@ -661,8 +652,88 @@ class AiriaAsyncClient(AiriaBaseClient):
         request_data = self._pre_get_pipeline_config(
             pipeline_id=pipeline_id,
             correlation_id=correlation_id,
-            api_version=api_version,
+            api_version=ApiVersion.V1.value,
         )
         resp = await self._make_request("GET", request_data)
 
         return GetPipelineConfigResponse(**resp)
+
+    async def create_conversation(
+        self,
+        user_id: str,
+        title: Optional[str] = None,
+        deployment_id: Optional[str] = None,
+        data_source_files: Dict[str, Any] = {},
+        is_bookmarked: bool = False,
+        correlation_id: Optional[str] = None,
+    ) -> CreateConversationResponse:
+        """
+        Create a new conversation.
+
+        Args:
+            user_id (str): The unique identifier of the user creating the conversation.
+            title (str, optional): The title for the conversation. If not provided,
+                the conversation will be created without a title.
+            deployment_id (str, optional): The unique identifier of the deployment
+                to associate with the conversation. If not provided, the conversation
+                will not be associated with any specific deployment.
+            data_source_files (dict): Configuration for data source files
+                to be associated with the conversation. If not provided, no data
+                source files will be associated.
+            is_bookmarked (bool): Whether the conversation should be bookmarked.
+                Defaults to False.
+            correlation_id (str, optional): A unique identifier for request tracing
+                and logging. If not provided, one will be automatically generated.
+
+        Returns:
+            CreateConversationResponse: A response object containing the created
+                conversation details including its ID, creation timestamp, and
+                all provided parameters.
+
+        Raises:
+            AiriaAPIError: If the API request fails, including cases where:
+                - The user_id doesn't exist (404)
+                - The deployment_id is invalid (404)
+                - Authentication fails (401)
+                - Access is forbidden (403)
+                - Server errors (5xx)
+
+        Example:
+            ```python
+            from airia import AiriaAsyncClient
+
+            client = AiriaAsyncClient(api_key="your_api_key")
+
+            # Create a basic conversation
+            conversation = await client.create_conversation(
+                user_id="user_123"
+            )
+            print(f"Created conversation: {conversation.conversation_id}")
+
+            # Create a conversation with all options
+            conversation = await client.create_conversation(
+                user_id="user_123",
+                title="My Research Session",
+                deployment_id="deployment_456",
+                data_source_files={"documents": ["doc1.pdf", "doc2.txt"]},
+                is_bookmarked=True
+            )
+            print(f"Created bookmarked conversation: {conversation.conversation_id}")
+            ```
+
+        Note:
+            The user_id is required and must correspond to a valid user in the system.
+            All other parameters are optional and can be set to None or their default values.
+        """
+        request_data = self._pre_create_conversation(
+            user_id=user_id,
+            title=title,
+            deployment_id=deployment_id,
+            data_source_files=data_source_files,
+            is_bookmarked=is_bookmarked,
+            correlation_id=correlation_id,
+            api_version=ApiVersion.V1.value,
+        )
+        resp = await self._make_request("POST", request_data)
+
+        return CreateConversationResponse(**resp)

airia/client/base_client.py

@@ -5,8 +5,10 @@ from urllib.parse import urljoin
 
 import loguru
 
+from ..constants import DEFAULT_BASE_URL, DEFAULT_TIMEOUT
 from ..logs import configure_logging, set_correlation_id
-from ..types import ApiVersion, RequestData
+from ..types._api_version import ApiVersion
+from ..types._request_data import RequestData
 
 
 class AiriaBaseClient:
@@ -17,9 +19,10 @@ class AiriaBaseClient:
 
     def __init__(
         self,
-        base_url: str = "https://api.airia.ai/",
+        base_url: str = DEFAULT_BASE_URL,
         api_key: Optional[str] = None,
-        timeout: float = 30.0,
+        bearer_token: Optional[str] = None,
+        timeout: float = DEFAULT_TIMEOUT,
         log_requests: bool = False,
         custom_logger: Optional["loguru.Logger"] = None,
     ):
@@ -28,12 +31,15 @@ class AiriaBaseClient:
 
         Args:
             api_key: API key for authentication. If not provided, will attempt to use AIRIA_API_KEY environment variable.
+            bearer_token: Bearer token for authentication. Must be provided explicitly (no environment variable fallback).
             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)
+        # Resolve authentication credentials
+        self.api_key, self.bearer_token = self.__class__._resolve_auth_credentials(
+            api_key, bearer_token
+        )
 
         # Store configuration
         self.base_url = base_url
@@ -44,27 +50,45 @@ class AiriaBaseClient:
         self.logger = configure_logging() if custom_logger is None else custom_logger
 
     @staticmethod
-    def _get_api_key(api_key: Optional[str] = None):
+    def _resolve_auth_credentials(
+        api_key: Optional[str] = None, bearer_token: Optional[str] = None
+    ):
         """
-        Get the API key from either the provided parameter or environment variable.
+        Resolve authentication credentials from parameters and environment variables.
 
         Args:
             api_key (Optional[str]): The API key provided as a parameter. Defaults to None.
+            bearer_token (Optional[str]): The bearer token provided as a parameter. Defaults to None.
 
         Returns:
-            str: The resolved API key.
+            tuple: (api_key, bearer_token) - exactly one will be non-None
 
         Raises:
-            ValueError: If no API key is provided through either method.
+            ValueError: If no authentication method is provided or if both are provided.
         """
-        api_key = api_key or os.environ.get("AIRIA_API_KEY")
-
-        if not api_key:
+        # Check for explicit conflict first
+        if api_key and bearer_token:
             raise ValueError(
-                "API key must be provided either as a parameter or through the AIRIA_API_KEY environment variable."
+                "Cannot provide both api_key and bearer_token. Please use only one authentication method."
             )
 
-        return api_key
+        # If bearer token is explicitly provided, use it exclusively
+        if bearer_token:
+            return None, bearer_token
+
+        # If API key is explicitly provided, use it exclusively
+        if api_key:
+            return api_key, None
+
+        # If neither is provided explicitly, fall back to environment variable
+        resolved_api_key = os.environ.get("AIRIA_API_KEY")
+        if resolved_api_key:
+            return resolved_api_key, None
+
+        # No authentication method found
+        raise ValueError(
+            "Authentication required. Provide either api_key (or set AIRIA_API_KEY environment variable) or bearer_token."
+        )
 
     def _prepare_request(
         self,
@@ -76,13 +100,18 @@ class AiriaBaseClient:
         # 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
+        # Set up base headers
         headers = {
-            "X-API-KEY": self.api_key,
             "X-Correlation-ID": correlation_id,
             "Content-Type": "application/json",
         }
 
+        # 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
@@ -92,6 +121,8 @@ class AiriaBaseClient:
             # 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 {}
@@ -218,6 +249,36 @@ class AiriaBaseClient:
 
         return request_data
 
+    def _pre_create_conversation(
+        self,
+        user_id: str,
+        title: Optional[str] = None,
+        deployment_id: Optional[str] = None,
+        data_source_files: Dict[str, Any] = {},
+        is_bookmarked: bool = False,
+        correlation_id: Optional[str] = None,
+        api_version: str = ApiVersion.V1.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}/Conversations")
+
+        payload = {
+            "userId": user_id,
+            "title": title,
+            "deploymentId": deployment_id,
+            "dataSourceFiles": data_source_files,
+            "isBookmarked": is_bookmarked,
+        }
+
+        request_data = self._prepare_request(
+            url=url, payload=payload, correlation_id=correlation_id
+        )
+
+        return request_data
+
     def _pre_get_projects(
         self,
         correlation_id: Optional[str] = None,