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

airia/types/sse/sse_messages.py

@@ -1,3 +1,11 @@
+"""
+Server-Sent Event (SSE) message types and models for the Airia API.
+
+This module defines all possible SSE message types that can be received during
+pipeline execution, including agent lifecycle events, processing steps, model
+streaming, and tool execution messages.
+"""
+
 from datetime import datetime, time
 from enum import Enum
 from typing import Optional, Union
@@ -6,6 +14,13 @@ from pydantic import BaseModel, ConfigDict
 
 
 class MessageType(str, Enum):
+    """
+    Enumeration of all possible SSE message types from the Airia API.
+
+    These message types correspond to different events that occur during
+    pipeline execution, agent processing, and streaming responses.
+    """
+
     AGENT_PING = "AgentPingMessage"
     AGENT_START = "AgentStartMessage"
     AGENT_INPUT = "AgentInputMessage"
@@ -35,11 +50,25 @@ class MessageType(str, Enum):
 
 
 class BaseSSEMessage(BaseModel):
+    """
+    Base class for all Server-Sent Event (SSE) messages from the Airia API.
+
+    All SSE messages include a message_type field that identifies the specific
+    type of event being reported.
+    """
+
     model_config = ConfigDict(use_enum_values=True)
     message_type: MessageType
 
 
 class AgentPingMessage(BaseSSEMessage):
+    """
+    Ping message sent periodically to maintain connection health.
+
+    These messages help verify that the connection is still active during
+    long-running pipeline executions.
+    """
+
     message_type: MessageType = MessageType.AGENT_PING
     timestamp: datetime
 
@@ -48,19 +77,38 @@ class AgentPingMessage(BaseSSEMessage):
 
 
 class BaseAgentMessage(BaseSSEMessage):
+    """
+    Base class for messages related to agent execution.
+
+    All agent messages include identifiers for the specific agent
+    and execution session.
+    """
+
     agent_id: str
     execution_id: str
 
 
 class AgentStartMessage(BaseAgentMessage):
+    """
+    Message indicating that an agent has started processing.
+    """
+
     message_type: MessageType = MessageType.AGENT_START
 
 
 class AgentInputMessage(BaseAgentMessage):
+    """
+    Message indicating that an agent has received input to process.
+    """
+
     message_type: MessageType = MessageType.AGENT_INPUT
 
 
 class AgentEndMessage(BaseAgentMessage):
+    """
+    Message indicating that an agent has finished processing.
+    """
+
     message_type: MessageType = MessageType.AGENT_END
 
 
@@ -68,22 +116,46 @@ class AgentEndMessage(BaseAgentMessage):
 
 
 class BaseStepMessage(BaseAgentMessage):
+    """
+    Base class for messages related to individual processing steps within an agent.
+
+    Steps represent discrete operations or tasks that an agent performs
+    as part of its overall processing workflow.
+    """
+
     step_id: str
     step_type: str
     step_title: Optional[str] = None
 
 
 class AgentStepStartMessage(BaseStepMessage):
+    """
+    Message indicating that a processing step has started.
+    """
+
     message_type: MessageType = MessageType.AGENT_STEP_START
     start_time: datetime
 
 
 class AgentStepHaltMessage(BaseStepMessage):
+    """
+    Message indicating that a step has been halted pending approval.
+
+    This occurs when human approval is required before proceeding
+    with potentially sensitive or high-impact operations.
+    """
+
     message_type: MessageType = MessageType.AGENT_STEP_HALT
     approval_id: str
 
 
 class AgentStepEndMessage(BaseStepMessage):
+    """
+    Message indicating that a processing step has completed.
+
+    Includes timing information and the final status of the step.
+    """
+
     message_type: MessageType = MessageType.AGENT_STEP_END
     end_time: datetime
     duration: time
@@ -91,6 +163,10 @@ class AgentStepEndMessage(BaseStepMessage):
 
 
 class AgentOutputMessage(BaseStepMessage):
