File size: 6,428 Bytes
a1e520f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 |
# Standard library imports
import os
import subprocess
import tempfile
import traceback
from pathlib import Path
from typing import Dict, Any, Optional
# Third-party imports
import gradio as gr
from fastapi import FastAPI, File, Form, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from jsonschema import validate, ValidationError
from schema import MCP_CONTEXT_SCHEMA, MCP_PREDICT_RESPONSE_SCHEMA
# Local imports
from utils import validate_mime_type, format_results_to_mcp, get_mcp_context, ALLOWED_MIME_TYPES
app = FastAPI(title="PoseBusters MCP API")
@app.post("/mcp/predict")
async def predict(
action: str = Form(...),
ligand_input: UploadFile = File(...),
protein_input: UploadFile = File(...),
crystal_input: Optional[UploadFile] = File(None),
) -> Dict[str, Any]:
"""
MCP-compliant prediction endpoint.
Validates file types and runs PoseBusters validation.
Response is validated against MCP_PREDICT_RESPONSE_SCHEMA.
"""
try:
# Validate MIME types
validate_mime_type(ligand_input, ".sdf")
validate_mime_type(protein_input, ".pdb")
if crystal_input:
validate_mime_type(crystal_input, ".sdf")
async def save_upload(upload: UploadFile, suffix: str) -> str:
"""Safely save uploaded file with secure naming."""
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
content = await upload.read()
tmp.write(content)
return tmp.name
# Create a temporary directory for isolated execution
with tempfile.TemporaryDirectory() as tmpdir:
try:
# Save files in the temporary directory
ligand_path = await save_upload(ligand_input, ".sdf")
protein_path = await save_upload(protein_input, ".pdb")
crystal_path = await save_upload(crystal_input, ".sdf") if crystal_input else None
# Move to temporary directory
os.chdir(tmpdir)
if action == "validate_pose":
cmd = ["bust", ligand_path, "-p", protein_path, "--outfmt", "csv"]
elif action == "redocking_validation":
if not crystal_path:
raise HTTPException(status_code=400, detail="Missing crystal ligand file")
cmd = ["bust", ligand_path, "-l", crystal_path, "-p", protein_path, "--outfmt", "csv"]
else:
raise HTTPException(status_code=400, detail=f"Unknown action: {action}")
# Use list form of subprocess.run to avoid shell injection
# Add timeout of 5 minutes
result = subprocess.run(
cmd,
capture_output=True,
text=True,
check=False, # Don't raise on non-zero exit
timeout=300 # 5 minutes timeout
)
response = format_results_to_mcp(
result.stdout,
result.stderr,
Path(ligand_input.filename).stem
)
try:
validate(instance=response, schema=MCP_PREDICT_RESPONSE_SCHEMA)
return response
except ValidationError as e:
return JSONResponse(
status_code=500,
content={
"object_id": "validation_results",
"data": {
"columns": ["ligand_id", "status", "passed/total", "details"],
"rows": [[
Path(ligand_input.filename).stem,
"❌",
"0/0",
f"Schema validation error: {str(e)}"
]]
}
}
)
finally:
# Clean up temporary files
for path in [ligand_path, protein_path, crystal_path]:
if path and os.path.exists(path):
try:
os.unlink(path)
except OSError:
pass
except subprocess.TimeoutExpired as e:
return JSONResponse(
status_code=500,
content={
"object_id": "validation_results",
"data": {
"columns": ["ligand_id", "status", "passed/total", "details"],
"rows": [[
Path(ligand_input.filename).stem,
"❌",
"0/0",
f"Process timed out after {e.timeout} seconds"
]]
}
}
)
except HTTPException:
raise
except Exception as e:
return JSONResponse(
status_code=500,
content={
"object_id": "validation_results",
"data": {
"columns": ["ligand_id", "status", "passed/total", "details"],
"rows": [[
Path(ligand_input.filename).stem,
"❌",
"0/0",
str(e)
]]
},
"error": str(e),
"traceback": traceback.format_exc()
}
)
@app.get("/mcp/context")
def context() -> Dict[str, Any]:
"""
Return MCP context with empty initial validation results.
Response is validated against MCP_CONTEXT_SCHEMA.
"""
try:
context = get_mcp_context()
validate(instance=context, schema=MCP_CONTEXT_SCHEMA)
return context
except ValidationError as e:
return JSONResponse(
status_code=500,
content={"error": f"Schema validation error: {str(e)}"}
)
def main():
"""Run the FastAPI server."""
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)
if __name__ == "__main__":
main() |