chore(retrosynthesis): modification tools

Author: AerDragon

ChemCoScientist/agents/agents.py

@@ -251,10 +251,23 @@ def chemist_node(state: dict, config: dict) -> Command:
             config["configurable"]["state"] = state
             agent_response = chem_agent.invoke({"messages": [("user", task_formatted)]})
 
+            def _parse_tool_content(content):
+                try:
+                    return json.loads(content)
+                except Exception:
+                    try:
+                        return ast.literal_eval(content)
+                    except Exception:
+                        return None
+
             updated_metadata = state.get("metadata", {}).copy()
             for message in agent_response["messages"]:
+                if isinstance(message, ToolMessage):
+                    print(f"TOOL MESSAGE: {message.name}")
                 if isinstance(message, ToolMessage) and message.name in ["detect_molecules", "detect_reactions"]:
-                    result = ast.literal_eval(message.content)
+                    result = _parse_tool_content(message.content)
+                    if result is None:
+                        continue
                     ocr_metadata = {"chem_ocr": result.get("metadata", None)}
                     if ocr_metadata["chem_ocr"]:
                         if "chem_ocr" in updated_metadata.keys():
@@ -263,7 +276,9 @@ def chemist_node(state: dict, config: dict) -> Command:
                             updated_metadata.update(ocr_metadata)
 
                 elif isinstance(message, ToolMessage) and message.name in ["calculate_docking_score"]:
-                    result = ast.literal_eval(message.content)
+                    result = _parse_tool_content(message.content)
+                    if result is None:
+                        continue
                     docking_metadata = {"docking": result.get("metadata", None)}
                     if docking_metadata["docking"]:
                         if "docking" in updated_metadata.keys():
@@ -272,23 +287,22 @@ def chemist_node(state: dict, config: dict) -> Command:
                             updated_metadata.update(docking_metadata)
 
                 elif isinstance(message, ToolMessage) and message.name in ["retrosynthesis_tree_search"]:
-                    try:
-                        result = ast.literal_eval(message.content)
-                    except Exception:
+                    result = _parse_tool_content(message.content)
+                    if result is None:
                         continue
+                    if isinstance(result, dict):
+                        print(f"RETRO RESULT KEYS: {list(result.keys())[:10]}")
                     updated_metadata.update({"retrosynthesis": result})
 
                 elif isinstance(message, ToolMessage) and message.name in ["classify_reaction"]:
-                    try:
-                        result = ast.literal_eval(message.content)
-                    except Exception:
+                    result = _parse_tool_content(message.content)
+                    if result is None:
                         continue
                     updated_metadata.update({"reaction_classification": result})
 
                 elif isinstance(message, ToolMessage) and message.name in ["forward_predict"]:
-                    try:
-                        result = ast.literal_eval(message.content)
-                    except Exception:
+                    result = _parse_tool_content(message.content)
+                    if result is None:
                         continue
                     updated_metadata.update({"forward_prediction": result})
 

ChemCoScientist/agents/agents_prompts.py

@@ -127,4 +127,4 @@
 You must detect and output every plausible chemical structure present in the image, even if the image is low-quality,
 sketchy, partial, or ambiguous. When uncertain, infer the most likely structure based on visible atoms, bonds, and geometry.
 Never return ‘no molecules detected’—instead describe all candidate structures with confidence scores.
-"""
+"""

ChemCoScientist/chemical_utils/retrosynthesis.py

