Daniel/AIcreat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 纯生产脚本更新

Browse Source
This commit is contained in:
Daniel
2026-04-28 12:10:27 +08:00
parent 2dc7f2e19c
commit 04f26bdaaf
6 changed files with 226 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

144
app/services/user_store.py
View File

@@ -78,6 +78,7 @@ class UserStore:
api_key TEXT NOT NULL,
base_url TEXT NOT NULL DEFAULT '',
model TEXT NOT NULL,
image_model TEXT NOT NULL DEFAULT '',
timeout_sec REAL NOT NULL DEFAULT 120.0,
max_output_tokens INTEGER NOT NULL DEFAULT 8192,
max_retries INTEGER NOT NULL DEFAULT 0,
@@ -87,6 +88,21 @@ class UserStore:
)
"""
)
c.execute(
"""
CREATE TABLE IF NOT EXISTS user_wallets (
user_id INTEGER PRIMARY KEY,
vip_enabled INTEGER NOT NULL DEFAULT 0,
token_balance INTEGER NOT NULL DEFAULT 0,
total_consumed_tokens INTEGER NOT NULL DEFAULT 0,
updated_at INTEGER NOT NULL,
FOREIGN KEY(user_id) REFERENCES users(id)
)
"""
)
ai_cols = self._table_columns(c, "ai_models")
if "image_model" not in ai_cols:
c.execute("ALTER TABLE ai_models ADD COLUMN image_model TEXT NOT NULL DEFAULT ''")
# 兼容历史单绑定结构，自动迁移为默认账号
rows = c.execute(
"SELECT user_id, appid, secret, author, thumb_media_id, thumb_image_path, updated_at FROM wechat_bindings"
@@ -299,6 +315,13 @@ class UserStore:
(username, pwd_hash, salt, reset_hash, reset_salt, now),
)
uid = int(cur.lastrowid)
c.execute(
"""
INSERT OR IGNORE INTO user_wallets(user_id, vip_enabled, token_balance, total_consumed_tokens, updated_at)
VALUES (?, 0, 0, 0, ?)
""",
(uid, now),
)
return {"id": uid, "username": username, "reset_code": reset_code}
except sqlite3.IntegrityError:
return None
@@ -440,12 +463,112 @@ class UserStore:
c.execute("DELETE FROM ai_models WHERE user_id=?", (user_id,))
c.execute("DELETE FROM user_prefs WHERE user_id=?", (user_id,))
c.execute("DELETE FROM wechat_bindings WHERE user_id=?", (user_id,))
c.execute("DELETE FROM user_wallets WHERE user_id=?", (user_id,))
c.execute(
"UPDATE users SET deleted_at=?, username=username || '#deleted' || ? WHERE id=?",
(now, str(now), user_id),
)
return True
def _ensure_wallet_row(self, c: sqlite3.Connection, user_id: int) -> None:
now = int(time.time())
c.execute(
"""
INSERT OR IGNORE INTO user_wallets(user_id, vip_enabled, token_balance, total_consumed_tokens, updated_at)
VALUES (?, 0, 0, 0, ?)
""",
(user_id, now),
)
def ensure_trial_tokens(self, user_id: int, trial_tokens: int) -> dict:
amount = max(0, int(trial_tokens))
now = int(time.time())
with self._conn() as c:
self._ensure_wallet_row(c, user_id)
row = c.execute(
"SELECT token_balance, total_consumed_tokens FROM user_wallets WHERE user_id=?",
(user_id,),
).fetchone()
current = int(row["token_balance"] or 0) if row else 0
if current <= 0 and amount > 0:
c.execute(
"""
UPDATE user_wallets
SET vip_enabled=1, token_balance=?, updated_at=?
WHERE user_id=?
""",
(amount, now, user_id),
)
return self.get_vip_status(user_id)
def get_vip_status(self, user_id: int) -> dict:
with self._conn() as c:
self._ensure_wallet_row(c, user_id)
row = c.execute(
"""
SELECT vip_enabled, token_balance, total_consumed_tokens, updated_at
FROM user_wallets
WHERE user_id=?
""",
(user_id,),
).fetchone()
return {
"vip_enabled": bool(int(row["vip_enabled"] or 0)) if row else False,
"token_balance": int(row["token_balance"] or 0) if row else 0,
"total_consumed_tokens": int(row["total_consumed_tokens"] or 0) if row else 0,
"updated_at": int(row["updated_at"] or 0) if row else 0,
}
def set_vip_enabled(self, user_id: int, enabled: bool) -> dict:
now = int(time.time())
with self._conn() as c:
self._ensure_wallet_row(c, user_id)
c.execute(
"UPDATE user_wallets SET vip_enabled=?, updated_at=? WHERE user_id=?",
(1 if enabled else 0, now, user_id),
)
return self.get_vip_status(user_id)
def recharge_tokens(self, user_id: int, tokens: int) -> dict:
add = max(0, int(tokens))
now = int(time.time())
with self._conn() as c:
self._ensure_wallet_row(c, user_id)
c.execute(
"""
UPDATE user_wallets
SET token_balance=token_balance + ?, vip_enabled=1, updated_at=?
WHERE user_id=?
""",
(add, now, user_id),
)
return self.get_vip_status(user_id)
def consume_tokens(self, user_id: int, tokens: int) -> tuple[bool, int]:
cost = max(0, int(tokens))
now = int(time.time())
with self._conn() as c:
self._ensure_wallet_row(c, user_id)
row = c.execute(
"SELECT token_balance FROM user_wallets WHERE user_id=?",
(user_id,),
).fetchone()
balance = int(row["token_balance"] or 0) if row else 0
if cost <= 0:
return True, balance
if balance < cost:
return False, balance
new_balance = balance - cost
c.execute(
"""
UPDATE user_wallets
SET token_balance=?, total_consumed_tokens=total_consumed_tokens + ?, updated_at=?
WHERE user_id=?
""",
(new_balance, cost, now, user_id),
)
return True, new_balance
def save_wechat_binding(
self,
user_id: int,
@@ -629,7 +752,7 @@ class UserStore:
with self._conn() as c:
rows = c.execute(
"""
SELECT id, model_name, base_url, model, timeout_sec, max_output_tokens, max_retries, updated_at
SELECT id, model_name, base_url, model, image_model, timeout_sec, max_output_tokens, max_retries, updated_at
FROM ai_models
WHERE user_id=?
ORDER BY updated_at DESC, id DESC
@@ -649,6 +772,7 @@ class UserStore:
"model_name": r["model_name"] or "",
"base_url": r["base_url"] or "",
"model": r["model"] or "",
"image_model": r["image_model"] or "",
"timeout_sec": float(r["timeout_sec"] or 120.0),
"max_output_tokens": int(r["max_output_tokens"] or 8192),
"max_retries": int(r["max_retries"] or 0),
@@ -665,6 +789,7 @@ class UserStore:
api_key: str,
base_url: str,
model: str,
image_model: str = "",
timeout_sec: float = 120.0,
max_output_tokens: int = 8192,
max_retries: int = 0,
@@ -676,20 +801,20 @@ class UserStore:
cur = c.execute(
"""
INSERT INTO ai_models
(user_id, model_name, api_key, base_url, model, timeout_sec, max_output_tokens, max_retries, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
(user_id, model_name, api_key, base_url, model, image_model, timeout_sec, max_output_tokens, max_retries, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(user_id, name, api_key, base_url, model, timeout_sec, max_output_tokens, max_retries, now),
(user_id, name, api_key, base_url, model, image_model, timeout_sec, max_output_tokens, max_retries, now),
)
except sqlite3.IntegrityError:
name = f"{name}-{now % 1000}"
cur = c.execute(
"""
INSERT INTO ai_models
(user_id, model_name, api_key, base_url, model, timeout_sec, max_output_tokens, max_retries, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
(user_id, model_name, api_key, base_url, model, image_model, timeout_sec, max_output_tokens, max_retries, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(user_id, name, api_key, base_url, model, timeout_sec, max_output_tokens, max_retries, now),
(user_id, name, api_key, base_url, model, image_model, timeout_sec, max_output_tokens, max_retries, now),
)
aid = int(cur.lastrowid)
c.execute(
@@ -773,7 +898,7 @@ class UserStore:
if aid:
row = c.execute(
"""
SELECT id, model_name, api_key, base_url, model, timeout_sec, max_output_tokens, max_retries, updated_at
SELECT id, model_name, api_key, base_url, model, image_model, timeout_sec, max_output_tokens, max_retries, updated_at
FROM ai_models
WHERE id=? AND user_id=?
""",
@@ -782,7 +907,7 @@ class UserStore:
if not row:
row = c.execute(
"""
SELECT id, model_name, api_key, base_url, model, timeout_sec, max_output_tokens, max_retries, updated_at
SELECT id, model_name, api_key, base_url, model, image_model, timeout_sec, max_output_tokens, max_retries, updated_at
FROM ai_models
WHERE user_id=?
ORDER BY updated_at DESC, id DESC
@@ -809,6 +934,7 @@ class UserStore:
"api_key": row["api_key"] or "",
"base_url": row["base_url"] or "",
"model": row["model"] or "",
"image_model": row["image_model"] or "",
"timeout_sec": float(row["timeout_sec"] or 120.0),
"max_output_tokens": int(row["max_output_tokens"] or 8192),
"max_retries": int(row["max_retries"] or 0),