+    """
+    Message containing the output result from a completed step.
+    """
+
     message_type: MessageType = MessageType.AGENT_OUTPUT
     step_result: str
 
@@ -99,15 +175,36 @@ class AgentOutputMessage(BaseStepMessage):
 
 
 class BaseStatusMessage(BaseStepMessage):
+    """
+    Base class for status update messages within processing steps.
+
+    Status messages provide real-time updates about what operations
+    are being performed during step execution.
+    """
+
     pass
 
 
 class AgentAgentCardMessage(BaseStatusMessage):
+    """
+    Message indicating that an agent card step is being processed.
+
+    Agent cards represent interactive UI components or displays
+    that provide rich information to users during pipeline execution.
+    """
+
     message_type: MessageType = MessageType.AGENT_AGENT_CARD
     step_name: str
 
 
 class AgentDatasearchMessage(BaseStatusMessage):
+    """
+    Message indicating that data source search is being performed.
+
+    This message is sent when an agent is querying or searching
+    through configured data sources to retrieve relevant information.
+    """
+
     message_type: MessageType = MessageType.AGENT_DATASEARCH
     datastore_id: str
     datastore_type: str
@@ -115,21 +212,49 @@ class AgentDatasearchMessage(BaseStatusMessage):
 
 
 class AgentInvocationMessage(BaseStatusMessage):
+    """
+    Message indicating that another agent is being invoked.
+
+    This occurs when the current agent calls or delegates work
+    to another specialized agent in the pipeline.
+    """
+
     message_type: MessageType = MessageType.AGENT_INVOCATION
     agent_name: str
 
 
 class AgentModelMessage(BaseStatusMessage):
+    """
+    Message indicating that a language model is being called.
+
+    This message is sent when an agent begins interacting with
+    a language model for text generation or processing.
+    """
+
     message_type: MessageType = MessageType.AGENT_MODEL
     model_name: str
 
 
 class AgentPythonCodeMessage(BaseStatusMessage):
+    """
+    Message indicating that Python code execution is taking place.
+
+    This message is sent when an agent executes custom Python code
+    blocks as part of its processing workflow.
+    """
+
     message_type: MessageType = MessageType.AGENT_PYTHON_CODE
     step_name: str
 
 
 class AgentToolActionMessage(BaseStatusMessage):
+    """
+    Message indicating that a tool or external service is being called.
+
+    This message is sent when an agent invokes an external tool,
+    API, or service to perform a specific action or retrieve data.
+    """
+
     message_type: MessageType = MessageType.AGENT_TOOL_ACTION
     step_name: str
     tool_name: str
@@ -139,33 +264,63 @@ class AgentToolActionMessage(BaseStatusMessage):
 
 
 class BaseModelStreamMessage(BaseAgentMessage):
+    """
+    Base class for language model streaming messages.
+
+    Model streaming allows real-time display of text generation
+    as it occurs, providing better user experience for long responses.
+    """
+
     step_id: str
     stream_id: str
 
 
 class AgentModelStreamStartMessage(BaseModelStreamMessage):
+    """
+    Message indicating that model text streaming has begun.
+    """
+
     message_type: MessageType = MessageType.AGENT_MODEL_STREAM_START
     model_name: str
 
 
 class AgentModelStreamErrorMessage(BaseModelStreamMessage):
+    """
+    Message indicating that an error occurred during model streaming.
+    """
+
     message_type: MessageType = MessageType.AGENT_MODEL_STREAM_ERROR
     error_message: str
 
 
 class AgentModelStreamFragmentMessage(BaseModelStreamMessage):
+    """
+    Fragment of streaming text content from a language model.
+
+    These messages contain individual chunks of text as they are generated
+    by the model, allowing for real-time display of results.
+    """
+
     message_type: MessageType = MessageType.AGENT_MODEL_STREAM_FRAGMENT
     index: int
     content: Optional[str] = None
 
 
 class AgentModelStreamEndMessage(BaseModelStreamMessage):
