mirror of
https://github.com/FuQuan233/nonebot-plugin-llmchat.git
synced 2026-08-13 10:09:27 +00:00
88 lines
3.5 KiB
Python
88 lines
3.5 KiB
Python
import asyncio
|
|
from collections import deque
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import aiofiles
|
|
from nonebot import logger
|
|
|
|
from .config import ScopedConfig
|
|
from .state import ChatState, StateStore
|
|
|
|
|
|
class StatePersistence:
|
|
def __init__(self, config: ScopedConfig, states: StateStore, group_file: Path, private_file: Path):
|
|
self.config = config
|
|
self.states = states
|
|
self.group_file = group_file
|
|
self.private_file = private_file
|
|
self._lock = asyncio.Lock()
|
|
|
|
@staticmethod
|
|
def _serialize(state: ChatState, include_probability: bool) -> dict:
|
|
data = {
|
|
"preset": state.preset_name,
|
|
"history": list(state.history),
|
|
"last_active": state.last_active,
|
|
"group_prompt": state.prompt,
|
|
"output_reasoning_content": state.output_reasoning_content,
|
|
}
|
|
if include_probability:
|
|
data["random_trigger_prob"] = state.random_trigger_prob
|
|
return data
|
|
|
|
async def _write_json(self, path: Path, data: dict) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.with_suffix(path.suffix + ".tmp")
|
|
async with aiofiles.open(temporary, "w", encoding="utf8") as file:
|
|
await file.write(json.dumps(data, ensure_ascii=False))
|
|
os.replace(temporary, path)
|
|
|
|
async def save(self) -> None:
|
|
async with self._lock:
|
|
await self._write_json(
|
|
self.group_file,
|
|
{key: self._serialize(state, True) for key, state in self.states.group_states.items()},
|
|
)
|
|
if self.config.enable_private_chat:
|
|
await self._write_json(
|
|
self.private_file,
|
|
{key: self._serialize(state, False) for key, state in self.states.private_states.items()},
|
|
)
|
|
|
|
async def _read_json(self, path: Path) -> dict:
|
|
if not path.exists():
|
|
return {}
|
|
try:
|
|
async with aiofiles.open(path, encoding="utf8") as file:
|
|
value = json.loads(await file.read())
|
|
return value if isinstance(value, dict) else {}
|
|
except (OSError, json.JSONDecodeError):
|
|
logger.exception(f"读取llmchat状态失败: {path}")
|
|
return {}
|
|
|
|
def _restore(self, raw: dict, *, is_group: bool) -> ChatState:
|
|
state = ChatState(
|
|
preset_name=raw.get(
|
|
"preset",
|
|
self.config.default_preset if is_group else self.config.private_chat_preset,
|
|
),
|
|
history_size=self.config.history_size,
|
|
past_events_size=self.config.past_events_size,
|
|
random_trigger_prob=raw.get("random_trigger_prob", self.config.random_trigger_prob) if is_group else 0.0,
|
|
)
|
|
state.history = deque(raw.get("history", []), maxlen=self.config.history_size * 2)
|
|
state.last_active = raw.get("last_active", state.last_active)
|
|
state.prompt = raw.get("group_prompt", raw.get("prompt"))
|
|
state.output_reasoning_content = raw.get("output_reasoning_content", False)
|
|
return state
|
|
|
|
async def load(self) -> None:
|
|
groups = await self._read_json(self.group_file)
|
|
for context_id, raw in groups.items():
|
|
self.states.group_states[int(context_id)] = self._restore(raw, is_group=True)
|
|
if self.config.enable_private_chat:
|
|
private = await self._read_json(self.private_file)
|
|
for context_id, raw in private.items():
|
|
self.states.private_states[int(context_id)] = self._restore(raw, is_group=False)
|