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

airia/logs.py

@@ -2,19 +2,19 @@ import os
 import sys
 import uuid
 from contextvars import ContextVar
-from typing import BinaryIO, Optional, TextIO, Union, overload
+from typing import BinaryIO, Optional, TextIO, Union
 
 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
@@ -48,28 +48,6 @@ def correlation_id_filter(record):
     return record
 
 
-@overload
-def configure_logging(
-    format_string: str = "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
-    level: str = "INFO",
-    sink: Union[TextIO, BinaryIO] = ...,
-    rotation: None = None,
-    retention: None = None,
-    include_correlation_id: bool = True,
-) -> "loguru.Logger": ...
-
-
-@overload
-def configure_logging(
-    format_string: str = "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
-    level: str = "INFO",
-    sink: Union[str, os.PathLike[str]] = ...,
-    rotation: Optional[str] = None,
-    retention: Optional[str] = None,
-    include_correlation_id: bool = True,
-) -> "loguru.Logger": ...
-
-
 def configure_logging(
     format_string: str = "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
     level: str = "INFO",
@@ -120,7 +98,7 @@ def configure_logging(
 
     # Add the new handler
     logger.add(sink, **kwargs)
-    
+
     return logger
 
 
@@ -129,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

@@ -1,11 +1,38 @@
-from typing import Any, Dict, Optional
+"""
+Internal data structures for HTTP request preparation.
 
-from pydantic import BaseModel
+This module defines the data models used internally by the SDK clients
+to organize and pass request information between methods.
+"""
+
+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,19 +0,0 @@
-from .get_projects import ProjectItem
-from .get_pipeline_config import GetPipelineConfigResponse
-from .pipeline_execution import (
-    PipelineExecutionDebugResponse,
-    PipelineExecutionResponse,
-    PipelineExecutionAsyncStreamedResponse,
-    PipelineExecutionStreamedResponse,
-)
-from .conversations import CreateConversationResponse
-
-__all__ = [
-    "PipelineExecutionDebugResponse",
-    "PipelineExecutionResponse",
-    "PipelineExecutionStreamedResponse",
-    "PipelineExecutionAsyncStreamedResponse",
-    "GetPipelineConfigResponse",
-    "ProjectItem",
-    "CreateConversationResponse",
-]

airia/types/api/conversations/__init__.py

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

airia/types/api/conversations/_conversations.py

@@ -0,0 +1,115 @@
+"""
+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
+
+from pydantic import BaseModel, Field
+
+
+class PolicyRedaction(BaseModel):
+    """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.
+
+    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
+    created_at: datetime = Field(alias="createdAt")
+    updated_at: datetime = Field(alias="updatedAt")
+    role: str
+    images: Optional[List[str]] = None
+
+
+class GetConversationResponse(BaseModel):
+    """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]
+    title: Optional[str] = None
+    websocket_url: Optional[str] = Field(None, alias="websocketUrl")
+    deployment_id: Optional[str] = Field(None, alias="deploymentId")
+    data_source_files: Dict[str, List[str]] = Field(alias="dataSourceFiles")
+    is_bookmarked: bool = Field(alias="isBookmarked")
+    policy_redactions: Optional[Dict[str, PolicyRedaction]] = Field(
+        None, alias="policyRedactions"
+    )
+    last_execution_status: Optional[str] = Field(None, alias="lastExecutionStatus")
+    last_execution_id: Optional[str] = Field(None, alias="lastExecutionId")
+
+
+class CreateConversationResponse(BaseModel):
+    """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")
+    deployment_id: str = Field(alias="deploymentId")
+    icon_id: Optional[str] = Field(None, alias="iconId")
+    icon_url: Optional[str] = Field(None, alias="iconUrl")
+    description: Optional[str] = None
+    space_name: Optional[str] = Field(None, alias="spaceName")

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/_pipeline_execution.py

@@ -0,0 +1,76 @@
+"""
+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
+
+
+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).
+
+    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.
+
+    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).
+
+    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).
+
+    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"]