+    """
+    Message indicating that model text streaming has completed.
+    """
+
     message_type: MessageType = MessageType.AGENT_MODEL_STREAM_END
     content_id: str
     duration: Optional[float] = None
 
 
 class AgentModelStreamUsageMessage(BaseModelStreamMessage):
+    """
+    Message containing token usage and cost information for model calls.
+    """
+
     message_type: MessageType = MessageType.AGENT_MODEL_STREAM_USAGE
     token: Optional[int] = None
     tokens_cost: Optional[float] = None
@@ -175,27 +330,53 @@ class AgentModelStreamUsageMessage(BaseModelStreamMessage):
 
 
 class BaseAgentAgentCardStreamMessage(BaseAgentMessage):
+    """
+    Base class for agent card streaming messages.
+
+    Agent card streaming allows real-time updates to interactive
+    UI components during their generation or processing.
+    """
+
     step_id: str
     stream_id: str
 
 
 class AgentAgentCardStreamStartMessage(BaseAgentAgentCardStreamMessage):
+    """
+    Message indicating that agent card streaming has begun.
+    """
+
     message_type: MessageType = MessageType.AGENT_AGENT_CARD_STREAM_START
     content: Optional[str] = None
 
 
 class AgentAgentCardStreamErrorMessage(BaseAgentAgentCardStreamMessage):
+    """
+    Message indicating that an error occurred during agent card streaming.
+    """
+
     message_type: MessageType = MessageType.AGENT_AGENT_CARD_STREAM_ERROR
     error_message: str
 
 
 class AgentAgentCardStreamFragmentMessage(BaseAgentAgentCardStreamMessage):
+    """
+    Fragment of streaming agent card content.
+
+    These messages contain individual chunks of agent card data
+    as they are generated, allowing for real-time UI updates.
+    """
+
     message_type: MessageType = MessageType.AGENT_AGENT_CARD_STREAM_FRAGMENT
     index: int
     content: Optional[str]
 
 
 class AgentAgentCardStreamEndMessage(BaseAgentAgentCardStreamMessage):
+    """
+    Message indicating that agent card streaming has completed.
+    """
+
     message_type: MessageType = MessageType.AGENT_AGENT_CARD_STREAM_END
     content: Optional[str] = None
 
@@ -204,15 +385,36 @@ class AgentAgentCardStreamEndMessage(BaseAgentAgentCardStreamMessage):
 
 
 class BaseAgentToolMessage(BaseStepMessage):
+    """
+    Base class for tool execution messages.
+
+    Tool messages track the lifecycle of external tool or service
+    calls made by agents during pipeline execution.
+    """
+
     id: str
     name: str
 
 
 class AgentToolRequestMessage(BaseAgentToolMessage):
+    """
+    Message indicating that a tool request has been initiated.
+
+    This message is sent when an agent begins calling an external
+    tool or service to perform a specific operation.
+    """
+
     message_type: MessageType = MessageType.AGENT_TOOL_REQUEST
 
 
 class AgentToolResponseMessage(BaseAgentToolMessage):
+    """
+    Message indicating that a tool request has completed.
+
+    This message contains the results and timing information
+    from a completed tool or service call.
+    """
+
     message_type: MessageType = MessageType.AGENT_TOOL_RESPONSE
     duration: time
     success: bool
@@ -246,6 +448,7 @@ SSEMessage = Union[
     AgentToolRequestMessage,
     AgentToolResponseMessage,
 ]
