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)