aegra-api 0.1.0__py3-none-any.whl

aegra_api/models/__init__.py

@@ -0,0 +1,76 @@
+"""Agent Protocol Pydantic models"""
+
+from aegra_api.models.assistants import (
+    AgentSchemas,
+    Assistant,
+    AssistantCreate,
+    AssistantList,
+    AssistantSearchRequest,
+    AssistantUpdate,
+)
+from aegra_api.models.auth import AuthContext, TokenPayload, User
+from aegra_api.models.errors import AgentProtocolError, get_error_type
+from aegra_api.models.runs import Run, RunCreate, RunStatus
+from aegra_api.models.store import (
+    StoreDeleteRequest,
+    StoreGetResponse,
+    StoreItem,
+    StorePutRequest,
+    StoreSearchRequest,
+    StoreSearchResponse,
+)
+from aegra_api.models.threads import (
+    Thread,
+    ThreadCheckpoint,
+    ThreadCheckpointPostRequest,
+    ThreadCreate,
+    ThreadHistoryRequest,
+    ThreadList,
+    ThreadSearchRequest,
+    ThreadSearchResponse,
+    ThreadState,
+    ThreadStateUpdate,
+    ThreadStateUpdateResponse,
+    ThreadUpdate,
+)
+
+__all__ = [
+    # Assistants
+    "Assistant",
+    "AssistantCreate",
+    "AssistantList",
+    "AssistantSearchRequest",
+    "AssistantUpdate",
+    "AgentSchemas",
+    # Threads
+    "Thread",
+    "ThreadCreate",
+    "ThreadList",
+    "ThreadSearchRequest",
+    "ThreadSearchResponse",
+    "ThreadState",
+    "ThreadStateUpdate",
+    "ThreadStateUpdateResponse",
+    "ThreadCheckpoint",
+    "ThreadCheckpointPostRequest",
+    "ThreadHistoryRequest",
+    # Runs
+    "Run",
+    "RunCreate",
+    "RunStatus",
+    # Store
+    "StorePutRequest",
+    "StoreGetResponse",
+    "StoreSearchRequest",
+    "StoreSearchResponse",
+    "StoreItem",
+    "StoreDeleteRequest",
+    # Errors
+    "AgentProtocolError",
+    "get_error_type",
+    # Auth
+    "User",
+    "AuthContext",
+    "TokenPayload",
+    "ThreadUpdate",
+]

aegra_api/models/assistants.py

@@ -0,0 +1,81 @@
+"""Assistant-related Pydantic models for Agent Protocol"""
+
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class AssistantCreate(BaseModel):
+    """Request model for creating assistants"""
+
+    assistant_id: str | None = Field(None, description="Unique assistant identifier (auto-generated if not provided)")
+    name: str | None = Field(
+        None,
+        description="Human-readable assistant name (auto-generated if not provided)",
+    )
+    description: str | None = Field(None, description="Assistant description")
+    config: dict[str, Any] | None = Field({}, description="Assistant configuration")
+    context: dict[str, Any] | None = Field({}, description="Assistant context")
+    graph_id: str = Field(..., description="LangGraph graph ID from aegra.json")
+    metadata: dict[str, Any] | None = Field({}, description="Metadata to use for searching and filtering assistants.")
+    if_exists: str | None = Field("error", description="What to do if assistant exists: error or do_nothing")
+
+
+class Assistant(BaseModel):
+    """Assistant entity model"""
+
+    assistant_id: str
+    name: str
+    description: str | None = None
+    config: dict[str, Any] = Field(default_factory=dict)
+    context: dict[str, Any] = Field(default_factory=dict)
+    graph_id: str
+    user_id: str
+    version: int = Field(..., description="The version of the assistant.")
+    metadata: dict[str, Any] = Field(default_factory=dict, alias="metadata_dict")
+    created_at: datetime
+    updated_at: datetime
+
+    model_config = ConfigDict(from_attributes=True, populate_by_name=True)
+
+
+class AssistantUpdate(BaseModel):
+    """Request model for creating assistants"""
+
+    name: str | None = Field(None, description="The name of the assistant (auto-generated if not provided)")
+    description: str | None = Field(None, description="The description of the assistant. Defaults to null.")
+    config: dict[str, Any] | None = Field({}, description="Configuration to use for the graph.")
+    graph_id: str = Field("agent", description="The ID of the graph")
+    context: dict[str, Any] | None = Field(
+        {},
+        description="The context to use for the graph. Useful when graph is configurable.",
+    )
+    metadata: dict[str, Any] | None = Field({}, description="Metadata to use for searching and filtering assistants.")
+
+
+class AssistantList(BaseModel):
+    """Response model for listing assistants"""
+
+    assistants: list[Assistant]
+    total: int
+
+
+class AssistantSearchRequest(BaseModel):
+    """Request model for assistant search"""
+
+    name: str | None = Field(None, description="Filter by assistant name")
+    description: str | None = Field(None, description="Filter by assistant description")
+    graph_id: str | None = Field(None, description="Filter by graph ID")
+    limit: int | None = Field(20, le=100, ge=1, description="Maximum results")
+    offset: int | None = Field(0, ge=0, description="Results offset")
+    metadata: dict[str, Any] | None = Field({}, description="Metadata to use for searching and filtering assistants.")
+
+
+class AgentSchemas(BaseModel):
+    """Agent schema definitions for client integration"""
+
+    input_schema: dict[str, Any] = Field(..., description="JSON Schema for agent inputs")
+    output_schema: dict[str, Any] = Field(..., description="JSON Schema for agent outputs")
+    state_schema: dict[str, Any] = Field(..., description="JSON Schema for agent state")
+    config_schema: dict[str, Any] = Field(..., description="JSON Schema for agent config")

