mirror of
http://88.130.71.182:3000/BlitTech/contexta_be.git
synced 2026-06-13 10:14:58 +00:00
107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
"""
|
|
Models router - serves available AI models based on user subscription plan.
|
|
|
|
Single source of truth flow:
|
|
config.py (PLAN_LIMITS + MODEL_CATALOG) → this router → frontend
|
|
|
|
To add/remove/rename a model, only edit config.py.
|
|
"""
|
|
from fastapi import APIRouter, Depends
|
|
from app.dependencies import get_current_user
|
|
from app.database import get_supabase
|
|
from app.config import PLAN_LIMITS, MODEL_CATALOG, DEFAULT_MODELS
|
|
from typing import List, Optional
|
|
from pydantic import BaseModel
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
router = APIRouter(prefix="/models", tags=["Models"])
|
|
|
|
|
|
# ─── Response Models ───────────────────────────────────────────────────────────
|
|
|
|
class AvailableModel(BaseModel):
|
|
id: str
|
|
name: str
|
|
provider: str
|
|
badge: str
|
|
description: Optional[str] = None
|
|
is_default: bool = False
|
|
|
|
|
|
class ModelsResponse(BaseModel):
|
|
models: List[AvailableModel]
|
|
plan: str
|
|
default_model: Optional[str] = None
|
|
has_premium_access: bool
|
|
upgrade_label: Optional[str] = None
|
|
|
|
|
|
# ─── Helpers ───────────────────────────────────────────────────────────────────
|
|
|
|
def _get_user_plan(user_id: str) -> str:
|
|
"""Get user's current subscription plan."""
|
|
supabase = get_supabase()
|
|
result = supabase.table("subscriptions") \
|
|
.select("plan") \
|
|
.eq("user_id", user_id) \
|
|
.eq("status", "active") \
|
|
.execute()
|
|
return result.data[0]["plan"] if result.data else "free"
|
|
|
|
|
|
# ─── Endpoint ─────────────────────────────────────────────────────────────────
|
|
|
|
@router.get("/available", response_model=ModelsResponse)
|
|
async def get_available_models(user=Depends(get_current_user)):
|
|
"""
|
|
Returns the list of AI models the user can access based on their plan.
|
|
Frontend uses this to populate model selection dynamically.
|
|
|
|
- free: gets a default model for preview/testing
|
|
- starter: Fireworks AI models
|
|
- pro: Fireworks + OpenAI + Anthropic + Google
|
|
- enterprise: all models
|
|
"""
|
|
plan = _get_user_plan(user.id)
|
|
plan_config = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"])
|
|
allowed_model_ids = plan_config.get("models", [])
|
|
|
|
# Enterprise wildcard → resolve to all catalog models
|
|
if "*" in allowed_model_ids:
|
|
allowed_model_ids = list(MODEL_CATALOG.keys())
|
|
|
|
# Build model list from catalog metadata
|
|
default_model_id = DEFAULT_MODELS.get(plan)
|
|
models: List[AvailableModel] = []
|
|
|
|
for model_id in allowed_model_ids:
|
|
meta = MODEL_CATALOG.get(model_id)
|
|
if not meta:
|
|
logger.warning(f"Model '{model_id}' in PLAN_LIMITS[{plan}] but not in MODEL_CATALOG — skipping")
|
|
continue
|
|
|
|
models.append(AvailableModel(
|
|
id=model_id,
|
|
name=meta["name"],
|
|
provider=meta["provider"],
|
|
badge=meta["badge"],
|
|
description=meta.get("description"),
|
|
is_default=(model_id == default_model_id),
|
|
))
|
|
|
|
# Determine upgrade messaging
|
|
has_premium = plan in ("pro", "enterprise")
|
|
upgrade_label = None
|
|
if plan == "free":
|
|
upgrade_label = "Upgrade to Starter for more models and publishing"
|
|
elif plan == "starter":
|
|
upgrade_label = "Upgrade to Pro for GPT-4o, Claude, Gemini"
|
|
|
|
return ModelsResponse(
|
|
models=models,
|
|
plan=plan,
|
|
default_model=default_model_id,
|
|
has_premium_access=has_premium,
|
|
upgrade_label=upgrade_label,
|
|
) |