413 lines
17 KiB
Python
413 lines
17 KiB
Python
import asyncio
|
||
import time
|
||
from typing import Dict, List, Optional, Tuple
|
||
from sqlalchemy import select, update, insert, and_, delete
|
||
from database.models import ConversationHistory, ConversationMessage
|
||
from services.database_service_base import DatabaseServiceBase
|
||
from utils.logger import session_logger
|
||
from config import settings
|
||
|
||
class DeviceConversationHistory:
|
||
def __init__(self):
|
||
self.history = []
|
||
self.last_interaction_time = asyncio.get_running_loop().time()
|
||
self.role_key = None # 角色标识
|
||
self.conversation_id = None # 会话ID，用于关联消息
|
||
|
||
class ConversationHistoryManager(DatabaseServiceBase):
|
||
def __init__(self):
|
||
super().__init__(service_name="conversation_history")
|
||
self.histories: Dict[str, DeviceConversationHistory] = {} # 内存缓存
|
||
self.lock = asyncio.Lock()
|
||
self.role_histories: Dict[Tuple[str, str], DeviceConversationHistory] = {} # (device_id, role_key) -> history
|
||
self.max_cache_size = 500 # 最大缓存条目数
|
||
|
||
async def get_history(self, device_id: str, role_key: Optional[str] = None) -> DeviceConversationHistory:
|
||
"""获取设备指定角色的对话历史"""
|
||
await self._init_database()
|
||
if not role_key:
|
||
from services.device_config import device_config_manager
|
||
device_config = await device_config_manager.get_config(device_id)
|
||
if device_config:
|
||
role_key = device_config.selected_role_key
|
||
else:
|
||
from config import settings
|
||
role_key = settings.selected_role_key
|
||
|
||
cache_key = (device_id, role_key)
|
||
if cache_key in self.role_histories:
|
||
return self.role_histories[cache_key]
|
||
|
||
db_session = await self.get_session()
|
||
try:
|
||
query = select(ConversationHistory).where(
|
||
and_(
|
||
ConversationHistory.device_id == device_id,
|
||
ConversationHistory.role_key == role_key
|
||
)
|
||
)
|
||
result = await db_session.execute(query)
|
||
db_history = result.scalar_one_or_none()
|
||
|
||
history = DeviceConversationHistory()
|
||
history.role_key = role_key
|
||
|
||
if db_history:
|
||
messages_query = select(ConversationMessage).where(
|
||
ConversationMessage.conversation_id == db_history.id
|
||
).order_by(ConversationMessage.timestamp)
|
||
|
||
messages_result = await db_session.execute(messages_query)
|
||
messages = messages_result.scalars().all()
|
||
|
||
message_pairs = []
|
||
user_msg = None
|
||
|
||
for msg in messages:
|
||
if msg.is_user:
|
||
user_msg = msg.content
|
||
elif user_msg is not None:
|
||
message_pairs.append({"user": user_msg, "assistant": msg.content})
|
||
user_msg = None
|
||
|
||
if user_msg is not None:
|
||
message_pairs.append({"user": user_msg, "assistant": ""})
|
||
|
||
history.history = message_pairs
|
||
history.last_interaction_time = db_history.last_interaction_time
|
||
history.conversation_id = db_history.id
|
||
|
||
self._add_to_cache(cache_key, history)
|
||
return history
|
||
else:
|
||
history = DeviceConversationHistory()
|
||
history.history = []
|
||
history.role_key = role_key
|
||
self._add_to_cache(cache_key, history)
|
||
return history
|
||
except Exception as e:
|
||
session_logger.error(device_id, "conversation_history", f"从数据库获取对话历史失败: {str(e)}")
|
||
history = DeviceConversationHistory()
|
||
history.role_key = role_key
|
||
return history
|
||
finally:
|
||
await db_session.close()
|
||
|
||
async def set_history(self, device_id: str, history: DeviceConversationHistory, role_key: Optional[str] = None):
|
||
"""设置设备指定角色的对话历史"""
|
||
await self._init_database()
|
||
if not role_key and not history.role_key:
|
||
from services.device_config import device_config_manager
|
||
device_config = await device_config_manager.get_config(device_id)
|
||
role_key = device_config.selected_role_key if device_config else settings.selected_role_key
|
||
|
||
effective_role_key = role_key or history.role_key
|
||
history.role_key = effective_role_key
|
||
cache_key = (device_id, effective_role_key)
|
||
self._add_to_cache(cache_key, history)
|
||
|
||
db_session = await self.get_session()
|
||
try:
|
||
query = select(ConversationHistory).where(
|
||
and_(
|
||
ConversationHistory.device_id == device_id,
|
||
ConversationHistory.role_key == effective_role_key
|
||
)
|
||
)
|
||
result = await db_session.execute(query)
|
||
db_history = result.scalar_one_or_none()
|
||
current_time = time.time()
|
||
|
||
if db_history:
|
||
stmt = update(ConversationHistory).where(
|
||
ConversationHistory.id == db_history.id
|
||
).values(
|
||
last_interaction_time=current_time
|
||
)
|
||
await db_session.execute(stmt)
|
||
conversation_id = db_history.id
|
||
else:
|
||
stmt = insert(ConversationHistory).values(
|
||
device_id=device_id,
|
||
role_key=effective_role_key,
|
||
last_interaction_time=current_time
|
||
)
|
||
result = await db_session.execute(stmt)
|
||
await db_session.commit()
|
||
|
||
query = select(ConversationHistory).where(
|
||
and_(
|
||
ConversationHistory.device_id == device_id,
|
||
ConversationHistory.role_key == effective_role_key
|
||
)
|
||
)
|
||
result = await db_session.execute(query)
|
||
db_history = result.scalar_one()
|
||
conversation_id = db_history.id
|
||
history.conversation_id = conversation_id
|
||
|
||
if history.history and len(history.history) > 0:
|
||
latest_msg = history.history[-1]
|
||
|
||
latest_msg_query = select(ConversationMessage).where(
|
||
and_(
|
||
ConversationMessage.conversation_id == conversation_id,
|
||
ConversationMessage.content == latest_msg["user"],
|
||
ConversationMessage.is_user == True
|
||
)
|
||
).order_by(ConversationMessage.timestamp.desc())
|
||
|
||
latest_msg_result = await db_session.execute(latest_msg_query)
|
||
existing_user_msg = latest_msg_result.scalar_one_or_none()
|
||
|
||
if not existing_user_msg:
|
||
user_stmt = insert(ConversationMessage).values(
|
||
conversation_id=conversation_id,
|
||
is_user=True,
|
||
content=latest_msg["user"],
|
||
timestamp=current_time - 0.1 # 确保用户消息在助手消息之前
|
||
)
|
||
await db_session.execute(user_stmt)
|
||
|
||
assistant_stmt = insert(ConversationMessage).values(
|
||
conversation_id=conversation_id,
|
||
is_user=False,
|
||
content=latest_msg["assistant"],
|
||
timestamp=current_time
|
||
)
|
||
await db_session.execute(assistant_stmt)
|
||
|
||
await db_session.commit()
|
||
session_logger.info(
|
||
device_id,
|
||
"conversation_history",
|
||
f"已保存设备 {device_id} 的角色 {effective_role_key} 对话历史到数据库"
|
||
)
|
||
except Exception as e:
|
||
await db_session.rollback()
|
||
session_logger.error(
|
||
device_id,
|
||
"conversation_history",
|
||
f"保存对话历史到数据库失败: {str(e)}"
|
||
)
|
||
finally:
|
||
await db_session.close()
|
||
|
||
async def remove_history(self, device_id: str, role_key: Optional[str] = None):
|
||
"""删除设备的对话历史"""
|
||
await self._init_database()
|
||
if not role_key:
|
||
from services.device_config import device_config_manager
|
||
device_config = await device_config_manager.get_config(device_id)
|
||
role_key = device_config.selected_role_key if device_config else settings.selected_role_key
|
||
|
||
cache_key = (device_id, role_key)
|
||
if cache_key in self.role_histories:
|
||
del self.role_histories[cache_key]
|
||
|
||
db_session = await self.get_session()
|
||
try:
|
||
query = select(ConversationHistory.id).where(
|
||
and_(
|
||
ConversationHistory.device_id == device_id,
|
||
ConversationHistory.role_key == role_key
|
||
)
|
||
)
|
||
result = await db_session.execute(query)
|
||
conversation_id = result.scalar_one_or_none()
|
||
|
||
if conversation_id:
|
||
msg_stmt = delete(ConversationMessage).where(
|
||
ConversationMessage.conversation_id == conversation_id
|
||
)
|
||
await db_session.execute(msg_stmt)
|
||
|
||
hist_stmt = delete(ConversationHistory).where(
|
||
ConversationHistory.id == conversation_id
|
||
)
|
||
await db_session.execute(hist_stmt)
|
||
|
||
await db_session.commit()
|
||
session_logger.info(
|
||
device_id,
|
||
"conversation_history",
|
||
f"已删除设备 {device_id} 的角色 {role_key} 对话历史"
|
||
)
|
||
except Exception as e:
|
||
await db_session.rollback()
|
||
session_logger.error(
|
||
device_id,
|
||
"conversation_history",
|
||
f"删除对话历史失败: {str(e)}"
|
||
)
|
||
finally:
|
||
await db_session.close()
|
||
|
||
async def clear_all_histories(self, device_id: str):
|
||
"""清除设备的所有对话历史（所有角色）"""
|
||
await self._init_database()
|
||
keys_to_remove = []
|
||
for (dev_id, _) in self.role_histories.keys():
|
||
if dev_id == device_id:
|
||
keys_to_remove.append((dev_id, _))
|
||
|
||
for key in keys_to_remove:
|
||
del self.role_histories[key]
|
||
|
||
db_session = await self.get_session()
|
||
try:
|
||
query = select(ConversationHistory.id).where(
|
||
ConversationHistory.device_id == device_id
|
||
)
|
||
result = await db_session.execute(query)
|
||
conversation_ids = result.scalars().all()
|
||
|
||
if conversation_ids:
|
||
for cid in conversation_ids:
|
||
msg_stmt = delete(ConversationMessage).where(
|
||
ConversationMessage.conversation_id == cid
|
||
)
|
||
await db_session.execute(msg_stmt)
|
||
|
||
stmt = delete(ConversationHistory).where(
|
||
ConversationHistory.device_id == device_id
|
||
)
|
||
await db_session.execute(stmt)
|
||
await db_session.commit()
|
||
|
||
session_logger.info(
|
||
device_id,
|
||
"conversation_history",
|
||
f"已清除设备 {device_id} 的所有对话历史"
|
||
)
|
||
except Exception as e:
|
||
await db_session.rollback()
|
||
session_logger.error(
|
||
device_id,
|
||
"conversation_history",
|
||
f"清除所有对话历史失败: {str(e)}"
|
||
)
|
||
finally:
|
||
await db_session.close()
|
||
|
||
async def get_all_histories(self):
|
||
"""获取所有历史记录，主要用于清理过期记录"""
|
||
await self._init_database()
|
||
db_session = await self.get_session()
|
||
try:
|
||
query = select(ConversationHistory)
|
||
result = await db_session.execute(query)
|
||
db_histories = result.scalars().all()
|
||
histories = []
|
||
|
||
for db_history in db_histories:
|
||
history = DeviceConversationHistory()
|
||
history.last_interaction_time = db_history.last_interaction_time
|
||
history.role_key = db_history.role_key
|
||
history.conversation_id = db_history.id
|
||
|
||
histories.append(((db_history.device_id, db_history.role_key), history))
|
||
|
||
return histories
|
||
except Exception as e:
|
||
session_logger.error("system", "conversation_history", f"获取所有对话历史失败: {str(e)}")
|
||
return []
|
||
finally:
|
||
await db_session.close()
|
||
|
||
async def get_device_role_histories(self, device_id: str) -> Dict[str, DeviceConversationHistory]:
|
||
"""获取设备的所有角色会话历史"""
|
||
await self._init_database()
|
||
db_session = await self.get_session()
|
||
try:
|
||
query = select(ConversationHistory).where(
|
||
ConversationHistory.device_id == device_id
|
||
)
|
||
result = await db_session.execute(query)
|
||
db_histories = result.scalars().all()
|
||
|
||
histories = {}
|
||
for db_history in db_histories:
|
||
role_key = db_history.role_key
|
||
history = DeviceConversationHistory()
|
||
history.role_key = role_key
|
||
history.last_interaction_time = db_history.last_interaction_time
|
||
history.conversation_id = db_history.id
|
||
|
||
messages_query = select(ConversationMessage).where(
|
||
ConversationMessage.conversation_id == db_history.id
|
||
).order_by(ConversationMessage.timestamp)
|
||
|
||
messages_result = await db_session.execute(messages_query)
|
||
messages = messages_result.scalars().all()
|
||
|
||
for msg in messages:
|
||
if msg.is_user:
|
||
history.history.append({
|
||
"user": msg.content,
|
||
"assistant": "",
|
||
"timestamp": msg.timestamp
|
||
})
|
||
else:
|
||
if history.history and "assistant" in history.history[-1]:
|
||
history.history[-1]["assistant"] = msg.content
|
||
|
||
histories[role_key] = history
|
||
|
||
return histories
|
||
except Exception as e:
|
||
session_logger.error(device_id, "conversation_history", f"获取设备所有角色历史失败: {str(e)}")
|
||
return {}
|
||
finally:
|
||
await db_session.close()
|
||
|
||
async def clear_cache(self, device_id: str = None, role_key: str = None):
|
||
"""清除指定设备和角色的缓存，或全部缓存"""
|
||
if device_id and role_key:
|
||
cache_key = (device_id, role_key)
|
||
if cache_key in self.role_histories:
|
||
del self.role_histories[cache_key]
|
||
elif device_id:
|
||
keys_to_remove = []
|
||
for (dev_id, _), _ in self.role_histories.items():
|
||
if dev_id == device_id:
|
||
keys_to_remove.append((dev_id, _))
|
||
|
||
for key in keys_to_remove:
|
||
if key in self.role_histories:
|
||
del self.role_histories[key]
|
||
else:
|
||
self.role_histories.clear()
|
||
|
||
def _add_to_cache(self, cache_key: Tuple[str, str], history: DeviceConversationHistory):
|
||
"""添加到缓存，检查大小限制"""
|
||
# 检查缓存大小限制
|
||
if len(self.role_histories) >= self.max_cache_size:
|
||
self._cleanup_old_cache_entries()
|
||
|
||
self.role_histories[cache_key] = history
|
||
|
||
def _cleanup_old_cache_entries(self):
|
||
"""清理最旧的缓存条目"""
|
||
if not self.role_histories:
|
||
return
|
||
|
||
# 按最后交互时间排序，删除最旧的25%条目
|
||
sorted_items = sorted(
|
||
self.role_histories.items(),
|
||
key=lambda x: x[1].last_interaction_time
|
||
)
|
||
cleanup_count = max(1, len(sorted_items) // 4)
|
||
|
||
for cache_key, _ in sorted_items[:cleanup_count]:
|
||
if cache_key in self.role_histories:
|
||
del self.role_histories[cache_key]
|
||
|
||
session_logger.info(
|
||
"system", "cache_cleanup",
|
||
f"对话历史缓存清理完成，删除了 {cleanup_count} 个条目，剩余 {len(self.role_histories)} 个"
|
||
)
|
||
|
||
|
||
conversation_history_manager = ConversationHistoryManager()
|