aegra_api/models/auth.py

@@ -0,0 +1,62 @@
+"""Authentication and user context models"""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+
+
+class User(BaseModel):
+    """User model that accepts any auth fields.
+
+    This model uses ConfigDict(extra="allow") to accept any additional fields
+    from auth handlers (e.g., subscription_tier, team_id) while maintaining
+    type hints for common fields.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    # Required
+    identity: str
+
+    # Optional with defaults
+    is_authenticated: bool = True
+    permissions: list[str] = []
+    display_name: str | None = None
+
+    # Common optional fields (for IDE hints)
+    org_id: str | None = None
+    email: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dict including all extra fields."""
+        return self.model_dump()
+
+    def __getattr__(self, name: str) -> Any:
+        """Allow attribute access to extra fields."""
+        try:
+            extra = object.__getattribute__(self, "__pydantic_extra__") or {}
+        except AttributeError:
+            extra = {}
+        if name in extra:
+            return extra[name]
+        raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")
+
+
+class AuthContext(BaseModel):
+    """Authentication context for request processing"""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    user: User
+    request_id: str | None = None
+
+
+class TokenPayload(BaseModel):
+    """JWT token payload structure"""
+
+    sub: str  # subject (user ID)
+    name: str | None = None
+    scopes: list[str] = []
+    org: str | None = None
+    exp: int | None = None
+    iat: int | None = None

aegra_api/models/enums.py

@@ -0,0 +1,29 @@
+"""Status enums for Aegra API specification."""
+
+from typing import Literal
+
+# Run status enum
+RunStatus = Literal[
+    "pending",
+    "running",
+    "error",
+    "success",
+    "timeout",
+    "interrupted",
+]
+
+# Thread status enum
+ThreadStatus = Literal[
+    "idle",
+    "busy",
+    "interrupted",
+    "error",
+]
+
+# Multitask strategy enum
+MultitaskStrategy = Literal[
+    "reject",
+    "rollback",
+    "interrupt",
+    "enqueue",
+]

aegra_api/models/errors.py

