refactor: migrate AI tree generator to provider abstraction
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -10,7 +10,6 @@
|
||||
import logging
|
||||
from typing import Annotated
|
||||
|
||||
import anthropic
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
@@ -52,7 +51,7 @@ def _require_ai_enabled() -> None:
|
||||
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.",
|
||||
detail="AI flow builder is not configured. Set GOOGLE_AI_API_KEY or ANTHROPIC_API_KEY.",
|
||||
)
|
||||
|
||||
|
||||
@@ -174,27 +173,6 @@ async def scaffold(
|
||||
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,
|
||||
@@ -216,6 +194,27 @@ async def scaffold(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||
detail=f"AI returned invalid output: {e}",
|
||||
)
|
||||
except Exception 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.",
|
||||
)
|
||||
|
||||
# Record successful usage
|
||||
await record_ai_usage(
|
||||
@@ -293,27 +292,6 @@ async def branch_detail(
|
||||
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,
|
||||
@@ -335,6 +313,27 @@ async def branch_detail(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||
detail=f"AI returned invalid output: {e}",
|
||||
)
|
||||
except Exception 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.",
|
||||
)
|
||||
|
||||
# Record successful usage
|
||||
await record_ai_usage(
|
||||
|
||||
Reference in New Issue
Block a user