ComfyUI-Ranking-API / chat_store.py
ZHIWEI666's picture
账号、私信与消息提醒配套更新
52590d6 verified
Raw History Blame Contribute Delete
7.34 kB
"""私信按真实参与者读取,所有消息修改使用锁内最新记录。"""
import copy
import logging
import math
import re
import time
import uuid
from fastapi import HTTPException
import 数据库连接 as db
logger = logging.getLogger(__name__)
HISTORY_TTL = 7 * 24 * 3600
MAX_MESSAGE_LENGTH = 1000
MESSAGE_FIELDS = frozenset({"id", "sender", "receiver", "content", "created_at", "is_read"})
def _unavailable():
return HTTPException(status_code=503, detail="聊天数据暂不可用")
def _timestamp(message):
value = message["created_at"]
return value / 1000 if value >= 1e12 else value
def _validate(data):
if not isinstance(data, dict):
raise _unavailable()
for key, messages in data.items():
if not isinstance(key, str) or not isinstance(messages, list):
raise _unavailable()
for message in messages:
if not isinstance(message, dict):
raise _unavailable()
if any(not isinstance(message.get(field), str) or not message[field]
for field in ("id", "sender", "receiver")):
raise _unavailable()
timestamp = message.get("created_at")
if not isinstance(message.get("content"), str) or type(timestamp) not in (int, float):
raise _unavailable()
try:
valid_timestamp = math.isfinite(timestamp) and timestamp >= 0
except (OverflowError, ValueError):
valid_timestamp = False
if not valid_timestamp or type(message.get("is_read", False)) is not bool:
raise _unavailable()
return data
def _read(name):
try:
return db.load_data(name, default_data={}, skip_cache=True, strict=True)
except Exception as error:
logger.warning("聊天读取失败: %s", type(error).__name__)
raise _unavailable() from None
def _mutate(updater):
try:
return db.atomic_update("chats.json", lambda data: updater(_validate(data)), default_data={})
except HTTPException:
raise
except Exception as error:
logger.warning("聊天写入失败: %s", type(error).__name__)
raise _unavailable() from None
def _pair(message, account, target):
return ((message["sender"] == account and message["receiver"] == target)
or (message["sender"] == target and message["receiver"] == account))
def _project(message):
result = copy.deepcopy({key: value for key, value in message.items() if key in MESSAGE_FIELDS})
result.setdefault("is_read", False)
return result
def send_message(sender, receiver, content, client_request_id=None):
if client_request_id is not None and (not isinstance(client_request_id, str)
or re.fullmatch(r"[0-9a-f]{32}", client_request_id) is None):
raise HTTPException(status_code=400, detail="发送请求标识不正确")
if not isinstance(receiver, str) or not receiver or not isinstance(content, str) or not content.strip() or len(content) > MAX_MESSAGE_LENGTH:
raise HTTPException(status_code=400, detail="私信内容或接收者不正确")
users = _read("users.json")
if not isinstance(users, dict):
raise _unavailable()
if receiver not in users:
raise HTTPException(status_code=404, detail="接收者不存在")
if not isinstance(users[receiver], dict) or users[receiver].get("account", receiver) != receiver:
raise _unavailable()
message = {"id": f"chat_{int(time.time())}_{uuid.uuid4().hex}", "sender": sender,
"receiver": receiver, "content": content, "created_at": int(time.time()), "is_read": False}
key = f"{min(sender, receiver)}_{max(sender, receiver)}"
def append(data):
if client_request_id is not None:
for messages in data.values():
for existing in messages:
if existing["sender"] == sender and existing.get("client_request_id") == client_request_id:
if existing["receiver"] != receiver or existing["content"] != content:
raise HTTPException(status_code=409, detail="发送请求标识已用于其它消息")
return {"id": existing["id"], "client_request_id": client_request_id}
message["client_request_id"] = client_request_id
data.setdefault(key, []).append(message)
if client_request_id is not None:
return {"id": message["id"], "client_request_id": client_request_id}
return _mutate(append)
def chat_list(account):
data = _validate(_read("chats.json"))
users = _read("users.json")
if not isinstance(users, dict):
raise _unavailable()
conversations = {}
# 旧键的下划线编码可能相撞,账号子串也不能代表参与者。
for messages in data.values():
for message in messages:
if message["sender"] == account:
target = message["receiver"]
elif message["receiver"] == account:
target = message["sender"]
else:
continue
conversations.setdefault(target, []).append(message)
result = []
for target, messages in conversations.items():
user = users.get(target, {})
if not isinstance(user, dict):
raise _unavailable()
last = max(messages, key=_timestamp)
result.append({"account": target, "name": user.get("name", target), "avatar": user.get("avatarDataUrl", ""),
"last_message": last["content"], "last_time": last["created_at"],
"unread_count": sum(message["receiver"] == account and not message.get("is_read", False)
for message in messages)})
result.sort(key=lambda item: item["last_time"] / 1000 if item["last_time"] >= 1e12 else item["last_time"], reverse=True)
return result
def chat_history(account, target):
now = time.time()
data = _validate(_read("chats.json"))
def expires(message):
return message.get("is_read", False) and now - _timestamp(message) >= HISTORY_TTL
def matching(message):
return _pair(message, account, target)
def needs_update(message):
return matching(message) and (expires(message) or (message["receiver"] == account and not message.get("is_read", False)))
def snapshot(data):
return sorted((_project(message) for messages in data.values() for message in messages
if matching(message)), key=_timestamp)
if not any(needs_update(message) for messages in data.values() for message in messages):
return snapshot(data)
def update(data):
for key, messages in data.items():
kept = []
for message in messages:
if matching(message):
if expires(message):
continue
if message["receiver"] == account:
message["is_read"] = True
kept.append(message)
data[key] = kept
return snapshot(data)
return _mutate(update)