"""Pydantic schemas for the AI Flow Builder wizard.""" from typing import Any, Literal, Optional from uuid import UUID from pydantic import BaseModel, Field, field_validator # ── Requests ── class AIStartRequest(BaseModel): """Stage 1: Foundation — engineer provides flow metadata.""" flow_type: Literal["troubleshooting", "procedural"] = Field( ..., description="Type of flow to generate" ) category_id: Optional[UUID] = None name: str = Field(..., min_length=1, max_length=255) description: str = Field("", max_length=2000) environment_tags: list[str] = Field(default_factory=list, max_length=20) @field_validator("environment_tags") @classmethod def validate_tags(cls, v: list[str]) -> list[str]: for tag in v: if len(tag) > 100: raise ValueError("Each environment tag must be 100 characters or fewer") if not tag.strip(): raise ValueError("Environment tags must not be empty") return v class AIScaffoldRequest(BaseModel): """Stage 2: Request AI-generated branch suggestions.""" conversation_id: UUID class AIBranchDetailRequest(BaseModel): """Stage 3: Request AI-generated detail for one branch.""" conversation_id: UUID branch_name: str = Field(..., min_length=1, max_length=255) class AIBranchUpdate(BaseModel): """A branch with optional user edits for assembly.""" name: str description: str = "" steps: Optional[dict[str, Any]] = None class AIAssembleRequest(BaseModel): """Stage 4: Assemble selected branches into a complete tree.""" conversation_id: UUID selected_branches: list[AIBranchUpdate] = Field(..., min_length=2) # ── Responses ── class AIStartResponse(BaseModel): """Response after creating a conversation.""" conversation_id: UUID status: str class AIBranchSuggestion(BaseModel): """A single branch suggestion from the AI.""" name: str description: str class AIScaffoldResponse(BaseModel): """Response with AI-suggested branches.""" conversation_id: UUID branches: list[AIBranchSuggestion] status: str class AIBranchDetailResponse(BaseModel): """Response with AI-generated detail for one branch.""" conversation_id: UUID branch_name: str steps: dict[str, Any] status: str class AITreeSummary(BaseModel): """Summary statistics for an assembled tree.""" node_count: int decision_count: int action_count: int solution_count: int depth: int class AIAssembleResponse(BaseModel): """Response with the fully assembled tree.""" tree_structure: dict[str, Any] suggested_name: str suggested_description: str summary: AITreeSummary status: str class AIQuotaStatusResponse(BaseModel): """Current user's AI quota status.""" plan: str monthly_used: int monthly_limit: Optional[int] monthly_reset_at: str daily_used: int daily_limit: Optional[int] daily_reset_at: str allowed: bool ai_enabled: bool