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

airia/constants.py

@@ -17,4 +17,4 @@ DEFAULT_ANTHROPIC_GATEWAY_URL = "https://gateway.airia.ai/anthropic"
 
 # Default timeouts
 DEFAULT_TIMEOUT = 30.0
-"""Default timeout in seconds for API requests."""
+"""Default timeout in seconds for API requests."""

airia/exceptions.py

@@ -1,20 +1,20 @@
 class AiriaAPIError(Exception):
-    '''
+    """
     Custom exception for Airia API errors.
-    
+
     This exception is raised when an API request to the Airia service fails.
     It contains both the HTTP status code and error message to help with
     debugging and proper error handling.
-    
+
     Attributes:
         status_code (int): The HTTP status code returned by the API
         message (str): The error message describing what went wrong
-    '''
-    
+    """
+
     def __init__(self, status_code: int, message: str):
         """
         Initialize the exception with a status code and error message.
-        
+
         Args:
             status_code (int): The HTTP status code of the failed request
             message (str): A descriptive error message
@@ -22,11 +22,11 @@ class AiriaAPIError(Exception):
         super().__init__(f"{status_code}: {message}")
         self.status_code = status_code
         self.message = message
-    
+
     def __str__(self) -> str:
         """
         Return a string representation of the exception.
-        
+
         Returns:
             str: A formatted string containing the status code and message
         """

airia/logs.py

@@ -8,13 +8,13 @@ import loguru
 from loguru import logger
 
 # Create a context variable to store correlation ID
-correlation_id_context: ContextVar[str] = ContextVar('correlation_id', default='')
+correlation_id_context: ContextVar[str] = ContextVar("correlation_id", default="")
 
 
 def get_correlation_id() -> str:
     """
     Get the current correlation ID from context or return empty string if not set.
-    
+
     Returns:
         str: The current correlation ID
     """
@@ -24,10 +24,10 @@ def get_correlation_id() -> str:
 def set_correlation_id(correlation_id: Optional[str] = None) -> str:
     """
     Set a correlation ID in the current context.
-    
+
     Args:
         correlation_id (Optional[str]): The correlation ID to set. If None, a new UUID will be generated.
-        
+
     Returns:
         str: The correlation ID that was set
     """
@@ -39,7 +39,7 @@ def set_correlation_id(correlation_id: Optional[str] = None) -> str:
 
 def clear_correlation_id() -> None:
     """Clear the correlation ID from the current context."""
-    correlation_id_context.set('')
+    correlation_id_context.set("")
 
 
 # Define a function to be used as a filter to inject correlation_id
