airia 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

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
 
 
@@ -46,8 +50,31 @@ class AiriaAsyncClient(AiriaBaseClient):
         )
 
         # 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(
@@ -117,24 +144,6 @@ class AiriaAsyncClient(AiriaBaseClient):
 
         return cls(base_url, api_key, timeout, log_requests, custom_logger)
 
-    async def __aenter__(self):
-        """Async context manager entry point."""
-        self.session = aiohttp.ClientSession(headers=self.headers)
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        """Async context manager exit point."""
-        if self.session:
-            await self.session.close()
-            self.session = None
-
-    def _check_session(self):
-        """Check if the client session is initialized."""
-        if not self.session:
-            raise RuntimeError(
-                "Client session not initialized. Use async with AiriaAsyncClient() as client: ..."
-            )
-
     def _handle_exception(
         self, e: aiohttp.ClientResponseError, url: str, correlation_id: str
     ):
@@ -257,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)
@@ -415,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,
@@ -462,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)

airia/client/base_client.py

@@ -11,6 +11,7 @@ from ..types import ApiVersion, RequestData
 
 class AiriaBaseClient:
     """Base client containing shared functionality for Airia API clients."""
+
     openai = None
     anthropic = None
 
@@ -186,3 +187,32 @@ class AiriaBaseClient:
         request_data = self._prepare_request(url, payload, correlation_id)
 
         return request_data
+
+    def _pre_get_active_pipelines_ids(
+        self,
+        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}/PipelinesConfig")
+        request_data = self._prepare_request(url, correlation_id=correlation_id)
+
+        return request_data
+
+    def _pre_get_pipeline_config(
+        self,
+        pipeline_id: str,
+        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}/PipelinesConfig/{pipeline_id}")
+        request_data = self._prepare_request(url, correlation_id=correlation_id)
+
+        return request_data

airia/client/sync_client.py

@@ -7,12 +7,14 @@ import requests
 from ..exceptions import AiriaAPIError
 from ..types import (
     ApiVersion,
+    GetPipelineConfigResponse,
     PipelineExecutionDebugResponse,
     PipelineExecutionResponse,
     PipelineExecutionV1StreamedResponse,
     PipelineExecutionV2StreamedResponse,
     RequestData,
 )
+from ..utils.sse_parser import parse_sse_stream_chunked
 from .base_client import AiriaBaseClient
 
 
@@ -150,7 +152,7 @@ class AiriaClient(AiriaBaseClient):
             status_code=e.response.status_code, message=sanitized_message
         ) from e
 
