Initial: integrated 2025 LawGraph (graphrag_pipeline) + 2026 kg_project

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-17 10:13:39 +08:00
commit 6c1a69af0d
83 changed files with 14295 additions and 0 deletions
@@ -0,0 +1,236 @@
+"""
+关系抽取模块
+使用多个LLM（GPT-4o, Doubao-pro, GLM-4）并行进行关系抽取
+"""
+
+import json
+import logging
+from typing import List, Dict, Any, Optional
+
+from ..prompts.re_prompts import REPromptBuilder
+from ..utils.llm_client import LLMClient, LLMProvider
+from ..utils.config import Config
+
+logger = logging.getLogger(__name__)
+
+
+class REExtractor:
+    """关系抽取器"""
+
+    def __init__(
+        self,
+        models: Optional[List[str]] = None,
+        config: Optional[Config] = None,
+    ):
+        """
+        初始化RE提取器
+
+        Args:
+            models: 使用的模型列表，默认使用多个模型并行
+            config: 配置对象
+        """
+        self.config = config or Config()
+        
+        # 默认使用多个模型（优先使用硅基流动）
+        if models is None:
+            models = []
+            # 优先使用硅基流动（使用两个模型）
+            if self.config.SILICONFLOW_API_KEY:
+                models.append(("siliconflow", "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"))
+                models.append(("siliconflow", "Qwen/Qwen2.5-7B-Instruct"))
+            # 如果有其他API密钥，也添加
+            if self.config.OPENAI_API_KEY:
+                models.append(("openai", "gpt-4o"))
+            if self.config.VOLCENGINE_ACCESS_KEY and self.config.VOLCENGINE_SECRET_KEY:
+                models.append(("doubao", "doubao-pro-32k"))
+            if self.config.ZHIPUAI_API_KEY:
+                models.append(("glm", "glm-4-airx"))
+            
+            # 如果都没有配置，至少使用硅基流动（即使没有密钥也会报错）
+            if not models:
+                models = [
+                    ("siliconflow", "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"),
+                    ("siliconflow", "Qwen/Qwen2.5-7B-Instruct")
+                ]
+        
+        self.clients = []
+        for provider_str, model in models:
+            try:
+                provider = LLMProvider(provider_str)
+                self.clients.append(LLMClient(provider, model, self.config))
+            except Exception as e:
+                logger.warning(f"初始化{provider_str}客户端失败: {e}，跳过该模型")
+        
+        self.prompt_builder = REPromptBuilder()
+
+    def extract(
+        self,
+        sentence: str,
+        entities: List[Dict[str, str]],
+        syntax_info: Dict = None,
+        context: str = None,
+        use_verification: bool = True,
+    ) -> List[Dict[str, Any]]:
+        """
+        从句子中提取关系（多模型并行，支持二次对话验证）
+
+        Args:
+            sentence: 输入句子
+            entities: 已识别的实体列表，格式：[{"text": "实体文本", "type": "实体类型"}, ...]
+            syntax_info: 句法分析结果
+            context: 上下文段落
+            use_verification: 是否使用二次对话验证
+
+        Returns:
+            结果列表，每个包含模型名称和三元组列表
+        """
+        # 格式化实体列表（用于Prompt）
+        formatted_entities = self._format_entities_for_prompt(entities)
+        
+        # 格式化句法信息
+        formatted_syntax = self._format_syntax_info(syntax_info) if syntax_info else "无句法信息"
+        
+        # 构建Prompt模板
+        template = self.prompt_builder.build_template(
+            entities=entities,
+            syntax_info=formatted_syntax,
+            context=context or ""
+        )
+        
+        # 构建第一阶段消息
+        system_msg = template.task_description["system"]
+        user_msg = template.task_description["user"]
+        
+        # 构建完整的用户消息
+        input_data = f"""已识别实体：
+{formatted_entities}
+
+句子文本：
+{sentence}
+
+句法分析结果：
+{formatted_syntax}
+
+上下文段落：
+{context or "无上下文"}"""
+        
+        full_user_msg = user_msg
+        if template.candidate_targets:
+            full_user_msg += "\n\n" + template.build_candidate_targets_section()
+        if template.task_examples:
+            full_user_msg += "\n\n" + template.build_task_examples_section()
+        if template.task_emphasis:
+            full_user_msg += "\n\n" + template.build_task_emphasis_section()
+        full_user_msg += f"\n\n输入数据：\n{input_data}"
+        
+        messages = [
+            {"role": "system", "content": system_msg},
+            {"role": "user", "content": full_user_msg}
+        ]
+        
+        # 并行调用多个模型
+        results = []
+        for i, client in enumerate(self.clients):
+            try:
+                model_name = f"model_{i+1}"
+                logger.info(f"使用模型 {model_name} 进行关系抽取...")
+                
+                # 第一阶段：初始提取
+                first_result = client.chat(messages)
+                
+                # 解析第一阶段结果
+                first_triplets = self._parse_triplets(first_result)
+                
+                # 第二阶段：验证和优化
+                if use_verification and first_triplets:
+                    verification_prompt = self.prompt_builder.build_verification_prompt()
+                    
+                    # 移除assistant消息，将第一阶段结果整合到user消息中（符合硅基流动API要求）
+                    second_messages = [
+                        {"role": "system", "content": "你需要验证和优化之前提取的三元组结果。请特别关注语义逻辑的正确性、遗漏关系的补全，以及利用上下文补全省略的成分。"},
+                        {"role": "user", "content": f"{verification_prompt}\n\n原始输入：\n{input_data}\n\n第一阶段结果：\n{first_result}\n\n请验证并优化这个结果，确保输出为有效的JSON格式。"}
+                    ]
+                    
+                    second_result = client.chat(second_messages)
+                    final_triplets = self._parse_triplets(second_result)
+                else:
+                    final_triplets = first_triplets
+                
+                results.append({
+                    "model": model_name,
+                    "triplets": final_triplets,
+                    "first_triplets": first_triplets,
+                })
+                
+            except Exception as e:
+                logger.error(f"模型 {i+1} 调用失败: {e}")
+                results.append({
+                    "model": f"model_{i+1}",
+                    "triplets": [],
+                    "error": str(e)
+                })
+        
+        return results
+
+    def _format_entities_for_prompt(self, entities: List[Dict[str, str]]) -> str:
+        """格式化实体列表用于Prompt"""
+        if not entities:
+            return "无识别实体"
+        
+        lines = []
+        for entity in entities:
+            entity_type = entity.get("type", "未知类型")
+            entity_text = entity.get("text", "")
+            lines.append(f"- {entity_type}: {entity_text}")
+        
+        return "\n".join(lines)
+
+    def _format_syntax_info(self, syntax_info: Dict) -> str:
+        """格式化句法信息"""
+        if not syntax_info:
+            return "无句法信息"
+        
+        formatted = []
+        if "dependency" in syntax_info:
+            deps = syntax_info["dependency"]
+            formatted.append("依存关系：")
+            for dep in deps[:5]:  # 只显示前5个
+                formatted.append(f"  {dep.get('word', '')} <-{dep.get('deprel', '')}- {dep.get('head', '')}")
+        
+        return "\n".join(formatted)
+
+    def _parse_triplets(self, result: str) -> List[Dict[str, str]]:
+        """解析三元组结果"""
+        # 尝试解析JSON
+        try:
+            triplets = json.loads(result)
+            if isinstance(triplets, list):
+                # 验证三元组格式
+                valid_triplets = []
+                for t in triplets:
+                    if isinstance(t, dict) and "head" in t and "relation" in t and "tail" in t:
+                        valid_triplets.append({
+                            "head": str(t.get("head", "")),
+                            "relation": str(t.get("relation", "")),
+                            "tail": str(t.get("tail", ""))
+                        })
+                return valid_triplets
+            elif isinstance(triplets, dict):
+                # 可能是嵌套结构，尝试提取
+                return []
+            return []
+        except json.JSONDecodeError:
+            # 尝试从文本中提取JSON
+            import re
+            json_match = re.search(r'\[.*\]', result, re.DOTALL)
+            if json_match:
+                try:
+                    triplets = json.loads(json_match.group())
+                    if isinstance(triplets, list):
+                        return triplets
+                except:
+                    pass
+            logger.warning(f"无法解析三元组结果: {result[:100]}")
+            return []
+
+