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

airia/client/sync_client.py

@@ -1,7 +1,6 @@
-from typing import Any, Dict, List, Literal, Optional, overload
+from typing import Optional
 
 import loguru
-import requests
 
 from ..constants import (
     DEFAULT_ANTHROPIC_GATEWAY_URL,
@@ -9,20 +8,13 @@ from ..constants import (
     DEFAULT_OPENAI_GATEWAY_URL,
     DEFAULT_TIMEOUT,
 )
-from ..exceptions import AiriaAPIError
-from ..types._api_version import ApiVersion
-from ..types._request_data import RequestData
-from ..types.api import (
-    CreateConversationResponse,
-    GetConversationResponse,
-    GetPipelineConfigResponse,
-    PipelineExecutionDebugResponse,
-    PipelineExecutionResponse,
-    PipelineExecutionStreamedResponse,
-    ProjectItem,
-)
-from ..utils.sse_parser import parse_sse_stream_chunked
+from ._request_handler import RequestHandler
 from .base_client import AiriaBaseClient
+from .conversations import Conversations
+from .pipeline_execution import PipelineExecution
+from .pipelines_config import PipelinesConfig
+from .project import Project
+from .store import Store
 
 
 class AiriaClient(AiriaBaseClient):
@@ -57,9 +49,19 @@ class AiriaClient(AiriaBaseClient):
             custom_logger=custom_logger,
         )
 
