feat(esm mcp): add get_prediction_pdb tool to download saved PDB (supports latest)

esm/mcp_output/mcp_plugin/mcp_service.py

@@ -150,20 +150,8 @@ def predict_structure(sequence: str):
             pdb_io = io.StringIO(response.text)
             structure = parser.get_structure("esmfold_prediction", pdb_io)
 
-            # Prefer persistent storage in Space (e.g., /data), fallback to local package dir
-            default_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "predictions")
-            preferred_dir = "/data/predictions"
-            predictions_dir = preferred_dir
-            try:
-                os.makedirs(predictions_dir, exist_ok=True)
-                # Test write permission with a temp touch
-                _probe_path = os.path.join(predictions_dir, ".write_probe")
-                with open(_probe_path, "w") as _probe:
-                    _probe.write("ok")
-                os.remove(_probe_path)
-            except Exception:
-                predictions_dir = default_dir
-                os.makedirs(predictions_dir, exist_ok=True)
+            predictions_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "predictions")
+            os.makedirs(predictions_dir, exist_ok=True)
 
             timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
             pdb_filename = f"prediction_{timestamp}.pdb"
@@ -265,6 +253,52 @@ def validate_protein_sequence(sequence: str):
         }
 
 
+@mcp.tool(name="get_prediction_pdb", description="Download saved PDB by filename or latest")
+def get_prediction_pdb(filename: str):
+    """Return PDB content saved under predictions directory.
+
+    Parameters:
+        filename (str): Filename in predictions directory. Use "latest" to fetch the most recent file.
+
+    Returns:
+        dict: success/result/error. result contains { filename, pdb_content, path }
+    """
+    try:
+        import glob
+        predictions_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "predictions")
+        if not os.path.isdir(predictions_dir):
+            return {"success": False, "result": None, "error": f"Predictions dir not found: {predictions_dir}"}
+
+        target_path: str
+        if filename.strip().lower() == "latest":
+            files = sorted(
+                glob.glob(os.path.join(predictions_dir, "*.pdb")),
+                key=lambda p: os.path.getmtime(p),
+                reverse=True,
+            )
+            if not files:
+                return {"success": False, "result": None, "error": "No PDB files found"}
+            target_path = files[0]
+            filename = os.path.basename(target_path)
+        else:
+            # prevent path traversal
+            safe_name = os.path.basename(filename)
+            target_path = os.path.join(predictions_dir, safe_name)
+            if not os.path.isfile(target_path):
+                return {"success": False, "result": None, "error": f"File not found: {safe_name}"}
+
+        with open(target_path, "r") as f:
+            content = f.read()
+
+        return {
+            "success": True,
+            "result": {"filename": filename, "pdb_content": content, "path": target_path},
+            "error": None,
+        }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
 def create_app():
     """
     Create and return a FastMCP instance.