feat: AI-assisted flow builder with 4-stage wizard
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>
This commit is contained in:
chihlasm
427
backend/app/api/endpoints/ai_builder.py
Normal file
427
backend/app/api/endpoints/ai_builder.py
Normal file
@@ -0,0 +1,427 @@
|
||||
"""AI Flow Builder wizard endpoints.
|
||||
|
||||
4-stage wizard:
|
||||
POST /ai/start — Stage 1: create conversation with metadata
|
||||
POST /ai/scaffold — Stage 2: AI suggests branches
|
||||
POST /ai/branch-detail — Stage 3: AI generates detail for one branch
|
||||
POST /ai/assemble — Stage 4: assemble branches into tree (no AI)
|
||||
GET /ai/quota — quota status
|
||||
"""
|
||||
import logging
|
||||
from typing import Annotated
|
||||
|
||||
import anthropic
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.api.deps import get_current_active_user, get_db, require_engineer_or_admin
|
||||
from app.core.config import settings
|
||||
from app.core.ai_conversation_store import (
|
||||
create_conversation,
|
||||
get_conversation,
|
||||
update_conversation,
|
||||
)
|
||||
from app.core.ai_quota_service import check_ai_quota, record_ai_usage, get_user_plan
|
||||
from app.core.ai_tree_generator_service import (
|
||||
scaffold_branches,
|
||||
generate_branch_detail,
|
||||
assemble_tree,
|
||||
)
|
||||
from app.models.user import User
|
||||
from app.schemas.ai_builder import (
|
||||
AIStartRequest,
|
||||
AIStartResponse,
|
||||
AIScaffoldRequest,
|
||||
AIScaffoldResponse,
|
||||
AIBranchDetailRequest,
|
||||
AIBranchDetailResponse,
|
||||
AIAssembleRequest,
|
||||
AIAssembleResponse,
|
||||
AIQuotaStatusResponse,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/ai", tags=["ai-builder"])
|
||||
|
||||
|
||||
def _require_ai_enabled() -> None:
|
||||
"""Raise 503 if AI is not configured."""
|
||||
if not settings.ai_enabled:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail="AI flow builder is not configured. Set ANTHROPIC_API_KEY.",
|
||||
)
|
||||
|
||||
|
||||
@router.get("/quota", response_model=AIQuotaStatusResponse)
|
||||
async def get_quota(
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
db: Annotated[AsyncSession, Depends(get_db)],
|
||||
):
|
||||
"""Get current user's AI quota status."""
|
||||
if not settings.ai_enabled:
|
||||
return AIQuotaStatusResponse(
|
||||
plan="free",
|
||||
monthly_used=0,
|
||||
monthly_limit=None,
|
||||
monthly_reset_at="",
|
||||
daily_used=0,
|
||||
daily_limit=None,
|
||||
daily_reset_at="",
|
||||
allowed=False,
|
||||
ai_enabled=False,
|
||||
)
|
||||
|
||||
_, quota_status = await check_ai_quota(
|
||||
user_id=current_user.id,
|
||||
account_id=current_user.account_id,
|
||||
db=db,
|
||||
billing_anchor=current_user.ai_billing_cycle_anchor_at,
|
||||
)
|
||||
return AIQuotaStatusResponse(
|
||||
**quota_status,
|
||||
ai_enabled=True,
|
||||
)
|
||||
|
||||
|
||||
@router.post("/start", response_model=AIStartResponse, status_code=201)
|
||||
async def start_conversation(
|
||||
data: AIStartRequest,
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
db: Annotated[AsyncSession, Depends(get_db)],
|
||||
_: None = Depends(require_engineer_or_admin),
|
||||
):
|
||||
"""Stage 1: Create a new AI wizard conversation with foundation metadata."""
|
||||
_require_ai_enabled()
|
||||
|
||||
# Check daily quota (anti-abuse)
|
||||
allowed, quota_status = await check_ai_quota(
|
||||
user_id=current_user.id,
|
||||
account_id=current_user.account_id,
|
||||
db=db,
|
||||
billing_anchor=current_user.ai_billing_cycle_anchor_at,
|
||||
)
|
||||
if not allowed:
|
||||
reset_key = (
|
||||
"daily_reset_at"
|
||||
if quota_status.get("deny_reason") == "daily"
|
||||
else "monthly_reset_at"
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||||
detail={
|
||||
"message": f"AI build limit exceeded ({quota_status['deny_reason']})",
|
||||
"reset_at": quota_status.get(reset_key),
|
||||
"quota": quota_status,
|
||||
},
|
||||
)
|
||||
|
||||
wizard_state = {
|
||||
"flow_type": data.flow_type,
|
||||
"name": data.name,
|
||||
"description": data.description,
|
||||
"environment_tags": data.environment_tags,
|
||||
"category_id": str(data.category_id) if data.category_id else None,
|
||||
}
|
||||
|
||||
conversation = await create_conversation(
|
||||
user_id=current_user.id,
|
||||
account_id=current_user.account_id,
|
||||
wizard_state=wizard_state,
|
||||
db=db,
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
return AIStartResponse(
|
||||
conversation_id=conversation.id,
|
||||
status=conversation.status,
|
||||
)
|
||||
|
||||
|
||||
@router.post("/scaffold", response_model=AIScaffoldResponse)
|
||||
async def scaffold(
|
||||
data: AIScaffoldRequest,
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
db: Annotated[AsyncSession, Depends(get_db)],
|
||||
_: None = Depends(require_engineer_or_admin),
|
||||
):
|
||||
"""Stage 2: AI suggests top-level branches."""
|
||||
_require_ai_enabled()
|
||||
|
||||
conversation = await get_conversation(
|
||||
data.conversation_id, current_user.id, db
|
||||
)
|
||||
|
||||
# Check per-flow call limit
|
||||
if conversation.question_rounds >= settings.AI_MAX_CALLS_PER_FLOW:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||||
detail="Maximum AI calls per flow exceeded",
|
||||
)
|
||||
|
||||
plan = await get_user_plan(current_user.account_id, db)
|
||||
|
||||
try:
|
||||
branches, input_tokens, output_tokens, cost = await scaffold_branches(
|
||||
conversation.wizard_state,
|
||||
)
|
||||
except anthropic.APIError as e:
|
||||
await record_ai_usage(
|
||||
user_id=current_user.id,
|
||||
account_id=current_user.account_id,
|
||||
conversation_id=conversation.id,
|
||||
generation_type="scaffold",
|
||||
tier=plan,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
estimated_cost=0,
|
||||
succeeded=False,
|
||||
counts_toward_quota=False,
|
||||
error_code=type(e).__name__,
|
||||
extra_data={"error": str(e)},
|
||||
db=db,
|
||||
)
|
||||
await db.commit()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_502_BAD_GATEWAY,
|
||||
detail="AI provider error. Please try again.",
|
||||
)
|
||||
except ValueError as e:
|
||||
await record_ai_usage(
|
||||
user_id=current_user.id,
|
||||
account_id=current_user.account_id,
|
||||
conversation_id=conversation.id,
|
||||
generation_type="scaffold",
|
||||
tier=plan,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
estimated_cost=0,
|
||||
succeeded=False,
|
||||
counts_toward_quota=False,
|
||||
error_code="invalid_output",
|
||||
extra_data={"error": str(e)},
|
||||
db=db,
|
||||
)
|
||||
await db.commit()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||
detail=f"AI returned invalid output: {e}",
|
||||
)
|
||||
|
||||
# Record successful usage
|
||||
await record_ai_usage(
|
||||
user_id=current_user.id,
|
||||
account_id=current_user.account_id,
|
||||
conversation_id=conversation.id,
|
||||
generation_type="scaffold",
|
||||
tier=plan,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
estimated_cost=cost,
|
||||
succeeded=True,
|
||||
counts_toward_quota=False,
|
||||
error_code=None,
|
||||
extra_data=None,
|
||||
db=db,
|
||||
)
|
||||
|
||||
# Update conversation state
|
||||
wizard_state = dict(conversation.wizard_state)
|
||||
wizard_state["branches"] = branches
|
||||
await update_conversation(
|
||||
conversation.id,
|
||||
current_user.id,
|
||||
{
|
||||
"status": "scaffolding",
|
||||
"wizard_state": wizard_state,
|
||||
"question_rounds": conversation.question_rounds + 1,
|
||||
},
|
||||
db,
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
return AIScaffoldResponse(
|
||||
conversation_id=conversation.id,
|
||||
branches=branches,
|
||||
status="scaffolding",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/branch-detail", response_model=AIBranchDetailResponse)
|
||||
async def branch_detail(
|
||||
data: AIBranchDetailRequest,
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
db: Annotated[AsyncSession, Depends(get_db)],
|
||||
_: None = Depends(require_engineer_or_admin),
|
||||
):
|
||||
"""Stage 3: AI generates detailed nodes for one branch."""
|
||||
_require_ai_enabled()
|
||||
|
||||
conversation = await get_conversation(
|
||||
data.conversation_id, current_user.id, db
|
||||
)
|
||||
|
||||
if conversation.question_rounds >= settings.AI_MAX_CALLS_PER_FLOW:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||||
detail="Maximum AI calls per flow exceeded",
|
||||
)
|
||||
|
||||
wizard_state = conversation.wizard_state
|
||||
existing_branches = [
|
||||
b.get("name", "") for b in wizard_state.get("branches", [])
|
||||
]
|
||||
|
||||
plan = await get_user_plan(current_user.account_id, db)
|
||||
|
||||
try:
|
||||
branch_tree, input_tokens, output_tokens, cost = (
|
||||
await generate_branch_detail(
|
||||
wizard_state,
|
||||
data.branch_name,
|
||||
existing_branches,
|
||||
)
|
||||
)
|
||||
except anthropic.APIError as e:
|
||||
await record_ai_usage(
|
||||
user_id=current_user.id,
|
||||
account_id=current_user.account_id,
|
||||
conversation_id=conversation.id,
|
||||
generation_type="branch_detail",
|
||||
tier=plan,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
estimated_cost=0,
|
||||
succeeded=False,
|
||||
counts_toward_quota=False,
|
||||
error_code=type(e).__name__,
|
||||
extra_data={"error": str(e), "branch_name": data.branch_name},
|
||||
db=db,
|
||||
)
|
||||
await db.commit()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_502_BAD_GATEWAY,
|
||||
detail="AI provider error. Please try again.",
|
||||
)
|
||||
except ValueError as e:
|
||||
await record_ai_usage(
|
||||
user_id=current_user.id,
|
||||
account_id=current_user.account_id,
|
||||
conversation_id=conversation.id,
|
||||
generation_type="branch_detail",
|
||||
tier=plan,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
estimated_cost=0,
|
||||
succeeded=False,
|
||||
counts_toward_quota=False,
|
||||
error_code="invalid_output",
|
||||
extra_data={"error": str(e), "branch_name": data.branch_name},
|
||||
db=db,
|
||||
)
|
||||
await db.commit()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||
detail=f"AI returned invalid output: {e}",
|
||||
)
|
||||
|
||||
# Record successful usage
|
||||
await record_ai_usage(
|
||||
user_id=current_user.id,
|
||||
account_id=current_user.account_id,
|
||||
conversation_id=conversation.id,
|
||||
generation_type="branch_detail",
|
||||
tier=plan,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
estimated_cost=cost,
|
||||
succeeded=True,
|
||||
counts_toward_quota=False,
|
||||
error_code=None,
|
||||
extra_data={"branch_name": data.branch_name},
|
||||
db=db,
|
||||
)
|
||||
|
||||
# Update conversation
|
||||
await update_conversation(
|
||||
conversation.id,
|
||||
current_user.id,
|
||||
{
|
||||
"status": "detailing",
|
||||
"question_rounds": conversation.question_rounds + 1,
|
||||
},
|
||||
db,
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
return AIBranchDetailResponse(
|
||||
conversation_id=conversation.id,
|
||||
branch_name=data.branch_name,
|
||||
steps=branch_tree,
|
||||
status="detailing",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/assemble", response_model=AIAssembleResponse)
|
||||
async def assemble(
|
||||
data: AIAssembleRequest,
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
db: Annotated[AsyncSession, Depends(get_db)],
|
||||
_: None = Depends(require_engineer_or_admin),
|
||||
):
|
||||
"""Stage 4: Assemble selected branches into a complete tree (no AI calls)."""
|
||||
conversation = await get_conversation(
|
||||
data.conversation_id, current_user.id, db
|
||||
)
|
||||
|
||||
wizard_state = conversation.wizard_state
|
||||
branches_for_assembly = [b.model_dump() for b in data.selected_branches]
|
||||
|
||||
try:
|
||||
tree_structure, name, description, stats = assemble_tree(
|
||||
wizard_state, branches_for_assembly
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||
detail=str(e),
|
||||
)
|
||||
|
||||
# Record quota-consuming usage on successful assembly
|
||||
plan = await get_user_plan(current_user.account_id, db)
|
||||
await record_ai_usage(
|
||||
user_id=current_user.id,
|
||||
account_id=current_user.account_id,
|
||||
conversation_id=conversation.id,
|
||||
generation_type="tree",
|
||||
tier=plan,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
estimated_cost=0,
|
||||
succeeded=True,
|
||||
counts_toward_quota=True,
|
||||
error_code=None,
|
||||
extra_data={"stats": stats},
|
||||
db=db,
|
||||
)
|
||||
|
||||
# Update conversation with assembled tree
|
||||
await update_conversation(
|
||||
conversation.id,
|
||||
current_user.id,
|
||||
{
|
||||
"status": "completed",
|
||||
"generated_tree": tree_structure,
|
||||
},
|
||||
db,
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
return AIAssembleResponse(
|
||||
tree_structure=tree_structure,
|
||||
suggested_name=name,
|
||||
suggested_description=description,
|
||||
summary=stats,
|
||||
status="completed",
|
||||
)
|
||||
@@ -5,6 +5,7 @@ from app.api.endpoints import ratings, analytics
|
||||
from app.api.endpoints import target_lists
|
||||
from app.api.endpoints import maintenance_schedules
|
||||
from app.api.endpoints import feedback
|
||||
from app.api.endpoints import ai_builder
|
||||
|
||||
api_router = APIRouter()
|
||||
|
||||
@@ -34,3 +35,4 @@ api_router.include_router(analytics.router)
|
||||
api_router.include_router(target_lists.router)
|
||||
api_router.include_router(maintenance_schedules.router)
|
||||
api_router.include_router(feedback.router)
|
||||
api_router.include_router(ai_builder.router)
|
||||
|
||||
Reference in New Issue
Block a user