entity_extractor.py

"""
实体和关系提取模块
使用LLM从文档中提取实体、关系和属性，构建知识图谱的基础
"""

from typing import List, Dict, Tuple
import time
import asyncio
import aiohttp
import json
try:
    from langchain_core.prompts import PromptTemplate
except ImportError:
    try:
        from langchain_core.prompts import PromptTemplate
    except ImportError:
        from langchain.prompts import PromptTemplate

from langchain_community.chat_models import ChatOllama
from langchain_core.output_parsers import JsonOutputParser
from config import LOCAL_LLM


class EntityExtractor:
    """实体提取器 - 使用LLM从文本中提取实体（支持异步批处理）"""
    
    def __init__(self, timeout: int = 180, max_retries: int = 3, enable_async: bool = True):
        """初始化实体提取器
        
        Args:
            timeout: LLM调用超时时间（秒）- 默认180秒以应对首次模型加载
            max_retries: 失败重试次数
            enable_async: 是否启用异步处理（默认启用）
        """
        self.llm = ChatOllama(
            model=LOCAL_LLM, 
            format="json", 
            temperature=0,
            timeout=timeout  # 添加超时设置
        )
        self.max_retries = max_retries
        self.enable_async = enable_async
        self.ollama_url = "http://localhost:11434/api/generate"
        self.timeout = timeout  # 保存超时设置供异步使用
        
        # 实体提取提示模板
        self.entity_prompt = PromptTemplate(
            template="""你是一个专业的实体识别专家。从以下文本中提取所有重要的实体。
            
实体类型包括:
- PERSON: 人物、作者、研究者
- ORGANIZATION: 组织、机构、公司
- CONCEPT: 技术概念、算法、方法论
- TECHNOLOGY: 具体技术、工具、框架
- PAPER: 论文、出版物
- EVENT: 事件、会议

文本内容:
{text}

请以JSON格式返回，包含以下字段:
{{
    "entities": [
        {{
            "name": "实体名称",
            "type": "实体类型",
            "description": "简短描述"
        }}
    ]
}}

不要包含前言或解释，只返回JSON。
""",
            input_variables=["text"]
        )
        
        # 关系提取提示模板
        self.relation_prompt = PromptTemplate(
            template="""你是一个关系抽取专家。从文本中识别实体之间的关系。

已识别的实体:
{entities}

文本内容:
{text}

请识别实体之间的关系，以JSON格式返回:
{{
    "relations": [
        {{
            "source": "源实体名称",
            "target": "目标实体名称",
            "relation_type": "关系类型",
            "description": "关系描述"
        }}
    ]
}}

关系类型包括: AUTHOR_OF, USES, BASED_ON, RELATED_TO, PART_OF, APPLIES_TO, IMPROVES, CITES

不要包含前言或解释，只返回JSON。
""",
            input_variables=["text", "entities"]
        )
        
        self.entity_chain = self.entity_prompt | self.llm | JsonOutputParser()
        self.relation_chain = self.relation_prompt | self.llm | JsonOutputParser()
    
    def extract_entities(self, text: str) -> List[Dict]:
        """
        从文本中提取实体（带重试机制）
        
        Args:
            text: 输入文本
            
        Returns:
            实体列表
        """
        for attempt in range(self.max_retries):
            try:
                print(f"   🔄 提取实体 (尝试 {attempt + 1}/{self.max_retries})...", end="")
                result = self.entity_chain.invoke({"text": text[:2000]})  # 限制长度
                entities = result.get("entities", [])
                print(f" ✅ 提取到 {len(entities)} 个实体")
                return entities
            except TimeoutError as e:
                print(f" ⏱️ 超时")
                if attempt < self.max_retries - 1:
                    wait_time = (attempt + 1) * 2
                    print(f"   ⏳ 等待 {wait_time} 秒后重试...")
                    time.sleep(wait_time)
                else:
                    print(f"   ❌ 实体提取最终失败: 超时")
                    return []
            except Exception as e:
                print(f" ❌ 错误: {str(e)[:100]}")
                if attempt < self.max_retries - 1:
                    time.sleep(1)
                else:
                    print(f"   ❌ 实体提取最终失败: {e}")
                    return []
        return []
    
    def extract_relations(self, text: str, entities: List[Dict]) -> List[Dict]:
        """
        从文本中提取实体关系（带重试机制）
        
        Args:
            text: 输入文本
            entities: 已识别的实体列表
            
        Returns:
            关系列表
        """
        if not entities:
            print("   ⚠️ 无实体，跳过关系提取")
            return []
        
        for attempt in range(self.max_retries):
            try:
                print(f"   🔄 提取关系 (尝试 {attempt + 1}/{self.max_retries})...", end="")
                entity_names = [e["name"] for e in entities]
                result = self.relation_chain.invoke({
                    "text": text[:2000],
                    "entities": ", ".join(entity_names)
                })
                relations = result.get("relations", [])
                print(f" ✅ 提取到 {len(relations)} 个关系")
                return relations
            except TimeoutError as e:
                print(f" ⏱️ 超时")
                if attempt < self.max_retries - 1:
                    wait_time = (attempt + 1) * 2
                    print(f"   ⏳ 等待 {wait_time} 秒后重试...")
                    time.sleep(wait_time)
                else:
                    print(f"   ❌ 关系提取最终失败: 超时")
                    return []
            except Exception as e:
                print(f" ❌ 错误: {str(e)[:100]}")
                if attempt < self.max_retries - 1:
                    time.sleep(1)
                else:
                    print(f"   ❌ 关系提取最终失败: {e}")
                    return []
        return []
    
    async def _async_llm_call(self, prompt: str, session: aiohttp.ClientSession, attempt: int = 0) -> Dict:
        """异步调用 Ollama API"""
        try:
            timeout = aiohttp.ClientTimeout(
                total=self.timeout,      # 总超时
                connect=30,               # 连接超时 30 秒
                sock_read=self.timeout    # 读取超时
            )
            
            async with session.post(
                self.ollama_url,
                json={
                    "model": LOCAL_LLM,
                    "prompt": prompt,
                    "format": "json",
                    "stream": False,
                    "options": {"temperature": 0}
                },
                timeout=timeout
            ) as response:
                if response.status == 200:
                    result = await response.json()
                    return json.loads(result.get('response', '{}'))
                else:
                    raise Exception(f"API返回错误: {response.status}")
        except (asyncio.TimeoutError, aiohttp.ClientError) as e:
            if attempt < self.max_retries - 1:
                wait_time = (attempt + 1) * 3
                await asyncio.sleep(wait_time)
                return await self._async_llm_call(prompt, session, attempt + 1)
            raise Exception(f"连接失败: {str(e)}")
        except Exception as e:
            if attempt < self.max_retries - 1:
                await asyncio.sleep(2)
                return await self._async_llm_call(prompt, session, attempt + 1)
            raise
    
    async def _extract_entities_async(self, text: str, doc_index: int, session: aiohttp.ClientSession) -> List[Dict]:
        """异步提取实体"""
        prompt = self.entity_prompt.format(text=text[:2000])
        
        for attempt in range(self.max_retries):
            try:
                print(f"   [文档 #{doc_index + 1}] 🔄 提取实体 (尝试 {attempt + 1}/{self.max_retries})...", end="")
                result = await self._async_llm_call(prompt, session, attempt)
                entities = result.get("entities", [])
                print(f" ✅ {len(entities)} 个实体")
                return entities
            except Exception as e:
                print(f" ❌ {str(e)[:50]}")
                if attempt == self.max_retries - 1:
                    return []
        return []
    
    async def _extract_relations_async(self, text: str, entities: List[Dict], doc_index: int, session: aiohttp.ClientSession) -> List[Dict]:
        """异步提取关系"""
        if not entities:
            return []
        
        entity_names = [e["name"] for e in entities]
        prompt = self.relation_prompt.format(
            text=text[:2000],
            entities=", ".join(entity_names)
        )
        
        for attempt in range(self.max_retries):
            try:
                print(f"   [文档 #{doc_index + 1}] 🔄 提取关系 (尝试 {attempt + 1}/{self.max_retries})...", end="")
                result = await self._async_llm_call(prompt, session, attempt)
                relations = result.get("relations", [])
                print(f" ✅ {len(relations)} 个关系")
                return relations
            except Exception as e:
                print(f" ❌ {str(e)[:50]}")
                if attempt == self.max_retries - 1:
                    return []
        return []
    
    async def _extract_from_document_async(self, document_text: str, doc_index: int, session: aiohttp.ClientSession) -> Dict:
        """异步处理单个文档"""
        print(f"\n🔍 [文档 #{doc_index + 1}] 开始异步提取...")
        
        # 并发提取实体和关系（先实体，再关系）
        entities = await self._extract_entities_async(document_text, doc_index, session)
        relations = await self._extract_relations_async(document_text, entities, doc_index, session)
        
        print(f"📊 [文档 #{doc_index + 1}] 完成: {len(entities)} 实体, {len(relations)} 关系")
        
        return {
            "entities": entities,
            "relations": relations
        }
    
    async def extract_batch_async(self, documents: List[Tuple[str, int]]) -> List[Dict]:
        """异步批量处理多个文档
        
        Args:
            documents: 文档列表，每个元素为 (document_text, doc_index) 元组
            
        Returns:
            提取结果列表
        """
        async with aiohttp.ClientSession() as session:
            tasks = [
                self._extract_from_document_async(doc_text, doc_idx, session)
                for doc_text, doc_idx in documents
            ]
            
            # 并发执行所有任务
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # 处理异常结果
            processed_results = []
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    print(f"⚠️ 文档 #{documents[i][1] + 1} 处理失败: {result}")
                    processed_results.append({"entities": [], "relations": []})
                else:
                    processed_results.append(result)
            
            return processed_results
    
    def extract_from_document(self, document_text: str, doc_index: int = 0) -> Dict:
        """
        从单个文档中提取实体和关系（同步接口，保持向后兼容）
        
        Args:
            document_text: 文档文本
            doc_index: 文档索引（用于日志）
            
        Returns:
            包含实体和关系的字典
        """
        # 同步方式调用（保持向后兼容）
        print(f"\n🔍 文档 #{doc_index + 1}: 开始提取...")
        
        entities = self.extract_entities(document_text)
        relations = self.extract_relations(document_text, entities)
        
        print(f"📊 文档 #{doc_index + 1} 完成: {len(entities)} 实体, {len(relations)} 关系")
        
        return {
            "entities": entities,
            "relations": relations
        }


