integrated_pipeline.py

import os
import json
import sys
import argparse
from typing import Dict, List, Optional
from pathlib import Path
from tqdm import tqdm

# 添加必要的路径
root_path = os.path.dirname(os.path.abspath(__file__))
print(root_path)
sys.path.append(root_path)
sys.path.append(os.path.join(root_path, "Models/ProTrek"))

# 导入所需模块
from interproscan import InterproScan
from Bio.Blast.Applications import NcbiblastpCommandline
from utils.utils import extract_interproscan_metrics, get_seqnid, extract_blast_metrics, rename_interproscan_keys
from go_integration_pipeline import GOIntegrationPipeline
from utils.generate_protein_prompt import generate_prompt, get_interpro_manager, get_lmdb_connection
from utils.openai_access import call_chatgpt

class IntegratedProteinPipeline:
    def __init__(self, 
                 blast_database: str = "uniprot_swissprot",
                 expect_value: float = 0.01,
                 interproscan_path: str = "interproscan/interproscan-5.75-106.0/interproscan.sh",
                 interproscan_libraries: List[str] = None,
                 go_topk: int = 2,
                 selected_info_types: List[str] = None,
                 pfam_descriptions_path: str = None,
                 go_info_path: str = None,
                 interpro_data_path: str = None,
                 lmdb_path: str = None,
                 args: argparse.Namespace = None):
        """
        整合蛋白质分析管道
        
        Args:
            blast_database: BLAST数据库名称
            expect_value: BLAST E-value阈值
            interproscan_path: InterProScan程序路径
            interproscan_libraries: InterProScan库列表
            go_topk: GO整合的topk参数
            selected_info_types: prompt生成时选择的信息类型
            pfam_descriptions_path: Pfam描述文件路径
            go_info_path: GO信息文件路径
            interpro_data_path: InterPro数据文件路径
            lmdb_path: LMDB数据库路径
        """
        self.blast_database = blast_database
        self.expect_value = expect_value
        self.interproscan_path = interproscan_path
        self.interproscan_libraries = interproscan_libraries or [
            "PFAM", "PIRSR", "PROSITE_PROFILES", "SUPERFAMILY", "PRINTS", 
            "PANTHER", "CDD", "GENE3D", "NCBIFAM", "SFLM", "MOBIDB_LITE", 
            "COILS", "PROSITE_PATTERNS", "FUNFAM", "SMART"
        ]
        self.go_topk = go_topk
        self.selected_info_types = selected_info_types or ['motif', 'go']
        
        # 文件路径配置
        self.pfam_descriptions_path = pfam_descriptions_path
        self.go_info_path = go_info_path
        self.interpro_data_path = interpro_data_path
        self.lmdb_path = lmdb_path
        self.interproscan_info_path = args.interproscan_info_path
        self.blast_info_path = args.blast_info_path
        
        # 初始化GO整合管道
        self.go_pipeline = GOIntegrationPipeline(topk=self.go_topk)
        
        # 初始化InterPro管理器（如果需要）
        self.interpro_manager = None
        other_types = [t for t in self.selected_info_types if t not in ['motif', 'go']]
        if other_types and self.interpro_data_path:
            self.interpro_manager = get_interpro_manager(self.interpro_data_path, None)
    
    def step1_run_blast_and_interproscan(self, input_fasta: str, temp_dir: str = "temp") -> tuple:
        """
        步骤1: 运行BLAST和InterProScan分析
        
        Args:
            input_fasta: 输入FASTA文件路径
            temp_dir: 临时文件目录
            
        Returns:
            tuple: (interproscan_info, blast_info)
        """
        print("步骤1: 运行BLAST和InterProScan分析...")
        
        # 创建临时目录
        os.makedirs(temp_dir, exist_ok=True)
        
        # 获取序列字典
        seq_dict = get_seqnid(input_fasta)
        print(f"读取到 {len(seq_dict)} 个序列")
        
        # 运行BLAST
        print("运行BLAST分析...")
        blast_xml = os.path.join(temp_dir, "blast_results.xml")
        blast_cmd = NcbiblastpCommandline(
            query=input_fasta,
            db=self.blast_database,
            out=blast_xml,
            outfmt=5,  # XML格式
            evalue=self.expect_value
        )
        blast_cmd()
        
        # 提取BLAST结果
        blast_results = extract_blast_metrics(blast_xml)
        blast_info = {}
        for uid, info in blast_results.items():
            blast_info[uid] = {"sequence": seq_dict[uid], "blast_results": info}
        
        # 运行InterProScan
        print("运行InterProScan分析...")
        interproscan_json = os.path.join(temp_dir, "interproscan_results.json")
        interproscan = InterproScan(self.interproscan_path)
        input_args = {
            "fasta_file": input_fasta,
            "goterms": True,
            "pathways": True,
            "save_dir": interproscan_json
        }
        interproscan.run(**input_args)
        
        # 提取InterProScan结果
        interproscan_results = extract_interproscan_metrics(
            interproscan_json, 
            librarys=self.interproscan_libraries
        )
        
        interproscan_info = {}
        for id, seq in seq_dict.items():
            info = interproscan_results[seq]
            info = rename_interproscan_keys(info)
            interproscan_info[id] = {"sequence": seq, "interproscan_results": info}
        
        # 清理临时文件
        if os.path.exists(blast_xml):
            os.remove(blast_xml)
        if os.path.exists(interproscan_json):
            os.remove(interproscan_json)
        
        print(f"步骤1完成: 处理了 {len(interproscan_info)} 个蛋白质")
        return interproscan_info, blast_info
    
    def step2_integrate_go_information(self, interproscan_info: Dict, blast_info: Dict) -> Dict:
        """
        步骤2: 整合GO信息
        
        Args:
            interproscan_info: InterProScan结果
            blast_info: BLAST结果
            
        Returns:
            Dict: 整合后的GO信息
        """
        print("步骤2: 整合GO信息...")
        
        # 使用GO整合管道进行第一层筛选
        protein_go_dict = self.go_pipeline.first_level_filtering(interproscan_info, blast_info)
        
        print(f"步骤2完成: 为 {len(protein_go_dict)} 个蛋白质整合了GO信息")
        return protein_go_dict
    
    def step3_generate_prompts(self, interproscan_info: Dict, blast_info: Dict, 
                              protein_go_dict: Dict) -> Dict:
        """
        步骤3: 生成蛋白质prompt
        
        Args:
            interproscan_info: InterProScan结果
            blast_info: BLAST结果
            protein_go_dict: 整合的GO信息
            
        Returns:
            Dict: 蛋白质ID到prompt的映射（如果有lmdb则包含QA对）
        """
        print("步骤3: 生成蛋白质prompt...")
        
        # 创建临时的GO整合文件格式（用于generate_prompt函数）
        temp_go_data = {}
        for protein_id, go_ids in protein_go_dict.items():
            temp_go_data[protein_id] = go_ids
        
        prompts_data = {}
        
        if self.lmdb_path:
            # 如果有lmdb路径，处理QA数据
            from utils.generate_protein_prompt import get_qa_data
            
            global_index = 0
            for protein_id in tqdm(interproscan_info.keys(), desc="生成prompts"):
                # 获取QA对
                qa_pairs = get_qa_data(protein_id, self.lmdb_path)
                
                for qa_pair in qa_pairs:
                    question = qa_pair['question']
                    ground_truth = qa_pair['ground_truth']
                    
                    # 生成prompt（需要修改generate_prompt函数以支持内存数据）
                    prompt = self._generate_prompt_from_memory(
                        protein_id, interproscan_info, temp_go_data, question
                    )
                    
                    if prompt:
                        prompts_data[global_index] = {
                            "index": global_index,
                            "protein_id": protein_id,
                            "prompt": prompt,
                            "question": question,
                            "ground_truth": ground_truth
                        }
                        global_index += 1
        else:
            # 如果没有lmdb路径，按原来的方式处理
            for protein_id in tqdm(interproscan_info.keys(), desc="生成prompts"):
                prompt = self._generate_prompt_from_memory(
                    protein_id, interproscan_info, temp_go_data
                )
                if prompt:
                    prompts_data[protein_id] = prompt
        
        print(f"步骤3完成: 生成了 {len(prompts_data)} 个prompt")
        return prompts_data
    
    def _generate_prompt_from_memory(self, protein_id: str, interproscan_info: Dict, 
                                   protein_go_dict: Dict, question: str = None) -> str:
        """
        从内存中的数据生成prompt，包含完整的motif和GO定义
        """
        try:
            from utils.protein_go_analysis import get_go_definition
            from jinja2 import Template
            from utils.generate_protein_prompt import get_prompt_template
            
            # 获取GO分析结果
            go_ids = protein_go_dict.get(protein_id, [])
            go_annotations = []
            all_related_definitions = {}
            
            if go_ids:
                for go_id in go_ids:
                    # 确保GO ID格式正确
                    clean_go_id = go_id.split(":")[-1] if ":" in go_id else go_id
                    go_annotations.append({"go_id": clean_go_id})
                    
                    # 获取GO定义
                    definition = get_go_definition(clean_go_id,self.go_info_path)
                    if definition:
                        all_related_definitions[clean_go_id] = definition
            
            # 获取motif信息
            motif_pfam = {}
            if self.pfam_descriptions_path:
                try:
                    # 从interproscan结果中提取pfam信息
                    interproscan_results = interproscan_info[protein_id].get('interproscan_results', {})
                    pfam_entries = interproscan_results.get('pfam_id', [])
                    
                    # 加载pfam描述
                    with open(self.pfam_descriptions_path, 'r') as f:
                        pfam_descriptions = json.load(f)
                    
                    # 构建motif_pfam字典
                    for entry in pfam_entries:
                        for pfam_id, ipr_id in entry.items():
                            if pfam_id and pfam_id in pfam_descriptions:
                                motif_pfam[pfam_id] = pfam_descriptions[pfam_id]['description']
                                
                except Exception as e:
                    print(f"获取motif信息时出错: {str(e)}")
            
            # 获取InterPro描述信息
            interpro_descriptions = {}
            other_types = [t for t in self.selected_info_types if t not in ['motif', 'go']]
            if other_types and self.interpro_manager:
                interpro_descriptions = self.interpro_manager.get_description(protein_id, other_types)
            
            # 准备模板数据
            template_data = {
                "protein_id": protein_id,
                "selected_info_types": self.selected_info_types,
                "go_data": {
                    "status": "success" if go_annotations else "no_data",
                    "go_annotations": go_annotations,
                    "all_related_definitions": all_related_definitions
                },
                "motif_pfam": motif_pfam,
                "interpro_descriptions": interpro_descriptions,
                "question": question
            }
            
            # 使用模板生成prompt
            PROMPT_TEMPLATE = get_prompt_template(self.selected_info_types, self.lmdb_path)
            template = Template(PROMPT_TEMPLATE)
            return template.render(**template_data)
            
        except Exception as e:
            print(f"生成prompt时出错 (protein_id: {protein_id}): {str(e)}")
            # 如果出错，返回简化版本的prompt
            return self._generate_simple_prompt(protein_id, interproscan_info, protein_go_dict, question)
    
    def _generate_simple_prompt(self, protein_id: str, interproscan_info: Dict, 
                               protein_go_dict: Dict, question: str = None) -> str:
        """
        生成简化版本的prompt（作为备用）
        """
        # 获取蛋白质序列
        sequence = interproscan_info[protein_id].get('sequence', '')
        
        # 获取GO信息
        go_ids = protein_go_dict.get(protein_id, [])
        
        # 获取motif信息
        interproscan_results = interproscan_info[protein_id].get('interproscan_results', {})
        pfam_entries = interproscan_results.get('pfam_id', [])
        
        # 简化的prompt生成逻辑
        prompt_parts = []
        
        if self.lmdb_path:
            from utils.prompts import FUNCTION_PROMPT
            prompt_parts.append(FUNCTION_PROMPT)
        else:
            from utils.prompts import ENZYME_PROMPT
            prompt_parts.append(ENZYME_PROMPT)
        
        prompt_parts.append("\ninput information:")
        
        # 添加motif信息
        if 'motif' in self.selected_info_types and pfam_entries:
            prompt_parts.append("\nmotif:")
            for entry in pfam_entries:
                for key, value in entry.items():
                    if value:
                        prompt_parts.append(f"{value}: 无详细描述")
        
        # 添加GO信息
        if 'go' in self.selected_info_types and go_ids:
            prompt_parts.append("\nGO:")
            for i, go_id in enumerate(go_ids[:10], 1):
                prompt_parts.append(f"▢ GO term{i}: {go_id}")
                prompt_parts.append(f"• definition: 无详细定义")
        
        if question:
            prompt_parts.append(f"\nquestion: \n{question}")
        
        return "\n".join(prompt_parts)
    
    def step4_generate_llm_answers(self, prompts_data: Dict, save_dir: str) -> None:
        """
        步骤4: 生成LLM答案
        
        Args:
            prompts_data: prompt数据
            save_dir: 保存目录
        """
        print("步骤4: 生成LLM答案...")
        
        # 创建保存目录
        os.makedirs(save_dir, exist_ok=True)
        
        if self.lmdb_path:
            # 如果有lmdb路径，处理QA数据
            for index, qa_item in tqdm(prompts_data.items(), desc="生成LLM答案"):
                try:
                    protein_id = qa_item['protein_id']
                    prompt = qa_item['prompt']
                    question = qa_item['question']
                    ground_truth = qa_item['ground_truth']
                    
                    # 调用LLM生成答案
                    llm_response = call_chatgpt(prompt)
                    
                    # 构建结果数据
                    result = {
                        'protein_id': protein_id,
                        'index': index,
                        'question': question,
                        'ground_truth': ground_truth,
                        'llm_answer': llm_response
                    }
                    
                    # 保存文件
                    save_path = os.path.join(save_dir, f"{protein_id}_{index}.json")
                    with open(save_path, 'w') as f:
                        json.dump(result, f, indent=2, ensure_ascii=False)
                        
                except Exception as e:
                    print(f"处理索引 {index} 时出错: {str(e)}")
        else:
            # 如果没有lmdb路径，按原来的方式处理
            for protein_id, prompt in tqdm(prompts_data.items(), desc="生成LLM答案"):
                try:
                    # 调用LLM生成答案
                    llm_response = call_chatgpt(prompt)
                    
                    # 构建结果数据
                    result = {
                        'protein_id': protein_id,
                        'prompt': prompt,
                        'llm_answer': llm_response
                    }
                    
                    # 保存文件
                    save_path = os.path.join(save_dir, f"{protein_id}.json")
                    with open(save_path, 'w') as f:
                        json.dump(result, f, indent=2, ensure_ascii=False)
                        
                except Exception as e:
                    print(f"处理蛋白质 {protein_id} 时出错: {str(e)}")
        
        print(f"步骤4完成: 结果已保存到 {save_dir}")
    
    def run(self, input_fasta: str, output_dir: str, temp_dir: str = "temp"):
        """
        运行完整的工作流
        
        Args:
            input_fasta: 输入FASTA文件路径
            output_dir: 输出目录
            temp_dir: 临时文件目录
        """
        print(f"开始运行整合蛋白质分析管道...")
        print(f"输入文件: {input_fasta}")
        print(f"输出目录: {output_dir}")
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        try:
            # 步骤1: 运行BLAST和InterProScan
            if self.interproscan_info_path is None or self.blast_info_path is None:
                interproscan_info, blast_info = self.step1_run_blast_and_interproscan(
                    input_fasta, temp_dir
                )
            else:
                interproscan_info = json.load(open(self.interproscan_info_path))
                blast_info = json.load(open(self.blast_info_path))
            
            # 步骤2: 整合GO信息
            protein_go_dict = self.step2_integrate_go_information(
                interproscan_info, blast_info
            )
            
            # 步骤3: 生成prompt
            prompts_data = self.step3_generate_prompts(
                interproscan_info, blast_info, protein_go_dict
            )
            print(prompts_data)
            
            # 步骤4: 生成LLM答案
            self.step4_generate_llm_answers(prompts_data, output_dir)
            
            print("整合管道运行完成！")
            
        except Exception as e:
            print(f"管道运行出错: {str(e)}")
            raise
        finally:
            # 清理临时目录
            print(f"清理临时目录: {temp_dir}")
            if os.path.exists(temp_dir):
                import shutil
                shutil.rmtree(temp_dir)

