重构数据库交互以使用 Peewee ORM
- 更新数据库连接和模型定义,以便使用 Peewee for SQLite。 - 在消息存储和检索功能中,用 Peewee ORM 查询替换 MongoDB 查询。 - 为 Messages、ThinkingLog 和 OnlineTime 引入了新的模型,以方便结构化数据存储。 - 增强了数据库操作的错误处理和日志记录。 - 删除了过时的 MongoDB 集合管理代码。 - 通过利用 Peewee 内置的查询和数据操作方法来提升性能。
This commit is contained in:
@@ -1,9 +1,10 @@
|
||||
from src.config.config import global_config
|
||||
from src.chat.message_receive.message import MessageRecv, MessageSending, Message
|
||||
from common.database.database import db
|
||||
from src.common.database.database_model import Messages, ThinkingLog
|
||||
import time
|
||||
import traceback
|
||||
from typing import List
|
||||
import json
|
||||
|
||||
|
||||
class InfoCatcher:
|
||||
@@ -60,8 +61,6 @@ class InfoCatcher:
|
||||
def catch_after_observe(self, obs_duration: float): # 这里可以有更多信息
|
||||
self.timing_results["sub_heartflow_observe_time"] = obs_duration
|
||||
|
||||
# def catch_shf
|
||||
|
||||
def catch_afer_shf_step(self, step_duration: float, past_mind: str, current_mind: str):
|
||||
self.timing_results["sub_heartflow_step_time"] = step_duration
|
||||
if len(past_mind) > 1:
|
||||
@@ -72,25 +71,10 @@ class InfoCatcher:
|
||||
self.heartflow_data["sub_heartflow_now"] = current_mind
|
||||
|
||||
def catch_after_llm_generated(self, prompt: str, response: str, reasoning_content: str = "", model_name: str = ""):
|
||||
# if self.response_mode == "heart_flow": # 条件判断不需要了喵~
|
||||
# self.heartflow_data["prompt"] = prompt
|
||||
# self.heartflow_data["response"] = response
|
||||
# self.heartflow_data["model"] = model_name
|
||||
# elif self.response_mode == "reasoning": # 条件判断不需要了喵~
|
||||
# self.reasoning_data["thinking_log"] = reasoning_content
|
||||
# self.reasoning_data["prompt"] = prompt
|
||||
# self.reasoning_data["response"] = response
|
||||
# self.reasoning_data["model"] = model_name
|
||||
|
||||
# 直接记录信息喵~
|
||||
self.reasoning_data["thinking_log"] = reasoning_content
|
||||
self.reasoning_data["prompt"] = prompt
|
||||
self.reasoning_data["response"] = response
|
||||
self.reasoning_data["model"] = model_name
|
||||
# 如果 heartflow 数据也需要通用字段,可以取消下面的注释喵~
|
||||
# self.heartflow_data["prompt"] = prompt
|
||||
# self.heartflow_data["response"] = response
|
||||
# self.heartflow_data["model"] = model_name
|
||||
|
||||
self.response_text = response
|
||||
|
||||
@@ -102,6 +86,7 @@ class InfoCatcher:
|
||||
):
|
||||
self.timing_results["make_response_time"] = response_duration
|
||||
self.response_time = time.time()
|
||||
self.response_messages = []
|
||||
for msg in response_message:
|
||||
self.response_messages.append(msg)
|
||||
|
||||
@@ -112,107 +97,110 @@ class InfoCatcher:
|
||||
@staticmethod
|
||||
def get_message_from_db_between_msgs(message_start: Message, message_end: Message):
|
||||
try:
|
||||
# 从数据库中获取消息的时间戳
|
||||
time_start = message_start.message_info.time
|
||||
time_end = message_end.message_info.time
|
||||
chat_id = message_start.chat_stream.stream_id
|
||||
|
||||
print(f"查询参数: time_start={time_start}, time_end={time_end}, chat_id={chat_id}")
|
||||
|
||||
# 查询数据库,获取 chat_id 相同且时间在 start 和 end 之间的数据
|
||||
messages_between = db.messages.find(
|
||||
{"chat_id": chat_id, "time": {"$gt": time_start, "$lt": time_end}}
|
||||
).sort("time", -1)
|
||||
messages_between_query = Messages.select().where(
|
||||
(Messages.chat_id == chat_id) &
|
||||
(Messages.time > time_start) &
|
||||
(Messages.time < time_end)
|
||||
).order_by(Messages.time.desc())
|
||||
|
||||
result = list(messages_between)
|
||||
result = list(messages_between_query)
|
||||
print(f"查询结果数量: {len(result)}")
|
||||
if result:
|
||||
print(f"第一条消息时间: {result[0]['time']}")
|
||||
print(f"最后一条消息时间: {result[-1]['time']}")
|
||||
print(f"第一条消息时间: {result[0].time}")
|
||||
print(f"最后一条消息时间: {result[-1].time}")
|
||||
return result
|
||||
except Exception as e:
|
||||
print(f"获取消息时出错: {str(e)}")
|
||||
print(traceback.format_exc())
|
||||
return []
|
||||
|
||||
def get_message_from_db_before_msg(self, message: MessageRecv):
|
||||
# 从数据库中获取消息
|
||||
message_id = message.message_info.message_id
|
||||
chat_id = message.chat_stream.stream_id
|
||||
message_id_val = message.message_info.message_id
|
||||
chat_id_val = message.chat_stream.stream_id
|
||||
|
||||
# 查询数据库,获取 chat_id 相同且 message_id 小于当前消息的 30 条数据
|
||||
messages_before = (
|
||||
db.messages.find({"chat_id": chat_id, "message_id": {"$lt": message_id}})
|
||||
.sort("time", -1)
|
||||
.limit(self.context_length * 3)
|
||||
) # 获取更多历史信息
|
||||
messages_before_query = Messages.select().where(
|
||||
(Messages.chat_id == chat_id_val) &
|
||||
(Messages.message_id < message_id_val)
|
||||
).order_by(Messages.time.desc()).limit(self.context_length * 3)
|
||||
|
||||
return list(messages_before)
|
||||
return list(messages_before_query)
|
||||
|
||||
def message_list_to_dict(self, message_list):
|
||||
# 存储简化的聊天记录
|
||||
result = []
|
||||
for message in message_list:
|
||||
if not isinstance(message, dict):
|
||||
message = self.message_to_dict(message)
|
||||
# print(message)
|
||||
for msg_item in message_list:
|
||||
processed_msg_item = msg_item
|
||||
if not isinstance(msg_item, dict):
|
||||
processed_msg_item = self.message_to_dict(msg_item)
|
||||
|
||||
if not processed_msg_item:
|
||||
continue
|
||||
|
||||
lite_message = {
|
||||
"time": message["time"],
|
||||
"user_nickname": message["user_info"]["user_nickname"],
|
||||
"processed_plain_text": message["processed_plain_text"],
|
||||
"time": processed_msg_item.get("time"),
|
||||
"user_nickname": processed_msg_item.get("user_nickname"),
|
||||
"processed_plain_text": processed_msg_item.get("processed_plain_text"),
|
||||
}
|
||||
result.append(lite_message)
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def message_to_dict(message):
|
||||
if not message:
|
||||
def message_to_dict(msg_obj):
|
||||
if not msg_obj:
|
||||
return None
|
||||
if isinstance(message, dict):
|
||||
return message
|
||||
return {
|
||||
# "message_id": message.message_info.message_id,
|
||||
"time": message.message_info.time,
|
||||
"user_id": message.message_info.user_info.user_id,
|
||||
"user_nickname": message.message_info.user_info.user_nickname,
|
||||
"processed_plain_text": message.processed_plain_text,
|
||||
# "detailed_plain_text": message.detailed_plain_text
|
||||
}
|
||||
if isinstance(msg_obj, dict):
|
||||
return msg_obj
|
||||
|
||||
if isinstance(msg_obj, Messages):
|
||||
return {
|
||||
"time": msg_obj.time,
|
||||
"user_id": msg_obj.user_id,
|
||||
"user_nickname": msg_obj.user_nickname,
|
||||
"processed_plain_text": msg_obj.processed_plain_text,
|
||||
}
|
||||
|
||||
if hasattr(msg_obj, 'message_info') and hasattr(msg_obj.message_info, 'user_info'):
|
||||
return {
|
||||
"time": msg_obj.message_info.time,
|
||||
"user_id": msg_obj.message_info.user_info.user_id,
|
||||
"user_nickname": msg_obj.message_info.user_info.user_nickname,
|
||||
"processed_plain_text": msg_obj.processed_plain_text,
|
||||
}
|
||||
|
||||
print(f"Warning: message_to_dict received an unhandled type: {type(msg_obj)}")
|
||||
return {}
|
||||
|
||||
def done_catch(self):
|
||||
"""将收集到的信息存储到数据库的 thinking_log 集合中喵~"""
|
||||
"""将收集到的信息存储到数据库的 thinking_log 表中喵~"""
|
||||
try:
|
||||
# 将消息对象转换为可序列化的字典喵~
|
||||
|
||||
thinking_log_data = {
|
||||
"chat_id": self.chat_id,
|
||||
"trigger_text": self.trigger_response_text,
|
||||
"response_text": self.response_text,
|
||||
"trigger_info": {
|
||||
"time": self.trigger_response_time,
|
||||
"message": self.message_to_dict(self.trigger_response_message),
|
||||
},
|
||||
"response_info": {
|
||||
"time": self.response_time,
|
||||
"message": self.response_messages,
|
||||
},
|
||||
"timing_results": self.timing_results,
|
||||
"chat_history": self.message_list_to_dict(self.chat_history),
|
||||
"chat_history_in_thinking": self.message_list_to_dict(self.chat_history_in_thinking),
|
||||
"chat_history_after_response": self.message_list_to_dict(self.chat_history_after_response),
|
||||
"heartflow_data": self.heartflow_data,
|
||||
"reasoning_data": self.reasoning_data,
|
||||
trigger_info_dict = self.message_to_dict(self.trigger_response_message)
|
||||
response_info_dict = {
|
||||
"time": self.response_time,
|
||||
"message": self.response_messages,
|
||||
}
|
||||
chat_history_list = self.message_list_to_dict(self.chat_history)
|
||||
chat_history_in_thinking_list = self.message_list_to_dict(self.chat_history_in_thinking)
|
||||
chat_history_after_response_list = self.message_list_to_dict(self.chat_history_after_response)
|
||||
|
||||
# 根据不同的响应模式添加相应的数据喵~ # 现在直接都加上去好了喵~
|
||||
# if self.response_mode == "heart_flow":
|
||||
# thinking_log_data["mode_specific_data"] = self.heartflow_data
|
||||
# elif self.response_mode == "reasoning":
|
||||
# thinking_log_data["mode_specific_data"] = self.reasoning_data
|
||||
|
||||
# 将数据插入到 thinking_log 集合中喵~
|
||||
db.thinking_log.insert_one(thinking_log_data)
|
||||
log_entry = ThinkingLog(
|
||||
chat_id=self.chat_id,
|
||||
trigger_text=self.trigger_response_text,
|
||||
response_text=self.response_text,
|
||||
trigger_info_json=json.dumps(trigger_info_dict) if trigger_info_dict else None,
|
||||
response_info_json=json.dumps(response_info_dict),
|
||||
timing_results_json=json.dumps(self.timing_results),
|
||||
chat_history_json=json.dumps(chat_history_list),
|
||||
chat_history_in_thinking_json=json.dumps(chat_history_in_thinking_list),
|
||||
chat_history_after_response_json=json.dumps(chat_history_after_response_list),
|
||||
heartflow_data_json=json.dumps(self.heartflow_data),
|
||||
reasoning_data_json=json.dumps(self.reasoning_data)
|
||||
)
|
||||
log_entry.save()
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
|
||||
Reference in New Issue
Block a user