class EntityDeduplicator:
    """实体去重和合并"""
    
    def __init__(self):
        self.llm = ChatOllama(model=LOCAL_LLM, format="json", temperature=0)
        
        self.merge_prompt = PromptTemplate(
            template="""判断以下两个实体是否指向同一个对象:

实体1: {entity1_name} - {entity1_desc}
实体2: {entity2_name} - {entity2_desc}

如果是同一个对象，返回:
{{
    "is_same": true,
    "canonical_name": "标准名称",
    "reason": "原因"
}}

如果不是，返回:
{{
    "is_same": false,
    "reason": "原因"
}}

只返回JSON，不要其他内容。
""",
            input_variables=["entity1_name", "entity1_desc", "entity2_name", "entity2_desc"]
        )
        
        self.merge_chain = self.merge_prompt | self.llm | JsonOutputParser()
    
    def _is_same_entity(self, entity1: Dict, entity2: Dict) -> bool:
        """
        使用LLM判断两个实体是否指向同一个对象
        
        Args:
            entity1: 实体1字典
            entity2: 实体2字典
            
        Returns:
            bool: 是否相同
        """
        try:
            # 准备输入
            input_data = {
                "entity1_name": entity1["name"],
                "entity1_desc": entity1.get("description", "无描述"),
                "entity2_name": entity2["name"],
                "entity2_desc": entity2.get("description", "无描述")
            }
            
            # 调用LLM
            result = self.merge_chain.invoke(input_data)
            
            # 解析结果
            return result.get("is_same", False)
        except Exception as e:
            print(f"   ⚠️ LLM判重失败 ({entity1['name']} vs {entity2['name']}): {e}")
            return False

    def deduplicate_entities(self, entities: List[Dict]) -> Dict:
        """
        去重实体列表
        
        Args:
            entities: 实体列表
            
        Returns:
            包含entities和mapping的字典
        """
        if len(entities) <= 1:
            # 返回字典格式，保持一致性
            entity_mapping = {entity["name"]: entity["name"] for entity in entities} if entities else {}
            return {
                "entities": entities,
                "mapping": entity_mapping
            }
        
        print(f"🔄 开始去重 {len(entities)} 个实体...")
        
        # 基于名称和LLM的智能去重
        unique_entities = {}
        entity_mapping = {}  # 映射别名到标准名称
        
        for entity in entities:
            name = entity["name"].lower().strip()
            
            # 查找是否有相似实体
            merged = False
            for canonical_name, canonical_entity in unique_entities.items():
                # 1. 简单的字符串匹配（作为预筛选）
                # 如果名称完全相同，或者是子串关系，则考虑合并
                is_substring = name in canonical_name or canonical_name in name
                
                if name == canonical_name:
                    # 完全匹配（忽略大小写），直接合并
                    entity_mapping[entity["name"]] = canonical_entity["name"]
                    merged = True
                    break
                elif is_substring:
                    # 子串匹配，使用LLM进行智能确认
                    # 例如："Python" 和 "Python Programming Language" -> 合并
                    # "Java" 和 "JavaScript" -> 不合并
                    if self._is_same_entity(entity, canonical_entity):
                        print(f"   ✨ 合并: {entity['name']} -> {canonical_entity['name']}")
                        entity_mapping[entity["name"]] = canonical_entity["name"]
                        merged = True
                        break
            
            if not merged:
                unique_entities[name] = entity
                entity_mapping[entity["name"]] = entity["name"]
                
        print(f"✅ 去重完成，剩余 {len(unique_entities)} 个唯一实体")
        
        return {
            "entities": list(unique_entities.values()),
            "mapping": entity_mapping
        }


def initialize_entity_extractor():
    """初始化实体提取器"""
    return EntityExtractor()