+"""Union type representing all possible SSE message types from the Airia API."""
 
 SSEDict = {
     MessageType.AGENT_PING.value: AgentPingMessage,
@@ -274,3 +477,9 @@ SSEDict = {
     MessageType.AGENT_TOOL_REQUEST.value: AgentToolRequestMessage,
     MessageType.AGENT_TOOL_RESPONSE.value: AgentToolResponseMessage,
 }
+"""
+Mapping from message type strings to their corresponding Pydantic model classes.
+
+This dictionary is used by the SSE parser to instantiate the correct message
+type based on the 'event' field in incoming SSE data.
+"""

airia/utils/sse_parser.py

@@ -1,3 +1,11 @@
+"""
+Server-Sent Events (SSE) stream parsing utilities.
+
+This module provides functions for parsing SSE streams from both synchronous and
+asynchronous HTTP responses. It handles the SSE protocol format and converts
+raw stream data into typed message objects.
+"""
+
 import json
 import re
 from typing import AsyncIterable, AsyncIterator, Iterable, Iterator
@@ -5,19 +13,36 @@ from typing import AsyncIterable, AsyncIterator, Iterable, Iterator
 from ..types.sse import SSEDict, SSEMessage
 
 
-def _to_snake_case(name: str):
+def _to_snake_case(name: str) -> str:
+    """
+    Convert camelCase or PascalCase string to snake_case.
+
+    Args:
+        name: The string to convert to snake_case
+
+    Returns:
+        The string converted to snake_case format
+    """
     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).
+    Parse Server-Sent Events (SSE) stream from an iterable of byte chunks.
+
+    This function processes streaming data from HTTP responses that follow the SSE protocol,
+    parsing event blocks and converting them into typed SSE message objects. It handles
+    incomplete chunks by buffering data until complete events are received.
 
     Args:
-        stream_chunks: Iterable of string chunks
+        stream_chunks: An iterable of byte chunks from an HTTP streaming response
 
     Yields:
-        dict: Dictionary containing 'event' and 'data' keys
+        SSEMessage: Typed SSE message objects corresponding to the parsed events
+
+    Note:
+        Events are expected to follow the SSE format with 'event:' and 'data:' lines,
+        terminated by double newlines (\n\n). The data portion should contain valid JSON.
     """
     buffer = ""
 
@@ -53,13 +78,21 @@ 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).
+    Asynchronously parse Server-Sent Events (SSE) stream from an async iterable of byte chunks.
+
+    This is the async version of parse_sse_stream_chunked, designed for use with
+    asynchronous HTTP clients. It processes streaming data that follows the SSE protocol,
+    parsing event blocks and converting them into typed SSE message objects.
 
     Args:
-        stream_chunks: Iterable of string chunks
+        stream_chunks: An async iterable of byte chunks from an HTTP streaming response
 
     Yields:
-        dict: Dictionary containing 'event' and 'data' keys
+        SSEMessage: Typed SSE message objects corresponding to the parsed events
+
+    Note:
+        Events are expected to follow the SSE format with 'event:' and 'data:' lines,
+        terminated by double newlines (\n\n). The data portion should contain valid JSON.
     """
     buffer = ""
 

airia-0.1.14.dist-info/METADATA