@@ -0,0 +1,29 @@
+"""Error response models for Agent Protocol"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class AgentProtocolError(BaseModel):
+    """Standard Agent Protocol error response"""
+
+    error: str = Field(..., description="Error type")
+    message: str = Field(..., description="Human-readable error message")
+    details: dict[str, Any] | None = Field(None, description="Additional error details")
+
+
+def get_error_type(status_code: int) -> str:
+    """Map HTTP status codes to error types"""
+    error_map = {
+        400: "bad_request",
+        401: "unauthorized",
+        403: "forbidden",
+        404: "not_found",
+        409: "conflict",
+        422: "validation_error",
+        500: "internal_error",
+        501: "not_implemented",
+        503: "service_unavailable",
+    }
+    return error_map.get(status_code, "unknown_error")

aegra_api/models/runs.py

@@ -0,0 +1,124 @@
+"""Run-related Pydantic models for Agent Protocol"""
+
+from datetime import datetime
+from typing import Any, Self
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    field_validator,
+    model_validator,
+)
+
+from aegra_api.utils.status_compat import validate_run_status
+
+
+class RunCreate(BaseModel):
+    """Request model for creating runs"""
+
+    assistant_id: str = Field(..., description="Assistant to execute")
+    input: dict[str, Any] | None = Field(
+        None,
+        description="Input data for the run. Optional when resuming from a checkpoint.",
+    )
+    config: dict[str, Any] | None = Field({}, description="Execution config")
+    context: dict[str, Any] | None = Field({}, description="Execution context")
+    checkpoint: dict[str, Any] | None = Field(
+        None,
+        description="Checkpoint configuration (e.g., {'checkpoint_id': '...', 'checkpoint_ns': ''})",
+    )
+    stream: bool = Field(False, description="Enable streaming response")
+    stream_mode: str | list[str] | None = Field(None, description="Requested stream mode(s)")
+    on_disconnect: str | None = Field(
+        None,
+        description="Behavior on client disconnect: 'cancel' (default) or 'continue'.",
+    )
+
+    multitask_strategy: str | None = Field(
+        None,
+        description="Strategy for handling concurrent runs on same thread: 'reject', 'interrupt', 'rollback', or 'enqueue'.",
+    )
+
+    # Human-in-the-loop fields (core HITL functionality)
+    command: dict[str, Any] | None = Field(
+        None,
+        description="Command for resuming interrupted runs with state updates or navigation",
+    )
+    interrupt_before: str | list[str] | None = Field(
+        None,
+        description="Nodes to interrupt immediately before they get executed. Use '*' for all nodes.",
+    )
+    interrupt_after: str | list[str] | None = Field(
+        None,
+        description="Nodes to interrupt immediately after they get executed. Use '*' for all nodes.",
+    )
+
+    # Subgraph configuration
+    stream_subgraphs: bool | None = Field(
+        False,
+        description="Whether to include subgraph events in streaming. When True, includes events from all subgraphs. When False (default when None), excludes subgraph events. Defaults to False for backwards compatibility.",
+    )
+
+    # Request metadata (top-level in payload)
+    metadata: dict[str, Any] | None = Field(
+        None,
+        description="Request metadata (e.g., from_studio flag)",
+    )
+
+    @model_validator(mode="after")
+    def validate_input_command_exclusivity(self) -> Self:
+        """Ensure input and command are mutually exclusive"""
+        # Allow empty input dict when command is present (frontend compatibility)
+        if self.input is not None and self.command is not None:
+            # If input is just an empty dict, treat it as None for compatibility
+            if self.input == {}:
+                self.input = None
+            else:
+                raise ValueError("Cannot specify both 'input' and 'command' - they are mutually exclusive")
+        if self.input is None and self.command is None:
+            if self.checkpoint is not None:
+                # Allow checkpoint-only requests by treating input as empty dict
+                self.input = {}
+            else:
+                raise ValueError("Must specify either 'input' or 'command'")
+        return self
+
+
+class Run(BaseModel):
+    """Run entity model
+
+    Status values: pending, running, error, success, timeout, interrupted
+    """
+
+    model_config = ConfigDict(from_attributes=True)
+
+    run_id: str
+    thread_id: str
+    assistant_id: str
+    status: str = "pending"  # Valid values: pending, running, error, success, timeout, interrupted
+    input: dict[str, Any]
+    output: dict[str, Any] | None = None
+    error_message: str | None = None
+    config: dict[str, Any] | None = {}
+    context: dict[str, Any] | None = {}
+    user_id: str
+    created_at: datetime
+    updated_at: datetime
+
+    @field_validator("status", mode="before")
+    @classmethod
+    def validate_status(cls, v: str) -> str:
+        """Validate status conforms to API specification."""
+        if not isinstance(v, str):
+            raise ValueError(f"Status must be a string, got {type(v)}")
+        return validate_run_status(v)
+
+
+class RunStatus(BaseModel):
+    """Simple run status response"""
+
+    run_id: str
+    status: str  # Standard status value
+
+    message: str | None = None

aegra_api/models/store.py

@@ -0,0 +1,67 @@
+"""Store-related Pydantic models for Agent Protocol"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class StorePutRequest(BaseModel):
+    """Request model for storing items"""
+
+    namespace: list[str] = Field(..., description="Storage namespace")
+    key: str = Field(..., description="Item key")
+    value: dict[str, Any] = Field(..., description="Item value (must be a JSON object)")
+
+    @field_validator("value", mode="before")
+    @classmethod
+    def validate_value_is_dict(cls, v: Any) -> dict[str, Any]:
+        """Validate that value is a dictionary.
+
+        LangGraph store requires values to be dictionaries for proper
+        serialization and search functionality.
+        """
+        if not isinstance(v, dict):
+            raise ValueError(f"Value must be a dictionary (JSON object), got {type(v).__name__}")
+        return v
+
+
+class StoreGetResponse(BaseModel):
+    """Response model for getting items"""
+
+    key: str
+    value: Any
+    namespace: list[str]
+
+
+class StoreSearchRequest(BaseModel):
+    """Request model for searching store items"""
+
+    namespace_prefix: list[str] = Field(..., description="Namespace prefix to search")
+    filter: dict[str, Any] | None = Field(None, description="Optional dictionary of key-value pairs to filter results.")
+    query: str | None = Field(None, description="Search query")
+    limit: int | None = Field(20, le=100, ge=1, description="Maximum results")
+    offset: int | None = Field(0, ge=0, description="Results offset")
+
+
+class StoreItem(BaseModel):
+    """Store item model"""
+
+    key: str
+    value: Any
+    namespace: list[str]
+
+
+class StoreSearchResponse(BaseModel):
+    """Response model for store search"""
+
+    items: list[StoreItem]
+    total: int
+    limit: int
+    offset: int
+
+
+class StoreDeleteRequest(BaseModel):
+    """Request body for deleting store items (SDK-compatible)."""
+
+    namespace: list[str]
+    key: str