@@ -98,7 +98,7 @@ def configure_logging(
 
     # Add the new handler
     logger.add(sink, **kwargs)
-    
+
     return logger
 
 
@@ -107,15 +107,15 @@ if __name__ == "__main__":
     # Basic configuration (uses sys.stderr)
     log = configure_logging()
     log.info("Basic logging configured successfully")
-    
+
     # Set a correlation ID
     set_correlation_id("request-123")
     log.info("This log has a correlation ID")
-    
+
     # Change correlation ID
     set_correlation_id("request-456")
     log.info("This log has a different correlation ID")
-    
+
     # Use auto-generated correlation ID
     set_correlation_id()
     log.info("This log has an auto-generated correlation ID")

airia/types/_request_data.py

@@ -4,28 +4,35 @@ Internal data structures for HTTP request preparation.
 This module defines the data models used internally by the SDK clients
 to organize and pass request information between methods.
 """
-from typing import Any, Dict, Optional
 
-from pydantic import BaseModel
+from io import BufferedIOBase
+from typing import Any, Dict, Optional, Tuple
+
+from pydantic import BaseModel, ConfigDict
 
 
 class RequestData(BaseModel):
     """
     Structured container for HTTP request components.
-    
+
     This internal data structure organizes all the components needed to make
     an HTTP request, including the URL, headers, payload, query parameters,
     and correlation ID for tracing.
-    
+
     Attributes:
         url: The complete URL for the HTTP request
         payload: Optional JSON payload for the request body
         params: Optional query parameters to append to the URL
+        files: Optional file data to be uploaded in the request body
         headers: HTTP headers including authentication and content-type
         correlation_id: Unique identifier for request tracing and logging
     """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     url: str
     payload: Optional[Dict[str, Any]]
     params: Optional[Dict[str, Any]]
+    files: Optional[Dict[str, Tuple[str, BufferedIOBase, str]]]
     headers: Dict[str, Any]
     correlation_id: str

airia/types/api/__init__.py

@@ -1,27 +0,0 @@
-"""
-API response models for the Airia SDK.
-
-This package contains Pydantic models that define the structure of responses
-from various Airia API endpoints, including pipeline execution, project management,
-conversation handling, and configuration retrieval.
-"""
-from .get_projects import ProjectItem
-from .get_pipeline_config import GetPipelineConfigResponse
-from .pipeline_execution import (
-    PipelineExecutionDebugResponse,
-    PipelineExecutionResponse,
-    PipelineExecutionAsyncStreamedResponse,
-    PipelineExecutionStreamedResponse,
-)
-from .conversations import CreateConversationResponse, GetConversationResponse
-
-__all__ = [
-    "PipelineExecutionDebugResponse",
-    "PipelineExecutionResponse",
-    "PipelineExecutionStreamedResponse",
-    "PipelineExecutionAsyncStreamedResponse",
-    "GetPipelineConfigResponse",
-    "ProjectItem",
-    "CreateConversationResponse",
-    "GetConversationResponse",
-]

airia/types/api/conversations/__init__.py

@@ -0,0 +1,3 @@
+from ._conversations import CreateConversationResponse, GetConversationResponse
+
+__all__ = ["CreateConversationResponse", "GetConversationResponse"]

airia/types/api/{conversations.py → conversations/_conversations.py}

@@ -4,6 +4,7 @@ Pydantic models for conversation management API responses.
 This module defines data structures for conversation operations including
 creation, retrieval, and message management within the Airia platform.
 """
+
 from typing import Optional, List, Dict
 from datetime import datetime
 
@@ -11,24 +12,37 @@ from pydantic import BaseModel, Field
 
 
 class PolicyRedaction(BaseModel):
-    """
-    Information about content that was redacted due to policy violations.
-    
+    """Information about content that was redacted due to policy violations.
+
     When content in a conversation violates platform policies, this model
     tracks what was redacted and where it occurred.
+
+    Attributes:
+        violating_text: The text content that violated platform policies
+        violating_message_index: Index of the message containing the violation
     """
+
     violating_text: str = Field(alias="violatingText")
     violating_message_index: int = Field(alias="violatingMessageIndex")
 
 
 class ConversationMessage(BaseModel):
-    """
-    Individual message within a conversation.
-    
+    """Individual message within a conversation.
+
     Represents a single message exchange in a conversation, which can be
     from a user, assistant, or system. Messages may include text content
     and optional image attachments.
+
+    Attributes:
+        id: Unique identifier for the message
+        conversation_id: ID of the conversation this message belongs to
+        message: Optional text content of the message
+        created_at: Timestamp when the message was created
+        updated_at: Timestamp when the message was last updated
+        role: Role of the message sender (user, assistant, system)
+        images: Optional list of image URLs or identifiers
     """
+
     id: str
     conversation_id: str = Field(alias="conversationId")
     message: Optional[str] = None
@@ -39,13 +53,26 @@ class ConversationMessage(BaseModel):
 
 
 class GetConversationResponse(BaseModel):
-    """
-    Complete conversation data including messages and metadata.
-    
+    """Complete conversation data including messages and metadata.
+
     This response contains all information about a conversation including
     its message history, associated files, execution status, and any
     content moderation actions that have been applied.
+
+    Attributes:
+        user_id: ID of the user who owns the conversation
+        conversation_id: Unique identifier for the conversation
+        messages: List of messages in the conversation
+        title: Optional title for the conversation
+        websocket_url: Optional WebSocket URL for real-time updates
+        deployment_id: Optional ID of the deployment handling the conversation
+        data_source_files: Dictionary mapping data sources to their files
+        is_bookmarked: Whether the conversation is bookmarked by the user
+        policy_redactions: Optional dictionary of policy violations and redactions
+        last_execution_status: Optional status of the last execution
+        last_execution_id: Optional ID of the last execution
     """
+
     user_id: str = Field(alias="userId")
     conversation_id: str = Field(alias="conversationId")
     messages: List[ConversationMessage]
@@ -62,12 +89,22 @@ class GetConversationResponse(BaseModel):
 
 
 class CreateConversationResponse(BaseModel):
-    """
-    Response data for newly created conversations.
-    
+    """Response data for newly created conversations.
+
     Contains the essential information needed to begin interacting with
     a new conversation, including connection details and visual metadata.
+
+    Attributes:
+        user_id: ID of the user who created the conversation
+        conversation_id: Unique identifier for the new conversation
+        websocket_url: WebSocket URL for real-time conversation updates
+        deployment_id: ID of the deployment handling the conversation
+        icon_id: Optional ID of the conversation icon
+        icon_url: Optional URL of the conversation icon
+        description: Optional description of the conversation
+        space_name: Optional name of the space containing the conversation
     """
+
     user_id: str = Field(alias="userId")
     conversation_id: str = Field(alias="conversationId")
     websocket_url: str = Field(alias="websocketUrl")

airia/types/api/pipeline_execution/__init__.py

@@ -0,0 +1,13 @@
+from ._pipeline_execution import (
+    PipelineExecutionAsyncStreamedResponse,
+    PipelineExecutionDebugResponse,
+    PipelineExecutionResponse,
+    PipelineExecutionStreamedResponse,
+)
+
+__all__ = [
+    "PipelineExecutionDebugResponse",
+    "PipelineExecutionResponse",
+    "PipelineExecutionStreamedResponse",
+    "PipelineExecutionAsyncStreamedResponse",
+]

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

@@ -4,56 +4,73 @@ 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
 
-from ..sse import SSEMessage
+from ...sse import SSEMessage
 
 
 class PipelineExecutionResponse(BaseModel):
-    """
-    Response model for standard pipeline execution requests.
-    
+    """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).
+
+    Attributes:
+        result: The execution result as a string
+        report: Always None for standard executions
+        is_backup_pipeline: Whether a backup pipeline was used for execution
     """
+
     result: str
     report: None
     is_backup_pipeline: bool = Field(alias="isBackupPipeline")
 
 
 class PipelineExecutionDebugResponse(BaseModel):
-    """
-    Response model for pipeline execution requests in debug mode.
-    
+    """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.
+
+    Attributes:
+        result: The execution result as a string
+        report: Dictionary containing debugging information and execution details
+        is_backup_pipeline: Whether a backup pipeline was used for execution
     """
+
     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).
-    
+    """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.
+
+    Attributes:
+        stream: Iterator that yields SSEMessage objects 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).
-    
+    """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.
+
+    Attributes:
+        stream: Async iterator that yields SSEMessage objects from the streaming response
     """
+
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     stream: AsyncIterator[SSEMessage]

airia/types/api/pipelines_config/__init__.py

@@ -0,0 +1,3 @@
+from .get_pipeline_config import GetPipelineConfigResponse
+
+__all__ = ["GetPipelineConfigResponse"]