Files
SmartChat/chat_session.py
2026-01-08 01:06:27 +08:00

735 lines
27 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import datetime
import json
import os.path
from typing import List, Tuple
import logging
import os
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import pytz
from tools import DatabaseManager
from file import ObjectFileManager, LocalFileManager
TZ = pytz.timezone('Asia/Shanghai')
def convert_to_utc(time_str, time_format='%Y-%m-%d %H:%M:%S', source_timezone=None):
# 解析时间字符串为datetime对象
local_dt = datetime.datetime.strptime(time_str, time_format)
# 如果指定了源时区,则将本地时间转换为UTC时间
if source_timezone:
tz = pytz.timezone(source_timezone)
local_dt = tz.localize(local_dt)
utc_dt = local_dt.astimezone(pytz.UTC)
else:
# 如果没有指定源时区,则假设时间字符串已经是UTC时间
utc_dt = pytz.UTC.localize(local_dt)
return utc_dt
class AIAgent:
"""AI代理类,整合了各种AI功能"""
def __init__(self, deepseek_api_key: str = "sk-c2f04b4d635b4df0aace0432f3c6d6f4",
max_retries: int = 3, timeout: int = 30):
"""初始化AI助手
Args:
deepseek_api_key: DeepSeek API密钥
max_retries: 最大重试次数
timeout: 请求超时时间(秒)
"""
self.deepseek_api_key = deepseek_api_key
self.timeout = timeout
# 配置日志
self._setup_logging()
# 创建带重试机制的会话
self.session = self._create_session(max_retries)
# 配置常量
self.BASE_REPORT_URL = "https://a.cmdp.cn/basiceg/v1/json/tpldev"
self.CACHE_DIR = "cache"
def _setup_logging(self):
"""配置日志"""
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
self.logger = logging.getLogger(__name__)
@staticmethod
def _create_session(max_retries: int) -> requests.Session:
"""创建带重试机制的会话
Args:
max_retries: 最大重试次数
Returns:
配置好的会话对象
"""
session = requests.Session()
# 配置重试策略
retry_strategy = Retry(
total=max_retries,
status_forcelist=[429, 500, 502, 503, 504],
backoff_factor=1
)
# 为HTTP和HTTPS适配器添加重试机制
adapter = HTTPAdapter(max_retries=retry_strategy)
session.mount("http://", adapter)
session.mount("https://", adapter)
return session
def _make_request(self, method: str, url: str, **kwargs):
"""发送HTTP请求(带错误处理)
Args:
method: HTTP方法
url: 请求URL
**kwargs: 其他请求参数
Returns:
JSON响应或None
"""
try:
# 设置默认超时
if 'timeout' not in kwargs:
kwargs['timeout'] = self.timeout
response = self.session.request(method, url, **kwargs)
response.raise_for_status()
self.logger.info(f"请求成功: {url}")
return response.json()
except requests.exceptions.RequestException as e:
self.logger.error(f"请求失败: {url} - {str(e)}")
return {}
except json.JSONDecodeError as e:
self.logger.error(f"JSON解析失败: {url} - {str(e)}")
return {}
def chat_deepseek(self, message: str, role: str = "", model: str = "deepseek-chat") -> str:
"""与DeepSeek聊天
Args:
message: 用户消息
role: 系统角色
model: 使用的模型
Returns:
AI回复内容
"""
url = "https://api.deepseek.com/chat/completions"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.deepseek_api_key}"
}
messages = []
if role:
messages.append({"role": "system", "content": role})
messages.append({"role": "user", "content": message})
body = {
"model": model,
"messages": messages,
"stream": False
}
result = self._make_request("POST", url, json=body, headers=headers)
if result and "choices" in result:
return result.get('choices', [{}])[0]['message']['content']
self.logger.warning("DeepSeek API调用失败或返回空结果")
return ""
agent = AIAgent(max_retries=1, timeout=180)
class ChatSessionManager:
"""聊天会话管理器类"""
def __init__(self, db_manager: DatabaseManager):
"""
初始化聊天会话管理器
Args:
db_manager: 数据库管理工具
"""
self.files_dir = os.path.join("", "files")
self.messages_dir = os.path.join("", "messages")
self.cache_path = None
# 初始化文件管理器
self.obj = ObjectFileManager()
self.lfm = LocalFileManager(self.files_dir)
self.dbm = db_manager
self.users = []
self.groups = []
self.members = {}
self.messages = []
self.files = {}
# 确保目录存在
os.makedirs(self.files_dir, exist_ok=True)
os.makedirs(self.messages_dir, exist_ok=True)
def execute_query(self, query: str, params=None, fetch_all=True):
result = []
if fetch_all:
rows = self.dbm.execute_query(query, params)
for row in rows:
result.append(list(row))
else:
rows = self.dbm.execute_query(query, params, False)
result = list(rows)
return result
def execute_event(self, query: str, params=None):
return self.dbm.execute_event(query, params)
def status(self):
row = self.dbm.execute_query("SELECT datetime FROM information order by datetime desc limit 1", fetch_all=False)
if row:
return row[0]
else:
return None
def _download_file(self, file_info, files_dir, history, lock):
"""
下载单个文件的线程函数
Args:
file_info: 文件信息
files_dir: 文件保存目录
history: 聊天历史字典
lock: 线程锁
"""
try:
path = file_info["path"].replace("s3://", "")
file_content = self.obj.get_file(path)
if path.endswith(".json"):
data = json.loads(file_content)
with lock:
history[data["msgid"]] = data
else:
# 保存非JSON文件
filename = path.split("/")[-1]
with open(os.path.join(files_dir, filename), "wb") as f:
f.write(file_content)
return path, True
except Exception as e:
return path, False
def download_chat_history(self, year: int, month: int, day: int, hour: int, max_workers: int = 5) -> bool:
"""
下载指定时间的聊天历史(多线程版本)
Args:
year: 年份
month: 月份
day: 日期
hour: 小时
max_workers: 最大线程数
Returns:
下载是否成功
"""
message_history_name = f"{year}_{month:02d}_{day:02d}_{hour:02d}.json"
print(f"Downloading chat history for {year}-{month:02d}-{day:02d}-{hour:02d} (多线程模式)")
try:
file_list = self.obj.get_file_list(f"zysl-log/wework/chat_msg/{year}/{month}/{day}/{hour}/")
print(f"共计{len(file_list)}个文件,使用{max_workers}个线程下载")
history = {}
failed_files = []
if file_list:
lock = threading.Lock()
# 使用线程池下载文件
with ThreadPoolExecutor(max_workers=max_workers) as executor:
# 提交所有下载任务
future_to_file = {
executor.submit(self._download_file, file_info, self.files_dir, history, lock): file_info
for file_info in file_list
}
# 处理完成的任务
completed_count = 0
total_count = len(file_list)
for future in as_completed(future_to_file):
file_info = future_to_file[future]
completed_count += 1
try:
path, success = future.result()
if success:
print(f"[{completed_count}/{total_count}] 下载成功: {path}")
else:
print(f"[{completed_count}/{total_count}] 下载失败: {path}")
failed_files.append(path)
except Exception as e:
print(f"[{completed_count}/{total_count}] 处理异常: {file_info['path']} - {str(e)}")
failed_files.append(file_info['path'])
# 检查是否有失败的文件
if failed_files:
print(f"警告: {len(failed_files)}个文件下载失败")
print("失败文件列表:", failed_files)
# 保存聊天历史JSON文件
output_path = os.path.join(self.messages_dir, message_history_name)
with open(output_path, "w", encoding="utf-8") as f:
json.dump(history, f, ensure_ascii=False, indent=2)
print(f"Downloading chat history for {year}-{month:02d}-{day:02d}-{hour:02d} complete")
print(f"成功下载: {len(file_list) - len(failed_files)}/{len(file_list)} 个文件")
# 如果有失败的文件,返回False
return len(failed_files) == 0
except Exception as e:
error_file = os.path.join(self.messages_dir, f"{message_history_name}.unfinished")
with open(error_file, "w", encoding="utf-8") as f:
f.write(str(e))
print(
f"Downloading chat history for {year}-{month:02d}-{day:02d}-{hour:02d} failed: please check the {message_history_name}.unfinished")
return False
def _parse_message(self, data, is_mixed: bool = False, idx: int = 0):
msg_type = data["msgtype"]
if msg_type == "text":
content = data['text']["content"]
elif msg_type in ["file", "image", "video", "voice", "emotion", "meeting_voice_call"]:
file_data = data[msg_type]
suffixs = [f"_{idx}", ".png", ".jpg", ".gif", ".mp4", "_.png", "_.jpg", "_.gif", "_.mp4",
("_" + file_data["filename"]) if msg_type == "file" else ""]
file_path = ""
for suffix in suffixs:
temp_path = f"{data['msgid']}_mixed{suffix}" if is_mixed else f"{data['msgid']}_{msg_type}{suffix}"
if self.lfm.exists(temp_path):
file_path = temp_path
break
if file_path != "":
ext_mapping = {
"file": (file_data.get('fileext', ''), "filesize"),
"image": ("png", file_data.get("filesize", 0)),
"video": ("mp4", file_data.get("filesize", 0)),
"voice": ("wav", file_data.get("voice_size", 0)),
"emotion": ("gif" if file_data.get("type", 0) == 1 else "png", file_data.get("imagesize", 0)),
"meeting_voice_call": ("wav", self.lfm.get_size(file_path)),
}
file_info = {
"id": file_data.get("md5sum", self.lfm.get_md5(file_path)),
"size": ext_mapping[msg_type][1],
"extension": ext_mapping[msg_type][0],
"content": file_path,
"original": self.cache_path + file_path,
"type": msg_type
}
self.files[file_info["id"]] = file_info
content = file_info["id"]
else:
file_path = f"{data['msgid']}_{msg_type}_0"
file_info = {
"id": file_data.get("md5sum", file_path),
"size": file_data.get('filesize', 0),
"extension": file_data.get('fileext', "unknown"),
"content": file_path,
"original": self.cache_path + file_path,
"type": msg_type
}
self.files[file_info["id"]] = file_info
content = file_info["id"]
elif msg_type == "link":
data = data[msg_type]
content = f'{data["description"]}{data["title"]}({data["link_url"]})'
else:
print(f"未知消息类型: {msg_type}")
content = json.dumps(data, ensure_ascii=False)
return content
def _parse_mixed_message(self, param, idx: int):
tp = param["msgtype"]
content = self._parse_message(param, True, idx)
if tp in ["file", "image", "video", "voice", "emotion", "meeting_voice_call"]:
self.files[content]["content"] = f'{param["msgid"]}_mixed_{idx}'
idx += 1
return idx, content
def parse_chat_history(self, year: int, month: int, day: int, hour: int):
"""
解析聊天历史数据
Args:
year: 年份
month: 月份
day: 日期
hour: 小时
Returns:
(用户列表, 群组列表, 成员映射, 消息列表, 文件映射)
"""
message_history_name = f"{year}_{month:02d}_{day:02d}_{hour:02d}.json"
history_path = os.path.join(self.messages_dir, message_history_name)
# 如果文件不存在,先下载
if not os.path.exists(history_path):
self.download_chat_history(year, month, day, hour,10)
with open(history_path, "r", encoding="utf-8") as f:
message_history = json.load(f)
try:
for message_data in message_history.values():
# 只处理发送的消息
if message_data.get("action") != "send":
continue
# 提取基本信息
from_user = message_data["from"]
to_users = message_data["tolist"]
room_id = message_data.get("roomid", "")
# 更新用户列表
self.users.extend([from_user] + to_users)
self.users = list(set(self.users))
# 确定接收者类型
to_type = 1 if room_id else 0
receiver = room_id if to_type == 1 else to_users[0]
# 处理群组信息
if room_id:
if room_id not in self.members:
self.members[room_id] = to_users + [from_user]
else:
self.members[room_id].extend(to_users + [from_user])
self.members[room_id] = list(set(self.members[room_id]))
if room_id not in self.groups:
self.groups.append(room_id)
# 构建消息对象
message = {
"id": message_data["msgid"],
"type": message_data["msgtype"],
"time": message_data["msgtime"],
"toType": to_type,
"from": from_user,
"to": receiver
}
# 根据消息类型处理内容
msg_type = message_data["msgtype"]
if msg_type == "mixed":
temp = []
idx = 0
for t in message_data["mixed"]["item"]:
idx, content = self._parse_mixed_message({"msgid": message_data["msgid"],
"msgtype": t["type"],
t["type"]: json.loads(t["content"])}, idx)
temp.append([t["type"], content])
message["content"] = json.dumps(temp, ensure_ascii=False)
else:
message["content"] = self._parse_message(message_data)
self.messages.append(message)
except Exception as e:
print(f"解析聊天记录失败: {e}")
raise e
def save_to_sqlite(self, sql: str, data: List[Tuple]) -> None:
"""
批量保存数据到SQLite数据库
Args:
sql: SQL语句
data: 数据列表
"""
if not data:
return
print(f"开始插入 {len(data)} 条记录")
self.dbm.execute_event(sql, data)
print(f"成功插入 {len(data)} 条记录")
def save_information_to_sqlite(self, date_time: str, message_num: int, file_num: int, user_num: int,
group_num: int) -> None:
"""
保存会话信息到SQLite数据库
Args:
date_time: 会话时间
message_num: 消息数量
file_num: 文件数量
user_num: 用户数量
group_num: 群组数量
"""
self.save_to_sqlite('''
INSERT OR IGNORE INTO "information"
(datetime, message, file, user, "group")
VALUES (?, ?, ?, ?, ?)
''', [(date_time, message_num, file_num, user_num, group_num)])
def save_all_to_sqlite(self) -> None:
"""
保存所有聊天数据到SQLite数据库
"""
# 保存文件信息
file_data = []
for file_info in self.files.values():
file_data.append((
file_info.get('id'),
file_info.get('type'),
file_info.get('content'),
file_info.get('original', ''),
file_info.get('extension', ''),
file_info.get('size', 0)
))
self.save_to_sqlite('''
INSERT OR IGNORE INTO file
(id, type, content, original, extension, size)
VALUES (?, ?, ?, ?, ?, ?)
''', file_data)
# 保存用户信息
user_data = [(user, user) for user in self.users]
self.save_to_sqlite('''
INSERT OR IGNORE INTO user
(id, name)
VALUES (?, ?)
''', user_data)
# 保存消息信息
message_data = []
for message in self.messages:
message_data.append((
message.get('id'),
message.get('type'),
message.get('time'),
message.get('toType'),
message.get('from'),
message.get('to'),
message.get('content')
))
self.save_to_sqlite('''
INSERT OR IGNORE INTO message
(id, type, time, toType, sender, receiver, content)
VALUES (?, ?, ?, ?, ?, ?, ?)
''', message_data)
# 保存群组信息
group_data = [(group, group) for group in self.groups]
self.save_to_sqlite('''
INSERT OR IGNORE INTO "group"
(id, name)
VALUES (?, ?)
''', group_data)
# 保存成员信息
member_data = []
for group_id, group_users in self.members.items():
for user in group_users:
member_data.append((group_id, user))
self.save_to_sqlite('''
INSERT OR IGNORE INTO "member"
(groupId, userId)
VALUES (?, ?)
''', member_data)
def process_chat_session(self, year: int, month: int, day: int, hour: int) -> bool:
"""
处理完整的聊天会话流程
Args:
year: 年份
month: 月份
day: 日期
hour: 小时
Returns:
处理是否成功
"""
self.users = []
self.groups = []
self.members = {}
self.messages = []
self.files = {}
try:
self.cache_path = f"zysl-log/wework/chat_msg/{year}/{month}/{day}/{hour}/"
# 解析聊天历史
self.parse_chat_history(year, month, day, hour)
# 保存到数据库
self.save_all_to_sqlite()
self.save_information_to_sqlite(f"{year}-{month:02d}-{day:02d}-{hour:02d}", len(self.messages),
len(self.files), len(self.users), len(self.groups))
print(f"成功处理 {year}-{month:02d}-{day:02d}-{hour:02d} 的聊天会话")
return True
except Exception as e:
print(f"处理聊天会话失败: {e}")
return False
def get_chat_history(self, start_time_str: str, end_time_str: str) -> list:
"""
获取指定时间的聊天历史记录
Args:
start_time_str: 开始时间,格式为 "YYYY-MM-DD HH:MM:SS"
end_time_str: 结束时间,格式为 "YYYY-MM-DD HH:MM:SS"
Returns:
聊天历史记录列表
"""
start_time = convert_to_utc(start_time_str, source_timezone="Asia/Shanghai").timestamp()
end_time = convert_to_utc(end_time_str, source_timezone="Asia/Shanghai").timestamp()
res = self.dbm.execute_query('''
SELECT message.time, message.content
FROM message
WHERE (message.type = 'text'
or message.type = 'mixed')
AND message.time >= ?
AND message.time <= ?
''', (str(start_time).replace(".0", "000"), str(end_time).replace(".0", "000")))
return res
@staticmethod
def analyze_qa_content(group_messages) -> dict:
group_qa = {}
for group_id in group_messages:
print(f"群ID: {group_id}")
print(f"消息数量: {len(group_messages[group_id])}")
content = ""
if len(group_messages[group_id]) <= 2:
continue
else:
message_mapping = {}
for msg in sorted(group_messages[group_id], key=lambda i: i[1]):
message_mapping[msg[4]] = msg[2]
content += json.dumps({
"user": msg[1],
"time": msg[2],
"content": msg[3],
"message_id": msg[4]
}, ensure_ascii=False) + ",\n"
template = f"""
以下是群中的消息记录
```json
[
{content}
]
```
""" + """
找到其中的问题和回答,如果没有回答就只整理问题,问题是与产品填报系统,上市公司,证券政策相关的问题
问题可能不只由一句话构成,回答也可能不只由一句话构成
多个问题需要分开放置
不需要任何解释说明只返回json格式
输出格式
```json
[
{
"question": ["message_id",...],
"answer": ["message_id",...]
}
]
```
"""
answer = agent.chat_deepseek(template)
answer = answer.replace("```json", "").replace("```", "") if "```json" in answer else answer
try:
answer = json.loads(answer)
except:
print(answer)
answer = []
for ans in answer:
time_list = [message_mapping.get(a, "") for a in ans["answer"]] + [message_mapping.get(q, "") for q
in ans["question"]]
start_time = min(time_list)
end_time = max(time_list)
ans["start_time"] = start_time
ans["end_time"] = end_time
group_qa[group_id] = answer
return group_qa
def save_qa_to_sqlite(self, group_qa: dict):
"""
将问答数据保存到 SQLite 数据库
Args:
group_qa: 问答数据字典,键为群ID,值为问答列表
"""
data = []
index_id = self.execute_query(f"SELECT id FROM qa order by id desc limit 1", fetch_all=False)[0]
for r in group_qa:
qas = group_qa.get(r, {})
for qa in qas:
Q = qa["question"]
A = qa["answer"]
index_id += 1
data.append((index_id, r, json.dumps(Q, ensure_ascii=False), json.dumps(A, ensure_ascii=False),
qa["start_time"],
qa["end_time"]))
self.execute_event('''
insert into qa (id, groupId, question, answer, start_time, end_time)
values (?, ?, ?, ?, ?, ?)
''', data)
def update_qa_data(self, start_time_str: str, end_time_str: str):
"""
更新指定时间范围内的问答内容
Args:
start_time_str: 开始时间,格式为 "YYYY-MM-DD HH:MM:SS"
end_time_str: 结束时间,格式为 "YYYY-MM-DD HH:MM:SS"
"""
start_time = convert_to_utc(start_time_str, source_timezone="Asia/Shanghai").timestamp()
end_time = convert_to_utc(end_time_str, source_timezone="Asia/Shanghai").timestamp()
data = self.execute_query('''
SELECT receiver, sender, time, content, id
FROM message
WHERE message.time >= ?
AND message.time <= ?
AND toType = 1
AND type in ('text', 'mixed', 'chatrecord')
''', (str(start_time).replace(".0", "000"), str(end_time).replace(".0", "000")))
group_messages = {}
for i in data:
group_id = i[0]
if group_id not in group_messages:
group_messages[group_id] = []
group_messages[group_id].append(i)
group_qa = self.analyze_qa_content(group_messages)
self.save_qa_to_sqlite(group_qa)