@@ -129,8 +129,8 @@ def classify_reaction_smiles(smiles: List[str], num_results: int = 10) -> Dict[s
 
 def forward_predict_products(
     smiles: List[str],
-    backend: str,
-    model_name: str = "wldn5",
+    backend: str = "wldn5",
+    model_name: str = "pistachio",
     reagents: str = "",
     solvent: str = "",
 ) -> Dict[str, Any]:
@@ -178,6 +178,7 @@ def forward_predict_products(
             error_msg = "Forward Prediction API returned None JSON response"
             logger.error(error_msg)
             raise ValueError(error_msg)
+        logger.info(f"FORWARD PREDICTION JSON RESPONSE: {json_response}")
         return json_response
     except requests.exceptions.RequestException as e:
         error_msg = (

ChemCoScientist/frontend/chat.py

@@ -5,6 +5,8 @@
 import os
 import streamlit as st
 import threading
+import re
+from typing import Optional, List
 from typing import Optional
 
 from io import BytesIO
@@ -356,6 +358,8 @@ def message_handler(user_query: str, placeholder: st.delta_generator.DeltaGenera
                             #os.remove(file)
 
                     # Store metadata in the message for later display
+                    
+                    logger.info(f"RESULT METADATA: {result}")
                     if "paper_analysis" in result["metadata"].keys():
                         st.session_state.messages[-1]["paper_analysis"] = result["metadata"]["paper_analysis"]
                         # Display the metadata immediately after storing it
@@ -452,10 +456,13 @@ def _reaction_smiles_to_image(reaction_smiles: str):
 
 
 def _render_reaction_smiles(reaction_smiles: str, caption: Optional[str] = None):
+    logger.info(f"RENDERING REACTION SMILES: {reaction_smiles}")
     normalized = reaction_smiles
     if isinstance(normalized, str):
         normalized = normalized.replace(" -> ", ">>").replace(" → ", ">>")
         normalized = normalized.replace(" + ", ".")
+        normalized = re.sub(r">{3,}", ">>", normalized)
+        normalized = re.sub(r">>\s*>>", ">>", normalized)
     img = _reaction_smiles_to_image(normalized)
     if img is not None:
         st.image(img, caption=caption)
@@ -464,6 +471,7 @@ def _render_reaction_smiles(reaction_smiles: str, caption: Optional[str] = None)
 
 
 def display_retrosynthesis_metadata(message):
+    logger.info(f"DISPLAYING RETROSYNTHESIS METADATA: {message}")
     data = message.get("retrosynthesis") or {}
     if not data:
         return
@@ -522,10 +530,13 @@ def _route_sort_key(r):
                 continue
             for step_idx, step in enumerate(steps, start=1):
                 reaction_smiles = step.get("reaction_smiles") or step.get("mapped_smiles")
+                logger.info(f"REACTION SMILES: {reaction_smiles}")
+                logger.info(f"STEP: {step}")
                 caption = f"Step {step_idx}"
                 if step.get("plausibility") is not None:
                     caption += f" | plausibility={step.get('plausibility')}"
                 if reaction_smiles:
+                    logger.info(f"RENDERING REACTION SMILES: {reaction_smiles}")
                     _render_reaction_smiles(reaction_smiles, caption=caption)
                 else:
                     st.markdown(f"**Step {step_idx}**")
@@ -548,10 +559,6 @@ def display_forward_prediction_metadata(message):
     if backend or model_name:
         st.markdown(f"**Backend:** `{backend}`  **Model:** `{model_name}`")
     inputs = data.get("inputs") or []
-    if inputs:
-        st.markdown("**Inputs:**")
-        for item in inputs:
-            st.code(item)
     predictions = data.get("predictions") or []
     if predictions and isinstance(predictions, list):
         predictions = sorted(
@@ -564,16 +571,16 @@ def display_forward_prediction_metadata(message):
         return
     if len(inputs) == 1:
         base = inputs[0]
-        for idx, pred in enumerate(predictions, start=1):
+        pred = predictions[0] if predictions else None
+        if pred:
             prod = pred.get("smiles")
             score = pred.get("score")
-            if not prod:
-                continue
-            reaction_smiles = f"{base}>>{prod}"
-            caption = f"Prediction {idx}"
-            if score is not None:
-                caption += f" | score={score}"
-            _render_reaction_smiles(reaction_smiles, caption=caption)
+            if prod:
+                reaction_smiles = f"{base}>>{prod}"
+                caption = "Best prediction"
+                if score is not None:
+                    caption += f" | score={score}"
+                _render_reaction_smiles(reaction_smiles, caption=caption)
     else:
         rows = [{"smiles": p.get("smiles"), "score": p.get("score")} for p in predictions]
         st.dataframe(rows)