add docking score calculation in chemist agent

ChemCoScientist/agents/agents.py

@@ -1,4 +1,5 @@
 import ast
+import asyncio
 import os
 import time
 import json
@@ -7,6 +8,7 @@
 import operator
 import streamlit as st
 from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import ToolMessage
 
 from langgraph.types import Command
 from langgraph.graph import END
@@ -18,9 +20,9 @@
     automl_prompt,
     ds_builder_prompt,
     worker_prompt,
-    chem_ocr_prompt
+    chemist_prompt,
 )
-from ChemCoScientist.tools import chem_tools, nanoparticle_tools, paper_analysis_tools, data_tools, chem_ocr_tools
+from ChemCoScientist.tools import chem_tools, nanoparticle_tools, paper_analysis_tools, data_tools
 from ChemCoScientist.tools.ml_tools import agents_tools as automl_tools
 from ChemCoScientist.download_papers.functions import download_papers
 
@@ -241,8 +243,10 @@ def chemist_node(state: dict, config: dict) -> Command:
     plan = state["plan"]
     llm = config["configurable"]["llm"]
 
+    current_prompt = f'{chemist_prompt}\nPass {{"session_id": None}} as a parameter to the detect_molecules and detect_reactions tools'
+
     chem_agent = create_react_agent(
-        llm, chem_tools, state_modifier=worker_prompt + "admet = qed"
+        llm, chem_tools, state_modifier=current_prompt
     )
 
     task_formatted = f"""For the following plan:\n{str(plan)}\n\nYou are tasked with executing: {task}."""
@@ -251,6 +255,26 @@ def chemist_node(state: dict, config: dict) -> Command:
         try:
             config["configurable"]["state"] = state
             agent_response = chem_agent.invoke({"messages": [("user", task_formatted)]})
+
+            updated_metadata = state.get("metadata", {}).copy()
+            for message in agent_response["messages"]:
+                if isinstance(message, ToolMessage) and message.name in ["detect_molecules", "detect_reactions", "extract_reactions"]:
+                    result = ast.literal_eval(message.content)
+                    ocr_metadata = {"chem_ocr": result.get("metadata", None)}
+                    if ocr_metadata["chem_ocr"]:
+                        if "chem_ocr" in updated_metadata.keys():
+                            updated_metadata["chem_ocr"].update(ocr_metadata["chem_ocr"])
+                        else:
+                            updated_metadata.update(ocr_metadata)
+
+                elif isinstance(message, ToolMessage) and message.name in ["calculate_docking"]:
+                    result = json.loads(message.content)
+                    docking_metadata = {"docking": result.get("metadata", None)}
+                    if docking_metadata["docking"]:
+                        if "docking" in updated_metadata.keys():
+                            updated_metadata["docking"].update(docking_metadata["docking"])
+                        else:
+                            updated_metadata.update(docking_metadata)
 
             return Command(update={
                 "past_steps": Annotated[set, operator.or_](set([
@@ -262,6 +286,7 @@ def chemist_node(state: dict, config: dict) -> Command:
                         tuple((m.type, m.content) for m in agent_response["messages"])
                     )
                 ])),
+                "metadata": Annotated[dict, operator.or_](updated_metadata),
             })
 
         except Exception as e:
@@ -398,73 +423,6 @@ def paper_analysis_agent(state: dict, config: dict) -> Command:
     })
 
 