-        # Initialize session for synchronous requests
-        self.session = requests.Session()
-        self.session.headers.update({"Content-Type": "application/json"})
+        self._request_handler = RequestHandler(
+            logger=self.logger,
+            timeout=self.timeout,
+            base_url=self.base_url,
+            api_key=self.api_key,
+            bearer_token=self.bearer_token,
+            log_requests=self.log_requests,
+        )
+        self.pipeline_execution = PipelineExecution(self._request_handler)
+        self.pipelines_config = PipelinesConfig(self._request_handler)
+        self.project = Project(self._request_handler)
+        self.conversations = Conversations(self._request_handler)
+        self.store = Store(self._request_handler)
 
     @classmethod
     def with_openai_gateway(
@@ -168,644 +170,11 @@ class AiriaClient(AiriaBaseClient):
             custom_logger=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 sensitive auth information is not included in error messages
-        sanitized_message = error_message
-        if self.api_key and self.api_key in sanitized_message:
-            sanitized_message = sanitized_message.replace(self.api_key, "[REDACTED]")
-        if self.bearer_token and self.bearer_token in sanitized_message:
-            sanitized_message = sanitized_message.replace(
-                self.bearer_token, "[REDACTED]"
-            )
-
-        # Raise custom exception with status code and sanitized message
-        raise AiriaAPIError(
-            status_code=e.response.status_code, message=sanitized_message
-        ) from e
-
-    def _make_request(
-        self, method: str, request_data: RequestData, return_json: bool = True
-    ) -> Dict[str, Any]:
-        """
-        Makes a synchronous HTTP request to the Airia API.
-
-        Args:
-            method (str): The HTTP method (e.g., 'GET', 'POST')
-            request_data: A dictionary containing the following request information:
-                - url: The endpoint URL for the request
-                - headers: HTTP headers to include in the request
-                - payload: The JSON payload/body for the request
-                - correlation_id: Unique identifier for request tracing
-            return_json (bool): Whether to return the response as JSON. Default is True.
-
-        Returns:
-            resp (Dict[str, Any]): The JSON response from the API as a dictionary.
-
-        Raises:
-            AiriaAPIError: If the API returns an error response, with details about the error
-            requests.HTTPError: For HTTP-related errors
-
-        Note:
-            This is an internal method used by other client methods to make API requests.
-            It handles logging, error handling, and API key redaction in error messages.
-        """
-        try:
-            # Make the request
-            response = self.session.request(
-                method=method,
-                url=request_data.url,
-                json=request_data.payload,
-                params=request_data.params,
-                headers=request_data.headers,
-                timeout=self.timeout,
-            )
-
-            # Log the response if enabled
-            if self.log_requests:
-                self.logger.info(
-                    f"API Response: {response.status_code} {response.reason}\n"
-                    f"URL: {request_data.url}\n"
-                    f"Correlation ID: {request_data.correlation_id}\n"
-                )
-
-            # Check for HTTP errors
-            response.raise_for_status()
-
-            # Returns the JSON response
-            if return_json:
-                return response.json()
-
-        except requests.HTTPError as e:
-            self._handle_exception(e, request_data.url, request_data.correlation_id)
-
-    def _make_request_stream(self, method: str, request_data: RequestData):
-        """
-        Makes a synchronous HTTP request to the Airia API.
-
-        Args:
-            method (str): The HTTP method (e.g., 'GET', 'POST')
-            request_data: A dictionary containing the following request information:
-                - url: The endpoint URL for the request
-                - headers: HTTP headers to include in the request
-                - payload: The JSON payload/body for the request
-                - correlation_id: Unique identifier for request tracing
-            stream (bool): If True, the response will be streamed instead of downloaded all at once
-
-        Yields:
-            resp (Iterator[str]): Yields chunks of the response as they are received.
-
-        Raises:
-            AiriaAPIError: If the API returns an error response, with details about the error
-            requests.HTTPError: For HTTP-related errors
-
-        Note:
-            This is an internal method used by other client methods to make API requests.
-            It handles logging, error handling, and API key redaction in error messages.
-        """
-        try:
-            # Make the request
-            response = self.session.request(
-                method=method,
-                url=request_data.url,
-                params=request_data.params,
-                json=request_data.payload,
-                headers=request_data.headers,
-                timeout=self.timeout,
-                stream=True,
-            )
-
-            # Log the response if enabled
-            if self.log_requests:
-                self.logger.info(
-                    f"API Response: {response.status_code} {response.reason}\n"
-                    f"URL: {request_data.url}\n"
-                    f"Correlation ID: {request_data.correlation_id}\n"
-                )
-
-            # Check for HTTP errors
-            response.raise_for_status()
-
-            # Yields the response content as a stream
-            for message in parse_sse_stream_chunked(response.iter_content()):
-                yield message
-
-        except requests.HTTPError as e:
-            self._handle_exception(e, request_data.url, request_data.correlation_id)
-
-    @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,
-    ) -> 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,
-    ) -> 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,
-    ) -> PipelineExecutionStreamedResponse: ...
-
-    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,
-    ):
-        """
-        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.
-
-        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=ApiVersion.V2.value,
-        )
-        resp = (
-            self._make_request_stream("POST", request_data)
-            if async_output
-            else self._make_request("POST", request_data)
-        )
-
-        if not async_output:
-            if not debug:
-                return PipelineExecutionResponse(**resp)
-            return PipelineExecutionDebugResponse(**resp)
-
-        return PipelineExecutionStreamedResponse(stream=resp)
-
-    def get_projects(self, correlation_id: Optional[str] = None) -> List[ProjectItem]:
-        """
-        Retrieve a list of all projects accessible to the authenticated user.
-
-        This method fetches comprehensive information about all projects that the
-        current user has access to, including project metadata, creation details,
-        and status information.
-
-        Args:
-            correlation_id (str, optional): A unique identifier for request tracing
-                and logging. If not provided, one will be automatically generated.
-
-        Returns:
-            List[ProjectItem]: A list of ProjectItem objects containing project
-                information. Returns an empty list if no projects are accessible
-                or found.
-
-        Raises:
-            AiriaAPIError: If the API request fails, including cases where:
-                - Authentication fails (401)
-                - Access is forbidden (403)
-                - Server errors (5xx)
-
-        Example:
-            ```python
-            from airia import AiriaClient
-
-            client = AiriaClient(api_key="your_api_key")
-
-            # Get all accessible projects
-            projects = client.get_projects()
-
-            for project in projects:
-                print(f"Project: {project.name}")
-                print(f"ID: {project.id}")
-                print(f"Description: {project.description}")
-                print(f"Created: {project.created_at}")
-                print("---")
-            ```
-
-        Note:
-            The returned projects are filtered based on the authenticated user's
-            permissions. Users will only see projects they have been granted
-            access to.
+    def close(self):
         """
-        request_data = self._pre_get_projects(
-            correlation_id=correlation_id, api_version=ApiVersion.V1.value
-        )
-        resp = self._make_request("GET", request_data)
-
-        if "items" not in resp or len(resp["items"]) == 0:
-            return []
-
-        return [ProjectItem(**item) for item in resp["items"]]
+        Closes the requests session to free up system resources.
 
-    def get_active_pipelines_ids(
-        self, project_id: Optional[str] = None, correlation_id: Optional[str] = None
-    ) -> List[str]:
+        This method should be called when the RequestHandler is no longer needed to ensure
+        proper cleanup of the underlying session and its resources.
         """