@@ -0,0 +1,221 @@
+Metadata-Version: 2.4
+Name: airia
+Version: 0.1.14
+Summary: Python SDK for Airia API
+Author-email: Airia LLC <support@airia.com>
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: requests>=2.32.3
+Requires-Dist: aiohttp>=3.11.14
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: pydantic>=2.11.0
+Provides-Extra: anthropic
+Requires-Dist: anthropic>=0.49.0; extra == "anthropic"
+Provides-Extra: openai
+Requires-Dist: openai>=1.74.0; extra == "openai"
+Provides-Extra: all
+Requires-Dist: anthropic>=0.49.0; extra == "all"
+Requires-Dist: openai>=1.74.0; extra == "all"
+Dynamic: license-file
+
+# Airia Python API Library
+
+[![PyPI version](https://badge.fury.io/py/airia.svg)](https://badge.fury.io/py/airia)
+[![Python versions](https://img.shields.io/pypi/pyversions/airia.svg)](https://pypi.org/project/airia/)
+[![License](https://img.shields.io/pypi/l/airia.svg)](https://pypi.org/project/airia/)
+
+The Airia Python library provides a clean and intuitive interface to interact with Airia API. The library offers both synchronous and asynchronous clients for maximum flexibility in your applications.
+
+## Features
+
+- **Dual Client Support**: Choose between synchronous (`AiriaClient`) and asynchronous (`AiriaAsyncClient`) implementations
+- **Pipeline Execution**: Easily run AI pipelines with customizable parameters
+- **Gateway Support**: Seamlessly integrate with OpenAI and Anthropic services through Airia gateways
+- **Error Handling**: Comprehensive error handling with custom exceptions
+- **Logging**: Built-in configurable logging with correlation ID support for request tracing
+- **Flexible Authentication**: Support for both API keys and bearer tokens with flexible configuration
+
+## Documentation and Quick Start Guides
+
+Full documentation and quick start guides are available [here](https://airiallc.github.io/airia-python).
+
+You can also run the documentation page locally with `mkdocs`:
+
+1. [Install development dependencies](#install-with-development-dependencies)
+2. Run `mkdocs serve` to start the local server
+
+## Installation
+
+You can install the package using pip or uv:
+
+<table>
+<tr>
+<th>pip</th>
+<th>uv</th>
+</tr>
+<tr>
+<td>
+
+```bash
+pip install airia
+```
+
+</td>
+<td>
+
+```bash
+uv add airia
+```
+
+</td>
+</tr>
+</table>
+
+### Install with optional dependencies
+
+The package supports optional dependencies for gateway functionality:
+
+<table>
+<tr>
+<th>OpenAI Gateway</th>
+<th>Anthropic Gateway</th>
+<th>All Gateways</th>
+</tr>
+<tr>
+<td>
+
+```bash
+pip install "airia[openai]"
+```
+
+</td>
+<td>
+
+```bash
+pip install "airia[anthropic]"
+```
+
+</td>
+<td>
+
+```bash
+pip install "airia[all]"
+```
+
+</td>
+</tr>
+</table>
+
+### Install with development dependencies
+
+Clone the repository:
+
+```bash
+git clone https://github.com/AiriaLLC/airia-python.git
+cd airia-python
+```
+
+Then, run one of the following commands:
+
+<table>
+<tr>
+<th>pip</th>
+<th>uv</th>
+</tr>
+<tr>
+<td>
+
+```bash
+pip install dependency-groups
+dev=$(python -m dependency_groups dev)
+pip install -e .
+pip install $dev
+```
+
+</td>
+<td>
+
+```bash
+uv sync --frozen --group dev
+```
+
+</td>
+</tr>
+</table>
+
+## Building from Source
+
+First make sure you have already cloned the repository, then run one of the following commands:
+
+<table>
+<tr>
+<th>pip</th>
+<th>uv</th>
+</tr>
+<tr>
+<td>
+
+```bash
+pip install build
+python -m build
+```
+
+</td>
+<td>
+
+```bash
+uv build
+```
+
+</td>
+</tr>
+</table>
+
+This will create both wheel and source distribution in the `dist/` directory.
+
+## Requirements
+
+- Python 3.9 or higher
+- Core dependencies:
+  - requests
+  - aiohttp
+  - loguru
+  - pydantic
+
+- Optional dependencies:
+  - OpenAI gateway: `openai>=1.74.0`
+  - Anthropic gateway: `anthropic>=0.49.0`
+
+## Development
+
+To run tests (make sure you have development dependencies installed):
+
+```bash
+pytest
+```
+
+For testing gateway functionality, install the optional dependencies:
+
+```bash
+# For OpenAI gateway tests
+pip install -e .[openai]
+pytest tests/test_openai_gateway.py
+
+# For Anthropic gateway tests
+pip install -e .[anthropic]
+pytest tests/test_anthropic_gateway.py
+
+# For all tests
+pip install -e .[all]
+pytest
+```