def main():
    parser = argparse.ArgumentParser(description="整合蛋白质分析管道")
    parser.add_argument("--input_fasta", type=str, required=True, help="输入FASTA文件路径")
    parser.add_argument("--output_dir", type=str, required=True, help="输出目录")
    parser.add_argument("--temp_dir", type=str, default="temp", help="临时文件目录")
    parser.add_argument('--interproscan_info_path', type=str, default=None, help="InterProScan结果文件路径")
    parser.add_argument('--blast_info_path', type=str, default=None, help="BLAST结果文件路径")
    
    
    # BLAST参数
    parser.add_argument("--blast_database", type=str, default="uniprot_swissprot", help="BLAST数据库")
    parser.add_argument("--expect_value", type=float, default=0.01, help="BLAST E-value阈值")
    
    # InterProScan参数
    parser.add_argument("--interproscan_path", type=str, 
                       default="interproscan/interproscan-5.75-106.0/interproscan.sh",
                       help="InterProScan程序路径")
    
    # GO整合参数
    parser.add_argument("--go_topk", type=int, default=2, help="GO整合topk参数")
    
    # Prompt生成参数
    parser.add_argument("--selected_info_types", type=str, nargs='+', 
                       default=['motif', 'go'], help="选择的信息类型")
    parser.add_argument("--pfam_descriptions_path", type=str, default='data/raw_data/all_pfam_descriptions.json', help="Pfam描述文件路径")
    parser.add_argument("--go_info_path", type=str, default='data/raw_data/go.json', help="GO信息文件路径")
    parser.add_argument("--interpro_data_path", type=str, default='data/raw_data/interpro_data.json', help="InterPro数据文件路径")
    parser.add_argument("--lmdb_path", type=str, help="LMDB数据库路径")
    
    args = parser.parse_args()
    
    # 创建管道实例
    pipeline = IntegratedProteinPipeline(
        blast_database=args.blast_database,
        expect_value=args.expect_value,
        interproscan_path=args.interproscan_path,
        go_topk=args.go_topk,
        selected_info_types=args.selected_info_types,
        pfam_descriptions_path=args.pfam_descriptions_path,
        go_info_path=args.go_info_path,
        interpro_data_path=args.interpro_data_path,
        lmdb_path=args.lmdb_path,
        args=args
    )
    
    # 运行管道
    pipeline.run(args.input_fasta, args.output_dir, args.temp_dir)

if __name__ == "__main__":
    main()