Update biopython/mcp_output/mcp_plugin/mcp_service.py

biopython/mcp_output/mcp_plugin/mcp_service.py

@@ -36,6 +36,10 @@ mcp = FastMCP("biopython_gene_species_identification")
 # Configure Entrez email (required by NCBI)
 Entrez.email = os.getenv("BIOPYTHON_ENTREZ_EMAIL", "biopython-mcp@huggingface.co")
 
+# Global storage for async BLAST jobs
+blast_jobs: Dict[str, Dict[str, Any]] = {}
+blast_lock = threading.Lock()
+
 
 # ==============================================
 # 基础工具 - Basic Utilities
@@ -141,7 +145,215 @@ def calculate_gc_content(payload: dict):
 
 
 # ==============================================
-# BLAST 搜索工具 - BLAST Search Tools
+# 异步 BLAST 搜索工具 - Async BLAST Tools
+# ==============================================
+
+def _run_blast_async(job_id: str, sequence: str, database: str, program: str, 
+                     hitlist_size: int, expect: float):
+    """后台运行 BLAST 搜索"""
+    try:
+        print(f"🔄 [Job {job_id}] Starting BLAST search in background...")
+        
+        with blast_lock:
+            blast_jobs[job_id]["status"] = "running"
+            blast_jobs[job_id]["message"] = "BLAST search in progress..."
+        
+        # 执行 BLAST
+        handle = NCBIWWW.qblast(
+            program=program,
+            database=database,
+            sequence=sequence,
+            hitlist_size=hitlist_size,
+            expect=expect,
+            format_type="XML"
+        )
+        
+        blast_record = NCBIXML.read(handle)
+        handle.close()
+        
+        # 处理结果
+        hits = []
+        for idx, alignment in enumerate(blast_record.alignments[:hitlist_size]):
+            if not alignment.hsps:
+                continue
+            
+            hsp = alignment.hsps[0]
+            align_len = int(hsp.align_length)
+            identities = int(hsp.identities)
+            identity_pct = (identities / align_len * 100.0) if align_len else 0
+            
+            hits.append({
+                "rank": idx + 1,
+                "accession": alignment.accession,
+                "hit_id": alignment.hit_id,
+                "hit_def": alignment.hit_def,
+                "length": alignment.length,
+                "bit_score": float(hsp.bits),
+                "evalue": float(hsp.expect),
+                "identity_pct": round(identity_pct, 2),
+                "identities": identities,
+                "align_length": align_len
+            })
+        
+        result = {
+            "query_length": len(sequence),
+            "database": database,
+            "program": program,
+            "total_hits": len(hits),
+            "hits": hits
+        }
+        
+        with blast_lock:
+            blast_jobs[job_id]["status"] = "completed"
+            blast_jobs[job_id]["result"] = result
+            blast_jobs[job_id]["message"] = f"Completed: {len(hits)} hits found"
+            blast_jobs[job_id]["completed_at"] = time.time()
+        
+        print(f"✅ [Job {job_id}] BLAST completed: {len(hits)} hits")
+        
+    except Exception as e:
+        with blast_lock:
+            blast_jobs[job_id]["status"] = "failed"
+            blast_jobs[job_id]["error"] = str(e)
+            blast_jobs[job_id]["message"] = f"Failed: {str(e)}"
+        
+        print(f"❌ [Job {job_id}] BLAST failed: {e}")
+
+
+@mcp.tool(name="blast_search_async", description="Start BLAST search asynchronously (returns job_id immediately)")
+def blast_search_async(payload: dict):
+    """
+    异步启动 BLAST 搜索，立即返回 job_id
+    Required fields: sequence
+    Optional fields: database (refseq_rna/nt), program (blastn), hitlist_size (20), expect (1e-10)
+    """
+    try:
+        sequence = str(payload.get("sequence", "")).strip()
+        if not sequence:
+            return {"success": False, "result": None, "error": "sequence is required"}
+        
+        sequence = "".join(sequence.split()).upper()
+        
+        database = payload.get("database", "refseq_rna")
+        program = payload.get("program", "blastn")
+        hitlist_size = int(payload.get("hitlist_size", 20))
+        expect = float(payload.get("expect", 1e-10))
+        
+        # 生成唯一 job_id
+        job_id = f"blast_{int(time.time() * 1000)}"
+        
+        # 初始化 job 状态
+        with blast_lock:
+            blast_jobs[job_id] = {
+                "status": "pending",
+                "created_at": time.time(),
+                "sequence_length": len(sequence),
+                "database": database,
+                "message": "Job created, starting soon..."
+            }
+        
+        # 在后台线程中运行 BLAST
+        thread = threading.Thread(
+            target=_run_blast_async,
+            args=(job_id, sequence, database, program, hitlist_size, expect),
+            daemon=True
+        )
+        thread.start()
+        
+        print(f"🚀 BLAST job {job_id} started asynchronously")
+        
+        return {
+            "success": True,
+            "result": {
+                "job_id": job_id,
+                "status": "pending",
+                "message": "BLAST search started. Use blast_check_status to check progress.",
+                "estimated_time": "1-15 minutes depending on database"
+            },
+            "error": None
+        }
+        
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="blast_check_status", description="Check status of async BLAST job")
+def blast_check_status(payload: dict):
+    """
+    检查异步 BLAST 任务的状态
+    Required fields: job_id
+    """
+    try:
+        job_id = payload.get("job_id", "")
+        if not job_id:
+            return {"success": False, "result": None, "error": "job_id is required"}
+        
+        with blast_lock:
+            if job_id not in blast_jobs:
+                return {
+                    "success": False,
+                    "result": None,
+                    "error": f"Job {job_id} not found"
+                }
+            
+            job = blast_jobs[job_id].copy()
+        
+        # 计算运行时间
+        elapsed = time.time() - job["created_at"]
+        job["elapsed_seconds"] = round(elapsed, 1)
+        
+        return {
+            "success": True,
+            "result": job,
+            "error": None
+        }
+        
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="blast_get_result", description="Get result of completed BLAST job")
+def blast_get_result(payload: dict):
+    """
+    获取已完成的 BLAST 任务结果
+    Required fields: job_id
+    """
+    try:
+        job_id = payload.get("job_id", "")
+        if not job_id:
+            return {"success": False, "result": None, "error": "job_id is required"}
+        
+        with blast_lock:
+            if job_id not in blast_jobs:
+                return {
+                    "success": False,
+                    "result": None,
+                    "error": f"Job {job_id} not found"
+                }
+            
+            job = blast_jobs[job_id]
+            
+            if job["status"] != "completed":
+                return {
+                    "success": False,
+                    "result": None,
+                    "error": f"Job {job_id} is not completed yet. Status: {job['status']}"
+                }
+            
+            result = job.get("result")
+        
+        return {
+            "success": True,
+            "result": result,
+            "error": None
+        }
+        
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+# ==============================================
+# BLAST 搜索工具 - BLAST Search Tools (同步版本，保留用于小序列)
 # ==============================================
 
 @mcp.tool(name="blast_search", description="Search NCBI database using BLAST to find similar sequences")
