""" 要素提取器：使用NER和LLM从文档中提取实体 支持分章节提取和实体去重。 """ import json import asyncio from typing import Dict, List, Any, Optional from loguru import logger from .ner_service import ner_service class ElementExtractor: """ 要素提取器 使用NER服务识别文档中的实体，可选使用LLM进行智能提取。 不预定义要素结构，返回动态识别的实体。 """ def __init__(self): self._deepseek_service = None @property def deepseek_service(self): """延迟加载deepseek服务""" if self._deepseek_service is None: try: from .deepseek_service import deepseek_service self._deepseek_service = deepseek_service except ImportError: logger.warning("DeepSeek服务未配置，LLM提取将跳过") self._deepseek_service = None return self._deepseek_service async def extract_from_text( self, text: str, attachment_id: int = 0, use_llm: bool = True ) -> Dict[str, Any]: """ 从纯文本中提取实体（主接口） Args: text: 文档纯文本 attachment_id: 附件ID use_llm: 是否使用LLM提取 Returns: { "entities": [...], # NER识别的实体列表 "llm_extractions": [...], # LLM提取的内容（可选） "statistics": {...} } """ logger.info(f"开始提取实体: attachment_id={attachment_id}, " f"text_length={len(text)}, use_llm={use_llm}") # 1. 使用NER服务提取实体 ner_entities = await self._extract_by_ner(text) logger.info(f"NER提取完成: {len(ner_entities)} 个实体") # 2. LLM智能提取（可选） llm_extractions = [] if use_llm and self.deepseek_service: llm_extractions = await self._extract_by_llm(text) logger.info(f"LLM提取完成: {len(llm_extractions)} 个内容") return { "entities": ner_entities, "llm_extractions": llm_extractions, "statistics": { "ner_entity_count": len(ner_entities), "llm_extraction_count": len(llm_extractions), "text_length": len(text) } } async def _extract_by_ner(self, text: str) -> List[Dict]: """ 使用NER服务提取实体 返回实体列表，每个实体包含： - text: 实体文本 - type: 实体类型（DATE, ORG, PERSON, NUMBER, CODE等） - label: 实体标签 - confidence: 置信度 - position: 位置信息 """ try: # 调用现有的NER服务，返回EntityInfo对象列表 entities = await ner_service.extract_entities(text) # 格式化输出（EntityInfo是Pydantic模型，使用属性访问） result = [] for entity in entities: result.append({ "text": entity.name, "type": entity.type, "label": entity.type, "confidence": entity.confidence, "position": { "start": entity.position.char_start if entity.position else 0, "end": entity.position.char_end if entity.position else 0 } }) return result except Exception as e: logger.error(f"NER提取失败: {e}") return [] async def _extract_by_llm(self, text: str) -> List[Dict]: """ 使用LLM智能提取关键信息 让LLM自动识别文档中的重要信息，不预设要提取什么。 """ if not self.deepseek_service: return [] try: # 截取文档前部分进行分析 sample_text = text[:8000] if len(text) > 8000 else text prompt = f"""请分析以下文档，提取其中的关键信息。 要求： 1. 识别文档类型（如：报告、合同、通知等） 2. 提取关键实体（如：组织名称、日期、金额、编号等） 3. 提取关键数据（如：得分、级别、数量等） 4. 以JSON格式返回 返回格式： {{ "document_type": "文档类型", "key_entities": [ {{"name": "实体名称", "type": "实体类型", "value": "实体值"}} ], "key_data": [ {{"name": "数据名称", "value": "数据值", "unit": "单位"}} ], "summary": "文档摘要（50字以内）" }} 文档内容： {sample_text} 只返回JSON，不要其他内容。""" response = await self.deepseek_service.chat(prompt) if response: # 尝试解析JSON try: # 清理响应，提取JSON部分 json_str = response.strip() if json_str.startswith("```"): json_str = json_str.split("```")[1] if json_str.startswith("json"): json_str = json_str[4:] data = json.loads(json_str) extractions = [] # 文档类型 if data.get("document_type"): extractions.append({ "name": "文档类型", "value": data["document_type"], "source": "llm" }) # 关键实体 for entity in data.get("key_entities", []): extractions.append({ "name": entity.get("name", ""), "type": entity.get("type", ""), "value": entity.get("value", ""), "source": "llm" }) # 关键数据 for item in data.get("key_data", []): value = item.get("value", "") if item.get("unit"): value = f"{value}{item['unit']}" extractions.append({ "name": item.get("name", ""), "value": value, "source": "llm" }) # 摘要 if data.get("summary"): extractions.append({ "name": "文档摘要", "value": data["summary"], "source": "llm" }) return extractions except json.JSONDecodeError: logger.warning(f"LLM返回的不是有效JSON: {response[:200]}") return [] return [] except Exception as e: logger.error(f"LLM提取失败: {e}") return [] async def extract_from_chapters( self, chapters: List[Dict], attachment_id: int = 0, use_llm: bool = True, parallel: bool = True ) -> Dict[str, Any]: """ 分章节提取实体，最后去重合并 Args: chapters: 章节列表，每个章节包含 {chapter_id, title, text} attachment_id: 附件ID use_llm: 是否使用LLM提取 parallel: 是否并行处理章节 Returns: { "entities": [...], # 去重后的实体列表 "chapter_entities": {...}, # 按章节分组的实体 "llm_extractions": [...], "statistics": {...} } """ logger.info(f"开始分章节提取: {len(chapters)} 个章节, parallel={parallel}") chapter_results = {} all_entities = [] all_llm_extractions = [] if parallel and len(chapters) > 1: # 并行处理章节 tasks = [] for chapter in chapters: task = self._extract_chapter(chapter, attachment_id, use_llm) tasks.append(task) results = await asyncio.gather(*tasks, return_exceptions=True) for chapter, result in zip(chapters, results): if isinstance(result, Exception): logger.error(f"章节 {chapter['chapter_id']} 提取失败: {result}") continue chapter_results[chapter['chapter_id']] = result else: # 串行处理章节 for chapter in chapters: try: result = await self._extract_chapter(chapter, attachment_id, use_llm) chapter_results[chapter['chapter_id']] = result except Exception as e: logger.error(f"章节 {chapter['chapter_id']} 提取失败: {e}") # 合并所有章节的实体 for chapter_id, result in chapter_results.items(): for entity in result.get('entities', []): entity['chapter_id'] = chapter_id all_entities.append(entity) all_llm_extractions.extend(result.get('llm_extractions', [])) # 去重 unique_entities = self._deduplicate_entities(all_entities) unique_llm = self._deduplicate_llm_extractions(all_llm_extractions) logger.info(f"分章节提取完成: 原始 {len(all_entities)} 个实体, 去重后 {len(unique_entities)} 个") return { "entities": unique_entities, "chapter_entities": chapter_results, "llm_extractions": unique_llm, "statistics": { "chapter_count": len(chapters), "total_entities_before_dedup": len(all_entities), "unique_entity_count": len(unique_entities), "llm_extraction_count": len(unique_llm) } } async def _extract_chapter( self, chapter: Dict, attachment_id: int, use_llm: bool ) -> Dict[str, Any]: """提取单个章节的实体""" chapter_id = chapter.get('chapter_id', 'unknown') title = chapter.get('title', '') text = chapter.get('text', '') if not text or len(text.strip()) < 10: return {"entities": [], "llm_extractions": []} logger.debug(f"提取章节 {chapter_id}: {title[:30]}... (长度: {len(text)})") # NER提取 entities = await self._extract_by_ner(text) # 为每个实体添加章节信息 for entity in entities: entity['chapter_id'] = chapter_id entity['chapter_title'] = title # LLM提取（可选） llm_extractions = [] if use_llm and self.deepseek_service: llm_extractions = await self._extract_by_llm(text) for item in llm_extractions: item['chapter_id'] = chapter_id item['chapter_title'] = title return { "entities": entities, "llm_extractions": llm_extractions } def _deduplicate_entities(self, entities: List[Dict]) -> List[Dict]: """ 实体去重 去重规则： 1. 相同类型+相同文本 -> 保留第一个出现的 2. 包含关系 -> 保留更长的实体 """ if not entities: return [] # 按 (type, text) 去重 seen = {} for entity in entities: key = (entity.get('type', ''), entity.get('text', '')) if key not in seen: seen[key] = entity else: # 保留置信度更高的 if entity.get('confidence', 0) > seen[key].get('confidence', 0): seen[key] = entity unique = list(seen.values()) # 处理包含关系（可选，较复杂） # 例如："中国电建集团" 和 "中国电建集团成都勘测设计研究院有限公司" # 保留更长的 final = [] texts = set() # 按文本长度降序排序 unique.sort(key=lambda x: len(x.get('text', '')), reverse=True) for entity in unique: text = entity.get('text', '') # 检查是否被更长的实体包含 is_substring = False for existing_text in texts: if text in existing_text and text != existing_text: is_substring = True break if not is_substring: final.append(entity) texts.add(text) # 恢复原始顺序（按位置） final.sort(key=lambda x: x.get('position', {}).get('start', 0)) return final def _deduplicate_llm_extractions(self, extractions: List[Dict]) -> List[Dict]: """LLM提取结果去重""" if not extractions: return [] seen = {} for item in extractions: key = (item.get('name', ''), item.get('value', '')) if key not in seen: seen[key] = item return list(seen.values()) # 创建单例 element_extractor = ElementExtractor()