aegra_api/models/threads.py

@@ -0,0 +1,152 @@
+"""Thread-related Pydantic models for Agent Protocol"""
+
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+from aegra_api.utils.status_compat import validate_thread_status
+
+
+class ThreadCreate(BaseModel):
+    """Request model for creating threads"""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    metadata: dict[str, Any] | None = Field(None, description="Thread metadata")
+    initial_state: dict[str, Any] | None = Field(None, description="LangGraph initial state")
+    thread_id: str | None = Field(
+        None,
+        alias="threadId",
+        description="Optional client-provided thread ID for idempotent creation",
+    )
+    if_exists: str | None = Field(
+        "raise",
+        alias="ifExists",
+        description="Behavior when thread exists: 'raise' (default) or 'do_nothing'",
+    )
+
+
+class ThreadUpdate(BaseModel):
+    """Request model for updating threads"""
+
+    metadata: dict[str, Any] | None = Field(None, description="Thread metadata to update")
+
+
+class Thread(BaseModel):
+    """Thread entity model
+
+    Status values: idle, busy, interrupted, error
+    """
+
+    model_config = ConfigDict(from_attributes=True)
+
+    thread_id: str
+    status: str = "idle"  # Valid values: idle, busy, interrupted, error
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    user_id: str
+    created_at: datetime
+    updated_at: datetime
+
+    @field_validator("status", mode="before")
+    @classmethod
+    def validate_status(cls, v: str) -> str:
+        """Validate status conforms to API specification."""
+        if not isinstance(v, str):
+            raise ValueError(f"Status must be a string, got {type(v)}")
+        return validate_thread_status(v)
+
+
+class ThreadList(BaseModel):
+    """Response model for listing threads"""
+
+    threads: list[Thread]
+    total: int
+
+
+class ThreadSearchRequest(BaseModel):
+    """Request model for thread search"""
+
+    metadata: dict[str, Any] | None = Field(None, description="Metadata filters")
+    status: str | None = Field(None, description="Thread status filter (idle, busy, interrupted, error)")
+    limit: int | None = Field(20, le=100, ge=1, description="Maximum results")
+    offset: int | None = Field(0, ge=0, description="Results offset")
+    order_by: str | None = Field("created_at DESC", description="Sort order")
+
+    @field_validator("status")
+    @classmethod
+    def validate_status(cls, v: str | None) -> str | None:
+        """Validate status filter conforms to API specification."""
+        if v is not None:
+            return validate_thread_status(v)
+        return v
+
+
+class ThreadSearchResponse(BaseModel):
+    """Response model for thread search"""
+
+    threads: list[Thread]
+    total: int
+    limit: int
+    offset: int
+
+
+class ThreadCheckpoint(BaseModel):
+    """Checkpoint identifier for thread history"""
+
+    checkpoint_id: str | None = None
+    thread_id: str | None = None
+    checkpoint_ns: str | None = ""
+
+
+class ThreadCheckpointPostRequest(BaseModel):
+    """Request model for fetching thread checkpoint"""
+
+    checkpoint: ThreadCheckpoint = Field(description="Checkpoint to fetch")
+    subgraphs: bool | None = Field(False, description="Include subgraph states")
+
+
+class ThreadState(BaseModel):
+    """Thread state model for history endpoint"""
+
+    values: dict[str, Any] = Field(description="Channel values (messages, etc.)")
+    next: list[str] = Field(default_factory=list, description="Next nodes to execute")
+    tasks: list[dict[str, Any]] = Field(default_factory=list, description="Tasks to execute")
+    interrupts: list[dict[str, Any]] = Field(default_factory=list, description="Interrupt data")
+    metadata: dict[str, Any] = Field(default_factory=dict, description="Checkpoint metadata")
+    created_at: datetime | None = Field(None, description="Timestamp of state creation")
+    checkpoint: ThreadCheckpoint = Field(description="Current checkpoint")
+    parent_checkpoint: ThreadCheckpoint | None = Field(None, description="Parent checkpoint")
+    checkpoint_id: str | None = Field(None, description="Checkpoint ID (for backward compatibility)")
+    parent_checkpoint_id: str | None = Field(None, description="Parent checkpoint ID (for backward compatibility)")
+
+
+class ThreadStateUpdate(BaseModel):
+    """Request model for updating thread state"""
+
+    values: dict[str, Any] | list[dict[str, Any]] | None = Field(
+        None, description="The values to update the state with"
+    )
+    checkpoint: dict[str, Any] | None = Field(None, description="The checkpoint to update the state of")
+    checkpoint_id: str | None = Field(None, description="Optional checkpoint ID to update from")
+    as_node: str | None = Field(None, description="Update the state as if this node had just executed")
+    # Also support query-like parameters for GET-like behavior via POST
+    subgraphs: bool | None = Field(False, description="Include states from subgraphs")
+    checkpoint_ns: str | None = Field(None, description="Checkpoint namespace")
+
+
+class ThreadStateUpdateResponse(BaseModel):
+    """Response model for thread state update"""
+
+    checkpoint: dict[str, Any] = Field(description="The checkpoint that was created/updated")
+
+
+class ThreadHistoryRequest(BaseModel):
+    """Request model for thread history endpoint"""
+
+    limit: int | None = Field(10, ge=1, le=1000, description="Number of states to return")
+    before: str | None = Field(None, description="Return states before this checkpoint ID")
+    metadata: dict[str, Any] | None = Field(None, description="Filter by metadata")
+    checkpoint: dict[str, Any] | None = Field(None, description="Checkpoint for subgraph filtering")
+    subgraphs: bool | None = Field(False, description="Include states from subgraphs")
+    checkpoint_ns: str | None = Field(None, description="Checkpoint namespace")

aegra_api/observability/__init__.py

@@ -0,0 +1 @@
+"""Observability module for the agent server."""