ai-robot-core/ai-service/app/schemas/retrieval_strategy.py

"""
Retrieval Strategy Schemas for AI Service.
[AC-AISVC-RES-01~15] Request/Response models aligned with OpenAPI contract.
"""

from enum import Enum
from typing import Any

from pydantic import BaseModel, Field, model_validator


class StrategyType(str, Enum):
    DEFAULT = "default"
    ENHANCED = "enhanced"


class ReactMode(str, Enum):
    REACT = "react"
    NON_REACT = "non_react"


class RolloutMode(str, Enum):
    OFF = "off"
    PERCENTAGE = "percentage"
    ALLOWLIST = "allowlist"


class ValidationCheckType(str, Enum):
    METADATA_CONSISTENCY = "metadata_consistency"
    EMBEDDING_PREFIX = "embedding_prefix"
    RRF_CONFIG = "rrf_config"
    PERFORMANCE_BUDGET = "performance_budget"


class RolloutConfig(BaseModel):
    """
    [AC-AISVC-RES-03] Grayscale rollout configuration.
    """

    mode: RolloutMode = Field(..., description="Rollout mode: off, percentage, or allowlist")
    percentage: float | None = Field(
        default=None,
        ge=0,
        le=100,
        description="Percentage of traffic for grayscale (0-100)",
    )
    allowlist: list[str] | None = Field(
        default=None,
        description="List of tenant IDs in allowlist",
    )

    @model_validator(mode="after")
    def validate_rollout_config(self) -> "RolloutConfig":
        if self.mode == RolloutMode.PERCENTAGE and self.percentage is None:
            raise ValueError("percentage is required when mode is 'percentage'")
        if self.mode == RolloutMode.ALLOWLIST and (
            self.allowlist is None or len(self.allowlist) == 0
        ):
            raise ValueError("allowlist is required when mode is 'allowlist'")
        return self


class RetrievalStrategyStatus(BaseModel):
    """
    [AC-AISVC-RES-01] Current retrieval strategy status.
    """

    active_strategy: StrategyType = Field(
        ...,
        alias="active_strategy",
        description="Current active strategy: default or enhanced",
    )
    react_mode: ReactMode = Field(
        ...,
        alias="react_mode",
        description="ReAct mode: react or non_react",
    )
    rollout: RolloutConfig = Field(..., description="Grayscale rollout configuration")

    model_config = {"populate_by_name": True}


class RetrievalStrategySwitchRequest(BaseModel):
    """
    [AC-AISVC-RES-02, AC-AISVC-RES-03, AC-AISVC-RES-05] Request to switch retrieval strategy.
    """

    target_strategy: StrategyType = Field(
        ...,
        alias="target_strategy",
        description="Target strategy to switch to",
    )
    react_mode: ReactMode | None = Field(
        default=None,
        alias="react_mode",
        description="ReAct mode to use",
    )
    rollout: RolloutConfig | None = Field(
        default=None,
        description="Grayscale rollout configuration",
    )
    reason: str | None = Field(
        default=None,
        description="Reason for strategy switch",
    )

    model_config = {"populate_by_name": True}


class RetrievalStrategySwitchResponse(BaseModel):
    """
    [AC-AISVC-RES-02] Response after strategy switch.
    """

    previous: RetrievalStrategyStatus = Field(..., description="Previous strategy status")
    current: RetrievalStrategyStatus = Field(..., description="Current strategy status")


class RetrievalStrategyValidationRequest(BaseModel):
    """
    [AC-AISVC-RES-04, AC-AISVC-RES-06, AC-AISVC-RES-08] Request to validate strategy.
    """

    strategy: StrategyType = Field(..., description="Strategy to validate")
    react_mode: ReactMode | None = Field(
        default=None,
        description="ReAct mode to validate",
    )
    checks: list[ValidationCheckType] | None = Field(
        default=None,
        description="List of checks to perform",
    )

    model_config = {"populate_by_name": True}


class ValidationResult(BaseModel):
    """
    [AC-AISVC-RES-06] Single validation check result.
    """

    check: str = Field(..., description="Check name")
    passed: bool = Field(..., description="Whether the check passed")
    message: str | None = Field(default=None, description="Additional message")


class RetrievalStrategyValidationResponse(BaseModel):
    """
    [AC-AISVC-RES-06] Validation response with all check results.
    """

    passed: bool = Field(..., description="Whether all checks passed")
    results: list[ValidationResult] = Field(..., description="Individual check results")


class RetrievalStrategyRollbackResponse(BaseModel):
    """
    [AC-AISVC-RES-07] Response after strategy rollback.
    """

    current: RetrievalStrategyStatus = Field(..., description="Current strategy status before rollback")
    rollback_to: RetrievalStrategyStatus = Field(..., description="Strategy status after rollback")


class StrategyAuditLog(BaseModel):
    """
    [AC-AISVC-RES-07] Audit log entry for strategy operations.
    """

    timestamp: str = Field(..., description="ISO timestamp of the operation")
    operation: str = Field(..., description="Operation type: switch, rollback, validate")
    previous_strategy: str | None = Field(default=None, description="Previous strategy")
    new_strategy: str | None = Field(default=None, description="New strategy")
    previous_react_mode: str | None = Field(default=None, description="Previous react mode")
    new_react_mode: str | None = Field(default=None, description="New react mode")
    reason: str | None = Field(default=None, description="Reason for the operation")
    operator: str | None = Field(default=None, description="Operator who performed the operation")
    tenant_id: str | None = Field(default=None, description="Tenant ID if applicable")
    metadata: dict[str, Any] | None = Field(default=None, description="Additional metadata")


class StrategyMetrics(BaseModel):
    """
    [AC-AISVC-RES-03, AC-AISVC-RES-08] Metrics for strategy operations.
    """

    strategy: StrategyType = Field(..., description="Current strategy")
    react_mode: ReactMode = Field(..., description="Current react mode")
    total_requests: int = Field(default=0, description="Total requests count")
    successful_requests: int = Field(default=0, description="Successful requests count")
    failed_requests: int = Field(default=0, description="Failed requests count")
    avg_latency_ms: float = Field(default=0.0, description="Average latency in ms")
    p99_latency_ms: float = Field(default=0.0, description="P99 latency in ms")
    direct_route_count: int = Field(default=0, description="Direct route count")
    react_route_count: int = Field(default=0, description="React route count")
    auto_route_count: int = Field(default=0, description="Auto route count")
    fallback_count: int = Field(default=0, description="Fallback to default count")
    last_updated: str | None = Field(default=None, description="Last update timestamp")