-    def _make_request(self, method: str, request_data: RequestData):
+    def _make_request(self, method: str, request_data: RequestData) -> Dict[str, Any]:
         """
         Makes a synchronous HTTP request to the Airia API.
 
@@ -247,8 +249,8 @@ class AiriaClient(AiriaBaseClient):
             response.raise_for_status()
 
             # Yields the response content as a stream
-            for chunk in response.iter_content():
-                yield chunk.decode("utf-8")
+            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)
@@ -426,10 +428,11 @@ class AiriaClient(AiriaBaseClient):
             api_version=api_version,
         )
         stream = async_output and api_version == ApiVersion.V2.value
-        if stream:
-            resp = self._make_request_stream("POST", request_data)
-        else:
-            resp = self._make_request("POST", request_data)
+        resp = (
+            self._make_request_stream("POST", request_data)
+            if stream
+            else self._make_request("POST", request_data)
+        )
 
         if not async_output:
             if not debug:
@@ -450,3 +453,112 @@ class AiriaClient(AiriaBaseClient):
             return PipelineExecutionV1StreamedResponse(**resp)
 
         return PipelineExecutionV2StreamedResponse(stream=resp)
+
+    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 = 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,
+        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 = self._make_request("GET", request_data)
+
+        return GetPipelineConfigResponse(**resp)

airia/types/__init__.py

@@ -1,12 +1,39 @@
-from .api_version import ApiVersion
-from .pipeline_execution import (
+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 (
+    AgentAgentCardMessage,
+    AgentAgentCardStreamEndMessage,
+    AgentAgentCardStreamErrorMessage,
+    AgentAgentCardStreamFragmentMessage,
+    AgentAgentCardStreamStartMessage,
+    AgentDatasearchMessage,
+    AgentEndMessage,
+    AgentInvocationMessage,
+    AgentModelMessage,
+    AgentModelStreamEndMessage,
+    AgentModelStreamErrorMessage,
+    AgentModelStreamFragmentMessage,
+    AgentModelStreamStartMessage,
+    AgentModelStreamUsageMessage,
+    AgentOutputMessage,
+    AgentPingMessage,
+    AgentPythonCodeMessage,
+    AgentStartMessage,
+    AgentStepEndMessage,
+    AgentStepHaltMessage,
+    AgentStepStartMessage,
+    AgentToolActionMessage,
+    AgentToolRequestMessage,
+    AgentToolResponseMessage,
+)
 
 __all__ = [
     "ApiVersion",
@@ -15,5 +42,30 @@ __all__ = [
     "PipelineExecutionV1StreamedResponse",
     "PipelineExecutionV2AsyncStreamedResponse",
     "PipelineExecutionV2StreamedResponse",
+    "GetPipelineConfigResponse",
     "RequestData",
+    "AgentPingMessage",
+    "AgentStartMessage",
+    "AgentEndMessage",
+    "AgentStepStartMessage",
+    "AgentStepHaltMessage",
+    "AgentStepEndMessage",
+    "AgentOutputMessage",
+    "AgentAgentCardMessage",
+    "AgentDatasearchMessage",
+    "AgentInvocationMessage",
+    "AgentModelMessage",
+    "AgentPythonCodeMessage",
+    "AgentToolActionMessage",
+    "AgentModelStreamStartMessage",
+    "AgentModelStreamEndMessage",
+    "AgentModelStreamErrorMessage",
+    "AgentModelStreamUsageMessage",
+    "AgentModelStreamFragmentMessage",
+    "AgentAgentCardStreamStartMessage",
+    "AgentAgentCardStreamErrorMessage",
+    "AgentAgentCardStreamFragmentMessage",
+    "AgentAgentCardStreamEndMessage",
+    "AgentToolRequestMessage",
+    "AgentToolResponseMessage",
 ]

airia/types/api/get_pipeline_config.py

@@ -0,0 +1,65 @@
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class Version(BaseModel):
+    pipeline_id: str = Field(alias="pipelineId")
+    major_version: int = Field(alias="majorVersion")
+    minor_version: int = Field(alias="minorVersion")
+    version_number: str = Field(alias="versionNumber")
+    is_draft_version: bool = Field(alias="isDraftVersion")
+    is_latest: bool = Field(alias="isLatest")
+    steps: Optional[List[Dict[str, Any]]] = None
+    alignment: str
+    id: str
+    tenant_id: str = Field(alias="tenantId")
+    project_id: str = Field(alias="projectId")
+    created_at: datetime = Field(alias="createdAt")
+    updated_at: datetime = Field(alias="updatedAt")
+    user_id: str = Field(alias="userId")
+
+
+class ExecutionStats(BaseModel):
+    success_count: int = Field(alias="successCount")
+    failure_count: int = Field(alias="failureCount")
+
+
+class GetPipelineConfigResponse(BaseModel):
+    deployment_id: Optional[str] = Field(alias="deploymentId")
+    deployment_name: Optional[str] = Field(alias="deploymentName")
+    deployment_description: Optional[str] = Field(alias="deploymentDescription")
+    user_keys: Dict[str, Any] = Field(alias="userKeys")
+    group_keys: Dict[str, Any] = Field(alias="groupKeys")
+    agent_icon: Optional[str] = Field(alias="agentIcon")
+    external: bool
+    active_version_id: str = Field(alias="activeVersionId")
+    name: str
+    execution_name: str = Field(alias="executionName")
+    description: str
+    video_link: Optional[str] = Field(alias="videoLink")
+    agent_icon_id: Optional[str] = Field(alias="agentIconId")
+    versions: List[Version]
+    execution_stats: ExecutionStats = Field(alias="executionStats")
+    industry: Optional[str]
+    sub_industries: List[str] = Field(alias="subIndustries")
+    agent_details: Dict[str, Any] = Field(alias="agentDetails")
+    agent_details_tags: List[str] = Field(alias="agentDetailsTags")
+    active_version: Version = Field(alias="activeVersion")
+    backup_pipeline_id: Optional[str] = Field(alias="backupPipelineId")
+    deployment: Optional[Any]
+    library_agent_id: Optional[str] = Field(alias="libraryAgentId")
+    library_imported_hash: Optional[str] = Field(alias="libraryImportedHash")
+    library_imported_version: Optional[str] = Field(alias="libraryImportedVersion")
+    is_deleted: Optional[bool] = Field(alias="isDeleted")
+    agent_trigger: Optional[Any] = Field(alias="agentTrigger")
+    api_key_id: Optional[str] = Field(alias="apiKeyId")
+    is_seeded: bool = Field(alias="isSeeded")
+    behaviours: List[Any]
+    id: str
+    tenant_id: str = Field(alias="tenantId")
+    project_id: str = Field(alias="projectId")
+    created_at: datetime = Field(alias="createdAt")
+    updated_at: datetime = Field(alias="updatedAt")
+    user_id: str = Field(alias="userId")

airia/types/{pipeline_execution.py → api/pipeline_execution.py}

@@ -2,6 +2,8 @@ from typing import Any, AsyncIterator, Dict, Iterator
 
 from pydantic import BaseModel, ConfigDict
 
+from ..sse_messages import SSEMessage
+
 
 class PipelineExecutionResponse(BaseModel):
     result: str
@@ -22,10 +24,10 @@ class PipelineExecutionV1StreamedResponse(BaseModel):
 class PipelineExecutionV2StreamedResponse(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    stream: Iterator[str]
+    stream: Iterator[SSEMessage]
 
 
 class PipelineExecutionV2AsyncStreamedResponse(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    stream: AsyncIterator[str]
+    stream: AsyncIterator[SSEMessage]

airia/types/request_data.py

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

airia/types/sse_messages.py

@@ -0,0 +1,269 @@
+from datetime import datetime, time
+from enum import Enum
+from typing import Optional, Union
+
+from pydantic import BaseModel, ConfigDict
+
+
+class MessageType(str, Enum):
+    AGENT_PING = "AgentPingMessage"
+    AGENT_START = "AgentStartMessage"
+    AGENT_END = "AgentEndMessage"
+    AGENT_STEP_START = "AgentStepStartMessage"
+    AGENT_STEP_HALT = "AgentStepHaltMessage"
+    AGENT_STEP_END = "AgentStepEndMessage"
+    AGENT_OUTPUT = "AgentOutputMessage"
+    AGENT_AGENT_CARD = "AgentAgentCardMessage"
+    AGENT_DATASEARCH = "AgentDatasearchMessage"
+    AGENT_INVOCATION = "AgentInvocationMessage"
+    AGENT_MODEL = "AgentModelMessage"
+    AGENT_PYTHON_CODE = "AgentPythonCodeMessage"
+    AGENT_TOOL_ACTION = "AgentToolActionMessage"
+    AGENT_MODEL_STREAM_START = "AgentModelStreamStartMessage"
+    AGENT_MODEL_STREAM_END = "AgentModelStreamEndMessage"
+    AGENT_MODEL_STREAM_ERROR = "AgentModelStreamErrorMessage"
+    AGENT_MODEL_STREAM_USAGE = "AgentModelStreamUsageMessage"
+    AGENT_MODEL_STREAM_FRAGMENT = "AgentModelStreamFragmentMessage"
+    MODEL_STREAM_FRAGMENT = "ModelStreamFragment"
+    AGENT_AGENT_CARD_STREAM_START = "AgentAgentCardStreamStartMessage"
+    AGENT_AGENT_CARD_STREAM_ERROR = "AgentAgentCardStreamErrorMessage"
+    AGENT_AGENT_CARD_STREAM_FRAGMENT = "AgentAgentCardStreamFragmentMessage"
+    AGENT_AGENT_CARD_STREAM_END = "AgentAgentCardStreamEndMessage"
+    AGENT_TOOL_REQUEST = "AgentToolRequestMessage"
+    AGENT_TOOL_RESPONSE = "AgentToolResponseMessage"
+
+
+class BaseSSEMessage(BaseModel):
+    model_config = ConfigDict(use_enum_values=True)
+    message_type: MessageType
+
+
+class AgentPingMessage(BaseSSEMessage):
+    message_type: MessageType = MessageType.AGENT_PING
+    timestamp: datetime
+
+
+### Agent Messages ###
+
+
+class BaseAgentMessage(BaseSSEMessage):
+    agent_id: str
+    execution_id: str
+
+
+class AgentStartMessage(BaseAgentMessage):
+    message_type: MessageType = MessageType.AGENT_START
+
+
+class AgentEndMessage(BaseAgentMessage):
+    message_type: MessageType = MessageType.AGENT_END
+
+
+### Step Messages ###
+
+
+class BaseStepMessage(BaseAgentMessage):
+    step_id: str
+    step_type: str
+    step_title: Optional[str] = None
+
+
+class AgentStepStartMessage(BaseStepMessage):
+    message_type: MessageType = MessageType.AGENT_STEP_START
+    start_time: datetime
+
+
+class AgentStepHaltMessage(BaseStepMessage):
+    message_type: MessageType = MessageType.AGENT_STEP_HALT
+    approval_id: str
+
+
+class AgentStepEndMessage(BaseStepMessage):
+    message_type: MessageType = MessageType.AGENT_STEP_END
+    end_time: datetime
+    duration: time
+    status: str
+
+
+class AgentOutputMessage(BaseStepMessage):
+    message_type: MessageType = MessageType.AGENT_OUTPUT
+    step_result: str
+
+
+### Status Messages ###
+
+
+class BaseStatusMessage(BaseStepMessage):
+    pass
+
+
+class AgentAgentCardMessage(BaseStatusMessage):
+    message_type: MessageType = MessageType.AGENT_AGENT_CARD
+    step_name: str
+
+
+class AgentDatasearchMessage(BaseStatusMessage):
+    message_type: MessageType = MessageType.AGENT_DATASEARCH
+    datastore_id: str
+    datastore_type: str
+    datastore_name: str
+
+
+class AgentInvocationMessage(BaseStatusMessage):
+    message_type: MessageType = MessageType.AGENT_INVOCATION
+    agent_name: str
+
+
+class AgentModelMessage(BaseStatusMessage):
+    message_type: MessageType = MessageType.AGENT_MODEL
+    model_name: str
+
+
+class AgentPythonCodeMessage(BaseStatusMessage):
+    message_type: MessageType = MessageType.AGENT_PYTHON_CODE
+    step_name: str
+
+
+class AgentToolActionMessage(BaseStatusMessage):
+    message_type: MessageType = MessageType.AGENT_TOOL_ACTION
+    step_name: str
+    tool_name: str
+
+
+### Model Stream Messages ###
+
+
+class BaseModelStreamMessage(BaseAgentMessage):
+    step_id: str
+    stream_id: str
+
+
+class AgentModelStreamStartMessage(BaseModelStreamMessage):
+    message_type: MessageType = MessageType.AGENT_MODEL_STREAM_START
+    model_name: str
+
+
+class AgentModelStreamErrorMessage(BaseModelStreamMessage):
+    message_type: MessageType = MessageType.AGENT_MODEL_STREAM_ERROR
+    error_message: str
+
+
+class AgentModelStreamFragmentMessage(BaseModelStreamMessage):
+    message_type: MessageType = MessageType.AGENT_MODEL_STREAM_FRAGMENT
+    index: int
+    content: Optional[str] = None
+
+
+class AgentModelStreamEndMessage(BaseModelStreamMessage):
+    message_type: MessageType = MessageType.AGENT_MODEL_STREAM_END
+    content_id: str
+    duration: Optional[float] = None
+
+
+class AgentModelStreamUsageMessage(BaseModelStreamMessage):
+    message_type: MessageType = MessageType.AGENT_MODEL_STREAM_USAGE
+    token: Optional[int] = None
+    tokens_cost: Optional[float] = None
+
+
+### Agent Card Messages ###
+
+
+class BaseAgentAgentCardStreamMessage(BaseAgentMessage):
+    step_id: str
+    stream_id: str
+
+
+class AgentAgentCardStreamStartMessage(BaseAgentAgentCardStreamMessage):
+    message_type: MessageType = MessageType.AGENT_AGENT_CARD_STREAM_START
+    content: Optional[str] = None
+
+
+class AgentAgentCardStreamErrorMessage(BaseAgentAgentCardStreamMessage):
+    message_type: MessageType = MessageType.AGENT_AGENT_CARD_STREAM_ERROR
+    error_message: str
+
+
+class AgentAgentCardStreamFragmentMessage(BaseAgentAgentCardStreamMessage):
+    message_type: MessageType = MessageType.AGENT_AGENT_CARD_STREAM_FRAGMENT
+    index: int
+    content: Optional[str]
+
+
+class AgentAgentCardStreamEndMessage(BaseAgentAgentCardStreamMessage):
+    message_type: MessageType = MessageType.AGENT_AGENT_CARD_STREAM_END
+    content: Optional[str] = None
+
+
+### Tool Messages ###
+
+
+class BaseAgentToolMessage(BaseStepMessage):
+    id: str
+    name: str
+
+
+class AgentToolRequestMessage(BaseAgentToolMessage):
+    message_type: MessageType = MessageType.AGENT_TOOL_REQUEST
+
+
+class AgentToolResponseMessage(BaseAgentToolMessage):
+    message_type: MessageType = MessageType.AGENT_TOOL_RESPONSE
+    duration: time
+    success: bool
+
+
+# Union type for all possible messages
+SSEMessage = Union[
+    AgentPingMessage,
+    AgentStartMessage,
+    AgentEndMessage,
+    AgentStepStartMessage,
+    AgentStepHaltMessage,
+    AgentStepEndMessage,
+    AgentOutputMessage,
+    AgentAgentCardMessage,
+    AgentDatasearchMessage,
+    AgentInvocationMessage,
+    AgentModelMessage,
+    AgentPythonCodeMessage,
+    AgentToolActionMessage,
+    AgentModelStreamStartMessage,
+    AgentModelStreamEndMessage,
+    AgentModelStreamErrorMessage,
+    AgentModelStreamUsageMessage,
+    AgentModelStreamFragmentMessage,
+    AgentAgentCardStreamStartMessage,
+    AgentAgentCardStreamErrorMessage,
+    AgentAgentCardStreamFragmentMessage,
+    AgentAgentCardStreamEndMessage,
+    AgentToolRequestMessage,
+    AgentToolResponseMessage,
+]
+
+SSEDict = {
+    MessageType.AGENT_PING.value: AgentPingMessage,
+    MessageType.AGENT_START.value: AgentStartMessage,
+    MessageType.AGENT_END.value: AgentEndMessage,
+    MessageType.AGENT_STEP_START.value: AgentStepStartMessage,
+    MessageType.AGENT_STEP_HALT.value: AgentStepHaltMessage,
+    MessageType.AGENT_STEP_END.value: AgentStepEndMessage,
+    MessageType.AGENT_OUTPUT.value: AgentOutputMessage,
+    MessageType.AGENT_AGENT_CARD.value: AgentAgentCardMessage,
+    MessageType.AGENT_DATASEARCH.value: AgentDatasearchMessage,
+    MessageType.AGENT_INVOCATION.value: AgentInvocationMessage,
+    MessageType.AGENT_MODEL.value: AgentModelMessage,
+    MessageType.AGENT_PYTHON_CODE.value: AgentPythonCodeMessage,
+    MessageType.AGENT_TOOL_ACTION.value: AgentToolActionMessage,
+    MessageType.AGENT_MODEL_STREAM_START.value: AgentModelStreamStartMessage,
+    MessageType.AGENT_MODEL_STREAM_END.value: AgentModelStreamEndMessage,
+    MessageType.AGENT_MODEL_STREAM_ERROR.value: AgentModelStreamErrorMessage,
+    MessageType.AGENT_MODEL_STREAM_USAGE.value: AgentModelStreamUsageMessage,
+    MessageType.AGENT_MODEL_STREAM_FRAGMENT.value: AgentModelStreamFragmentMessage,
+    MessageType.AGENT_AGENT_CARD_STREAM_START.value: AgentAgentCardStreamStartMessage,
+    MessageType.AGENT_AGENT_CARD_STREAM_ERROR.value: AgentAgentCardStreamErrorMessage,
+    MessageType.AGENT_AGENT_CARD_STREAM_FRAGMENT.value: AgentAgentCardStreamFragmentMessage,
+    MessageType.AGENT_AGENT_CARD_STREAM_END.value: AgentAgentCardStreamEndMessage,
+    MessageType.AGENT_TOOL_REQUEST.value: AgentToolRequestMessage,
+    MessageType.AGENT_TOOL_RESPONSE.value: AgentToolResponseMessage,
+}

airia/utils/sse_parser.py

@@ -0,0 +1,91 @@
+import json
+import re
+from typing import AsyncIterable, AsyncIterator, Iterable, Iterator
+
+from ..types.sse_messages import SSEDict, SSEMessage
+
+
+def _to_snake_case(name: str):
+    return re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower()
+
+
+def parse_sse_stream_chunked(stream_chunks: Iterable[bytes]) -> Iterator[SSEMessage]:
+    """
+    Parse SSE stream from an iterable of chunks (e.g., from a streaming response).
+
+    Args:
+        stream_chunks: Iterable of string chunks
+
+    Yields:
+        dict: Dictionary containing 'event' and 'data' keys
+    """
+    buffer = ""
+
+    for chunk in stream_chunks:
+        buffer += chunk.decode("utf-8")
+
+        # Look for complete events (ending with double newline)
+        while "\n\n" in buffer:
+            event_block, buffer = buffer.split("\n\n", 1)
+
+            if not event_block.strip():
+                continue
+
+            event_name = None
+            event_data = None
+
+            # Parse each line in the event block
+            for line in event_block.strip().split("\n"):
+                line = line.strip()
+                if line.startswith("event:"):
+                    event_name = line[6:].strip()
+                elif line.startswith("data:"):
+                    data_json = line[5:].strip()
+                    event_data = json.loads(data_json)
+
+            if event_name and event_data:
+                yield SSEDict[event_name](
+                    **{_to_snake_case(k): v for k, v in event_data.items()}
+                )
+
+
+async def async_parse_sse_stream_chunked(
+    stream_chunks: AsyncIterable[bytes],
+) -> AsyncIterator[SSEMessage]:
+    """
+    Parse SSE stream from an iterable of chunks (e.g., from a streaming response).
+
+    Args:
+        stream_chunks: Iterable of string chunks
+
+    Yields:
+        dict: Dictionary containing 'event' and 'data' keys
+    """
+    buffer = ""
+
+    async for chunk in stream_chunks:
+        buffer += chunk.decode("utf-8")
+
+        # Look for complete events (ending with double newline)
+        while "\n\n" in buffer:
+            event_block, buffer = buffer.split("\n\n", 1)
+
+            if not event_block.strip():
+                continue
+
+            event_name = None
+            event_data = None
+
+            # Parse each line in the event block
+            for line in event_block.strip().split("\n"):
+                line = line.strip()
+                if line.startswith("event:"):
+                    event_name = line[6:].strip()
+                elif line.startswith("data:"):
+                    data_json = line[5:].strip()
+                    event_data = json.loads(data_json)
+
+            if event_name and event_data:
+                yield SSEDict[event_name](
+                    **{_to_snake_case(k): v for k, v in event_data.items()}
+                )

{airia-0.1.5.dist-info → airia-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airia
-Version: 0.1.5
+Version: 0.1.6
 Summary: Python SDK for Airia API
 Author-email: Airia LLC <support@airia.com>
 License: MIT
@@ -182,7 +182,7 @@ This will create both wheel and source distribution in the `dist/` directory.
 from airia import AiriaClient
 
 client = AiriaClient(
-    base_url="https://api.airia.com",        # Default: "https://api.airia.com"
+    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)
@@ -223,7 +223,7 @@ response = client.execute_pipeline(
 )
 
 for c in response.stream:
-    print(c, end="")
+    print(c)
 ```
 
 ### Asynchronous Usage
@@ -233,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())
 ```
@@ -251,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 response.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.

airia-0.1.6.dist-info/RECORD

@@ -0,0 +1,19 @@
+airia/__init__.py,sha256=T39gO8E5T5zxlw-JP78ruxOu7-LeKOJCJzz6t40kdQo,231
+airia/exceptions.py,sha256=4Z55n-cRJrtTa5-pZBIK2oZD4-Z99aUtKx_kfTFYY5o,1146
+airia/logs.py,sha256=17YZ4IuzOF0m5bgofj9-QYlJ2BYR2kRZbBVQfFSLFEk,5441
+airia/client/__init__.py,sha256=6gSQ9bl7j79q1HPE0o5py3IRdkwWWuU_7J4h05Dd2o8,127
+airia/client/async_client.py,sha256=QtnH1MNfF_ONx2YHsTz7Y4pFBTrggHWz0LLJiITREwc,23694
+airia/client/base_client.py,sha256=IQYksAw6xCFmAarPUGxyp2bHm95iTzdZUv8FNUq4fHk,7962
+airia/client/sync_client.py,sha256=iUqRn7Bv5hi-Qcg3Vgfc-Z1gx5T42f9Vq4pUTPsyG5M,22750
+airia/types/__init__.py,sha256=s6_uMrrKzhk2wg1lpaRRIl4AQ8hUxKbCNLfz-VK7tAs,2166
+airia/types/api_version.py,sha256=Uzom6O2ZG92HN_Z2h-lTydmO2XYX9RVs4Yi4DJmXytE,255
+airia/types/request_data.py,sha256=m8lBwFliK2_kZU2TgeLdhiO_7v826lVRhA5_8NYB_NM,206
+airia/types/sse_messages.py,sha256=wSdowY07AjEO8R73SJrFPJtkfIBS4satUpNytjKQq2U,8305
+airia/types/api/get_pipeline_config.py,sha256=hHyIfzpH-miayAAUtgf5ywXfqOg53dUMiEfDHW1I-tg,2837
+airia/types/api/pipeline_execution.py,sha256=4zGM5W4aXiAibAdgWQCIZBC_blW6QYAlAHoVUMbnCF0,752
+airia/utils/sse_parser.py,sha256=h3TcBvXqUIngTqgY6yYxZCGLnC1eI6meQzYr13aFAb8,2791
+airia-0.1.6.dist-info/licenses/LICENSE,sha256=R3ClUMMKPRItIcZ0svzyj2taZZnFYw568YDNzN9KQ1Q,1066
+airia-0.1.6.dist-info/METADATA,sha256=1gvnnQKQhkJR31GkL4RG9Vw91UTi9Tclcjv5d9nxpMk,13471
+airia-0.1.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+airia-0.1.6.dist-info/top_level.txt,sha256=qUQEKfs_hdOYTwjKj1JZbRhS5YeXDNaKQaVTrzabS6w,6
+airia-0.1.6.dist-info/RECORD,,

airia-0.1.5.dist-info/RECORD

@@ -1,16 +0,0 @@
-airia/__init__.py,sha256=T39gO8E5T5zxlw-JP78ruxOu7-LeKOJCJzz6t40kdQo,231
-airia/exceptions.py,sha256=4Z55n-cRJrtTa5-pZBIK2oZD4-Z99aUtKx_kfTFYY5o,1146
-airia/logs.py,sha256=17YZ4IuzOF0m5bgofj9-QYlJ2BYR2kRZbBVQfFSLFEk,5441
-airia/client/__init__.py,sha256=6gSQ9bl7j79q1HPE0o5py3IRdkwWWuU_7J4h05Dd2o8,127
-airia/client/async_client.py,sha256=4JU6RePXjznulnQXkfo9ZVX3zLeUyclOgxwlaM7-5Js,18816
-airia/client/base_client.py,sha256=IEVsskGpiEL3jw7YLc96YWUZ5rpUZula52IZHe2D8mE,6844
-airia/client/sync_client.py,sha256=mCLViWkW-z8QhCy8CmlEsT8dssBb1g6V2B3fAYgYU1A,18336
-airia/types/__init__.py,sha256=OOFtJ0VIbO98px89u7cq645iL7CYDOel43g85IAjRRw,562
-airia/types/api_version.py,sha256=Uzom6O2ZG92HN_Z2h-lTydmO2XYX9RVs4Yi4DJmXytE,255
-airia/types/pipeline_execution.py,sha256=obp8KOz-ShNxEciF1YQCSpHXPoJUyEa_9uPv-VHf6YA,699
-airia/types/request_data.py,sha256=RbVPWPRIYuT7FaolJKn19IVzuQKbLM4VhXM5r7UXTUY,186
-airia-0.1.5.dist-info/licenses/LICENSE,sha256=R3ClUMMKPRItIcZ0svzyj2taZZnFYw568YDNzN9KQ1Q,1066
-airia-0.1.5.dist-info/METADATA,sha256=WIPLholUrToVlvO4K269qOSYHT4fD7XHpVSrjowOF6A,10745
-airia-0.1.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-airia-0.1.5.dist-info/top_level.txt,sha256=qUQEKfs_hdOYTwjKj1JZbRhS5YeXDNaKQaVTrzabS6w,6
-airia-0.1.5.dist-info/RECORD,,