-        Retrieve a list of active pipeline IDs.
-
-        This method fetches the IDs of all active pipelines, optionally filtered by project.
-        Active pipelines are those that are currently deployed and available for execution.
-
-        Args:
-            project_id (str, optional): The unique identifier of the project to filter
-                pipelines by. If not provided, returns active pipelines from all projects
-                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.
-
-        Returns:
-            List[str]: A list of pipeline IDs that are currently active. Returns an
-                empty list if no active pipelines are found.
-
-        Raises:
-            AiriaAPIError: If the API request fails, including cases where:
-                - The project_id doesn't exist (404)
-                - Authentication fails (401)
-                - Access is forbidden (403)
-                - Server errors (5xx)
-
-        Example:
-            ```python
-            from airia import AiriaClient
-
-            client = AiriaClient(api_key="your_api_key")
-
-            # Get all active pipeline IDs
-            pipeline_ids = client.get_active_pipelines_ids()
-            print(f"Found {len(pipeline_ids)} active pipelines")
-
-            # Get active pipeline IDs for a specific project
-            project_pipelines = client.get_active_pipelines_ids(
-                project_id="your_project_id"
-            )
-            print(f"Project has {len(project_pipelines)} active pipelines")
-            ```
-
-        Note:
-            Only pipelines with active versions are returned. Inactive or archived
-            pipelines are not included in the results.
-        """
-        request_data = self._pre_get_active_pipelines_ids(
-            project_id=project_id,
-            correlation_id=correlation_id,
-            api_version=ApiVersion.V1.value,
-        )
-        resp = self._make_request("GET", request_data)
-
-        if "items" not in resp or len(resp["items"]) == 0:
-            return []
-
-        pipeline_ids = [r["activeVersion"]["pipelineId"] for r in resp["items"]]
-
-        return pipeline_ids
-
-    def get_pipeline_config(
-        self, pipeline_id: str, correlation_id: Optional[str] = None
-    ) -> GetPipelineConfigResponse:
-        """
-        Retrieve configuration details for a specific pipeline.
-
-        This method fetches comprehensive information about a pipeline including its
-        deployment details, execution statistics, version information, and metadata.
-
-        Args:
-            pipeline_id (str): The unique identifier of the pipeline to retrieve
-                configuration for.
-            correlation_id (str, optional): A unique identifier for request tracing
-                and logging. If not provided, one will be automatically generated.
-
-        Returns:
-            GetPipelineConfigResponse: A response object containing the pipeline
-                configuration.
-
-        Raises:
-            AiriaAPIError: If the API request fails, including cases where:
-                - The pipeline_id doesn't exist (404)
-                - Authentication fails (401)
-                - Access is forbidden (403)
-                - Server errors (5xx)
-
-        Example:
-            ```python
-            from airia import AiriaClient
-
-            client = AiriaClient(api_key="your_api_key")
-
-            # Get pipeline configuration
-            config = client.get_pipeline_config(
-                pipeline_id="your_pipeline_id"
-            )
-
-            print(f"Pipeline: {config.deployment_name}")
-            print(f"Description: {config.deployment_description}")
-            print(f"Success rate: {config.execution_stats.success_count}")
-            print(f"Active version: {config.active_version.version_number}")
-            ```
-
-        Note:
-            This method only retrieves configuration information and does not
-            execute the pipeline. Use execute_pipeline() to run the pipeline.
-        """
-        request_data = self._pre_get_pipeline_config(
-            pipeline_id=pipeline_id,
-            correlation_id=correlation_id,
-            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)
-
-    def get_conversation(
-        self, conversation_id: str, correlation_id: Optional[str] = None
-    ) -> GetConversationResponse:
-        """
-        Retrieve detailed information about a specific conversation by its ID.
-
-        This method fetches comprehensive information about a conversation including
-        all messages, metadata, policy redactions, and execution status.
-
-        Args:
-            conversation_id (str): The unique identifier of the conversation to retrieve.
-            correlation_id (str, optional): A unique identifier for request tracing
-                and logging. If not provided, one will be automatically generated.
-
-        Returns:
-            GetConversationResponse: A response object containing the conversation
-                details including user ID, messages, title, deployment information,
-                data source files, bookmark status, policy redactions, and execution status.
-
-        Raises:
-            AiriaAPIError: If the API request fails, including cases where:
-                - The conversation_id doesn't exist (404)
-                - Authentication fails (401)
-                - Access is forbidden (403)
-                - Server errors (5xx)
-
-        Example:
-            ```python
-            from airia import AiriaClient
-
-            client = AiriaClient(api_key="your_api_key")
-
-            # Get conversation details
-            conversation = client.get_conversation(
-                conversation_id="conversation_123"
-            )
-
-            print(f"Conversation: {conversation.title}")
-            print(f"User: {conversation.user_id}")
-            print(f"Messages: {len(conversation.messages)}")
-            print(f"Bookmarked: {conversation.is_bookmarked}")
-
-            # Access individual messages
-            for message in conversation.messages:
-                print(f"[{message.role}]: {message.message}")
-            ```
-
-        Note:
-            This method only retrieves conversation information and does not
-            modify or execute any operations on the conversation.
-        """
-        request_data = self._pre_get_conversation(
-            conversation_id=conversation_id,
-            correlation_id=correlation_id,
-            api_version=ApiVersion.V1.value,
-        )
-        resp = self._make_request("GET", request_data)
-
-        return GetConversationResponse(**resp)
-
-    def delete_conversation(
-        self,
-        conversation_id: str,
-        correlation_id: Optional[str] = None,
-    ) -> None:
-        """
-        Delete a conversation by its ID.
-
-        This method permanently removes a conversation and all associated data
-        from the Airia platform. This action cannot be undone.
-
-        Args:
-            conversation_id: The unique identifier of the conversation to delete
-            correlation_id: Optional correlation ID for request tracing
-
-        Returns:
-            None: This method returns nothing upon successful deletion
-
-        Raises:
-            AiriaAPIError: If the API request fails or the conversation doesn't exist
-        """
-        request_data = self._pre_delete_conversation(
-            conversation_id=conversation_id,
-            correlation_id=correlation_id,
-            api_version=ApiVersion.V1.value,
-        )
-        self._make_request("DELETE", request_data, return_json=False)
+        self._request_handler.close()