feat: add token usage tracking for AI recommendations and subscription management

app/api/v1/recommendations.py

@@ -106,6 +106,7 @@ def recommend_outfit(
             role="ai",
             content=recommendation.get("recommendation", ""),
             selected_item_ids=[str(sid) for sid in selected_ids] or None,
+            tokens_used=recommendation.get("tokens_used") or 0,
         ))
         conv.updated_at = datetime.utcnow()  # đẩy cuộc trò chuyện lên đầu danh sách
         db.commit()

app/api/v1/subscription.py

@@ -1,4 +1,5 @@
 from fastapi import APIRouter, Depends
+from sqlalchemy import func
 from sqlalchemy.orm import Session
 from datetime import datetime, timedelta, timezone
 from app.core.database import get_db
@@ -13,6 +14,9 @@ router = APIRouter()
 # Free-tier caps advertised on the landing page's pricing table.
 FREE_WARDROBE_ITEM_LIMIT = 30
 FREE_AI_RECOMMENDATIONS_PER_WEEK = 5
+# AI Stylist usage is metered like an LLM API quota: total tokens consumed
+# by "ai" responses over a rolling 30-day window.
+FREE_AI_TOKEN_LIMIT_PER_30D = 50_000
 
 
 def _build_response(db: Session, user: User) -> SubscriptionResponse:
@@ -33,6 +37,17 @@ def _build_response(db: Session, user: User) -> SubscriptionResponse:
         .count()
     )
 
+    month_ago = datetime.now(timezone.utc) - timedelta(days=30)
+    tokens_used = (
+        db.query(func.coalesce(func.sum(ChatMessage.tokens_used), 0))
+        .filter(
+            ChatMessage.user_id == user.id,
+            ChatMessage.role == "ai",
+            ChatMessage.created_at >= month_ago,
+        )
+        .scalar()
+    )
+
     return SubscriptionResponse(
         plan=user.plan,
         renews_at=user.premium_renews_at,
@@ -40,6 +55,8 @@ def _build_response(db: Session, user: User) -> SubscriptionResponse:
         wardrobe_items_limit=None if is_premium else FREE_WARDROBE_ITEM_LIMIT,
         ai_recommendations_used=ai_used,
         ai_recommendations_limit=None if is_premium else FREE_AI_RECOMMENDATIONS_PER_WEEK,
+        ai_tokens_used=int(tokens_used or 0),
+        ai_tokens_limit=None if is_premium else FREE_AI_TOKEN_LIMIT_PER_30D,
     )
 
 

app/main.py

@@ -31,6 +31,11 @@ try:
             conn.execute(text("ALTER TABLE chat_messages ADD COLUMN IF NOT EXISTS conversation_id UUID;"))
         except Exception:
             pass
+        # Ensure chat_messages has tokens_used (added after the table first shipped)
+        try:
+            conn.execute(text("ALTER TABLE chat_messages ADD COLUMN IF NOT EXISTS tokens_used INTEGER NOT NULL DEFAULT 0;"))
+        except Exception:
+            pass
         # Ensure users has plan/premium_renews_at (added after the table first shipped)
         try:
             conn.execute(text("ALTER TABLE users ADD COLUMN IF NOT EXISTS plan VARCHAR(20) NOT NULL DEFAULT 'free';"))

app/models/chat.py

@@ -1,4 +1,4 @@
-from sqlalchemy import Column, String, DateTime, func, ForeignKey, Text
+from sqlalchemy import Column, String, DateTime, func, ForeignKey, Text, Integer
 from sqlalchemy.dialects.postgresql import UUID, ARRAY
 from app.core.database import Base
 import uuid
@@ -28,4 +28,6 @@ class ChatMessage(Base):
     content = Column(Text, nullable=False)
     # Với tin nhắn "ai" là gợi ý phối đồ: các ID sản phẩm được chọn (để dựng lại thumbnail).
     selected_item_ids = Column(ARRAY(UUID(as_uuid=True)), nullable=True)
+    # Số token LLM tiêu tốn cho tin nhắn "ai" này (dùng để tính quota trợ lý AI theo token).
+    tokens_used = Column(Integer, nullable=False, server_default="0", default=0)
     created_at = Column(DateTime(timezone=True), server_default=func.now())

app/schemas/subscription.py

@@ -9,6 +9,8 @@ class SubscriptionResponse(BaseModel):
     wardrobe_items_limit: Optional[int]  # None = unlimited
     ai_recommendations_used: int
     ai_recommendations_limit: Optional[int]  # None = unlimited
+    ai_tokens_used: int
+    ai_tokens_limit: Optional[int]  # None = unlimited
 
     class Config:
         from_attributes = True

app/services/llm_service.py

@@ -76,17 +76,19 @@ class LLMService:
         prompt = self._build_prompt(user_profile, wardrobe_items, message, history, weather, event)
 
         try:
-            content = self._chat_completion(prompt)
+            content, tokens_used = self._chat_completion(prompt)
             data = self._parse_json(content)
             if "is_recommendation" not in data:
                 data["is_recommendation"] = len(data.get("selected_item_ids", [])) > 0
+            data["tokens_used"] = tokens_used
             return data
         except Exception as e:
             logging.error(f"Error calling LLM provider '{self.provider}': {e}")
             return {
                 "is_recommendation": False,
                 "recommendation": "Hệ thống gặp lỗi khi liên kết với AI gợi ý. Hãy thử lại sau.",
-                "selected_item_ids": []
+                "selected_item_ids": [],
+                "tokens_used": 0,
             }
 
     # ------------------------------------------------------------------ prompt
@@ -159,7 +161,7 @@ class LLMService:
         """
 
     # --------------------------------------------------------------- providers
-    def _chat_completion(self, prompt: str) -> str:
+    def _chat_completion(self, prompt: str) -> tuple[str, int]:
         if self.provider == "openai":
             url = "https://api.openai.com/v1/chat/completions"
             headers = {
@@ -192,7 +194,11 @@ class LLMService:
 
         response = requests.post(url, headers=headers, json=payload, timeout=timeout)
         response.raise_for_status()
-        return response.json()["choices"][0]["message"]["content"]
+        body = response.json()
+        content = body["choices"][0]["message"]["content"]
+        # Both OpenAI and NVIDIA's OpenAI-compatible endpoint return a "usage" object.
+        tokens_used = int((body.get("usage") or {}).get("total_tokens") or 0)
+        return content, tokens_used
 
     # ------------------------------------------------------------------ parsing
     @staticmethod
@@ -219,7 +225,8 @@ class LLMService:
         return {
             "is_recommendation": is_rec,
             "recommendation": "Mock Outfit: Áo sơ mi trắng phối quần tây đen nhã nhặn phù hợp cho môi trường văn phòng/phỏng vấn." if is_rec else f"Tôi hiểu bạn đang nói về: '{message}'. Trong vai trò trợ lý thời trang, tôi khuyên bạn nên tự tin thể hiện phong cách của mình!",
-            "selected_item_ids": [str(item["id"]) for item in wardrobe_items[:2]] if (wardrobe_items and is_rec) else []
+            "selected_item_ids": [str(item["id"]) for item in wardrobe_items[:2]] if (wardrobe_items and is_rec) else [],
+            "tokens_used": 0,
         }