diff --git a/backend/python/app/api/chat_service.py b/backend/python/app/api/chat_service.py index 8a90165..7a83622 100644 --- a/backend/python/app/api/chat_service.py +++ b/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}