xml-pipeline/xml_pipeline/server/models.py

"""
models.py — Pydantic models for AgentServer API.

These models define the JSON structure for API responses.
Uses camelCase for JSON keys (JavaScript convention).
"""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, ConfigDict, Field


def to_camel(string: str) -> str:
    """Convert snake_case to camelCase."""
    components = string.split("_")
    return components[0] + "".join(x.title() for x in components[1:])


class CamelModel(BaseModel):
    """Base model with camelCase JSON serialization."""

    model_config = ConfigDict(
        alias_generator=to_camel,
        populate_by_name=True,
    )


# =============================================================================
# Enums
# =============================================================================


class AgentState(str, Enum):
    """Agent processing state."""

    IDLE = "idle"
    PROCESSING = "processing"
    WAITING = "waiting"
    ERROR = "error"
    PAUSED = "paused"


class ThreadStatus(str, Enum):
    """Thread lifecycle status."""

    ACTIVE = "active"
    COMPLETED = "completed"
    ERROR = "error"
    KILLED = "killed"


class OrganismStatus(str, Enum):
    """Organism running status."""

    STARTING = "starting"
    RUNNING = "running"
    STOPPING = "stopping"
    STOPPED = "stopped"


# =============================================================================
# API Response Models
# =============================================================================


class AgentInfo(CamelModel):
    """Agent information for API response."""

    name: str
    description: str
    is_agent: bool = Field(alias="isAgent")
    peers: List[str] = Field(default_factory=list)
    payload_class: str = Field(alias="payloadClass")
    state: AgentState = AgentState.IDLE
    current_thread: Optional[str] = Field(None, alias="currentThread")
    queue_depth: int = Field(0, alias="queueDepth")
    last_activity: Optional[datetime] = Field(None, alias="lastActivity")
    message_count: int = Field(0, alias="messageCount")


class MessageInfo(CamelModel):
    """Message information for API response."""

    id: str
    thread_id: str = Field(alias="threadId")
    from_id: str = Field(alias="from")
    to_id: str = Field(alias="to")
    payload_type: str = Field(alias="payloadType")
    payload: Dict[str, Any] = Field(default_factory=dict)
    timestamp: datetime
    slot_index: Optional[int] = Field(None, alias="slotIndex")


class ThreadInfo(CamelModel):
    """Thread information for API response."""

    id: str
    status: ThreadStatus = ThreadStatus.ACTIVE
    participants: List[str] = Field(default_factory=list)
    message_count: int = Field(0, alias="messageCount")
    created_at: datetime = Field(alias="createdAt")
    last_activity: Optional[datetime] = Field(None, alias="lastActivity")
    error: Optional[str] = None


class OrganismInfo(CamelModel):
    """Organism overview for API response."""

    name: str
    status: OrganismStatus = OrganismStatus.RUNNING
    uptime_seconds: float = Field(0.0, alias="uptimeSeconds")
    agent_count: int = Field(0, alias="agentCount")
    active_threads: int = Field(0, alias="activeThreads")
    total_messages: int = Field(0, alias="totalMessages")
    identity_configured: bool = Field(False, alias="identityConfigured")


class OrganismConfig(CamelModel):
    """Sanitized organism configuration for API response."""

    name: str
    port: int = 8765
    thread_scheduling: str = Field("breadth-first", alias="threadScheduling")
    max_concurrent_pipelines: int = Field(50, alias="maxConcurrentPipelines")
    max_concurrent_handlers: int = Field(20, alias="maxConcurrentHandlers")
    listeners: List[str] = Field(default_factory=list)
    # Note: Secrets like API keys are never exposed


# =============================================================================
# Request Models
# =============================================================================


class InjectRequest(CamelModel):
    """Request body for POST /inject."""

    to: str
    payload: Dict[str, Any]
    thread_id: Optional[str] = Field(None, alias="threadId")