-def chem_ocr_agent(state: dict, config: dict) -> Command:
-    """
-    Extracts molecular structures and reaction information from images.
-
-    This agent processes user-provided chemical images—such as reaction schemes, 
-    drawn molecules, or figures from papers—and converts them into machine-readable 
-    formats. It attempts to identify molecular structures, reaction components, 
-    and other depicted chemical entities, returning standardized SMILES.
-
-    Args:
-        state (dict): The current state of the interaction, including images or PDFs provided by user.
-        config (dict): Configuration settings, including the OCR pipeline to use.
-
-    Returns:
-        Command: An object containing the next step in the process ('replan' or `END`) 
-        and updates to the state, including extracted SMILES, user images with detected chemical entities
-        and any error produced during parsing.
-    """
-    print("--------------------------------")
-    print("ChemOCR agent called")
-    print("Current task:")
-    print(state["task"])
-    print("--------------------------------")
-
-    llm: BaseChatModel = config["configurable"]["llm"]
-
-    task = state["task"]
-
-    chem_ocr_agent = create_react_agent(
-        llm, chem_ocr_tools, state_modifier=chem_ocr_prompt
-    )
-
-    for attempt in range(3):
-        try:
-            response = chem_ocr_agent.invoke({"messages": [("user", task)]})
-
-            result = ast.literal_eval(response["messages"][2].content)
-
-            answer_serialized = json.dumps(result["answer"], sort_keys=True)
-
-            updated_metadata = state.get("metadata", {}).copy()
-            ocr_metadata = {"chem_ocr": result.get("metadata", None)}
-            if ocr_metadata["chem_ocr"]:
-                if "chem_ocr" in updated_metadata.keys():
-                    updated_metadata["chem_ocr"].update(ocr_metadata["chem_ocr"])
-                else:
-                    updated_metadata.update(ocr_metadata)
-
-            return Command(update={
-                "past_steps": Annotated[set, operator.or_](set([
-                    (task, answer_serialized)
-                ])),
-                "nodes_calls": Annotated[set, operator.or_](set([
-                    ("chem_ocr_agent", (("text", answer_serialized),))
-                ])),
-                "metadata": Annotated[dict, operator.or_](updated_metadata),
-            })
-        except Exception as e:
-            print(f"ChemOCR agent error: {str(e)}. Retrying ({attempt + 1}/3)")
-            time.sleep(1.2 ** attempt)
-
-    return Command(goto=END, update={
-        "response": "I cannot extract molecules or reactions right now."
-                    "Can I help with something else?"
-    })
-
-
 def papers_search_agent(state: dict, config: dict) -> Command:
     """
     Searches for entity IDs or scientific papers based on user query and downloads papers' PDFs.

ChemCoScientist/agents/agents_prompts.py

@@ -73,6 +73,19 @@
 worker_prompt = "You are a helpful assistant. You can use provided tools. \
     If there is no appropriate tool, or you can't use one, answer yourself"
 
+
+chemist_prompt = f"""You are a specialized chemistry assistant with access to powerful molecular analysis tools. 
+Your role is to help users with chemical structure analysis, property calculations, and molecular data processing.
+
+WORKFLOW GUIDELINES:
+- Always start by understanding what the user needs: structure conversion, property calculation, visualization, or image analysis.
+- For property calculations, use smiles2prop to get comprehensive molecular descriptors.
+- For ADMET properties, note that admet = qed (QED - Quantitative Estimate of Drug-likeness is a key ADMET metric).
+- If a tool is not available or fails, provide helpful guidance based on your chemistry knowledge.
+- Always explain your reasoning and the results you obtain from tool calls.
+
+Remember: You are a chemistry expert. Use the tools effectively, but also leverage your knowledge to provide comprehensive answers and guidance."""
+
 paper_agent_prompt = """
 You are a helpful assistant. You can use provided tools. If there is no appropriate tool, or you can't use anyone, 
 answer yourself.

ChemCoScientist/chemical_utils/chemical_functions.py

