109 lines
3.1 KiB
Python
109 lines
3.1 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import time
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from yuxi.utils.logging_config import logger
|
||
|
|
|
||
|
|
DEFAULT_TTL_SECONDS = 3600
|
||
|
|
MAX_CACHE_ENTRIES = 10000
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class SentMessageEntry:
|
||
|
|
message_id: str
|
||
|
|
channel_user_id: str
|
||
|
|
chat_id: str
|
||
|
|
content_hash: str
|
||
|
|
mode: str
|
||
|
|
sent_at: float
|
||
|
|
expires_at: float
|
||
|
|
metadata: dict[str, Any] | None = None
|
||
|
|
|
||
|
|
|
||
|
|
class WeChatSendCache:
|
||
|
|
def __init__(self, ttl_seconds: int = DEFAULT_TTL_SECONDS):
|
||
|
|
self._cache: dict[str, SentMessageEntry] = {}
|
||
|
|
self._ttl = ttl_seconds
|
||
|
|
self._id_index: dict[str, str] = {}
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _content_key(content: str) -> str:
|
||
|
|
import hashlib
|
||
|
|
|
||
|
|
return hashlib.sha256(content.encode()).hexdigest()[:12]
|
||
|
|
|
||
|
|
def store(
|
||
|
|
self,
|
||
|
|
message_id: str,
|
||
|
|
channel_user_id: str,
|
||
|
|
chat_id: str,
|
||
|
|
content: str,
|
||
|
|
mode: str = "wecom",
|
||
|
|
metadata: dict[str, Any] | None = None,
|
||
|
|
) -> None:
|
||
|
|
self._prune()
|
||
|
|
now = time.time()
|
||
|
|
entry = SentMessageEntry(
|
||
|
|
message_id=message_id,
|
||
|
|
channel_user_id=channel_user_id,
|
||
|
|
chat_id=chat_id,
|
||
|
|
content_hash=self._content_key(content),
|
||
|
|
mode=mode,
|
||
|
|
sent_at=now,
|
||
|
|
expires_at=now + self._ttl,
|
||
|
|
metadata=metadata,
|
||
|
|
)
|
||
|
|
self._cache[message_id] = entry
|
||
|
|
self._id_index[message_id] = message_id
|
||
|
|
logger.debug(f"[WeChat/SendCache] Stored message_id={message_id}")
|
||
|
|
|
||
|
|
def get(self, message_id: str) -> SentMessageEntry | None:
|
||
|
|
entry = self._cache.get(message_id)
|
||
|
|
if entry and entry.expires_at < time.time():
|
||
|
|
self._cache.pop(message_id, None)
|
||
|
|
self._id_index.pop(message_id, None)
|
||
|
|
return None
|
||
|
|
return entry
|
||
|
|
|
||
|
|
def get_attached_results(self, message_id: str) -> dict[str, Any]:
|
||
|
|
entry = self.get(message_id)
|
||
|
|
if entry is None:
|
||
|
|
return {}
|
||
|
|
return {
|
||
|
|
"message_id": entry.message_id,
|
||
|
|
"channel_user_id": entry.channel_user_id,
|
||
|
|
"chat_id": entry.chat_id,
|
||
|
|
"sent_at": entry.sent_at,
|
||
|
|
"mode": entry.mode,
|
||
|
|
**(entry.metadata or {}),
|
||
|
|
}
|
||
|
|
|
||
|
|
def remove(self, message_id: str) -> None:
|
||
|
|
self._cache.pop(message_id, None)
|
||
|
|
self._id_index.pop(message_id, None)
|
||
|
|
|
||
|
|
def hit_rate(self) -> float:
|
||
|
|
total = len(self._cache)
|
||
|
|
if total == 0:
|
||
|
|
return 1.0
|
||
|
|
now = time.time()
|
||
|
|
valid = sum(1 for e in self._cache.values() if e.expires_at > now)
|
||
|
|
return valid / total
|
||
|
|
|
||
|
|
def _prune(self) -> None:
|
||
|
|
if len(self._cache) < MAX_CACHE_ENTRIES:
|
||
|
|
return
|
||
|
|
now = time.time()
|
||
|
|
expired = [mid for mid, e in self._cache.items() if e.expires_at < now]
|
||
|
|
for mid in expired:
|
||
|
|
self._cache.pop(mid, None)
|
||
|
|
self._id_index.pop(mid, None)
|
||
|
|
if expired:
|
||
|
|
logger.info(f"[WeChat/SendCache] Pruned {len(expired)} expired entries")
|
||
|
|
|
||
|
|
def clear(self) -> None:
|
||
|
|
self._cache.clear()
|
||
|
|
self._id_index.clear()
|