class InjectResponse(CamelModel):
    """Response body for POST /inject."""

    thread_id: str = Field(alias="threadId")
    message_id: str = Field(alias="messageId")


class SubscribeRequest(CamelModel):
    """WebSocket subscription filter."""

    threads: Optional[List[str]] = None
    agents: Optional[List[str]] = None
    payload_types: Optional[List[str]] = Field(None, alias="payloadTypes")
    events: Optional[List[str]] = None


# =============================================================================
# WebSocket Event Models
# =============================================================================


class WSEvent(CamelModel):
    """Base WebSocket event."""

    event: str


class WSConnectedEvent(WSEvent):
    """Sent on WebSocket connection with full state snapshot."""

    event: str = "connected"
    organism: OrganismInfo
    agents: List[AgentInfo]
    threads: List[ThreadInfo]


class WSAgentStateEvent(WSEvent):
    """Agent state changed."""

    event: str = "agent_state"
    agent: str
    state: AgentState
    current_thread: Optional[str] = Field(None, alias="currentThread")


class WSMessageEvent(WSEvent):
    """New message in the system."""

    event: str = "message"
    message: MessageInfo


class WSThreadCreatedEvent(WSEvent):
    """New thread started."""

    event: str = "thread_created"
    thread: ThreadInfo


class WSThreadUpdatedEvent(WSEvent):
    """Thread status changed."""

    event: str = "thread_updated"
    thread_id: str = Field(alias="threadId")
    status: ThreadStatus
    message_count: int = Field(alias="messageCount")


class WSErrorEvent(WSEvent):
    """Error occurred."""

    event: str = "error"
    thread_id: Optional[str] = Field(None, alias="threadId")
    agent: Optional[str] = None
    error: str
    timestamp: datetime


# =============================================================================
# List Response Models
# =============================================================================


class AgentListResponse(CamelModel):
    """Response for GET /agents."""

    agents: List[AgentInfo]
    count: int


class ThreadListResponse(CamelModel):
    """Response for GET /threads."""

    threads: List[ThreadInfo]
    count: int
    total: int
    offset: int
    limit: int


class MessageListResponse(CamelModel):
    """Response for GET /messages or /threads/{id}/messages."""

    messages: List[MessageInfo]
    count: int
    total: int
    offset: int
    limit: int


class ErrorResponse(CamelModel):
    """Error response."""

    error: str
    detail: Optional[str] = None


# =============================================================================
# Capability/Introspection Models
# =============================================================================


class CapabilityInfo(CamelModel):
    """Basic capability info for listing."""

    name: str
    description: str
    is_agent: bool = Field(alias="isAgent")
    peers: List[str] = Field(default_factory=list)
    root_tag: str = Field(alias="rootTag")


class CapabilityDetail(CamelModel):
    """Detailed capability info including schema."""

    name: str
    description: str
    is_agent: bool = Field(alias="isAgent")
    peers: List[str] = Field(default_factory=list)
    root_tag: str = Field(alias="rootTag")
    payload_class: str = Field(alias="payloadClass")
    schema_xsd: Optional[str] = Field(None, alias="schemaXsd")
    example_xml: Optional[str] = Field(None, alias="exampleXml")


class CapabilityListResponse(CamelModel):
    """Response for GET /capabilities."""

    capabilities: List[CapabilityInfo]
    count: int


# =============================================================================
# Usage/Gas Tracking Models
# =============================================================================


class UsageTotals(CamelModel):
    """Aggregate usage statistics."""

    total_tokens: int = Field(0, alias="totalTokens")
    prompt_tokens: int = Field(0, alias="promptTokens")
    completion_tokens: int = Field(0, alias="completionTokens")
    request_count: int = Field(0, alias="requestCount")
    total_cost: float = Field(0.0, alias="totalCost")
    avg_latency_ms: float = Field(0.0, alias="avgLatencyMs")


