Spaces:
Running
Running
Download chat_store.py from ZHIWEI666/ComfyUI-Ranking-API: direct link, hf CLI and curl.
- Browser
- Download file 7.34 kB
-
https://huggingface.co/spaces/ZHIWEI666/ComfyUI-Ranking-API/resolve/main/chat_store.py
- Command line
-
hf download hf://spaces/ZHIWEI666/ComfyUI-Ranking-API/chat_store.py
-
curl -L -o chat_store.py https://huggingface.co/spaces/ZHIWEI666/ComfyUI-Ranking-API/resolve/main/chat_store.py
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) | |