#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 内容结构化处理模块 主要功能： 1. 从数据库中拉取需要结构化的数据 2. 调用Gemini API进行内容结构化 3. 将结构化结果更新到数据库 """ import os import json import time import sys import threading from typing import Dict, Any, List, Optional, Tuple # 导入自定义模块 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from utils.mysql_db import MysqlHelper from gemini import GeminiProcessor from utils.file import File from utils.logging_config import get_logger class StructureProcessor: def __init__(self): # 设置日志 self.logger = get_logger('StructureProcessor') # 初始化处理器 self.processor = GeminiProcessor() self.system_prompt = File.read_file('../prompt/structure.md') self.logger.info("系统提示词加载完成") self.logger.debug(f"系统提示词: {self.system_prompt}") # 线程控制 self.lock = threading.Lock() self.stop_event = threading.Event() self.threads = [] def build_query_conditions(self, query_word: Optional[str], source_type: Optional[str], source_channel: Optional[str]) -> Tuple[str, Tuple]: """构建查询条件和参数""" conditions = ["multimodal_recognition is not null", "structured_data is null"] params = [] if query_word is not None: conditions.append("query_word = %s") params.append(query_word) if source_type is not None: conditions.append("source_type = %s") params.append(source_type) if source_channel is not None: conditions.append("source_channel = %s") params.append(source_channel) where_clause = " AND ".join(conditions) return where_clause, tuple(params) def process_single_record(self, query_word: Optional[str], source_type: Optional[str], source_channel: Optional[str]) -> bool: """处理单条记录""" try: with self.lock: # 构建查询条件和参数 where_clause, params = self.build_query_conditions(query_word, source_type, source_channel) # 先查询一条需要处理的记录 select_sql = f""" SELECT id, multimodal_recognition FROM knowledge_search_content WHERE {where_clause} LIMIT 1 """ records = MysqlHelper.get_values(select_sql, params) if not records: self.logger.warning("没有找到需要处理的记录") return False row = records[0] record_id = row[0] # 标记为处理中，防止其他线程取到重复处理 mark_sql = """ UPDATE knowledge_search_content SET structured_data = '{}' WHERE id = %s """ MysqlHelper.update_values(mark_sql, (record_id,)) self.logger.info(f"开始处理记录 ID: {record_id}") # 处理内容 result = self.processor.process(row[1], self.system_prompt) self.logger.info(f"处理完成，结果长度: {len(str(result))}") self.logger.debug(f"处理结果: {result}") # 更新数据库为实际结果 update_sql = """ UPDATE knowledge_search_content SET structured_data = %s WHERE id = %s """ MysqlHelper.update_values(update_sql, (result, record_id)) self.logger.info(f"记录 {record_id} 处理完成并更新数据库") return True except Exception as e: self.logger.error(f"处理记录失败: {str(e)}", exc_info=True) return False def worker_thread(self, thread_id: int, query_word: Optional[str], source_type: Optional[str], source_channel: Optional[str]): """工作线程函数""" thread_logger = get_logger(f'WorkerThread-{thread_id}') thread_logger.info(f"线程 {thread_id} 启动") while not self.stop_event.is_set(): try: # 尝试处理一条记录 success = self.process_single_record(query_word, source_type, source_channel) if not success: thread_logger.info(f"没有找到需要处理的记录，等待5秒后重试") # 等待时也要检查停止信号 if self.stop_event.wait(5): break continue # 处理成功后等待5秒再处理下一条 thread_logger.info(f"处理完成，等待5秒后处理下一条") # 等待时也要检查停止信号 if self.stop_event.wait(5): break except Exception as e: thread_logger.error(f"发生错误: {str(e)}", exc_info=True) # 等待时也要检查停止信号 if self.stop_event.wait(5): break thread_logger.info(f"线程 {thread_id} 已停止") def start_multi_thread_processing(self, query_word: Optional[str], source_type: Optional[str], source_channel: Optional[str]): """启动多线程处理""" self.threads = [] self.logger.info("启动多线程处理...") self.logger.info(f"查询条件: query_word={query_word}, source_type={source_type}, source_channel={source_channel}") # 创建5个线程，间隔5秒启动 for i in range(5): thread = threading.Thread( target=self.worker_thread, args=(i + 1, query_word, source_type, source_channel) ) self.threads.append(thread) # 启动线程 thread.start() self.logger.info(f"线程 {i + 1} 已启动") # 等待5秒后启动下一个线程 if i < 4: # 最后一个线程不需要等待 self.logger.info("等待5秒后启动下一个线程...") time.sleep(5) self.logger.info("所有线程已启动，使用 ./start_structure.sh stop 停止") try: # 等待所有线程完成 for thread in self.threads: thread.join() except KeyboardInterrupt: self.logger.info("收到停止信号，正在停止所有线程...") self.stop_all_threads() def stop_all_threads(self): """停止所有线程""" self.logger.info("正在停止所有线程...") self.stop_event.set() # 等待所有线程结束 for i, thread in enumerate(self.threads): if thread.is_alive(): self.logger.info(f"等待线程 {i + 1} 结束...") thread.join(timeout=10) # 最多等待10秒 if thread.is_alive(): self.logger.warning(f"线程 {i + 1} 未能正常结束") else: self.logger.info(f"线程 {i + 1} 已正常结束") self.logger.info("所有线程已停止") def main(): """主函数""" import argparse parser = argparse.ArgumentParser(description='内容结构化处理脚本') parser.add_argument('--query_word', default=None, help='query词') parser.add_argument('--source_type', default=None, help='数据源类型') parser.add_argument('--source_channel', default=None, help='数据源渠道') args = parser.parse_args() try: processor = StructureProcessor() processor.start_multi_thread_processing( query_word=args.query_word, source_type=args.source_type, source_channel=args.source_channel ) except Exception as e: print(f"程序执行失败: {str(e)}") sys.exit(1) if __name__ == "__main__": # 测试单条记录处理 processor = StructureProcessor() processor.process_single_record(query_word=None, source_type=None, source_channel=None)