backend_update_260131

app/main.py

@@ -7,7 +7,7 @@ from passlib.context import CryptContext
 from sqlalchemy.exc import IntegrityError
 from contextlib import asynccontextmanager
 import httpx
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import List
 
 @asynccontextmanager
@@ -115,9 +115,12 @@ async def create_asset(
     
     new_asset = models.Asset(
         title=asset.title,
+        type=asset.type,
         content_outer_encrypted=encrypted_content,
         private_key_shard=asset.private_key_shard,
-        author_id=current_user.id
+        author_id=current_user.id,
+        created_at=datetime.utcnow(),
+        updated_at=datetime.utcnow()
     )
     db.add(new_asset)
     await db.commit()
@@ -199,7 +202,7 @@ async def assign_asset(
     
     heir_result = await db.execute(
         select(models.User).where(
-            models.User.username == assignment.heir_name
+            models.User.email == assignment.heir_email
         )
     )
     heir_user = heir_result.scalars().first()
@@ -216,7 +219,21 @@ async def assign_asset(
     asset.heir = heir_user
     await db.commit()
 
-    return {"message": f"Asset assigned to {assignment.heir_name}"}
+    return {"message": f"Asset assigned to {assignment.heir_email}"}
+
+
+@app.get("/assets/designated", response_model=List[schemas.AssetOut])
+async def get_designated_assets(
+    current_user: models.User = Depends(auth.get_current_user),
+    db: AsyncSession = Depends(database.get_db)
+):
+    """
+    Query assets where the current user is the designated heir.
+    """
+    result = await db.execute(
+        select(models.Asset).where(models.Asset.heir_id == current_user.id)
+    )
+    return result.scalars().all()
 
 
 
@@ -253,14 +270,34 @@ async def declare_user_guale(
         "guale": target_user.guale
     }
 
-# 用于测试热加载
-@app.post("/post1")
-async def test1():
-    a=2
-    b=3
-    c = a+b
-    return {"msg": f"this is a msg {c}"}
 
+async def get_or_create_token_usage(user_id: int, db: AsyncSession):
+    # Get current week start (Monday)
+    now = datetime.utcnow()
+    monday = now - timedelta(days=now.weekday())
+    week_start = monday.replace(hour=0, minute=0, second=0, microsecond=0)
+
+    result = await db.execute(
+        select(models.UserTokenUsage).where(models.UserTokenUsage.user_id == user_id)
+    )
+    usage = result.scalars().first()
+
+    if not usage:
+        usage = models.UserTokenUsage(
+            user_id=user_id,
+            tokens_used=0,
+            last_reset_at=week_start
+        )
+        db.add(usage)
+        await db.commit()
+        await db.refresh(usage)
+    #每周重置token使用情况
+    elif usage.last_reset_at < week_start:
+        usage.tokens_used = 0
+        usage.last_reset_at = week_start
+        await db.commit()
+    
+    return usage
 
 @app.post("/ai/proxy", response_model=schemas.AIResponse)
 async def ai_proxy(
@@ -272,6 +309,43 @@ async def ai_proxy(
     Proxy relay for AI requests.
     Fetches AI configuration from the database.
     """
+    def get_quota_exceeded_response():
+        return {
+            "id": f"chatcmpl-{int(datetime.utcnow().timestamp())}",
+            "object": "chat.completion",
+            "created": int(datetime.utcnow().timestamp()),
+            "model": "quota-manager",
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": "quota exceeded, please upgrade plan"
+                    },
+                    "finish_reason": "stop"
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0
+            }
+        }
+
+    # 1. 检查 Tier 是否允许使用 AI
+    result = await db.execute(
+        select(models.SubscriptionPlans).where(models.SubscriptionPlans.name == current_user.tier)
+    )
+    tier_plan = result.scalars().first()
+    
+    if not tier_plan or not tier_plan.can_use_ai_proxy:
+        return get_quota_exceeded_response()
+
+    # 2. 检查本周 Token 使用是否超过限制
+    usage_record = await get_or_create_token_usage(current_user.id, db)
+    if usage_record.tokens_used >= tier_plan.weekly_token_limit:
+        return get_quota_exceeded_response()
+
     # Fetch active AI config
     result = await db.execute(
         select(models.AIConfig).where(models.AIConfig.is_active == True)
@@ -290,6 +364,9 @@ async def ai_proxy(
     payload = ai_request.model_dump()
     payload["model"] = config.default_model
     
+    current_user.last_active_at = datetime.utcnow()
+    await db.commit()
+    
     async with httpx.AsyncClient() as client:
         try:
             response = await client.post(
@@ -299,7 +376,15 @@ async def ai_proxy(
                 timeout=30.0
             )
             response.raise_for_status()
-            return response.json()
+            ai_data = response.json()
+            
+            # 3. 记录使用的 Token
+            total_tokens = ai_data.get("usage", {}).get("total_tokens", 0)
+            if total_tokens > 0:
+                usage_record.tokens_used += total_tokens
+                await db.commit()
+                
+            return ai_data
         except httpx.HTTPStatusError as e:
             raise HTTPException(
                 status_code=e.response.status_code,
@@ -311,3 +396,12 @@ async def ai_proxy(
                 detail=f"An error occurred while requesting AI provider: {str(e)}"
             )
 
+# 用于测试热加载
+@app.post("/post1")
+async def test1():
+    a=2
+    b=3
+    c = a+b
+    return {"msg": f"this is a msg {c}"}
+
+