phase1-billing-meilisearch

backend/python/app/api/chat_service.py

@@ -8598,3 +8598,26 @@ def _launch_media_task_bg(ai_msg_id, ai_model_id, prompt, user_id):
     _t = threading.Thread(target=_run, daemon=True)
     _t.start()
 
+# Phase 1: Token billing
+async def _bill_tokens(user_id, prompt_tokens, completion_tokens, model_name, conv_id):
+    try:
+        from app.database import SessionLocal
+        from sqlalchemy import text
+        db = SessionLocal()
+        total = prompt_tokens + completion_tokens
+        cost = round(total * 0.002 / 1000, 6)
+        db.execute(text("UPDATE user_balances SET balance = balance - :c, total_used = total_used + :c, updated_at = NOW() WHERE user_id = :u AND balance >= :c"), {"c": cost, "u": user_id})
+        db.execute(text("INSERT INTO token_usage_logs (id, user_id, model_name, conversation_id, prompt_tokens, completion_tokens, total_tokens, cost, created_at) VALUES (gen_random_uuid(), :u, :m, :cv, :pt, :ct, :tt, :c, NOW())"), {"u": user_id, "m": model_name, "cv": conv_id, "pt": prompt_tokens, "ct": completion_tokens, "tt": total, "c": cost})
+        db.commit()
+        db.close()
+    except Exception:
+        pass
+
+async def get_user_balance(user_id):
+    from app.database import SessionLocal
+    from sqlalchemy import text
+    db = SessionLocal()
+    r = db.execute(text("SELECT balance, total_used FROM user_balances WHERE user_id = :u"), {"u": user_id}).fetchone()
+    db.close()
+    if r: return {"balance": float(r[0]), "total_used": float(r[1])}
+    return {"balance": 10.0, "total_used": 0}