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

airia/client/sync_client.py

@@ -1,19 +1,24 @@
 from typing import Any, Dict, List, Literal, Optional, overload
-from urllib.parse import urljoin
 
 import loguru
 import requests
 
+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,
     PipelineExecutionDebugResponse,
     PipelineExecutionResponse,
-    PipelineExecutionV1StreamedResponse,
-    PipelineExecutionV2StreamedResponse,
+    PipelineExecutionStreamedResponse,
     ProjectItem,
-    RequestData,
 )
 from ..utils.sse_parser import parse_sse_stream_chunked
 from .base_client import AiriaBaseClient
@@ -24,9 +29,10 @@ class AiriaClient(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,
     ):
@@ -36,6 +42,7 @@ class AiriaClient(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.
@@ -43,6 +50,7 @@ class AiriaClient(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,
@@ -55,10 +63,10 @@ class AiriaClient(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,
@@ -77,22 +85,28 @@ class AiriaClient(AiriaBaseClient):
         """
         from openai import OpenAI
 
-        api_key = cls._get_api_key(api_key)
-        cls.openai = OpenAI(
+        client = cls(
+            base_url=base_url,
             api_key=api_key,
+            timeout=timeout,
+            log_requests=log_requests,
+            custom_logger=custom_logger,
+        )
+        cls.openai = OpenAI(
+            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,
@@ -111,14 +125,47 @@ class AiriaClient(AiriaBaseClient):
         """
         from anthropic import Anthropic
 
-        api_key = cls._get_api_key(api_key)
-        cls.anthropic = Anthropic(
+        client = cls(
+            base_url=base_url,
             api_key=api_key,
+            timeout=timeout,
+            log_requests=log_requests,
+            custom_logger=custom_logger,
+        )
+        cls.anthropic = Anthropic(
+            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 synchronous 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: requests.HTTPError, url: str, correlation_id: str):
         # Log the error response if enabled
@@ -141,12 +188,14 @@ class AiriaClient(AiriaBaseClient):
             # 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
-        )
+        # 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(
@@ -278,7 +327,6 @@ class AiriaClient(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
@@ -301,7 +349,6 @@ class AiriaClient(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
@@ -324,31 +371,7 @@ class AiriaClient(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,
-    ) -> 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: ...
+    ) -> PipelineExecutionStreamedResponse: ...
 
     def execute_pipeline(
         self,
@@ -369,7 +392,6 @@ class AiriaClient(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,
     ):
         """
         Execute a pipeline with the provided input.
@@ -393,7 +415,6 @@ class AiriaClient(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.
@@ -428,12 +449,11 @@ class AiriaClient(AiriaBaseClient):
             additional_info=additional_info,
             prompt_variables=prompt_variables,
             correlation_id=correlation_id,
-            api_version=api_version,
+            api_version=ApiVersion.V2.value,
         )
-        stream = async_output and api_version == ApiVersion.V2.value
         resp = (
             self._make_request_stream("POST", request_data)
-            if stream
+            if async_output
             else self._make_request("POST", request_data)
         )
 
@@ -442,26 +462,9 @@ class AiriaClient(AiriaBaseClient):
                 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,
-                payload={"socketIdentifier": resp},
-                correlation_id=request_data.headers["X-Correlation-ID"],
-            )
-            resp = self._make_request("POST", request_data)
-
-            return PipelineExecutionV1StreamedResponse(**resp)
-
-        return PipelineExecutionV2StreamedResponse(stream=resp)
+        return PipelineExecutionStreamedResponse(stream=resp)
 
-    def get_projects(
-        self,
-        correlation_id: Optional[str] = None,
-        api_version: str = ApiVersion.V1.value,
-    ) -> List[ProjectItem]:
+    def get_projects(self, correlation_id: Optional[str] = None) -> List[ProjectItem]:
         """
         Retrieve a list of all projects accessible to the authenticated user.
 
@@ -472,8 +475,6 @@ class AiriaClient(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
@@ -481,7 +482,6 @@ class AiriaClient(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)
@@ -510,7 +510,7 @@ class AiriaClient(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 = self._make_request("GET", request_data)
 
@@ -520,10 +520,7 @@ class AiriaClient(AiriaBaseClient):
         return [ProjectItem(**item) for item in resp["items"]]
 
     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.
@@ -537,15 +534,12 @@ class AiriaClient(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)
@@ -576,7 +570,7 @@ class AiriaClient(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 = self._make_request("GET", request_data)
 
@@ -588,10 +582,7 @@ class AiriaClient(AiriaBaseClient):
         return pipeline_ids
 
     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.
@@ -602,8 +593,6 @@ class AiriaClient(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.
 
@@ -612,7 +601,6 @@ class AiriaClient(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)
@@ -643,8 +631,88 @@ class AiriaClient(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 = self._make_request("GET", request_data)
 
         return GetPipelineConfigResponse(**resp)
+
+    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 AiriaClient
+
+            client = AiriaClient(api_key="your_api_key")
+
+            # Create a basic conversation
+            conversation = client.create_conversation(
+                user_id="user_123"
+            )
+            print(f"Created conversation: {conversation.id}")
+
+            # Create a conversation with all options
+            conversation = 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.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 = self._make_request("POST", request_data)
+
+        return CreateConversationResponse(**resp)

airia/constants.py

@@ -0,0 +1,9 @@
+"""Constants used throughout the Airia SDK."""
+
+# Default API endpoints
+DEFAULT_BASE_URL = "https://api.airia.ai/"
+DEFAULT_OPENAI_GATEWAY_URL = "https://gateway.airia.ai/openai/v1"
+DEFAULT_ANTHROPIC_GATEWAY_URL = "https://gateway.airia.ai/anthropic"
+
+# Default timeouts
+DEFAULT_TIMEOUT = 30.0

airia/types/api/__init__.py

@@ -0,0 +1,19 @@
+from .get_projects import ProjectItem
+from .get_pipeline_config import GetPipelineConfigResponse
+from .pipeline_execution import (
+    PipelineExecutionDebugResponse,
+    PipelineExecutionResponse,
+    PipelineExecutionAsyncStreamedResponse,
+    PipelineExecutionStreamedResponse,
+)
+from .conversations import CreateConversationResponse
+
+__all__ = [
+    "PipelineExecutionDebugResponse",
+    "PipelineExecutionResponse",
+    "PipelineExecutionStreamedResponse",
+    "PipelineExecutionAsyncStreamedResponse",
+    "GetPipelineConfigResponse",
+    "ProjectItem",
+    "CreateConversationResponse",
+]

airia/types/api/conversations.py

@@ -0,0 +1,14 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+
+class CreateConversationResponse(BaseModel):
+    user_id: str = Field(alias="userId")
+    conversation_id: str = Field(alias="conversationId")
+    websocket_url: str = Field(alias="websocketUrl")
+    deployment_id: str = Field(alias="deploymentId")
+    icon_id: Optional[str] = Field(None, alias="iconId")
+    icon_url: Optional[str] = Field(None, alias="iconUrl")
+    description: Optional[str] = None
+    space_name: Optional[str] = Field(None, alias="spaceName")

airia/types/api/pipeline_execution.py

@@ -1,33 +1,29 @@
 from typing import Any, AsyncIterator, Dict, Iterator
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field
 
-from ..sse_messages import SSEMessage
+from ..sse import SSEMessage
 
 
 class PipelineExecutionResponse(BaseModel):
     result: str
     report: None
-    isBackupPipeline: bool
+    is_backup_pipeline: bool = Field(alias="isBackupPipeline")
 
 
 class PipelineExecutionDebugResponse(BaseModel):
     result: str
     report: Dict[str, Any]
-    isBackupPipeline: bool
+    is_backup_pipeline: bool = Field(alias="isBackupPipeline")
 
 
-class PipelineExecutionV1StreamedResponse(BaseModel):
-    webSocketUrl: str
-
-
-class PipelineExecutionV2StreamedResponse(BaseModel):
+class PipelineExecutionStreamedResponse(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     stream: Iterator[SSEMessage]
 
 
-class PipelineExecutionV2AsyncStreamedResponse(BaseModel):
+class PipelineExecutionAsyncStreamedResponse(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     stream: AsyncIterator[SSEMessage]

airia/types/{__init__.py → sse/__init__.py}

@@ -1,15 +1,6 @@
-from .api.get_projects import ProjectItem
-from .api.get_pipeline_config import GetPipelineConfigResponse
-from .api.pipeline_execution import (
-    PipelineExecutionDebugResponse,
-    PipelineExecutionResponse,
-    PipelineExecutionV1StreamedResponse,
-    PipelineExecutionV2AsyncStreamedResponse,
-    PipelineExecutionV2StreamedResponse,
-)
-from .api_version import ApiVersion
-from .request_data import RequestData
 from .sse_messages import (
+    SSEDict,
+    SSEMessage,
     AgentAgentCardMessage,
     AgentAgentCardStreamEndMessage,
     AgentAgentCardStreamErrorMessage,
@@ -37,15 +28,8 @@ from .sse_messages import (
 )
 
 __all__ = [
-    "ApiVersion",
-    "PipelineExecutionDebugResponse",
-    "PipelineExecutionResponse",
-    "PipelineExecutionV1StreamedResponse",
-    "PipelineExecutionV2AsyncStreamedResponse",
-    "PipelineExecutionV2StreamedResponse",
-    "GetPipelineConfigResponse",
-    "ProjectItem",
-    "RequestData",
+    "SSEDict",
+    "SSEMessage",
     "AgentPingMessage",
     "AgentStartMessage",
     "AgentEndMessage",

airia/utils/sse_parser.py

@@ -2,7 +2,7 @@ import json
 import re
 from typing import AsyncIterable, AsyncIterator, Iterable, Iterator
 
-from ..types.sse_messages import SSEDict, SSEMessage
+from ..types.sse import SSEDict, SSEMessage
 
 
 def _to_snake_case(name: str):