import json from datetime import datetime from typing import Dict, Any import os from dotenv import load_dotenv from minio import Minio from parser.simple_object_parser import SimpleObjectParser from embedding import EmbeddingTool from milvus_api import MilvusAPI, MilvusDocument from logger.logging import setup_logging load_dotenv() logger = setup_logging() # Minio 配置信息。 MINIO_BUCKET = os.getenv('MINIO_BUCKET') MINIO_ENDPOINT = os.getenv('MINIO_ENDPOINT') MINIO_ACCESS_KEY = os.getenv('MINIO_ACCESS_KEY') MINIO_SECRET_KEY = os.getenv('MINIO_SECRET_KEY') MINIO_REGION = os.getenv('MINIO_REGION') MINIO_USE_VIRTUAL_HOST = os.getenv('MINIO_USE_VIRTUAL_HOST', 'false').lower() == 'true' MINIO_USE_SSL = os.getenv('MINIO_USE_SSL', 'false').lower() == 'true' ## 获取MinIO客户端 def _get_minio_client() -> Minio: client = Minio( endpoint=MINIO_ENDPOINT, access_key=MINIO_ACCESS_KEY, secret_key=MINIO_SECRET_KEY, secure=MINIO_USE_SSL, region=MINIO_REGION, ) if MINIO_USE_VIRTUAL_HOST: client.enable_virtual_style_endpoint() return client class MinioEventHandler: """Minio事件监听器""" def __init__(self, bucket_name: str = None): """ 初始化监听器 Args: bucket_name (str): 要监听的桶名称，默认从环境变量获取 """ self.bucket_name = bucket_name or os.getenv("MINIO_BUCKET", "mildoc") # 初始化各个组件 self.minio_client = _get_minio_client() # 初始化解析器 logger.info("初始化解析器...") self.parser: SimpleObjectParser = SimpleObjectParser(minio_client=self.minio_client) # 初始化Milvus logger.info("初始化Milvus...") self.milvus_api: MilvusAPI = MilvusAPI() # 测试embedding工具 logger.info("测试embedding工具...") self.embedding_tool: EmbeddingTool = EmbeddingTool() logger.info("所有组件初始化完成！") def _extract_event_info(self, event_data: Dict[str, Any]) -> Dict[str, Any]: """ 从事件数据中提取关键信息 Args: event_data (Dict[str, Any]): 事件数据 Returns: Dict[str, Any]: 提取的信息 """ try: #logger.info(f"event_data: {event_data}") logger.info(f"数据: {json.dumps(event_data, ensure_ascii=False, indent=2)}") record = event_data.get('Records', [{}])[0] s3_info = record.get('s3', {}) return { 'event_name': record.get('eventName', ''), 'event_time': record.get('eventTime', ''), 'bucket_name': s3_info.get('bucket', {}).get('name', ''), 'object_name': s3_info.get('object', {}).get('key', ''), 'object_size': s3_info.get('object', {}).get('size', 0), 'content_type': s3_info.get('object', {}).get('contentType', ''), 'etag': s3_info.get('object', {}).get('eTag', ''), } except Exception as e: logger.error(f"提取事件信息失败: {e}") return {} def _handle_object_created(self, event_info: Dict[str, Any]): """ 处理对象创建事件 Args: event_info (Dict[str, Any]): 事件信息 """ try: bucket_name = event_info['bucket_name'] object_name = event_info['object_name'] logger.info(f"\n=== 处理新增对象: {bucket_name}/{object_name} ===") logger.info(f"对象大小: {event_info['object_size']} 字节") logger.info(f"内容类型: {event_info['content_type']}") # 直接调用_process_single_object方法处理 self._process_single_object(bucket_name, object_name, force_update=True) except Exception as e: logger.error(f"处理对象创建事件失败: {e}") def _handle_object_deleted(self, event_info: Dict[str, Any]): """ 处理对象删除事件 Args: event_info (Dict[str, Any]): 事件信息 """ try: bucket_name = event_info['bucket_name'] object_name = event_info['object_name'] doc_path_name = object_name # 不再包含bucket_name前缀 logger.info(f"\n=== 处理删除对象: {bucket_name}/{object_name} ===") # 从Milvus中删除相关记录 logger.info("从Milvus中查找并删除相关记录...") # 使用MilvusAPI的删除方法 if self.milvus_api.delete_existing_document(doc_path_name): logger.info(f"成功删除文档记录: {doc_path_name}") else: logger.error(f"删除文档记录失败: {doc_path_name}") except Exception as e: logger.error(f"处理对象删除事件失败: {e}") def _process_event(self, event_data: Dict[str, Any]): """ 处理单个事件 Args: event_data (Dict[str, Any]): 事件数据 """ try: # 提取事件信息 event_info = self._extract_event_info(event_data) if not event_info: logger.error("无法提取事件信息，跳过处理") return event_name = event_info['event_name'] timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") logger.info(f"\n[{timestamp}] 收到事件: {event_name}") logger.info(f"对象: {event_info['bucket_name']}/{event_info['object_name']}") # 根据事件类型进行处理 if 'ObjectCreated' in event_name: self._handle_object_created(event_info) elif 'ObjectRemoved' in event_name: self._handle_object_deleted(event_info) else: logger.error(f"未处理的事件类型: {event_name}") except Exception as e: logger.error(f"处理事件时出错: {e}") def _process_single_object(self, bucket_name: str, object_name: str, force_update: bool = False): """ 处理单个对象（用于全量刷新和排查补漏） Args: bucket_name (str): 桶名称 object_name (str): 对象名称 force_update (bool): 是否强制更新（True=全量刷新，False=排查补漏） Returns: bool: 处理是否成功 """ try: doc_path_name = object_name # 不再包含bucket_name前缀 # 如果是排查补漏模式，先检查是否已存在 if not force_update: if self.milvus_api.check_document_exists(doc_path_name): logger.info(f" 文档已存在，跳过: {object_name}") return True logger.info(f" 处理文档: {object_name}") # 解析对象内容 parse_result = self.parser.parse_object(bucket_name, object_name) if 'error' in parse_result: logger.error(f" 解析失败: {parse_result['error']}") return False if not parse_result['contents']: logger.error(f" 未提取到文本内容，跳过") return True logger.info(f" 解析成功，获得 {len(parse_result['contents'])} 个文本片段") # 如果是强制更新，先删除已存在的记录 if force_update: self.milvus_api.delete_existing_document(doc_path_name) # 为每个文本片段生成embedding并存储到Milvus success_count = 0 for i, content in enumerate(parse_result['contents']): try: # 生成embedding向量 embedding_vector = self.embedding_tool.get_embedding(content) if not embedding_vector: logger.error(f" 片段 {i+1} embedding生成失败，跳过") continue # 准备文档数据 doc_data = MilvusDocument( doc_name=parse_result['doc_name'], doc_path_name=parse_result['doc_path_name'], doc_type=parse_result['doc_type'], doc_md5=parse_result['doc_md5'], doc_length=parse_result['doc_length'], content=content, content_vector=embedding_vector, embedding_model=self.embedding_tool.model ) # 存储到Milvus（允许重复，因为我们已经处理了去重逻辑） if self.milvus_api.insert_document(doc_data): success_count += 1 else: logger.error(f" 片段 {i+1} 存储失败") except Exception as e: logger.error(f" 处理片段 {i+1} 时出错: {e}") continue logger.info(f" 完成！成功存储 {success_count}/{len(parse_result['contents'])} 个片段") return success_count > 0 except Exception as e: logger.error(f" 处理对象失败: {e}") return False def full_update(self): """ 模式1：全量刷新 - 遍历Minio桶中的所有数据并更新到Milvus """ logger.info(f"\n=== 模式1：全量刷新 ===") logger.info(f"正在遍历桶 '{self.bucket_name}' 中的所有对象...") try: # 获取桶中的所有对象 objects = self.minio_client.list_objects(self.bucket_name, recursive=True) total_objects = 0 processed_objects = 0 for obj in objects: object_name = obj.object_name # 跳过文件夹 if object_name.endswith('/'): continue total_objects += 1 logger.info(f"\n[{total_objects}] 处理对象: {object_name}") if self._process_single_object(self.bucket_name, object_name, force_update=True): processed_objects += 1 self.milvus_api.flush_collection() logger.info(f"\n=== 全量刷新完成 ===") logger.info(f"总对象数: {total_objects}") logger.info(f"成功处理: {processed_objects}") logger.info(f"失败数量: {total_objects - processed_objects}") except Exception as e: logger.error(f"全量刷新失败: {e}") def backfill_update(self): """ 模式2：排查补漏 - 检查Milvus中不存在的文档并新增 """ logger.info(f"\n=== 模式2：排查补漏 ===") logger.info(f"正在检查桶 '{self.bucket_name}' 中缺失的文档...") try: # 获取桶中的所有对象 objects = self.minio_client.list_objects(self.bucket_name, recursive=True) total_objects = 0 new_objects = 0 existing_objects = 0 for obj in objects: object_name = obj.object_name # 跳过文件夹 if object_name.endswith('/'): continue total_objects += 1 logger.info(f"\n[{total_objects}] 检查对象: {object_name}") # 检查是否已存在 if self.milvus_api.check_document_exists(object_name): logger.info(f" 已存在，跳过") existing_objects += 1 else: logger.info(f" 不存在，开始处理...") if self._process_single_object(self.bucket_name, object_name, force_update=False): new_objects += 1 self.milvus_api.flush_collection() logger.info(f"\n=== 排查补漏完成 ===") logger.info(f"总对象数: {total_objects}") logger.info(f"已存在: {existing_objects}") logger.info(f"新增: {new_objects}") logger.info(f"失败数量: {total_objects - existing_objects - new_objects}") except Exception as e: logger.error(f"排查补漏失败: {e}") def start_listening(self): """ 模式3：增量更新 - 根据消息通知进行增量更新 """ logger.info(f"\n=== 模式3：增量更新 ===") logger.info(f"开始监听桶 '{self.bucket_name}' 的事件...") logger.info("按 Ctrl+C 停止监听") try: # 监听桶事件 events = self.minio_client.listen_bucket_notification( bucket_name=self.bucket_name, events=['s3:ObjectCreated:*', 's3:ObjectRemoved:*'] ) for event in events: try: if event: # 解析事件数据 if isinstance(event, bytes): event_data = json.loads(event.decode('utf-8')) elif isinstance(event, str): event_data = json.loads(event) elif isinstance(event, dict): event_data = event else: logger.error(f"未知的事件数据类型: {type(event)}") continue # 处理事件 self._process_event(event_data) except json.JSONDecodeError as e: logger.error(f"解析事件数据失败: {e}") except Exception as e: logger.error(f"处理事件失败: {e}") except KeyboardInterrupt: logger.info("\n监听已停止") except Exception as e: logger.error(f"监听过程中出错: {e}")