NousResearch/atropos
1
0
Fork 0
mirror of https://github.com/NousResearch/atropos.git synced 2026-04-19 12:57:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

additional fixes to alphafold2_multimer and tool_executor

Browse source
This commit is contained in:
based-tachikoma 2025-05-21 21:28:30 -07:00
parent 6783a077cc
commit b01023ad3a
2 changed files with 30 additions and 56 deletions
Download patch file Download diff file Expand all files Collapse all files

38
environments/hack0/protein_design_env/tool_executor.py
View file

@ -97,7 +97,6 @@ class ToolExecutor:
if self.debug_protein_design_calls:
logger.warning(f"DEBUG MODE: Bypassing AlphaFold2 API call for workflow {item_id}.")
# Use a relative path within the package
module_dir = Path(__file__).parent
fixed_pdb_path = module_dir / "debug_target.pdb"
@ -352,8 +351,7 @@ class ToolExecutor:
self._debug_af2m_call_count += 1
mock_plddt = 87.5 if self._debug_af2m_call_count % 2 == 1 else 45.2
success_message = f"DEBUG MODE: Returning {'high' if mock_plddt > 50 else 'low'}-quality mock results (call #{self._debug_af2m_call_count})"
# In debug mode, ToolExecutor still handles file saving
debug_pdb_filename = f"complex_{item_id}_s{current_internal_step}_af2m_DEBUG_pLDDT{mock_plddt:.2f}.pdb"
debug_pdb_path = self.output_dir / debug_pdb_filename
try:
@ -377,7 +375,6 @@ class ToolExecutor:
}
return {"tool_output": tool_output, "state_updates": state_updates}
# Call AF2-Multimer - no output_dir passed here
api_result = await call_alphafold2_multimer(
sequences=all_input_sequences_for_multimer,
api_key=self.nim_api_key,
@ -386,34 +383,28 @@ class ToolExecutor:
polling_interval=self.polling_interval
)
# Check for explicit failure from the call_alphafold2_multimer function
if api_result is None or (isinstance(api_result, dict) and api_result.get("success") is False):
error_detail = "AF2-Multimer call failed or returned None."
if isinstance(api_result, dict):
error_detail = api_result.get("error", "AF2-Multimer call failed with unspecified error.")
detail_info = api_result.get("detail", "")
if detail_info: error_detail += f" Details: {detail_info}"
logger.error(f"Workflow {item_id}: AF2-Multimer call failed: {error_detail}. API Result: {api_result}")
tool_output = {"success": False, "error": error_detail}
state_updates["complex_evaluated"] = False
return {"tool_output": tool_output, "state_updates": state_updates}
# api_result should now be like: {"structures": [{"model_index": ..., "pdb_content": "...", "average_plddt": ...}, ...]}
# or {"success": True, "structures": [...]}
all_structures_info = api_result.get("structures")
if not all_structures_info or not isinstance(all_structures_info, list):
# This case covers if _process_pdb_and_scores_from_api returned success but empty structures
# or if the structure of api_result is unexpected
message = api_result.get("message", "No structures returned from AF2-Multimer process.")
logger.warning(f"Workflow {item_id}: {message}. API Result: {api_result}")
if not all_structures_info and isinstance(all_structures_info, list): # Empty list of structures
if not all_structures_info and isinstance(all_structures_info, list):
tool_output = {"success": True, "message": "AF2-Multimer ran, but no PDB structures were produced by the API.", "plddt": 0.0, "complex_file_path": None}
state_updates["af2_multimer_plddt"] = 0.0
state_updates["complex_evaluated"] = True # Evaluated, but with no result
state_updates["complex_evaluated"] = True
state_updates["complex_pdb_content_path"] = None
else: # Malformed result
else:
tool_output = {"success": False, "error": "AF2-Multimer returned unexpected data or no structures."}
state_updates["complex_evaluated"] = False
return {"tool_output": tool_output, "state_updates": state_updates}
@ -434,19 +425,17 @@ class ToolExecutor:
state_updates["complex_evaluated"] = False
return {"tool_output": tool_output, "state_updates": state_updates}
# Now, save the PDB content of the best structure
best_pdb_content = best_structure_info.get("pdb_content")
best_plddt = best_structure_info.get("average_plddt", 0.0) # Should be same as highest_plddt
best_model_idx = best_structure_info.get("model_index", "NA") # Use NA if not found
best_plddt = best_structure_info.get("average_plddt", 0.0)
best_model_idx = best_structure_info.get("model_index", "NA")
if not best_pdb_content:
logger.error(f"Workflow {item_id}: Best AF2-Multimer structure (Model {best_model_idx}, pLDDT {best_plddt:.2f}) found, but PDB content is missing.")
tool_output = {"success": False, "error": f"Best model (pLDDT {best_plddt:.2f}) has no PDB content."}
state_updates["complex_evaluated"] = False # Or True with pLDDT, but no path
state_updates["complex_evaluated"] = False
state_updates["af2_multimer_plddt"] = best_plddt
return {"tool_output": tool_output, "state_updates": state_updates}
# Construct filename and save
complex_pdb_filename = f"complex_{item_id}_s{current_internal_step}_af2m_model{best_model_idx}_pLDDT{best_plddt:.2f}.pdb"
complex_pdb_path = self.output_dir / complex_pdb_filename
@ -454,7 +443,7 @@ class ToolExecutor:
with open(complex_pdb_path, "w", encoding='utf-8') as f:
f.write(best_pdb_content)
logger.info(f"Workflow {item_id}: AlphaFold2-Multimer complete. Saved best model (Index {best_model_idx}) with pLDDT: {best_plddt:.2f} from {len(all_structures_info)} models to {complex_pdb_path}")
state_updates["complex_pdb_content_path"] = str(complex_pdb_path)
state_updates["af2_multimer_plddt"] = best_plddt
state_updates["complex_evaluated"] = True
@ -471,11 +460,10 @@ class ToolExecutor:
except IOError as e:
logger.error(f"Workflow {item_id}: Failed to save best AF2-Multimer PDB (Model {best_model_idx}, pLDDT {best_plddt:.2f}) to {complex_pdb_path}: {e}")
tool_output = {"success": False, "error": f"Failed to save best complex PDB: {e}"}
# Still record the pLDDT for reward, even if saving failed
state_updates["af2_multimer_plddt"] = best_plddt
state_updates["complex_pdb_content_path"] = None # Path is not valid
state_updates["complex_evaluated"] = True # It was evaluated, saving failed
state_updates["complex_pdb_content_path"] = None
state_updates["complex_evaluated"] = True
return {"tool_output": tool_output, "state_updates": state_updates}
@ -504,4 +492,4 @@ class ToolExecutor:
return {
"tool_output": {"success": False, "error": f"Unknown tool name: {tool_name}"},
"state_updates": {}
}
}