""" PlanningAgent V2 - 增强版任务规划 支持根据用户意图生成合理的分析流程 集成Knowledge Agent获取参数推荐 """ import json import re from typing import Dict, List, Optional from .base import AgentMessage, AgentState, BaseAgent from gfa_core.mas.llm import LLMManager class PlanningAgent(BaseAgent): """ 任务规划Agent V2 根据用户意图和可用资源，生成最优的分析执行计划 """ name = "planning" role = "Workflow Planning Specialist" # 标准分析流程定义 STANDARD_FLOWS = { "prepare": ["genome_prepare"], "basic": ["identify"], "identify_refine": ["identify_refine"], "phylogeny": ["identify", "phylogeny"], "structure": ["identify", "phylogeny", "motif", "domain", "structure"], "promoter": ["identify", "phylogeny", "promoter"], "synteny": ["identify", "phylogeny", "synteny"], "karyotype": ["identify", "phylogeny", "karyotype"], "interspecies": ["identify", "phylogeny", "interspecies_synteny"], "expression": ["identify", "phylogeny", "expression"], "complete": [ "identify", "phylogeny", "motif", "domain", "structure", "promoter", "synteny", "expression" ], "rnaseq": [ "rnaseq_qc", "rnaseq_align", "rnaseq_quant", "rnaseq_sample_qc", "rnaseq_de", "rnaseq_go_enrich", "rnaseq_kegg_enrich", "rnaseq_fgsea", "rnaseq_wgcna", ], "evolution": ["identify", "phylogeny", "kaks"], "interaction": ["identify", "ppi"] } # 阶段依赖关系 STAGE_DEPENDENCIES = { "genome_prepare": [], "identify_refine": [], "phylogeny": ["identify"], "motif": ["identify"], "domain": ["identify"], "structure": ["identify"], "promoter": ["identify", "phylogeny"], "synteny": ["identify"], "karyotype": ["identify"], "interspecies_synteny": ["identify"], "expression": ["identify"], "kaks": ["identify", "phylogeny"], "ppi": ["identify"], "rnaseq_align": ["rnaseq_qc"], "rnaseq_quant": ["rnaseq_align"], "rnaseq_sample_qc": ["rnaseq_quant"], "rnaseq_wgcna": ["rnaseq_sample_qc"], "rnaseq_de": ["rnaseq_sample_qc"], "rnaseq_go_enrich": ["rnaseq_de"], "rnaseq_kegg_enrich": ["rnaseq_de"], "rnaseq_fgsea": ["rnaseq_de"], "rnaseq_enrich": ["rnaseq_de"], "rnaseq_candidate_family": ["rnaseq_de", "rnaseq_go_enrich", "rnaseq_kegg_enrich", "rnaseq_fgsea"] } # 阶段描述 STAGE_DESCRIPTIONS = { "genome_prepare": "数据准备总流程 - 多物种文件整理、GFF/FASTA标准化、最长转录本和参考序列合并", "identify_refine": "修订鉴定成员 - 根据人工修订ID列表重建final_members.id和final_members.pep.fa", "identify": "基因家族鉴定 - 使用HMM和BLAST识别基因家族成员", "phylogeny": "系统发育分析 - 构建进化树和亚家族分类", "motif": "Motif分析 - 识别保守的蛋白基序", "domain": "结构域分析 - 分析蛋白结构域组成", "structure": "结构预测 - 预测蛋白三维结构", "promoter": "启动子分析 - 分析顺式作用元件", "synteny": "共线性分析 - 分析基因组共线性关系", "karyotype": "核型分析 - 生成染色体分布图", "interspecies_synteny": "物种间共线性 - 跨物种比较分析", "expression": "表达分析 - 分析基因表达模式", "kaks": "Ka/Ks分析 - 计算进化速率和选择压力", "ppi": "PPI分析 - 蛋白互作网络分析", "rnaseq_qc": "RNA-seq质控 - FastQC质量评估", "rnaseq_align": "RNA-seq比对 - reads比对到基因组", "rnaseq_quant": "RNA-seq定量 - 基因表达量计算", "rnaseq_sample_qc": "RNA-seq样本质量评估 - VST、PCA、样本相关性与聚类", "rnaseq_wgcna": "RNA-seq WGCNA - 加权共表达网络与模块性状相关性", "rnaseq_de": "差异表达 - DESeq2差异分析", "rnaseq_go_enrich": "RNA-seq GO富集分析 - 独立使用 gene2go 背景", "rnaseq_kegg_enrich": "RNA-seq KEGG富集分析 - 独立使用 gene2ko/KO 背景", "rnaseq_fgsea": "RNA-seq fgsea - 基于排序基因和 KEGG 基因集的 GSEA", "rnaseq_enrich": "RNA-seq富集分析 - 兼容旧入口；规划中优先使用 GO/KEGG 子阶段" } def __init__(self, llm_manager: LLMManager): super().__init__(llm_manager) async def execute(self, state: AgentState, message: AgentMessage) -> AgentState: """ 执行任务规划 根据用户意图生成分析执行计划 """ query = message.payload.get("query", "") intent = state.results.get("results", {}).get("intent", {}) or state.results.get("intent", {}) # 提取关键信息 gene_family = intent.get("gene_family", "") or self._extract_gene_family(query) species = intent.get("species", "") or self._extract_species(query) stage = intent.get("stage", "") # 确定分析流程 flow = self._determine_flow(query, stage) # 生成执行计划 plan = await self._generate_plan(query, gene_family, species, flow) # 更新状态 state.results["task_plan"] = plan state.results["planning_info"] = { "gene_family": gene_family, "species": species, "flow_type": flow, "total_stages": len(plan.get("stages", [])) } state.messages.append({ "role": "planning", "content": f"📋 已生成分析计划: {len(plan.get('stages', []))}个阶段", }) return state def _extract_gene_family(self, query: str) -> str: """从查询中提取基因家族名称""" # 常见基因家族模式 patterns = [ r'(\w+)\s*基因家族', r'基因家族\s*(\w+)', r'\b(SPL|NAC|WRKY|MYB|bHLH|WD40|C2H2|C3H|NBS|AP2|ARF|AUX|B3|BZIP|C2C2|CCAAT|CO-like|CPP|CSD|CTF|CYP|Dof|E2F|ERF|FAR1|G2-like|GATA|GeBP|GRAS|GRF|HB|HMA|HMG|HRT|HSF|JUMONJI|KNOX|LBD|LFY|LIM|LOB|LSD|MADS|MULE|MYB|NAC|ND|Nin|NLP|NPR|OFP|Orphan|Ots|PHD|PLATZ|PPR|Pseudo|RWP|S1Fa|SAP|SBP|Sigma|SMC|SRS|TALE|TCP|Tify|TRAF|Trihelix|TST|TUB|ULT|VOZ|Whirly|WOX|WRKY|XAL2|YABBY|ZF)\b' ] for pattern in patterns: match = re.search(pattern, query, re.IGNORECASE) if match: return match.group(1).upper() return "SPL" # 默认 def _extract_species(self, query: str) -> str: """从查询中提取物种信息""" species_map = { 'ft': 'Ft', 'tartary': 'Ft', '苦荞': 'Ft', 'at': 'At', 'arabidopsis': 'At', '拟南芥': 'At', 'os': 'Os', 'rice': 'Os', '水稻': 'Os', 'zm': 'Zm', 'maize': 'Zm', '玉米': 'Zm', 'gm': 'Gm', 'soybean': 'Gm', '大豆': 'Gm', 'sl': 'Sl', 'tomato': 'Sl', '番茄': 'Sl' } query_lower = query.lower() for key, value in species_map.items(): if key in query_lower: return value return "Ft" # 默认 def _determine_flow(self, query: str, stage: str) -> str: """确定分析流程类型""" query_lower = query.lower() # RNA-seq关键词 if any(kw in query_lower for kw in ['rnaseq', 'rna-seq', '转录组', '表达谱']): return "rnaseq" # 进化分析关键词 if any(kw in query_lower for kw in ['kaks', 'ka/ks', '进化', '选择压力']): return "evolution" # PPI关键词 if any(kw in query_lower for kw in ['ppi', '互作', 'interaction', 'string']): return "interaction" # 根据明确指定的stage if stage: flow_map = { "identify": "basic", "phylogeny": "phylogeny", "motif": "structure", "domain": "structure", "structure": "structure", "promoter": "promoter", "synteny": "synteny", "karyotype": "karyotype", "interspecies_synteny": "interspecies", "expression": "expression", "kaks": "evolution", "ppi": "interaction" } return flow_map.get(stage, "complete") # 完整分析关键词 if any(kw in query_lower for kw in ['完整', '全套', '全部', 'complete', 'full']): return "complete" return "complete" # 默认完整分析 async def _generate_plan(self, query: str, gene_family: str, species: str, flow: str) -> dict: """生成详细执行计划""" # 获取流程阶段列表 stages = self.STANDARD_FLOWS.get(flow, self.STANDARD_FLOWS["complete"]) # 构建阶段详情 stage_details = [] for i, stage_id in enumerate(stages): stage_info = { "id": stage_id, "name": self.STAGE_DESCRIPTIONS.get(stage_id, stage_id), "description": self.STAGE_DESCRIPTIONS.get(stage_id, ""), "order": i + 1, "dependencies": self.STAGE_DEPENDENCIES.get(stage_id, []), "params": await self._get_default_params(stage_id, gene_family, species) } stage_details.append(stage_info) # 构建执行计划 plan = { "query": query, "gene_family": gene_family, "species": species, "flow_type": flow, "prefix": f"{species}.{gene_family}", "stages": stage_details, "total_stages": len(stages), "estimated_time": self._estimate_time(stages), "required_inputs": self._get_required_inputs(stages), "outputs": self._get_outputs(stages, gene_family, species) } return plan async def _get_default_params(self, stage_id: str, gene_family: str, species: str) -> dict: """获取阶段的默认参数，并尝试从Knowledge Agent获取推荐""" default_params = { "identify": {"cpu": 8, "dom_i_evalue_max": 0.001, "hmm_cov_min": 0.3, "blast_evalue": "1e-05"}, "phylogeny": {"cpu": 8, "bootstrap": 1000, "model": "MFP"}, "motif": {"cpu": 8, "nmotifs": 10, "min_width": 6, "max_width": 200}, "domain": {"cpu": 8, "e_value": 1e-5}, "structure": {"cpu": 8}, "promoter": {"cpu": 8}, "synteny": {"cpu": 8}, "karyotype": {"cpu": 8}, "interspecies_synteny": {"cpu": 8}, "expression": {"cpu": 8}, "kaks": {"cpu": 8, "model": "YN00"}, "ppi": {"species": "ath", "score": 200}, "rnaseq_qc": {"cpu": 8}, "rnaseq_align": {"cpu": 8}, "rnaseq_quant": {"cpu": 8}, "rnaseq_de": {"cpu": 8}, "rnaseq_enrich": {"cpu": 8, "top_pathways": 20} } params = default_params.get(stage_id, {"cpu": 8}).copy() params["prefix"] = f"{species}.{gene_family}" # 尝试从 Knowledge Agent 获取参数推荐 try: recommended = await self._get_params_from_knowledge(stage_id, gene_family, species) if recommended: # 合并推荐参数（推荐参数优先） params.update(recommended) print(f"[PlanningAgent] 从Knowledge获取到{stage_id}阶段参数推荐") except Exception as e: print(f"[PlanningAgent] Knowledge参数推荐失败: {e}，使用默认参数") return params async def _get_params_from_knowledge(self, stage_id: str, gene_family: str, species: str) -> dict: """从Knowledge Agent获取参数推荐""" from .knowledge import create_knowledge_agent from gfa_core.mas.llm import LLMManager # 创建 Knowledge Agent llm_manager = LLMManager() knowledge_agent = create_knowledge_agent(llm_manager) # 构建查询 query = f"推荐{stage_id}阶段的分析参数，基因家族：{gene_family}，物种：{species}" # 构建消息 from .base import AgentMessage, AgentState import uuid message = AgentMessage( sender="planning", receiver="knowledge", message_type="query", payload={"query": query, "mode": "consultation"} ) state = AgentState( session_id=f"planning_{uuid.uuid4().hex[:8]}", results={}, messages=[] ) # 执行查询 result_state = await knowledge_agent.execute(state, message) # 解析推荐参数 answer = result_state.results.get("knowledge_response", "") # 尝试从回答中提取参数（简化版） recommended = {} if "cpu" in answer.lower(): # 提取CPU推荐 import re cpu_match = re.search(r'(\d+)\s*核', answer) if cpu_match: recommended["cpu"] = int(cpu_match.group(1)) return recommended def _estimate_time(self, stages: list) -> str: """估算总执行时间""" time_map = { "identify": 5, "phylogeny": 30, "motif": 20, "domain": 15, "structure": 30, "promoter": 20, "synteny": 30, "karyotype": 10, "interspecies_synteny": 40, "expression": 15, "kaks": 20, "ppi": 10, "rnaseq_qc": 30, "rnaseq_align": 60, "rnaseq_quant": 30, "rnaseq_sample_qc": 10, "rnaseq_de": 20, "rnaseq_enrich": 15, "rnaseq_go_enrich": 15, "rnaseq_kegg_enrich": 15, "rnaseq_fgsea": 10, "rnaseq_wgcna": 30 } total_minutes = sum(time_map.get(s, 10) for s in stages) if total_minutes < 60: return f"约{total_minutes}分钟" else: hours = total_minutes // 60 mins = total_minutes % 60 return f"约{hours}小时{mins}分钟" if mins > 0 else f"约{hours}小时" def _get_required_inputs(self, stages: list) -> list: """获取所需输入文件""" inputs = [] if "identify" in stages: inputs.extend([ {"name": "HMM模型文件", "type": "hmm", "required": True, "example": "SPL_Ath.hmm"}, {"name": "蛋白质序列", "type": "fasta", "required": True, "example": "Ft.pep.fasta"} ]) if any(s.startswith("rnaseq") for s in stages): inputs.extend([ {"name": "FASTQ文件", "type": "fastq", "required": True, "example": "*.fastq.gz"}, {"name": "样本配置", "type": "yaml", "required": True, "example": "samples.yaml"} ]) return inputs def _get_outputs(self, stages: list, gene_family: str, species: str) -> list: """获取预期输出""" prefix = f"{species}.{gene_family}" outputs = [] output_map = { "identify": [{"name": "最终成员列表", "file": f"{prefix}.final_members.id"}], "phylogeny": [{"name": "进化树", "file": f"{prefix}.tree.nwk"}], "kaks": [{"name": "Ka/Ks结果", "file": f"{prefix}.kaks_results.csv"}], "ppi": [{"name": "PPI网络", "file": f"{prefix}.ppi_network.tsv"}] } for stage in stages: if stage in output_map: outputs.extend(output_map[stage]) return outputs def create_planning_agent(llm_manager): """创建PlanningAgent实例""" return PlanningAgent(llm_manager)