Implements the complete AI flow builder feature using a guided 4-stage wizard (Foundation → Scaffold → Branch Detail → Review & Assemble). AI assists at bounded points using Claude Haiku for cost-efficient structured JSON generation (~$0.01-0.03/flow). Backend: new models (ai_conversations, ai_usage), Alembic migration, quota enforcement with billing anchor, Anthropic API integration with prompt caching, tree validation, conversation CRUD with 24h TTL, APScheduler cleanup job, 5 API endpoints, Pydantic schemas. Frontend: TypeScript types, API client, Zustand store for wizard state, 7 components (modal, step indicator, foundation form, branch selector, branch detail view, tree preview, quota display), MyTreesPage integration with "Build with AI" button (hidden when AI not configured). Tests: 14 validator unit tests + 11 endpoint integration tests with mocked Anthropic (zero real API spend). All 25 tests passing. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
117 lines
2.6 KiB
Python
117 lines
2.6 KiB
Python
"""Pydantic schemas for the AI Flow Builder wizard."""
|
|
from typing import Any, Literal, Optional
|
|
from uuid import UUID
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
# ── 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)
|
|
|
|
|
|
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
|