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

airia/types/api/get_pipeline_config.py

@@ -1,9 +1,21 @@
+"""
+Pydantic models for pipeline configuration API responses.
+
+This module defines comprehensive data structures for pipeline configuration exports,
+including all components like agents, models, tools, data sources, and deployment settings.
+"""
 from typing import Any, Dict, List, Optional
 
 from pydantic import BaseModel, Field
 
 
 class Metadata(BaseModel):
+    """
+    Pipeline metadata and export configuration.
+    
+    Contains version information, export settings, and descriptive metadata
+    about the pipeline configuration.
+    """
     id: str
     export_version: str = Field(alias="exportVersion")
     tagline: Optional[str] = None
@@ -17,6 +29,12 @@ class Metadata(BaseModel):
 
 
 class Agent(BaseModel):
+    """
+    AI agent configuration and workflow definition.
+    
+    Represents the core agent that executes the pipeline, including its
+    identity, industry specialization, and step-by-step workflow configuration.
+    """
     name: str
     execution_name: str = Field(alias="executionName")
     agent_description: Optional[str] = Field(alias="agentDescription", default=None)
@@ -91,6 +109,12 @@ class Tool(BaseModel):
 
 
 class Model(BaseModel):
+    """
+    Language model configuration and deployment settings.
+    
+    Defines an AI model used in the pipeline, including its deployment details,
+    pricing configuration, authentication settings, and capabilities.
+    """
     id: str
     display_name: str = Field(alias="displayName")
     model_name: str = Field(alias="modelName")
@@ -167,6 +191,13 @@ class DataSource(BaseModel):
 
 
 class GetPipelineConfigResponse(BaseModel):
