airia 0.1.4__tar.gz → 0.1.6__tar.gz

{airia-0.1.4 → airia-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airia
-Version: 0.1.4
+Version: 0.1.6
 Summary: Python SDK for Airia API
 Author-email: Airia LLC <support@airia.com>
 License: MIT
@@ -176,6 +176,20 @@ This will create both wheel and source distribution in the `dist/` directory.
 
 ## Quick Start
 
+### Client Instantiation
+
+```python
+from airia import AiriaClient
+
+client = AiriaClient(
+    base_url="https://api.airia.ai",        # Default: "https://api.airia.ai"
+    api_key=None,                  # Or set AIRIA_API_KEY environment variable
+    timeout=30.0,                            # Request timeout in seconds (default: 30.0)
+    log_requests=False,                       # Enable request/response logging (default: False)
+    custom_logger=None                       # Use custom logger (default: None - uses built-in)
+)
+```
+
 ### Synchronous Usage
 
 ```python
@@ -208,8 +222,8 @@ response = client.execute_pipeline(
     async_output=True
 )
 
-for c in resp.stream:
-    print(c, end="")
+for c in response.stream:
+    print(c)
 ```
 
 ### Asynchronous Usage
@@ -219,13 +233,12 @@ import asyncio
 from airia import AiriaAsyncClient
 
 async def main():
-    # Use async context manager for proper session handling
-    async with AiriaAsyncClient(api_key="your_api_key") as client:
-        response = await client.execute_pipeline(
-            pipeline_id="your_pipeline_id",
-            user_input="Tell me about quantum computing"
-        )
-        print(response.result)
+    client = AiriaAsyncClient(api_key="your_api_key")
+    response = await client.execute_pipeline(
+        pipeline_id="your_pipeline_id",
+        user_input="Tell me about quantum computing"
+    )
+    print(response.result)
 
 asyncio.run(main())
 ```
@@ -237,19 +250,111 @@ import asyncio
 from airia import AiriaAsyncClient
 
 async def main():
-    # Use async context manager for proper session handling
-    async with AiriaAsyncClient(api_key="your_api_key") as client:
-        response = await client.execute_pipeline(
-            pipeline_id="your_pipeline_id",
-            user_input="Tell me about quantum computing",
-            async_output=True
-        )
-        async for c in resp.stream:
-            print(c, end="")
+    client = AiriaAsyncClient(api_key="your_api_key")
+    response = await client.execute_pipeline(
+        pipeline_id="your_pipeline_id",
+        user_input="Tell me about quantum computing",
+        async_output=True
+    )
+    async for c in response.stream:
+        print(c)
 
 asyncio.run(main())
 ```
 
+## Streaming Event Parsing
+
+When using streaming mode (`async_output=True`), the API returns Server-Sent Events (SSE) that contain different types of messages throughout the pipeline execution. You can parse and filter these events to extract specific information.
+
+### Available Message Types
+
+The streaming response includes various message types defined in `airia.types`. Here are the key ones:
+
+- `AgentModelStreamFragmentMessage` - Contains actual LLM output chunks
+- `AgentModelStreamStartMessage` - Indicates LLM streaming has started
+- `AgentModelStreamEndMessage` - Indicates LLM streaming has ended
+- `AgentStepStartMessage` - Indicates a pipeline step has started
+- `AgentStepEndMessage` - Indicates a pipeline step has ended
+- `AgentOutputMessage` - Contains step output
+
+<details>
+<summary>Click to expand the full list of message types</summary>
+
+```python
+[
+    AgentPingMessage,
+    AgentStartMessage,
+    AgentEndMessage,
+    AgentStepStartMessage,
+    AgentStepHaltMessage,
+    AgentStepEndMessage,
+    AgentOutputMessage,
+    AgentAgentCardMessage,
+    AgentDatasearchMessage,
+    AgentInvocationMessage,
+    AgentModelMessage,
+    AgentPythonCodeMessage,
+    AgentToolActionMessage,
+    AgentModelStreamStartMessage,
+    AgentModelStreamEndMessage,
+    AgentModelStreamErrorMessage,
+    AgentModelStreamUsageMessage,
+    AgentModelStreamFragmentMessage,
+    AgentAgentCardStreamStartMessage,
+    AgentAgentCardStreamErrorMessage,
+    AgentAgentCardStreamFragmentMessage,
+    AgentAgentCardStreamEndMessage,
+    AgentToolRequestMessage,
+    AgentToolResponseMessage,
+]
+```
+
+</details>
+
+### Filtering LLM Output
+
+To extract only the actual LLM output text from the stream:
+
+```python
+from airia import AiriaClient
+from airia.types import AgentModelStreamFragmentMessage
+
+client = AiriaClient(api_key="your_api_key")
+
+response = client.execute_pipeline(
+    pipeline_id="your_pipeline_id",
+    user_input="Tell me about quantum computing",
+    async_output=True
+)
+
+# Filter and display only LLM output
+for event in response.stream:
+    if isinstance(event, AgentModelStreamFragmentMessage) and event.index != -1:
+        print(event.content, end="", flush=True)
+```
+
+## Pipeline Configuration Retrieval
+
+You can retrieve detailed configuration information about a pipeline using the `get_pipeline_config` method:
+
+> To get a list of all active pipeline ids, run the `get_active_pipelines_ids` method.
+
+```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")
+
+# Access configuration details
+print(f"Pipeline Name: {config.deployment_name}")
+print(f"Description: {config.deployment_description}")
+print(f"Active Version: {config.active_version.version_number}")
+print(f"Success Count: {config.execution_stats.success_count}")
+print(f"Failure Count: {config.execution_stats.failure_count}")
+```
+
 ## Gateway Usage
 
 Airia provides gateway capabilities for popular AI services like OpenAI and Anthropic, allowing you to use your Airia API key with these services.
@@ -292,6 +397,8 @@ response = client.anthropic.messages.create(
 print(response.content[0].text)
 ```
 
+You can set the Gateway URL by passing the `gateway_url` parameter when using the gateway constructors. The default values are `https://gateway.airia.ai/openai/v1` for OpenAI and `https://gateway.airia.ai/anthropic` for Anthropic.
+
 ### Asynchronous Gateway Usage
 
 Both gateways also support asynchronous usage:

{airia-0.1.4 → airia-0.1.6}/README.md

@@ -146,6 +146,20 @@ This will create both wheel and source distribution in the `dist/` directory.
 
 ## Quick Start
 
+### Client Instantiation
+
+```python
+from airia import AiriaClient
+
+client = AiriaClient(
+    base_url="https://api.airia.ai",        # Default: "https://api.airia.ai"
+    api_key=None,                  # Or set AIRIA_API_KEY environment variable
+    timeout=30.0,                            # Request timeout in seconds (default: 30.0)
+    log_requests=False,                       # Enable request/response logging (default: False)
+    custom_logger=None                       # Use custom logger (default: None - uses built-in)
+)
+```
+
 ### Synchronous Usage
 
 ```python
@@ -178,8 +192,8 @@ response = client.execute_pipeline(
     async_output=True
 )
 
-for c in resp.stream:
-    print(c, end="")
+for c in response.stream:
+    print(c)
 ```
 
 ### Asynchronous Usage
@@ -189,13 +203,12 @@ import asyncio
 from airia import AiriaAsyncClient
 
 async def main():
-    # Use async context manager for proper session handling
-    async with AiriaAsyncClient(api_key="your_api_key") as client:
-        response = await client.execute_pipeline(
-            pipeline_id="your_pipeline_id",
-            user_input="Tell me about quantum computing"
-        )
-        print(response.result)
+    client = AiriaAsyncClient(api_key="your_api_key")
+    response = await client.execute_pipeline(
+        pipeline_id="your_pipeline_id",
+        user_input="Tell me about quantum computing"
+    )
+    print(response.result)
 
 asyncio.run(main())
 ```
@@ -207,19 +220,111 @@ import asyncio
 from airia import AiriaAsyncClient
 
 async def main():
-    # Use async context manager for proper session handling
-    async with AiriaAsyncClient(api_key="your_api_key") as client:
-        response = await client.execute_pipeline(
-            pipeline_id="your_pipeline_id",
-            user_input="Tell me about quantum computing",
-            async_output=True
-        )
-        async for c in resp.stream:
-            print(c, end="")
+    client = AiriaAsyncClient(api_key="your_api_key")
+    response = await client.execute_pipeline(
+        pipeline_id="your_pipeline_id",
+        user_input="Tell me about quantum computing",
+        async_output=True
+    )
+    async for c in response.stream:
+        print(c)
 
 asyncio.run(main())
 ```
 
+## Streaming Event Parsing
+
+When using streaming mode (`async_output=True`), the API returns Server-Sent Events (SSE) that contain different types of messages throughout the pipeline execution. You can parse and filter these events to extract specific information.
+
+### Available Message Types
+
+The streaming response includes various message types defined in `airia.types`. Here are the key ones:
+
+- `AgentModelStreamFragmentMessage` - Contains actual LLM output chunks
+- `AgentModelStreamStartMessage` - Indicates LLM streaming has started
+- `AgentModelStreamEndMessage` - Indicates LLM streaming has ended
+- `AgentStepStartMessage` - Indicates a pipeline step has started
+- `AgentStepEndMessage` - Indicates a pipeline step has ended
+- `AgentOutputMessage` - Contains step output
+
+<details>
+<summary>Click to expand the full list of message types</summary>
+
+```python
+[
+    AgentPingMessage,
+    AgentStartMessage,
+    AgentEndMessage,
+    AgentStepStartMessage,
+    AgentStepHaltMessage,
+    AgentStepEndMessage,
+    AgentOutputMessage,
+    AgentAgentCardMessage,
+    AgentDatasearchMessage,
+    AgentInvocationMessage,
+    AgentModelMessage,
+    AgentPythonCodeMessage,
+    AgentToolActionMessage,
+    AgentModelStreamStartMessage,
+    AgentModelStreamEndMessage,
+    AgentModelStreamErrorMessage,
+    AgentModelStreamUsageMessage,
+    AgentModelStreamFragmentMessage,
+    AgentAgentCardStreamStartMessage,
+    AgentAgentCardStreamErrorMessage,
+    AgentAgentCardStreamFragmentMessage,
+    AgentAgentCardStreamEndMessage,
+    AgentToolRequestMessage,
+    AgentToolResponseMessage,
+]
+```
+
+</details>
+
+### Filtering LLM Output
+
+To extract only the actual LLM output text from the stream:
+
+```python
+from airia import AiriaClient
+from airia.types import AgentModelStreamFragmentMessage
+
+client = AiriaClient(api_key="your_api_key")
+
+response = client.execute_pipeline(
+    pipeline_id="your_pipeline_id",
+    user_input="Tell me about quantum computing",
+    async_output=True
+)
+
+# Filter and display only LLM output
+for event in response.stream:
+    if isinstance(event, AgentModelStreamFragmentMessage) and event.index != -1:
+        print(event.content, end="", flush=True)
+```
+
+## Pipeline Configuration Retrieval
+
+You can retrieve detailed configuration information about a pipeline using the `get_pipeline_config` method:
+
+> To get a list of all active pipeline ids, run the `get_active_pipelines_ids` method.
+
+```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")
+
+# Access configuration details
+print(f"Pipeline Name: {config.deployment_name}")
+print(f"Description: {config.deployment_description}")
+print(f"Active Version: {config.active_version.version_number}")
+print(f"Success Count: {config.execution_stats.success_count}")
+print(f"Failure Count: {config.execution_stats.failure_count}")
+```
+
 ## Gateway Usage
 
 Airia provides gateway capabilities for popular AI services like OpenAI and Anthropic, allowing you to use your Airia API key with these services.
@@ -262,6 +367,8 @@ response = client.anthropic.messages.create(
 print(response.content[0].text)
 ```
 
+You can set the Gateway URL by passing the `gateway_url` parameter when using the gateway constructors. The default values are `https://gateway.airia.ai/openai/v1` for OpenAI and `https://gateway.airia.ai/anthropic` for Anthropic.
+
 ### Asynchronous Gateway Usage
 
 Both gateways also support asynchronous usage:

{airia-0.1.4 → airia-0.1.6}/airia/client/async_client.py

@@ -1,3 +1,5 @@
+import asyncio
+import weakref
 from typing import Any, AsyncIterator, Dict, List, Literal, Optional, overload
 from urllib.parse import urljoin
 
@@ -7,12 +9,14 @@ import loguru
 from ..exceptions import AiriaAPIError
 from ..types import (
     ApiVersion,
+    GetPipelineConfigResponse,
     PipelineExecutionDebugResponse,
     PipelineExecutionResponse,
     PipelineExecutionV1StreamedResponse,
     PipelineExecutionV2AsyncStreamedResponse,
     RequestData,
 )
+from ..utils.sse_parser import async_parse_sse_stream_chunked
 from .base_client import AiriaBaseClient
 
 
@@ -21,6 +25,7 @@ class AiriaAsyncClient(AiriaBaseClient):
 
     def __init__(
         self,
+        base_url: str = "https://api.airia.ai/",
         api_key: Optional[str] = None,
         timeout: float = 30.0,
         log_requests: bool = False,
@@ -30,20 +35,52 @@ class AiriaAsyncClient(AiriaBaseClient):
         Initialize the asynchronous Airia API client.
 
         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.
             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)
+        super().__init__(
+            base_url=base_url,
+            api_key=api_key,
+            timeout=timeout,
+            log_requests=log_requests,
+            custom_logger=custom_logger,
+        )
 
         # Session will be initialized in __aenter__
-        self.session = None
         self.headers = {"Content-Type": "application/json"}
-    
+        self.session = aiohttp.ClientSession(headers=self.headers)
+
+        # Register finalizer to clean up session when client is garbage collected
+        self._finalizer = weakref.finalize(self, self._cleanup_session, self.session)
+
+    @staticmethod
+    def _cleanup_session(session: aiohttp.ClientSession):
+        """Static method to clean up session - called by finalizer"""
+        if session and not session.closed:
+            # Create a new event loop if none exists
+            try:
+                loop = asyncio.get_event_loop()
+                if loop.is_closed():
+                    raise RuntimeError("Event loop is closed")
+            except RuntimeError:
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+
+            # Close the session
+            if not loop.is_running():
+                loop.run_until_complete(session.close())
+            else:
+                # If loop is running, schedule the close operation
+                asyncio.create_task(session.close())
+
     @classmethod
     def with_openai_gateway(
         cls,
+        base_url: str = "https://api.airia.ai/",
+        gateway_url: str = "https://gateway.airia.ai/openai/v1",
         api_key: Optional[str] = None,
         timeout: float = 30.0,
         log_requests: bool = False,
@@ -54,6 +91,8 @@ class AiriaAsyncClient(AiriaBaseClient):
         Initialize the asynchronous Airia API client with AsyncOpenAI gateway capabilities.
 
         Args:
+            base_url: Base URL of the Airia API.
+            gateway_url: Base URL of the Airia Gateway API.
             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.
@@ -65,15 +104,17 @@ class AiriaAsyncClient(AiriaBaseClient):
         api_key = cls._get_api_key(api_key)
         cls.openai = AsyncOpenAI(
             api_key=api_key,
-            base_url="https://gateway.airia.ai/openai/v1",
+            base_url=gateway_url,
             **kwargs,
         )
 
-        return cls(api_key, timeout, log_requests, custom_logger)
+        return cls(base_url, api_key, timeout, log_requests, custom_logger)
 
     @classmethod
     def with_anthropic_gateway(
         cls,
+        base_url: str = "https://api.airia.ai/",
+        gateway_url: str = "https://gateway.airia.ai/anthropic",
         api_key: Optional[str] = None,
         timeout: float = 30.0,
         log_requests: bool = False,
@@ -84,6 +125,8 @@ class AiriaAsyncClient(AiriaBaseClient):
         Initialize the asynchronous Airia API client with AsyncAnthropic gateway capabilities.
 
         Args:
+            base_url: Base URL of the Airia API.
+            gateway_url: Base URL of the Airia Gateway API.
             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.
@@ -95,29 +138,11 @@ class AiriaAsyncClient(AiriaBaseClient):
         api_key = cls._get_api_key(api_key)
         cls.anthropic = AsyncAnthropic(
             api_key=api_key,
-            base_url="https://gateway.airia.ai/anthropic",
+            base_url=gateway_url,
             **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: ..."
-            )
+        return cls(base_url, api_key, timeout, log_requests, custom_logger)
 
     def _handle_exception(
         self, e: aiohttp.ClientResponseError, url: str, correlation_id: str
@@ -241,8 +266,10 @@ class AiriaAsyncClient(AiriaBaseClient):
                 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")
+                async for message in async_parse_sse_stream_chunked(
+                    response.content.iter_any()
+                ):
+                    yield message
 
         except aiohttp.ClientResponseError as e:
             self._handle_exception(e, request_data.url, request_data.correlation_id)
@@ -399,8 +426,6 @@ class AiriaAsyncClient(AiriaBaseClient):
             ...     )
             >>> print(response.result)
         """
-        self._check_session()
-
         request_data = self._pre_execute_pipeline(
             pipeline_id=pipeline_id,
             user_input=user_input,
@@ -446,3 +471,112 @@ class AiriaAsyncClient(AiriaBaseClient):
             return PipelineExecutionV1StreamedResponse(**resp)
 
         return PipelineExecutionV2AsyncStreamedResponse(stream=resp)
+
+    async def get_active_pipelines_ids(
+        self,
+        correlation_id: Optional[str] = None,
+        api_version: str = ApiVersion.V1.value,
+    ) -> List[str]:
+        """
+        Retrieve a list of active pipeline IDs.
+        
+        This method fetches all currently active pipeline IDs from the Airia API.
+        These IDs can be used with other methods like execute_pipeline() or 
+        get_pipeline_config().
+        
+        Args:
+            api_version (str, optional): API version to use for the request. 
+                Must be one of the supported versions. Defaults to "v1".
+            correlation_id (str, optional): Unique identifier for request tracing
+                and logging. If not provided, a new UUID will be automatically 
+                generated.
+        
+        Returns:
+            List[str]: A list of active pipeline ID strings. Returns an empty list
+                if no active pipelines are found.
+        
+        Raises:
+            ValueError: If the provided API version is not supported.
+            AiriaAPIError: If the API request fails, including network errors,
+                authentication failures, or server errors.
+        
+        Example:
+            >>> client = AiriaClient(api_key="your_api_key")
+            >>> pipeline_ids = client.get_active_pipelines_ids()
+            >>> print(f"Found {len(pipeline_ids)} active pipelines")
+            >>> for pipeline_id in pipeline_ids:
+            ...     print(f"Pipeline ID: {pipeline_id}")
+        """
+        request_data = self._pre_get_active_pipelines_ids(
+            correlation_id=correlation_id, api_version=api_version
+        )
+        resp = await 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
+
+    async def get_pipeline_config(
+        self,
+        pipeline_id: str,
+        correlation_id: Optional[str] = None,
+        api_version: str = ApiVersion.V1.value,
+    ) -> 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.
+            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.
+
+        Returns:
+            GetPipelineConfigResponse: A response object containing the pipeline
+                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)
+                - 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=api_version,
+        )
+        resp = await self._make_request("GET", request_data)
+
+        return GetPipelineConfigResponse(**resp)