Update biopython/mcp_output/mcp_plugin/mcp_service.py

biopython/mcp_output/mcp_plugin/mcp_service.py

@@ -6,8 +6,14 @@ sys.path.insert(0, source_path)
 
 from fastmcp import FastMCP
 from Bio.SeqIO import parse, read, write
+from Bio.Seq import Seq
+from Bio.SeqRecord import SeqRecord
 from Bio.Blast.NCBIWWW import qblast
-from Bio.Entrez import efetch, esearch
+from Bio.Entrez import efetch, esearch, email
+import Bio.Entrez
+
+# 设置 NCBI Entrez email（必需）
+Bio.Entrez.email = "biopython-mcp@huggingface.co"
 
 mcp = FastMCP("biopython_service")
 
@@ -30,6 +36,25 @@ def seqrecord_to_dict(record):
         ] if record.features else []
     }
 
+def dict_to_seqrecord(seq_dict):
+    """Convert a dictionary back to a SeqRecord object."""
+    if isinstance(seq_dict, SeqRecord):
+        return seq_dict
+    
+    # 从字典创建 SeqRecord
+    record = SeqRecord(
+        Seq(seq_dict.get("sequence", "")),
+        id=seq_dict.get("id", ""),
+        name=seq_dict.get("name", ""),
+        description=seq_dict.get("description", "")
+    )
+    
+    # 添加注释
+    if "annotations" in seq_dict:
+        record.annotations.update(seq_dict["annotations"])
+    
+    return record
+
 @mcp.tool(name="seqio_parse", description="Parse sequence data from a file.")
 def seqio_parse(file_path: str, format: str) -> dict:
     """
@@ -76,7 +101,7 @@ def seqio_write(sequences, file_path: str, format: str) -> dict:
     Writes sequences to a file.
 
     Parameters:
-    - sequences: List of sequences to write.
+    - sequences: List of sequences to write (can be SeqRecord objects or dicts).
     - file_path: Path to the output file.
     - format: Format of the output file (e.g., 'fasta').
 
@@ -84,7 +109,22 @@ def seqio_write(sequences, file_path: str, format: str) -> dict:
     - A dictionary with success status and number of records written or error message.
     """
     try:
-        count = write(sequences, file_path, format)
+        # 如果 sequences 是列表，将每个元素转换为 SeqRecord
+        if isinstance(sequences, list):
+            records = []
+            for seq in sequences:
+                if isinstance(seq, dict):
+                    records.append(dict_to_seqrecord(seq))
+                else:
+                    records.append(seq)
+        else:
+            # 如果是单个对象
+            if isinstance(sequences, dict):
+                records = [dict_to_seqrecord(sequences)]
+            else:
+                records = [sequences]
+        
+        count = write(records, file_path, format)
         return {"success": True, "result": count, "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
@@ -134,15 +174,32 @@ def entrez_esearch(db: str, term: str) -> dict:
     Searches NCBI's Entrez databases.
 
     Parameters:
-    - db: Database to search (e.g., 'pubmed').
-    - term: Search term.
+    - db: Database to search (e.g., 'nucleotide', 'protein', 'gene', 'pubmed').
+    - term: Search term. Use proper NCBI query syntax:
+        - gene_name[GENE] - Search by gene name
+        - organism[ORGN] - Search by organism
+        - "Homo sapiens"[ORGN] - Species filter
+        - "RefSeq"[Filter] - Only RefSeq records
 
     Returns:
     - A dictionary with success status and search results or error message.
+    
+    Examples:
+    - "TP53[gene] AND human[organism]"
+    - "BRCA1[gene]"
+    - "insulin[protein name]"
     """
     try:
-        handle = esearch(db=db, term=term)
-        return {"success": True, "result": handle.read(), "error": None}
+        # 确保 email 已设置
+        if not Bio.Entrez.email:
+            Bio.Entrez.email = "biopython-mcp@huggingface.co"
+        
+        # 使用 retmax 参数获取更多结果
+        handle = esearch(db=db, term=term, retmax=100)
+        result = handle.read()
+        handle.close()
+        
+        return {"success": True, "result": result, "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}