+    """
+    Complete pipeline configuration export response.
+    
+    This is the root response model containing all components of a pipeline
+    configuration, including the agent definition, associated resources,
+    and deployment settings.
+    """
     metadata: Metadata
     agent: Agent
     data_sources: Optional[List[DataSource]] = Field(

airia/types/api/get_projects.py

@@ -1,3 +1,9 @@
+"""
+Pydantic models for project management API responses.
+
+This module defines the data structures returned by project-related endpoints,
+including project listings and associated pipeline information.
+"""
 from datetime import datetime
 from typing import Any, List, Optional
 
@@ -5,11 +11,49 @@ from pydantic import BaseModel, Field
 
 
 class Pipeline(BaseModel):
+    """
+    Basic pipeline information associated with a project.
+    
+    Represents a simplified view of pipeline data within project contexts,
+    containing only essential identification information.
+    """
     id: str
     name: str
 
 
 class ProjectItem(BaseModel):
+    """
+    Comprehensive project information and metadata.
+    
+    This model represents a complete project entity with all associated resources,
+    budget information, security settings, and organizational details. Projects
+    serve as containers for pipelines, models, data sources, and other AI resources.
+    
+    Attributes:
+        tenant_id: Unique identifier for the tenant/organization
+        created_at: Timestamp when the project was created
+        require_classification: Whether data classification is required
+        budget_amount: Optional budget limit for the project
+        budget_period: Time period for budget calculations
+        budget_alert: Budget alert threshold configuration
+        budget_stop: Whether to stop operations when budget is exceeded
+        used_budget_amount: Amount of budget currently consumed
+        resume_ends_at: When the project resumption period ends
+        updated_at: Timestamp of last project modification
+        pipelines: List of pipelines associated with this project
+        models: AI models available in this project
+        data_sources: Data sources configured for this project
+        prompts: Prompt templates available in this project
+        api_keys: API key configurations for external services
+        memories: Memory/context storage configurations
+        project_icon: Base64 encoded project icon image
+        project_icon_id: Unique identifier for the project icon
+        description: Human-readable project description
+        project_type: Classification of project type
+        classifications: Data classification settings
+        id: Unique project identifier
+        name: Human-readable project name
+    """
     tenant_id: str = Field(alias="tenantId")
     created_at: datetime = Field(alias="createdAt")
     require_classification: bool = Field(alias="requireClassification")

airia/types/api/pipeline_execution.py

@@ -1,3 +1,9 @@
+"""
+Pydantic models for pipeline execution API responses.
+
+This module defines the response models returned by pipeline execution endpoints,
+including both synchronous and streaming response types.
+"""
 from typing import Any, AsyncIterator, Dict, Iterator
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -6,24 +12,48 @@ from ..sse import SSEMessage
 
 
 class PipelineExecutionResponse(BaseModel):
+    """
+    Response model for standard pipeline execution requests.
+    
+    This model represents the response when executing a pipeline in normal mode
+    (not debug mode and not streaming).
+    """
     result: str
     report: None
     is_backup_pipeline: bool = Field(alias="isBackupPipeline")
 
 
 class PipelineExecutionDebugResponse(BaseModel):
+    """
+    Response model for pipeline execution requests in debug mode.
+    
+    This model includes additional debugging information in the report field
+    that provides insights into the pipeline's execution process.
+    """
     result: str
     report: Dict[str, Any]
     is_backup_pipeline: bool = Field(alias="isBackupPipeline")
 
 
 class PipelineExecutionStreamedResponse(BaseModel):
+    """
+    Response model for streaming pipeline execution requests (synchronous client).
+    
+    This model contains an iterator that yields SSEMessage objects as they
+    are received from the streaming response.
+    """
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     stream: Iterator[SSEMessage]
 
 
 class PipelineExecutionAsyncStreamedResponse(BaseModel):
+    """
+    Response model for streaming pipeline execution requests (asynchronous client).
+    
+    This model contains an async iterator that yields SSEMessage objects as they
+    are received from the streaming response.
+    """
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     stream: AsyncIterator[SSEMessage]

airia/types/sse/__init__.py

@@ -1,3 +1,10 @@
+"""
+Server-Sent Event (SSE) message types for the Airia SDK.
+
+This package contains all the message types that can be received via SSE streams
+during pipeline execution, including agent lifecycle events, processing steps,
+model streaming fragments, and tool execution updates.
+"""
 from .sse_messages import (
     SSEDict,
     SSEMessage,

airia/types/sse/sse_messages.py

@@ -1,3 +1,10 @@
+"""
+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,8 +13,15 @@ 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"
     AGENT_END = "AgentEndMessage"
     AGENT_STEP_START = "AgentStepStartMessage"
     AGENT_STEP_HALT = "AgentStepHaltMessage"
@@ -34,11 +48,23 @@ 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
 
@@ -47,15 +73,34 @@ 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
 
 
@@ -63,22 +108,42 @@ 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
@@ -86,6 +151,9 @@ class AgentStepEndMessage(BaseStepMessage):
 
 
 class AgentOutputMessage(BaseStepMessage):
+    """
+    Message containing the output result from a completed step.
+    """
     message_type: MessageType = MessageType.AGENT_OUTPUT
     step_result: str
 
@@ -94,15 +162,33 @@ 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
@@ -110,21 +196,45 @@ 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
@@ -134,33 +244,57 @@ 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
@@ -170,27 +304,48 @@ 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
 
@@ -199,15 +354,33 @@ 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
@@ -217,6 +390,7 @@ class AgentToolResponseMessage(BaseAgentToolMessage):
 SSEMessage = Union[
     AgentPingMessage,
     AgentStartMessage,
+    AgentInputMessage,
     AgentEndMessage,
     AgentStepStartMessage,
     AgentStepHaltMessage,
@@ -240,10 +414,12 @@ SSEMessage = Union[
     AgentToolRequestMessage,
     AgentToolResponseMessage,
 ]
+"""Union type representing all possible SSE message types from the Airia API."""
 
 SSEDict = {
     MessageType.AGENT_PING.value: AgentPingMessage,
     MessageType.AGENT_START.value: AgentStartMessage,
+    MessageType.AGENT_INPUT.value: AgentInputMessage,
     MessageType.AGENT_END.value: AgentEndMessage,
     MessageType.AGENT_STEP_START.value: AgentStepStartMessage,
     MessageType.AGENT_STEP_HALT.value: AgentStepHaltMessage,
@@ -267,3 +443,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,10 @@
+"""
+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 +12,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 +77,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 = ""