kiro/dashboard_api/models.py

# -*- coding: utf-8 -*-
"""
Pydantic models for dashboard API.
"""

from datetime import datetime
from typing import Optional, List, Dict, Any
from pydantic import BaseModel, Field


class DashboardStats(BaseModel):
    """Current gateway statistics."""
    uptime_seconds: float
    total_requests: int
    total_errors: int
    error_rate: float
    active_requests: int
    
    # Token usage
    total_input_tokens: int
    total_output_tokens: int
    
    # Per-model breakdown
    requests_by_model: Dict[str, int]
    tokens_by_model: Dict[str, Dict[str, int]]
    
    # Latency stats
    avg_latency_ms: float
    p50_latency_ms: float
    p95_latency_ms: float
    p99_latency_ms: float
    
    # Account stats (if load balancing enabled)
    accounts_total: int
    accounts_healthy: int
    accounts_unhealthy: int


class MetricsQuery(BaseModel):
    """Query parameters for metrics endpoint."""
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None
    model: Optional[str] = None
    granularity: str = Field(default="1m", description="1m, 5m, 15m, 1h, 1d")
    limit: int = Field(default=100, ge=1, le=1000)


class MetricsDataPoint(BaseModel):
    """Single metrics data point."""
    timestamp: datetime
    requests: int
    errors: int
    input_tokens: int
    output_tokens: int
    avg_latency_ms: float


class MetricsResponse(BaseModel):
    """Metrics query response."""
    query: MetricsQuery
    data: List[MetricsDataPoint]
    total_requests: int
    total_errors: int
    total_input_tokens: int
    total_output_tokens: int


class AccountCreate(BaseModel):
    """Create account request."""
    id: str
    name: str
    enabled: bool = True
    weight: int = 1
    priority: int = 0
    
    # Credentials (one required)
    refresh_token: Optional[str] = None
    creds_file: Optional[str] = None
    sqlite_db: Optional[str] = None
    
    # Optional
    region: str = "us-east-1"
    profile_arn: Optional[str] = None


class AccountUpdate(BaseModel):
    """Update account request."""
    name: Optional[str] = None
    enabled: Optional[bool] = None
    weight: Optional[int] = None
    priority: Optional[int] = None
    
    # Credentials
    refresh_token: Optional[str] = None
    creds_file: Optional[str] = None
    sqlite_db: Optional[str] = None
    
    # Optional
    region: Optional[str] = None
    profile_arn: Optional[str] = None


class AccountResponse(BaseModel):
    """Account response (masks sensitive data)."""
    id: str
    name: str
    enabled: bool
    weight: int
    priority: int
    region: str
    
    # Health info
    health_status: str
    request_count: int
    error_count: int
    active_connections: int
    error_rate: float
    last_used: Optional[datetime]
    last_health_check: Optional[datetime]
    
    # Masked credentials info
    has_refresh_token: bool
    has_creds_file: bool
    has_sqlite_db: bool


class ConfigResponse(BaseModel):
    """Gateway configuration response."""
    version: str
    
    # Server settings
    server_host: str
    server_port: int
    
    # API settings
    region: str
    
    # Model settings
    hidden_models: Dict[str, str]
    model_aliases: Dict[str, str]
    hidden_from_list: List[str]
    
    # Feature flags
    fake_reasoning_enabled: bool
    truncation_recovery: bool
    metrics_enabled: bool
    load_balancer_enabled: bool
    load_balancer_strategy: str
    
    # Debug
    debug_mode: str
    log_level: str


class ConfigUpdate(BaseModel):
    """Update configuration request."""
    model_aliases: Optional[Dict[str, str]] = None
    hidden_from_list: Optional[List[str]] = None
    fake_reasoning_enabled: Optional[bool] = None
    truncation_recovery: Optional[bool] = None
    log_level: Optional[str] = None


class LogEntry(BaseModel):
    """Log entry for streaming."""
    timestamp: datetime
    level: str
    message: str
    module: Optional[str] = None
    function: Optional[str] = None
    line: Optional[int] = None


class HealthResponse(BaseModel):
    """Detailed health check response."""
    status: str  # healthy, degraded, unhealthy
    version: str
    uptime_seconds: float
    
    # Component health
    auth_healthy: bool
    api_reachable: bool
    
    # Account health (if load balancing)
    accounts_status: Optional[Dict[str, str]] = None
    
    # Metrics health
    metrics_enabled: bool
    prometheus_ready: bool