class ThreadBudgetInfo(CamelModel):
    """Token budget info for a thread."""

    thread_id: str = Field(alias="threadId")
    max_tokens: int = Field(alias="maxTokens")
    prompt_tokens: int = Field(alias="promptTokens")
    completion_tokens: int = Field(alias="completionTokens")
    total_tokens: int = Field(alias="totalTokens")
    remaining: int
    percent_used: float = Field(alias="percentUsed")
    is_exhausted: bool = Field(alias="isExhausted")


class AgentUsageInfo(CamelModel):
    """Usage info for a specific agent."""

    agent_id: str = Field(alias="agentId")
    total_tokens: int = Field(0, alias="totalTokens")
    prompt_tokens: int = Field(0, alias="promptTokens")
    completion_tokens: int = Field(0, alias="completionTokens")
    request_count: int = Field(0, alias="requestCount")
    total_cost: float = Field(0.0, alias="totalCost")


class ModelUsageInfo(CamelModel):
    """Usage info for a specific model."""

    model: str
    total_tokens: int = Field(0, alias="totalTokens")
    prompt_tokens: int = Field(0, alias="promptTokens")
    completion_tokens: int = Field(0, alias="completionTokens")
    request_count: int = Field(0, alias="requestCount")
    total_cost: float = Field(0.0, alias="totalCost")


class UsageOverview(CamelModel):
    """Complete usage overview (gas gauge)."""

    totals: UsageTotals
    by_agent: List[AgentUsageInfo] = Field(default_factory=list, alias="byAgent")
    by_model: List[ModelUsageInfo] = Field(default_factory=list, alias="byModel")
    active_threads: int = Field(0, alias="activeThreads")


class UsageResponse(CamelModel):
    """Response for GET /usage."""

    usage: UsageOverview


class ThreadBudgetListResponse(CamelModel):
    """Response for GET /usage/threads."""

    threads: List[ThreadBudgetInfo]
    count: int
    default_max_tokens: int = Field(alias="defaultMaxTokens")


# =============================================================================
# Usage History Models (Persistent)
# =============================================================================


class UsageEventInfo(CamelModel):
    """A single usage event from history."""

    id: int
    timestamp: str
    thread_id: str = Field(alias="threadId")
    agent_id: Optional[str] = Field(None, alias="agentId")
    model: str
    provider: str
    prompt_tokens: int = Field(alias="promptTokens")
    completion_tokens: int = Field(alias="completionTokens")
    total_tokens: int = Field(alias="totalTokens")
    latency_ms: float = Field(alias="latencyMs")
    estimated_cost: Optional[float] = Field(None, alias="estimatedCost")
    metadata: dict = Field(default_factory=dict)


class UsageHistoryResponse(CamelModel):
    """Response for GET /usage/history."""

    events: List[UsageEventInfo]
    count: int
    total: int
    offset: int
    limit: int


class BillingSummaryResponse(CamelModel):
    """Response for GET /usage/billing."""

    org_id: Optional[str] = Field(None, alias="orgId")
    start_time: str = Field(alias="startTime")
    end_time: str = Field(alias="endTime")
    total_tokens: int = Field(alias="totalTokens")
    prompt_tokens: int = Field(alias="promptTokens")
    completion_tokens: int = Field(alias="completionTokens")
    request_count: int = Field(alias="requestCount")
    total_cost: float = Field(alias="totalCost")
    by_model: dict = Field(default_factory=dict, alias="byModel")
    by_agent: dict = Field(default_factory=dict, alias="byAgent")


class DailyUsagePoint(CamelModel):
    """A single day's usage for charting."""

    date: str
    total_tokens: int = Field(alias="totalTokens")
    request_count: int = Field(alias="requestCount")
    total_cost: float = Field(alias="totalCost")


class DailyUsageResponse(CamelModel):
    """Response for GET /usage/daily."""

    days: List[DailyUsagePoint]
    count: int