""" =========================================== 路由标签识别模块 - RouteLabel =========================================== 功能：使用LLM识别用户查询中涉及的节点类型和实体 作为后续Neo4j查询的入口节点 =========================================== """ import json import logging import os import time import asyncio from typing import List, Dict, Any, Optional, Tuple from pydantic import BaseModel, Field from graph_search.init_prompts import get_extract_entities_prompt, get_rewrite_query_prompt logger = logging.getLogger(__name__) #logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) # 不依赖继承链 class RouteItem(BaseModel): """路由项：节点类型和实体""" label: str = Field(..., description="节点类型，比如'Drug'、'Disease'") entity: str = Field(..., description="实体文本") class RouteOutput(BaseModel): """路由输出：包含多个路由项""" outputs: List[RouteItem] class RouteLabel: """ 路由标签识别类 功能：使用LLM识别用户查询中涉及的节点类型和实体 """ def __init__(self, llm_api, optional_labels: list[str], schema: str = None): """ 初始化路由识别器 Args: llm_api: LLM API实例（OpenaiAPI） optional_labels: 可选节点标签列表 schema: Neo4j图谱schema信息（可选，用于帮助模型识别节点类型） """ self.llm_api = llm_api self.optional_labels = optional_labels self.schema = schema async def _extract_entities(self, query: str) -> List[Dict[str, str]]: """ 任务1：抽取主语实体 Args: query: 用户查询 Returns: List[Dict]: 路由结果列表，格式如 [{"label": "Drug", "entity": "xxx"}] """ # 使用 init_prompts 中的 prompt 生成函数 prompt = get_extract_entities_prompt( query=query, optional_labels=self.optional_labels, schema=self.schema ) try: llm_start = time.time() model = os.getenv("OPENAI_MODEL") response = await self.llm_api.open_api_chat_async_json(prompt, model, temperature=0) llm_elapsed = (time.time() - llm_start) * 1000 # 解析JSON响应 parse_start = time.time() content = response.strip() # 如果包含markdown代码块，提取JSON if "```json" in content: json_start = content.find("```json") + 7 json_end = content.find("```", json_start) if json_end > json_start: content = content[json_start:json_end].strip() elif "```" in content: json_start = content.find("```") + 3 json_end = content.find("```", json_start) if json_end > json_start: content = content[json_start:json_end].strip() # 尝试找到JSON数组或对象的开始和结束位置 if not content.startswith('[') and not content.startswith('{'): for start_char in ['[', '{']: start_idx = content.find(start_char) if start_idx != -1: content = content[start_idx:] break # 解析JSON try: outputs_data = json.loads(content) except json.JSONDecodeError as e: logger.warning(f"任务1 JSON解析失败，尝试修复: {e}") if content.rfind(']') > content.rfind('}'): end_idx = content.rfind(']') + 1 content = content[:end_idx] elif content.rfind('}') > -1: end_idx = content.rfind('}') + 1 content = content[:end_idx] try: outputs_data = json.loads(content) except json.JSONDecodeError: logger.error(f"任务1 JSON解析最终失败，原始内容: {content[:200]}...") return [] # 提取entities if isinstance(outputs_data, dict) and "entities" in outputs_data: route_results = [ {"label": item.get("label", ""), "entity": item.get("entity", "")} for item in outputs_data["entities"] ] elif isinstance(outputs_data, list): # 兼容旧格式：直接是列表 route_results = [ {"label": item.get("label", ""), "entity": item.get("entity", "")} for item in outputs_data ] elif isinstance(outputs_data, dict) and "outputs" in outputs_data: # 兼容旧格式：包含 outputs 字段 route_results = [ {"label": item.get("label", ""), "entity": item.get("entity", "")} for item in outputs_data["outputs"] ] else: route_results = [] return route_results except Exception as e: logger.error(f"任务1（实体抽取）失败: {e}") return [] async def _rewrite_query(self, query: str) -> Dict[str, Any]: """ 任务2：重写用户问题并进行分类 Args: query: 用户查询 Returns: Dict: 包含重写后的查询和分类信息 - rewritten_query: 重写后的查询字符串 - query_classification: 分类信息字典 - query_type: "aggregate"（统计类）或 "detail"（明细类） - has_time: 是否包含时间信息（bool） - time_type: 时间统计类型（"single"/"dual"/"multi"/None） """ # 使用 init_prompts 中的 prompt 生成函数 prompt = get_rewrite_query_prompt( query=query, schema=self.schema ) try: llm_start = time.time() model = os.getenv("OPENAI_MODEL") response = await self.llm_api.open_api_chat_async_json(prompt, model, temperature=0) # 解析JSON响应 content = response.strip() # 如果包含markdown代码块，提取JSON if "```json" in content: json_start = content.find("```json") + 7 json_end = content.find("```", json_start) if json_end > json_start: content = content[json_start:json_end].strip() elif "```" in content: json_start = content.find("```") + 3 json_end = content.find("```", json_start) if json_end > json_start: content = content[json_start:json_end].strip() # 尝试找到JSON数组或对象的开始和结束位置 if not content.startswith('[') and not content.startswith('{'): for start_char in ['[', '{']: start_idx = content.find(start_char) if start_idx != -1: content = content[start_idx:] break # 解析JSON try: outputs_data = json.loads(content) except json.JSONDecodeError as e: logger.warning(f"任务2 JSON解析失败，尝试修复: {e}") if content.rfind(']') > content.rfind('}'): end_idx = content.rfind(']') + 1 content = content[:end_idx] elif content.rfind('}') > -1: end_idx = content.rfind('}') + 1 content = content[:end_idx] try: outputs_data = json.loads(content) except json.JSONDecodeError: logger.error(f"任务2 JSON解析最终失败，原始内容: {content[:200]}...") return query # 提取rewritten_query和分类信息 if isinstance(outputs_data, dict) and "rewritten_query" in outputs_data: rewritten_query = outputs_data.get("rewritten_query", query) # 提取分类信息 query_classification = outputs_data.get("query_classification", {}) if not isinstance(query_classification, dict): query_classification = {} # 确保分类信息格式正确 result = { "rewritten_query": rewritten_query, "query_classification": { "query_type": query_classification.get("query_type", "detail"), "has_time": query_classification.get("has_time", False), "time_type": query_classification.get("time_type") } } else: # 如果解析失败，返回默认值 result = { "rewritten_query": query, "query_classification": { "query_type": "detail", "has_time": False, "time_type": None } } return result except Exception as e: logger.error(f"任务2（查询重写）失败: {e}") # 返回默认值 return { "rewritten_query": query, "query_classification": { "query_type": "detail", "has_time": False, "time_type": None } } async def route(self, query: str) -> Tuple[List[Dict[str, str]], str, Dict[str, Any]]: """ 路由标签识别：识别标签，抽取实体 功能：使用LLM识别用户查询中涉及的节点类型和实体。 将两个任务（实体抽取和查询重写）拆分为独立的异步调用，并行执行以提高效率。 Args: query: 用户查询 Returns: tuple: (路由结果列表, 重写后的查询, 分类信息) 路由结果列表格式如 [{"label": "Drug", "entity": "xxx"}] 重写后的查询字符串，如果未获取到则返回原始query 分类信息字典，包含 query_type, has_time, time_type """ try: # ========== 并行执行两个任务 ========== total_start = time.time() logger.info(f"[时间统计] RouteLabel: 开始并行执行两个任务（实体抽取 + 查询重写+分类）...") # 使用 asyncio.gather 并行执行两个任务 route_results, rewrite_result = await asyncio.gather( self._extract_entities(query), self._rewrite_query(query) ) logger.info(route_results) logger.info(111111111111111111111111111111111111111) # 从重写结果中提取信息 if isinstance(rewrite_result, dict): rewritten_query = rewrite_result.get("rewritten_query", query) query_classification = rewrite_result.get("query_classification", { "query_type": "detail", "has_time": False, "time_type": None }) else: # 兼容旧格式（如果返回的是字符串） rewritten_query = rewrite_result if isinstance(rewrite_result, str) else query query_classification = { "query_type": "detail", "has_time": False, "time_type": None } total_elapsed = (time.time() - total_start) * 1000 logger.info(f"[时间统计] RouteLabel: 两个任务并行执行完成，总耗时: {total_elapsed:.2f}ms") logger.info(f"最终结果 - 入口节点标签与实体: {route_results}") logger.info(f"最终结果 - 原始查询: '{query}' -> 重写后查询: '{rewritten_query}'") logger.info(f"最终结果 - 分类信息: {query_classification}") return route_results, rewritten_query, query_classification except Exception as e: logger.error(f"路由标签识别失败: {e}") return [], query, { "query_type": "detail", "has_time": False, "time_type": None }