384 lines
13 KiB
Python
384 lines
13 KiB
Python
"""
|
||
聊天服务 - 处理聊天记录的读写操作
|
||
|
||
基于 SillyTavern JSONL 格式的聊天记录管理
|
||
"""
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional, Any
|
||
import json
|
||
import logging
|
||
from datetime import datetime
|
||
import uuid
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class ChatService:
|
||
"""聊天服务类，处理聊天记录的CRUD操作"""
|
||
|
||
def __init__(self, data_path: Path):
|
||
"""
|
||
初始化聊天服务
|
||
|
||
Args:
|
||
data_path: 数据目录路径
|
||
"""
|
||
self.data_path = data_path
|
||
self.chat_dir = data_path / "chat"
|
||
self.chat_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
def list_all_chats(self) -> Dict[str, List[Dict]]:
|
||
"""
|
||
获取所有角色和聊天列表
|
||
|
||
Returns:
|
||
Dict[str, List[Dict]]: 字典结构，键是角色名称，值是该角色的聊天信息列表
|
||
"""
|
||
result = {}
|
||
|
||
if not self.chat_dir.exists():
|
||
logger.warning(f"聊天目录不存在: {self.chat_dir}")
|
||
return result
|
||
|
||
for role_dir in self.chat_dir.iterdir():
|
||
try:
|
||
if role_dir.is_dir():
|
||
chats = []
|
||
|
||
for chat_file in role_dir.glob("*.jsonl"):
|
||
chat_info = self._get_chat_summary(role_dir.name, chat_file.stem)
|
||
if chat_info:
|
||
chats.append(chat_info)
|
||
|
||
if chats:
|
||
result[role_dir.name] = chats
|
||
|
||
except Exception as e:
|
||
logger.error(f"处理角色目录 {role_dir.name} 时出错: {str(e)}")
|
||
continue
|
||
|
||
return {"chat": [{"role_name": role, **chat} for role, chats in result.items() for chat in chats]}
|
||
|
||
def _get_chat_summary(self, role_name: str, chat_name: str) -> Optional[Dict]:
|
||
"""
|
||
获取聊天摘要信息
|
||
|
||
Args:
|
||
role_name: 角色名称
|
||
chat_name: 聊天名称
|
||
|
||
Returns:
|
||
Dict: 聊天摘要信息，如果文件不存在则返回None
|
||
"""
|
||
chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
|
||
|
||
if not chat_file.exists():
|
||
return None
|
||
|
||
try:
|
||
with open(chat_file, 'r', encoding='utf-8') as f:
|
||
lines = f.readlines()
|
||
|
||
if not lines:
|
||
return None
|
||
|
||
# 第一行是header
|
||
header = json.loads(lines[0])
|
||
|
||
# 计算消息数量（排除header）
|
||
message_count = len(lines) - 1
|
||
|
||
# 获取最后修改时间
|
||
last_modified = datetime.fromtimestamp(
|
||
chat_file.stat().st_mtime
|
||
).isoformat()
|
||
|
||
# 获取最后一条消息预览
|
||
last_message = ""
|
||
if message_count > 0:
|
||
try:
|
||
last_msg_data = json.loads(lines[-1])
|
||
last_message = last_msg_data.get("mes", "")
|
||
except:
|
||
pass
|
||
|
||
return {
|
||
"chat_name": chat_name,
|
||
"user_name": header.get("user_name", "User"),
|
||
"character_name": header.get("character_name", ""),
|
||
"last_modified": last_modified,
|
||
"message_count": message_count,
|
||
"last_message": last_message
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"读取聊天摘要失败 {role_name}/{chat_name}: {str(e)}")
|
||
return None
|
||
|
||
def get_chat(self, role_name: str, chat_name: str) -> Dict[str, Any]:
|
||
"""
|
||
获取指定聊天的完整内容
|
||
|
||
Args:
|
||
role_name: 角色名称
|
||
chat_name: 聊天名称
|
||
|
||
Returns:
|
||
Dict: 包含metadata和messages的字典
|
||
|
||
Raises:
|
||
FileNotFoundError: 聊天文件不存在
|
||
"""
|
||
chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
|
||
|
||
if not chat_file.exists():
|
||
raise FileNotFoundError(f"Chat not found: {role_name}/{chat_name}")
|
||
|
||
try:
|
||
with open(chat_file, 'r', encoding='utf-8') as f:
|
||
lines = f.readlines()
|
||
|
||
if not lines:
|
||
raise ValueError(f"Empty chat file: {role_name}/{chat_name}")
|
||
|
||
# 第一行是header
|
||
header = json.loads(lines[0])
|
||
|
||
# 解析消息
|
||
messages = []
|
||
for i, line in enumerate(lines[1:], start=1):
|
||
if line.strip(): # 跳过空行
|
||
msg_data = json.loads(line)
|
||
# 确保有floor字段
|
||
if "floor" not in msg_data:
|
||
msg_data["floor"] = i
|
||
|
||
messages.append(msg_data)
|
||
|
||
return {
|
||
"metadata": {
|
||
"user_name": header.get("user_name", "User"),
|
||
"character_name": header.get("character_name", ""),
|
||
"chat_id": header.get("chat_id_hash", ""),
|
||
"integrity": header.get("integrity", "")
|
||
},
|
||
"messages": messages
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"读取聊天失败 {role_name}/{chat_name}: {str(e)}")
|
||
raise
|
||
|
||
def create_chat(self, role_name: str, chat_name: str, metadata: Dict = None) -> Dict:
|
||
"""
|
||
创建新聊天
|
||
|
||
Args:
|
||
role_name: 角色名称
|
||
chat_name: 聊天名称
|
||
metadata: 聊天元数据
|
||
|
||
Returns:
|
||
Dict: 创建的聊天信息
|
||
|
||
Raises:
|
||
FileExistsError: 聊天已存在
|
||
"""
|
||
chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
|
||
|
||
if chat_file.exists():
|
||
raise FileExistsError(f"Chat already exists: {role_name}/{chat_name}")
|
||
|
||
# 创建角色目录
|
||
chat_file.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 构建header
|
||
header = {
|
||
"user_name": metadata.get("user_name", "User") if metadata else "User",
|
||
"character_name": metadata.get("character_name", role_name) if metadata else role_name,
|
||
"integrity": str(uuid.uuid4()),
|
||
"chat_id_hash": str(uuid.uuid4()),
|
||
"note_prompt": "",
|
||
"note_interval": 0,
|
||
"note_position": 0,
|
||
"note_depth": 0,
|
||
"note_role": 0,
|
||
"extensions": {},
|
||
"timedWorldInfo": {},
|
||
"variables": {},
|
||
"tainted": False,
|
||
"lastInContextMessageId": -1
|
||
}
|
||
|
||
# 写入header
|
||
with open(chat_file, 'w', encoding='utf-8') as f:
|
||
f.write(json.dumps(header, ensure_ascii=False) + '\n')
|
||
|
||
return {
|
||
"role_name": role_name,
|
||
"chat_name": chat_name,
|
||
"metadata": {
|
||
"user_name": header["user_name"],
|
||
"character_name": header["character_name"]
|
||
}
|
||
}
|
||
|
||
def add_message(self, role_name: str, chat_name: str, message_data: Dict) -> Dict:
|
||
"""
|
||
向聊天添加新消息
|
||
|
||
Args:
|
||
role_name: 角色名称
|
||
chat_name: 聊天名称
|
||
message_data: 消息数据
|
||
|
||
Returns:
|
||
Dict: 添加的消息
|
||
|
||
Raises:
|
||
FileNotFoundError: 聊天不存在
|
||
"""
|
||
chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
|
||
|
||
if not chat_file.exists():
|
||
raise FileNotFoundError(f"Chat not found: {role_name}/{chat_name}")
|
||
|
||
try:
|
||
# 读取现有消息以确定floor
|
||
with open(chat_file, 'r', encoding='utf-8') as f:
|
||
lines = f.readlines()
|
||
|
||
# 计算下一个floor号
|
||
next_floor = len(lines) - 1 # 减去header行
|
||
|
||
# 构建完整的消息数据
|
||
full_message = {
|
||
"name": message_data.get("name", "User"),
|
||
"is_user": message_data.get("is_user", True),
|
||
"is_system": message_data.get("is_system", False),
|
||
"floor": next_floor,
|
||
"send_date": message_data.get("send_date", str(int(datetime.now().timestamp() * 1000))),
|
||
"mes": message_data.get("mes", ""),
|
||
"extra": message_data.get("extra", {}),
|
||
"swipes": message_data.get("swipes", []),
|
||
"swipe_id": message_data.get("swipe_id", 0),
|
||
"force_avatar": None,
|
||
"variables": [],
|
||
"variables_initialized": [],
|
||
"is_ejs_processed": []
|
||
}
|
||
|
||
# 追加消息到文件
|
||
with open(chat_file, 'a', encoding='utf-8') as f:
|
||
f.write(json.dumps(full_message, ensure_ascii=False) + '\n')
|
||
|
||
return full_message
|
||
|
||
except Exception as e:
|
||
logger.error(f"添加消息失败 {role_name}/{chat_name}: {str(e)}")
|
||
raise
|
||
|
||
def update_message(self, role_name: str, chat_name: str, floor: int, update_data: Dict) -> Dict:
|
||
"""
|
||
更新指定楼层的消息
|
||
|
||
Args:
|
||
role_name: 角色名称
|
||
chat_name: 聊天名称
|
||
floor: 楼层号
|
||
update_data: 更新的数据
|
||
|
||
Returns:
|
||
Dict: 更新后的消息
|
||
|
||
Raises:
|
||
FileNotFoundError: 聊天不存在
|
||
ValueError: 楼层不存在
|
||
"""
|
||
chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
|
||
|
||
if not chat_file.exists():
|
||
raise FileNotFoundError(f"Chat not found: {role_name}/{chat_name}")
|
||
|
||
try:
|
||
# 读取所有行
|
||
with open(chat_file, 'r', encoding='utf-8') as f:
|
||
lines = f.readlines()
|
||
|
||
# 找到对应的消息行（floor + 1，因为第0行是header）
|
||
message_line_index = floor + 1
|
||
|
||
if message_line_index >= len(lines):
|
||
raise ValueError(f"Floor {floor} not found in chat")
|
||
|
||
# 解析并更新消息
|
||
msg_data = json.loads(lines[message_line_index])
|
||
msg_data.update(update_data)
|
||
|
||
# 写回文件
|
||
lines[message_line_index] = json.dumps(msg_data, ensure_ascii=False) + '\n'
|
||
|
||
with open(chat_file, 'w', encoding='utf-8') as f:
|
||
f.writelines(lines)
|
||
|
||
return msg_data
|
||
|
||
except Exception as e:
|
||
logger.error(f"更新消息失败 {role_name}/{chat_name}/{floor}: {str(e)}")
|
||
raise
|
||
|
||
def delete_message(self, role_name: str, chat_name: str, floor: int) -> Dict:
|
||
"""
|
||
删除指定楼层的消息
|
||
|
||
Args:
|
||
role_name: 角色名称
|
||
chat_name: 聊天名称
|
||
floor: 楼层号
|
||
|
||
Returns:
|
||
Dict: 被删除的消息
|
||
|
||
Raises:
|
||
FileNotFoundError: 聊天不存在
|
||
ValueError: 楼层不存在
|
||
"""
|
||
chat_file = self.chat_dir / role_name / f"{chat_name}.jsonl"
|
||
|
||
if not chat_file.exists():
|
||
raise FileNotFoundError(f"Chat not found: {role_name}/{chat_name}")
|
||
|
||
try:
|
||
# 读取所有行
|
||
with open(chat_file, 'r', encoding='utf-8') as f:
|
||
lines = f.readlines()
|
||
|
||
# 找到对应的消息行
|
||
message_line_index = floor + 1
|
||
|
||
if message_line_index >= len(lines):
|
||
raise ValueError(f"Floor {floor} not found in chat")
|
||
|
||
# 保存被删除的消息
|
||
deleted_msg = json.loads(lines[message_line_index])
|
||
|
||
# 删除该行
|
||
del lines[message_line_index]
|
||
|
||
# 重新编号后续消息的floor
|
||
for i in range(message_line_index, len(lines)):
|
||
if lines[i].strip(): # 跳过空行
|
||
msg_data = json.loads(lines[i])
|
||
msg_data["floor"] = i - 1 # 重新计算floor
|
||
lines[i] = json.dumps(msg_data, ensure_ascii=False) + '\n'
|
||
|
||
# 写回文件
|
||
with open(chat_file, 'w', encoding='utf-8') as f:
|
||
f.writelines(lines)
|
||
|
||
return deleted_msg
|
||
|
||
except Exception as e:
|
||
logger.error(f"删除消息失败 {role_name}/{chat_name}/{floor}: {str(e)}")
|
||
raise
|