@@ -0,0 +1,210 @@
+import requests
+from typing import List, Dict, Any, Callable
+from functools import wraps
+from dotenv import load_dotenv
+import os
+import logging
+import inspect
+from definitions import CONFIG_PATH
+
+load_dotenv(CONFIG_PATH)
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+CHEM_SERVICES_HOST = os.environ.get("CHEM_SERVICES_HOST")
+CHEM_SERVICES_PORT = os.environ.get("CHEM_SERVICES_PORT")
+CHEM_SERVICES_URL = f"http://{CHEM_SERVICES_HOST}:{CHEM_SERVICES_PORT}"
+REQUEST_TIMEOUT = 60
+
+
+def handle_api_request(endpoint: str, file_param_name: str = None, ):
+    """
+    Decorator for handling requests to Chemical ToolsService API.
+
+    Args:
+        endpoint (str): API endpoint path (e.g., "/extract_molecules_from_figure/")
+        file_param_name (str): Name of the file parameter in multipart/form-data (e.g., "image" or "pdf_file")
+
+    Returns:
+        A decorator that wraps a function and performs all necessary checks.
+    """
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        def wrapper(*args, **kwargs) -> Any:
+            """
+            Wrapper that executes API request and handles all errors.
+
+            For file uploads:
+                Args: file_data (bytes): File data (image or PDF)
+
+            For parameter requests:
+                Args: Any parameters passed to the function (e.g., smiles, pdb_id)
+
+            Returns:
+                Data from the "data" field of the API response
+            """
+            try:
+                api_url = f"{CHEM_SERVICES_URL}{endpoint}"
+                logger.info(f"Calling ChemService API: {api_url}")
+
+                if file_param_name:
+                    if args:
+                        file_data = args[0]
+                    elif file_param_name in kwargs:
+                        file_data = kwargs.pop(file_param_name)
+                    else:
+                        raise ValueError(f"File data must be provided as first argument or '{file_param_name}' keyword")
+
+                    response = requests.post(
+                        api_url,
+                        files={file_param_name: file_data},
+                        timeout=REQUEST_TIMEOUT
+                    )
+                else:
+                    params = {}
+                    if args:
+                        sig = inspect.signature(func)
+                        param_names = list(sig.parameters.keys())
+                        for i, arg in enumerate(args):
+                            if i < len(param_names):
+                                params[param_names[i]] = arg
+                    params.update(kwargs)
+
+                    response = requests.post(
+                        api_url,
+                        params=params,
+                        timeout=REQUEST_TIMEOUT
+                    )
+                if response.status_code != 200:
+                    error_msg = f"ChemService API returned status {response.status_code}: {response.text[:500]}"
+                    logger.error(error_msg)
+                    return {'errors': error_msg}
+
+                json_response = response.json()
+                if json_response is None:
+                    error_msg = "ChemService API returned None JSON response"
+                    logger.error(error_msg)
+                    return {'errors': error_msg}
+
+                if "data" not in json_response:
+                    error_msg = f"ChemService API response missing 'data' field. Response: {json_response}"
+                    logger.error(error_msg)
+                    return {'errors': error_msg}
+
+                return json_response
+
+            except requests.exceptions.RequestException as e:
+                error_msg = f"Failed to connect to ChemService API at {CHEM_SERVICES_URL}: {str(e)}"
+                logger.error(error_msg)
+                return {'errors': error_msg}
+        return wrapper
+    return decorator
+
+
+@handle_api_request(endpoint="/extract_reactions_from_pdf/", file_param_name="pdf_file")
+def extract_reactions_from_pdf(file: bytes) -> List[Dict]:
+    """
+    Extract reactions information from a PDF file.
+    Response contains list of reactions for each page of the PDF.
+    Each reaction contains list of reactants, products and conditions.
+
+    Args:
+        file (bytes): PDF file to extract reactions from.
+    Returns:
+        response (List[Dict]): List of reactions in pdf file for each page.
+    Raises:
+        ConnectionError: If API is unavailable or connection fails.
+        ValueError: If API returns invalid response.
+        RuntimeError: For unexpected errors.
+    """
+    pass
+
+
+@handle_api_request(endpoint="/extract_reactions_from_figure/", file_param_name="image")
+def extract_reactions_from_figure(image: bytes) -> List[Dict]:
+    """
+    Extract reactions information from an image.
+
+    Response contains list of reactions on the image.
+    Each reaction contains list of reactants, products and conditions.
+    Args:
+        image (bytes): Image to extract reactions from.
+    Returns:
+        response (List[Dict]): List of reactions on the image.
+    """
+    pass
+
+
+@handle_api_request(endpoint="/extract_molecules_from_pdf/", file_param_name="pdf_file")
+def extract_molecules_from_pdf(file: bytes) -> List[Dict]:
+    """
+    Extract molecules information from a PDF file.
+    Response contains list of molecules for each page of the PDF.
+    Each molecule contains bbox and smiles.
+
+    Args:
+        file (bytes): PDF file to extract molecules from.
+    Returns:
+        response (List[Dict]): List of molecules in pdf file for each page.
+    """
+    pass
+
+
+@handle_api_request(endpoint="/extract_molecules_from_figure/", file_param_name="image")
+def extract_molecules_from_figure(image: bytes) -> List[Dict]:
+    """
+    Extract molecules information from an image.
+    Response contains list of molecules on the image.
+    Each molecule contains bbox and smiles.
+
+    Args:
+        image (bytes): Image to extract molecules from.
+    Returns:
+        response (List[Dict]): List of molecules on the image.
+    """
+    pass
+
+
+@handle_api_request(endpoint="/convert_image_to_smiles/", file_param_name="image")
+def convert_image_to_smiles(image: bytes) -> str:
+    """
+    Convert an image to a smiles string.
+    Response contains smiles string of the image.
+    Args:
+        image (bytes): Image to convert to smiles.
+    Returns:
+        response (str): SMILES string of the image.
+    """
+    pass
+
+@handle_api_request(endpoint="/docking/", file_param_name=None)
+def calculate_docking_score(smiles: str, pdb_id: str) -> str:
+    """
+    Calculate docking score for a molecule.
+    Response contains docking score for the molecule.
+    Args:
+        smiles (str): SMILES string of the molecule.
+        pdb_id (str): ID of the PDB file containing the receptor structure.
+    Returns:
+        response (str): Docking score for the molecule.
+    """
+    pass
+
+
+def remove_keys(obj: Any, keys_to_remove: set[str] = {"bbox", "score"}) -> Any:
+    """Processes ChemService json output to remove unnecessary keys like 'score' and 'bbox'."""
+    if isinstance(obj, dict):
+        for k in keys_to_remove:
+            obj.pop(k, None)
+        for v in obj.values():
+            remove_keys(v, keys_to_remove)
+    elif isinstance(obj, list):
+        for item in obj:
+            remove_keys(item, keys_to_remove)
+    return obj
+
+
+if __name__ == "__main__":
+    result = calculate_docking_score(smiles="C1CCCCC1", pdb_id="5vfi")
+    print(result)