@@ -467,6 +679,174 @@ def predict_species(payload: dict):
 # 完整流程工具 - Complete Workflow Tool
 # ==============================================
 
+@mcp.tool(name="identify_gene_species_async", 
+          description="Complete workflow with async BLAST - returns job_id immediately, check status separately")
+def identify_gene_species_async(payload: dict):
+    """
+    异步版本的完整基因物种鉴定流程
+    Required fields: sequence
+    Optional fields: min_identity (70), max_hits (20), database (refseq_rna)
+    
+    Returns job_id immediately. Use these tools to track progress:
+    1. blast_check_status(job_id) - Check if BLAST is done
+    2. continue_workflow_after_blast(job_id) - Complete remaining steps after BLAST finishes
+    """
+    try:
+        sequence = payload.get("sequence", "")
+        min_identity = float(payload.get("min_identity", 70.0))
+        max_hits = int(payload.get("max_hits", 20))
+        database = payload.get("database", "refseq_rna")
+        
+        # Step 1: Validate sequence
+        print("📋 Step 1: Validating sequence...")
+        val_result = validate_sequence({"sequence": sequence})
+        
+        if not val_result["success"]:
+            return {
+                "success": False,
+                "result": None,
+                "error": "Sequence validation failed"
+            }
+        
+        validated_seq = val_result["result"]["sequence"]
+        
+        # Step 2: Start async BLAST
+        print("🚀 Step 2: Starting async BLAST search...")
+        blast_async_result = blast_search_async({
+            "sequence": validated_seq,
+            "database": database,
+            "hitlist_size": max_hits
+        })
+        
+        if not blast_async_result["success"]:
+            return blast_async_result
+        
+        job_id = blast_async_result["result"]["job_id"]
+        
+        # Store workflow parameters for later continuation
+        with blast_lock:
+            if job_id in blast_jobs:
+                blast_jobs[job_id]["workflow_params"] = {
+                    "min_identity": min_identity,
+                    "max_hits": max_hits,
+                    "sequence_length": len(validated_seq)
+                }
+        
+        return {
+            "success": True,
+            "result": {
+                "job_id": job_id,
+                "status": "blast_running",
+                "message": "BLAST search started. Use 'blast_check_status' to monitor progress, then 'continue_workflow_after_blast' to complete analysis.",
+                "next_steps": [
+                    f"1. Check status: blast_check_status({{\"job_id\": \"{job_id}\"}})",
+                    f"2. When completed: continue_workflow_after_blast({{\"job_id\": \"{job_id}\"}})"
+                ]
+            },
+            "error": None
+        }
+        
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="continue_workflow_after_blast",
+          description="Continue species identification workflow after BLAST completes")
+def continue_workflow_after_blast(payload: dict):
+    """
+    在 BLAST 完成后继续工作流程
+    Required fields: job_id
+    """
+    try:
+        job_id = payload.get("job_id", "")
+        if not job_id:
+            return {"success": False, "result": None, "error": "job_id is required"}
+        
+        # Get BLAST result
+        blast_result_response = blast_get_result({"job_id": job_id})
+        if not blast_result_response["success"]:
+            return blast_result_response
+        
+        blast_result = blast_result_response["result"]
+        hits = blast_result["hits"]
+        
+        # Get workflow parameters
+        with blast_lock:
+            workflow_params = blast_jobs[job_id].get("workflow_params", {})
+        
+        min_identity = workflow_params.get("min_identity", 70.0)
+        
+        workflow_results = {"steps": []}
+        
+        # Step 3: Filter hits
+        print("🔬 Step 3: Filtering BLAST hits...")
+        filter_result = filter_blast_hits({
+            "hits": hits,
+            "min_identity": min_identity,
+            "max_hits": 10
+        })
+        workflow_results["steps"].append({
+            "step": 3,
+            "name": "filter_blast_hits",
+            "result": filter_result
+        })
+        
+        filtered_hits = filter_result["result"]["hits"]
+        
+        # Step 4: Extract species
+        print("🌍 Step 4: Extracting species information...")
+        species_result = extract_species_from_blast({"hits": filtered_hits})
+        workflow_results["steps"].append({
+            "step": 4,
+            "name": "extract_species_from_blast",
+            "result": species_result
+        })
+        
+        species_data = species_result["result"]["species_data"]
+        
+        # Step 5: Aggregate scores
+        print("📊 Step 5: Aggregating species scores...")
+        agg_result = aggregate_species_scores({"species_data": species_data})
+        workflow_results["steps"].append({
+            "step": 5,
+            "name": "aggregate_species_scores",
+            "result": agg_result
+        })
+        
+        species_rankings = agg_result["result"]["species_rankings"]
+        
+        # Step 6: Predict species
+        print("🎯 Step 6: Predicting species...")
+        pred_result = predict_species({"species_rankings": species_rankings})
+        workflow_results["steps"].append({
+            "step": 6,
+            "name": "predict_species",
+            "result": pred_result
+        })
+        
+        # Compile final result
+        final_result = {
+            "job_id": job_id,
+            "input_sequence_length": workflow_params.get("sequence_length"),
+            "total_blast_hits": len(hits),
+            "filtered_hits": len(filtered_hits),
+            "unique_species_found": len(species_rankings),
+            "prediction": pred_result["result"] if pred_result["success"] else None,
+            "workflow": workflow_results
+        }
+        
+        print("✅ Complete workflow finished successfully!")
+        
+        return {
+            "success": True,
+            "result": final_result,
+            "error": None
+        }
+        
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
 @mcp.tool(name="identify_gene_species_complete", 
           description="Complete workflow: input unknown gene sequence, identify species through BLAST search and analysis")
 def identify_gene_species_complete(payload: dict):