From bcf22697e3fb1f4cf29c56b53b54dceb5f697957 Mon Sep 17 00:00:00 2001 From: rodiongolovinsky Date: Tue, 13 Jan 2026 18:19:05 +0300 Subject: [PATCH 1/8] add docking score calculation in chemist agent --- ChemCoScientist/conf/create_conf.py | 2 ++ pyproject.toml | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/ChemCoScientist/conf/create_conf.py b/ChemCoScientist/conf/create_conf.py index 2ee137b7..32a5ea54 100644 --- a/ChemCoScientist/conf/create_conf.py +++ b/ChemCoScientist/conf/create_conf.py @@ -159,6 +159,8 @@ and an empty or absent `metadata.papers`. """ +======= +>>>>>>> beee5da (add docking score calculation in chemist agent) additional_agents_description = ( automl_agent_description diff --git a/pyproject.toml b/pyproject.toml index e899f554..e75ddc92 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,4 +53,4 @@ requires = ["poetry-core>=2.0.0,<3.0.0"] build-backend = "poetry.core.masonry.api" [tool.poetry] -package-mode = false \ No newline at end of file +package-mode = false From c57fe6558a4fd58a64e695ca9ca2576df46c5f84 Mon Sep 17 00:00:00 2001 From: rodiongolovinsky Date: Sun, 8 Feb 2026 19:12:07 +0300 Subject: [PATCH 2/8] add MCP server for chemical tools --- ChemCoScientist/agents/agents.py | 5 +- ChemCoScientist/mcp/chemical_server.py | 677 +++++++++++++++++++++++++ 2 files changed, 681 insertions(+), 1 deletion(-) create mode 100644 ChemCoScientist/mcp/chemical_server.py diff --git a/ChemCoScientist/agents/agents.py b/ChemCoScientist/agents/agents.py index a0b2566f..c39e81d6 100644 --- a/ChemCoScientist/agents/agents.py +++ b/ChemCoScientist/agents/agents.py @@ -9,6 +9,8 @@ import streamlit as st from langchain_core.language_models import BaseChatModel from langchain_core.messages import ToolMessage +import logging +from langchain_mcp_adapters.client import MultiServerMCPClient from langgraph.types import Command from langgraph.graph import END @@ -244,9 +246,10 @@ def chemist_node(state: dict, config: dict) -> Command: llm = config["configurable"]["llm"] current_prompt = f'{chemist_prompt}\nPass {{"session_id": None}} as a parameter to the detect_molecules and detect_reactions tools' + chemical_tools = _get_chemical_mcp_tools() chem_agent = create_react_agent( - llm, chem_tools, state_modifier=current_prompt + llm, chemical_tools, state_modifier=current_prompt ) task_formatted = f"""For the following plan:\n{str(plan)}\n\nYou are tasked with executing: {task}.""" diff --git a/ChemCoScientist/mcp/chemical_server.py b/ChemCoScientist/mcp/chemical_server.py new file mode 100644 index 00000000..f62fa80d --- /dev/null +++ b/ChemCoScientist/mcp/chemical_server.py @@ -0,0 +1,677 @@ +import base64 +from fastmcp import FastMCP +import logging +from pathlib import Path +import os +from ChemCoScientist.chemical_utils.ocr_pipeline import * +from ChemCoScientist.chemical_utils.chemical_functions import * +import os +from typing import Annotated, Optional, List, Dict +from urllib.parse import quote + +import pubchempy as pcp +import py3Dmol +import rdkit.Chem as Chem +import requests +from langchain_core.runnables.config import RunnableConfig +from langchain_experimental.utilities import PythonREPL +from rdkit.Chem import AllChem +from rdkit.Chem.Descriptors import CalcMolDescriptors +from typing import Dict, List, Optional +from definitions import CONFIG_PATH +from pathlib import Path +from dotenv import load_dotenv +from ChemCoScientist.chemical_utils.ocr_pipeline import molecules_ocr, reactions_ocr +from ChemCoScientist.chemical_utils.chemical_functions import calculate_docking_score + +import aiohttp +import base64 +import asyncio +import json +import re +import pandas as pd +from io import StringIO + +load_dotenv(CONFIG_PATH) + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +CHEMBL_BASE = "https://www.ebi.ac.uk/chembl/api/data" +VALID_AFFINITY_TYPES = {"Ki", "Kd", "IC50", "EC50"} +repl = PythonREPL() + + +mcp = FastMCP("ChemTools") + + +def _run_async(coro): + """Run async coroutine from both sync and async contexts.""" + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + + if loop and loop.is_running(): + # running in async env (e.g. LangChain) + return asyncio.ensure_future(coro) + else: + # safe to call asyncio.run() + return asyncio.run(coro) + + +async def fetch_uniprot_id( + session: aiohttp.ClientSession, + protein_name: str, + organism_id: int = 9606, + max_retries: int = 5, + delay: float = 0.5 +) -> Optional[str]: + """ + Asynchronously fetch UniProt ID for a given protein name. + Retries up to `max_retries` times in case of network or transient API errors. + """ + url = "https://rest.uniprot.org/uniprotkb/search" + params = { + "query": f"{protein_name} AND organism_id:{organism_id}", + "format": "json", + "size": 1, + "fields": "accession", + } + + for attempt in range(max_retries): + try: + async with session.get(url, params=params, timeout=aiohttp.ClientTimeout(total=15)) as resp: + if resp.status != 200: + await asyncio.sleep(delay * (1 + attempt * 0.5)) + continue + data = await resp.json() + results = data.get("results", []) + if results: + return results[0].get("primaryAccession") + return None + except asyncio.CancelledError: + raise + except Exception as e: + logger.error(f"[UniProt] Attempt {attempt+1} failed: {str(e)}") + await asyncio.sleep(delay * (1 + attempt * 0.5)) + return None + + +async def fetch_affinity_bindingdb( + session: aiohttp.ClientSession, + uniprot_id: str, + affinity_type: str, + cutoff: int, + max_retries: int = 5, + delay: float = 0.5 +) -> List[Dict]: + """ + Asynchronously retrieve affinity values from BindingDB for a given UniProt ID. + Retries on network errors or incomplete data. + """ + url = ( + f"http://bindingdb.org/rest/getLigandsByUniprot?" + f"uniprot={uniprot_id};{cutoff}&response=application/json" + ) + + get_smiles = lambda x: re.sub(r'\s*\|.*\|$', '', x) + + for attempt in range(max_retries): + try: + async with session.get(url, timeout=aiohttp.ClientTimeout(total=60)) as resp: + if resp.status != 200: + logger.error(f"[BindingDB] HTTP {resp.status} for {uniprot_id}, retrying...") + await asyncio.sleep(delay * (1 + attempt * 0.5)) + continue + data = json.loads(await resp.text()) + affinities = ( + data.get("getLindsByUniprotResponse", {}).get("bdb.affinities", []) + or data.get("bdb.affinities", []) + or [] + ) + + result = [{'monomerid': a.get('bdb.monomerid'), + 'smiles': get_smiles(a.get('bdb.smile')), + 'affinity_type': a.get('bdb.affinity_type'), + 'affinity': a.get('bdb.affinity')} for a in affinities if a.get("bdb.affinity_type") == affinity_type] + return result + except asyncio.CancelledError: + raise + except Exception as e: + logger.error(f"[BindingDB] Attempt {attempt+1} failed: {str(e)}") + await asyncio.sleep(delay * (1 + attempt * 0.5)) + return [] + + +async def _aio_fetch_json( + session: aiohttp.ClientSession, + url: str, + timeout: int = 30, + max_retries: int = 4, + retry_delay: float = 0.5, + semaphore: Optional[asyncio.Semaphore] = None +) -> dict: + for attempt in range(max_retries): + try: + if semaphore: + async with semaphore: + async with session.get(url, timeout=aiohttp.ClientTimeout(total=timeout)) as resp: + if resp.status == 200: + data = await resp.json() + if isinstance(data, dict): + return data + else: + async with session.get(url, timeout=aiohttp.ClientTimeout(total=timeout)) as resp: + if resp.status == 200: + data = await resp.json() + if isinstance(data, dict): + return data + except asyncio.CancelledError: + raise + except Exception: + pass + await asyncio.sleep(retry_delay * (1 + 0.5 * attempt)) + return {} + + +async def _resolve_chembl_target_id( + session: aiohttp.ClientSession, + target_name: str, + limit: int = 5, + max_retries: int = 3 +) -> str: + """ + Search ChEMBL target endpoint and return first target_chembl_id. + Retries until valid data is obtained. Returns empty string if none found. + """ + if not target_name: + return "" + + for attempt in range(max_retries): + url = f"{CHEMBL_BASE}/target/search?q={quote(target_name)}&format=json&limit={limit}" + data = await _aio_fetch_json(session, url) + targets = data.get("targets", []) + if targets and isinstance(targets, list): + chembl_ids = [(target.get('target_chembl_id'), target.get('organism')) for target in targets] + if chembl_ids: + return chembl_ids + await asyncio.sleep(0.5 * (1 + attempt * 0.5)) + + return "" + + +def _normalize_activities(activities, target_id, affinity_type): + out = [] + for act in activities: + val = act.get("standard_value") + out.append({ + "smiles": act.get("canonical_smiles") or "", + "affinity_type": affinity_type, + "affinity_value": float(val) if val not in (None, "", "NA") else None, + "affinity_units": act.get("standard_units") or "", + "source": "ChEMBL", + "target_id": target_id + }) + return out + + +async def _fetch_chembl_activity_async( + session: aiohttp.ClientSession, + target_id: str, + affinity_type: str = "Ki", + limit_per_page: int = 1000, + max_records: int = 100000, + semaphore: Optional[asyncio.Semaphore] = None, +) -> List[Dict]: + """Fetch all ChEMBL activity pages concurrently for a given target. + Correct handling of total_count vs max_records (no premature return).""" + if affinity_type not in VALID_AFFINITY_TYPES: + return [] + + # First page (offset=0) to determine total_count and collect initial activities + base_url = ( + f"{CHEMBL_BASE}/activity.json?" + f"target_chembl_id={quote(target_id)}&" + f"standard_type={quote(affinity_type)}&" + f"limit={limit_per_page}&offset=0&include=molecule" + ) + first_data = await _aio_fetch_json(session, base_url, semaphore=semaphore) + if not first_data: + return [] + + activities = first_data.get("activities", []) or [] + results = _normalize_activities(activities, target_id, affinity_type) + + page_meta = first_data.get("page_meta", {}) or {} + try: + total_count = int(page_meta.get("total_count", len(results))) + except Exception: + total_count = len(results) + + # Determine how many records we should fetch (cap by max_records) + desired_total = min(total_count, max_records) + already = len(results) + remaining = max(0, desired_total - already) + if remaining <= 0: + return results + + # Build offsets for the remaining pages (start at limit_per_page) + offsets = list(range(limit_per_page, limit_per_page + remaining, limit_per_page)) + # But ensure offsets do not exceed desired_total + offsets = [o for o in offsets if o < desired_total] + + async def fetch_page(offset: int) -> List[Dict]: + url = ( + f"{CHEMBL_BASE}/activity.json?" + f"target_chembl_id={quote(target_id)}&" + f"standard_type={quote(affinity_type)}&" + f"limit={limit_per_page}&offset={offset}&include=molecule" + ) + data = await _aio_fetch_json(session, url, semaphore=semaphore) + acts = data.get("activities", []) if data else [] + return _normalize_activities(acts, target_id, affinity_type) + + # Limit concurrency across all pages + other targets using provided semaphore + tasks = [fetch_page(off) for off in offsets] + # run and collect (exceptions are returned) + page_results = await asyncio.gather(*tasks, return_exceptions=True) + + for pr in page_results: + if isinstance(pr, list): + results.extend(pr) + # if pr is Exception or unexpected, skip (we keep previous retry logic in _aio_fetch_json) + + # Trim results to desired_total in case last page(s) overshot + if len(results) > desired_total: + results = results[:desired_total] + + return results + + +async def fetch_chembl_data( + target_name: str, + target_id: Optional[str] = None, + affinity_type: str = "Ki", + max_records: int = 10000, + concurrency_limit: int = 10, +) -> List[Dict]: + """ + High-performance concurrent ChEMBL data fetcher. + Fetches multiple targets and multiple pages concurrently with controlled concurrency. + """ + semaphore = asyncio.Semaphore(concurrency_limit) + results: List[Dict] = [] + + async with aiohttp.ClientSession() as session: + # Resolve targets + if not target_id: + chembl_targets = await _resolve_chembl_target_id(session, target_name) + if not chembl_targets: + return [] + else: + chembl_targets = [(target_id, "unknown")] + + # Concurrently fetch all targets + tasks = [ + _fetch_chembl_activity_async( + session=session, + target_id=tid, + affinity_type=affinity_type, + max_records=max_records, + semaphore=semaphore, + ) + for tid, _ in chembl_targets + ] + + all_data = await asyncio.gather(*tasks, return_exceptions=True) + + for (tid, organism), data in zip(chembl_targets, all_data): + if isinstance(data, Exception): + continue + for rec in data: + rec["target_id"] = tid + rec["organism"] = organism + results.extend(data) + + return results + + +@mcp.tool() +def fetch_activity_data( + source: str, + protein_name: str, + dir_to_save: str, + protein_id: Optional[str] = None, + affinity_type: str = "IC50", + cutoff: int = 10000, +) -> str: + """ + Unified data retrieval tool for biochemical databases. + + This function fetches protein-ligand interaction or activity data from supported sources + such as BindingDB and ChEMBL. It automatically handles protein ID resolution and + standardized affinity type filtering. + + Args: + source (str): Name of data source ("bindingdb" or "chembl"). + protein_name (str): Target protein name. + dir_to_save (str): directory to save parsed data in csv format + protein_id (str, optional): Target protein id. If passed, protein_name is ignored + affinity_type (str, optional): Type of affinity (Ki, Kd, IC50). Defaults to "Ki". + cutoff (int, optional): Optional threshold (nM) for BindingDB. Defaults to 10000. + + Returns: + str: Summary of results with path to file and some statistics + Returns error string if data not found or error occurs. + """ + source = source.lower().strip() + if affinity_type not in VALID_AFFINITY_TYPES: + return f"Invalid affinity type '{affinity_type}'. Must be one of {VALID_AFFINITY_TYPES}" + + async def _main(): + async with aiohttp.ClientSession() as session: + if source == "bindingdb": + target_id = protein_id # avoid shadowing outer var + if not target_id: + resolved_id = await fetch_uniprot_id(session, protein_name) + if not resolved_id: + return f"[BindingDB] Could not find UniProt ID for '{protein_name}'" + target_id = resolved_id + + entries = await fetch_affinity_bindingdb( + session, target_id, affinity_type, cutoff + ) + return entries + + elif source == "chembl": + entries = await fetch_chembl_data( + target_name=protein_name, + target_id=protein_id, + affinity_type=affinity_type + ) + return entries + + else: + return f"Unsupported data source '{source}'. Use 'bindingdb' or 'chembl'." + + try: + results = _run_async(_main()) + file_name = os.path.join(dir_to_save, f'{protein_name}_{affinity_type}_{source}.csv') + if isinstance(results, list): + os.makedirs(dir_to_save, exist_ok=True) + df = pd.DataFrame(results) + if len(df)>0: + df.to_csv(file_name, index=False) + buffer = StringIO() + df.info(buf=buffer) + info_str = buffer.getvalue() + return_str = f"The data was saved to {file_name}. Here is info about dataset: {info_str}" + del df + else: + return_str = f"The data was not saved because it is empty" + return return_str + else: + return results + except Exception as e: + return f"[fetch_activity_data] Error: {str(e)}" + + +@mcp.tool() +def python_repl_tool( + code: Annotated[str, "The python code to execute"], +): + """ + Use this tool to perform calculations or execute Python code. It provides a safe environment for code execution without access to external resources like files, networks, or external libraries. + + Args: + code (str): The Python code to execute. + + Returns: + str: The result of the execution, including the code and its standard output. If an error occurs during execution, the error message is returned instead. + """ + try: + result = repl.run(code) + except BaseException as e: + # logger.exception(f"'python_repl_tool' failed with error: {e}") + return f"Failed to execute. Error: {repr(e)}" + result_str = ( + f"Successfully executed:\n```python\n{code}\n```\nStdout: {result}" + ) + return result_str + + +@mcp.tool() +def name2smiles( + mol: Annotated[str, "Name of a molecule"], +): + """ + Convert a molecule name to its SMILES representation. + + This method attempts to retrieve the SMILES string for a given molecule name using a chemical database. It handles potential errors during the retrieval process and provides informative messages if the conversion fails. + + Args: + mol (str): The name of the molecule to convert. + + Returns: + str: The SMILES string representation of the molecule if successful, + an error message if the conversion fails after multiple attempts, + or a "couldn't obtain smiles" message if the name is invalid. + """ + max_attempts = 3 + for attempts in range(max_attempts): + try: + compound = pcp.get_compounds(mol, "name") + smiles = compound[0].canonical_smiles + return smiles + except BaseException as e: + # logger.exception(f"'name2smiles' failed with error: {e}") + return f"Failed to execute. Error: {repr(e)}" + return "I've couldn't obtain smiles, the name is wrong" + + +@mcp.tool() +def smiles2name(smiles: Annotated[str, "SMILES of a molecule"]): + """ + Converts a SMILES string representing a molecule into its IUPAC name. + + Args: + smiles (str): The SMILES string of the molecule. + + Returns: + str: The IUPAC name of the molecule, or an error message if the conversion fails. + """ + + url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/smiles/{smiles}/property/IUPACName/JSON" + max_attempts = 3 + for attempts in range(max_attempts): + try: + response = requests.get(url) + if response.status_code == 200: + data = response.json() + iupac_name = data["PropertyTable"]["Properties"][0]["IUPACName"] + return iupac_name + else: + return "I've couldn't get iupac name" + + except BaseException as e: + # logger.exception(f"'smiles2name' failed with error: {e}") + return f"Failed to execute. Error: {repr(e)}" + return "I've couldn't get iupac name" + + +@mcp.tool() +def smiles2prop( + smiles: Annotated[str, "SMILES of a molecule"], iupac: Optional[str] = None +): + """ + Calculate molecular properties from a SMILES string or IUPAC name. + + Args: + smiles (str): The SMILES string of the molecule. + iupac (str, optional): The IUPAC name of the molecule. If provided, the SMILES string will be derived from it. Defaults to None. + + Returns: + CalcMolDescriptors: An object containing calculated molecular properties. + Returns an error message as a string if the calculation fails. + """ + + try: + if iupac: + compound = pcp.get_compounds(iupac, "name") + if len(compound): + smiles = compound[0].canonical_smiles + + res = CalcMolDescriptors(Chem.MolFromSmiles(smiles)) + return res + except BaseException as e: + # logger.exception(f"'smiles2prop' failed with error: {e}") + return f"Failed to execute. Error: {repr(e)}" + + +@mcp.tool() +def visualize_molecule( + smiles: Annotated[str, "SMILES of a molecule"], + config: RunnableConfig, +): + """ + Visualizes a molecule from its SMILES representation and saves the 3D structure as an HTML file. + + Args: + smiles (str): The SMILES string representing the molecule to visualize. + config (RunnableConfig): Configuration object containing necessary settings, + including the path to save the visualization. + + Returns: + str: A message indicating success or failure of the visualization process. + On success, it confirms the molecule was visualized and saved. + On failure, it provides an error message. + """ + try: + mol = Chem.MolFromSmiles(smiles) + if mol: + mol = Chem.Mol(mol) + mol = AllChem.AddHs(mol, addCoords=True) + AllChem.EmbedMolecule(mol) + AllChem.MMFFOptimizeMolecule(mol) + + view = py3Dmol.view( + data=Chem.MolToMolBlock(mol), # Convert the RDKit molecule for py3Dmol + style={ + "stick": {}, + "sphere": {"scale": 0.3}, + }, + width=600, + height=400, + ) + view.setBackgroundColor("#b8bfcc") + view.zoomTo() + html_content = view.write_html() + + state = config["configurable"].get("state") + # tool_call_id: Annotated[str, InjectedToolCallId] = state['messages'][-1]["tool_calls"][0]['id'] + + path_to_results = os.path.join( + os.environ.get("PATH_TO_RESULTS"), "vis_mols" + ) + if not os.path.exists(path_to_results): + os.makedirs(path_to_results) + + with open( + os.path.join(path_to_results, "vis.html"), "w", encoding="utf-8" + ) as f: + f.write(html_content) + + answer = f"I've successfully generated images of {smiles} molecule" + return answer + else: + return f"I've couldn't visualize this molecule. Perhaps SMILES is invalid" + + except BaseException as e: + return f"Failed to execute. Error: {repr(e)}" + + +@mcp.tool() +def extract_reactions() -> Dict: + """Detects chemical reactions in uploaded images and converts them into structured reaction elements using the `reactions_ocr` pipeline. + + Image paths are read from the directory given by the `IMG_STORAGE_PATH` environment variable. + + Returns: + dict: On success: dictionary from `reactions_ocr` (image filenames to reactants, + conditions, products) and annotated images saved as _annotated.jpg. + On failure or no images: dict with an `"answer"` key and an explanatory message. + """ + logger.info('Running extract_reactions tool...') + try: + directory = Path(os.environ.get('IMG_STORAGE_PATH')) + images = [str(f.resolve()) for f in directory.iterdir() if f.is_file() and f.suffix.lower() in ['.jpg', '.png', '.jpeg']] + if not images: + return {'answer': 'No images provided for reactions recognition.'} + return reactions_ocr(images) + except Exception as e: + logger.error(f'reactions_recognition ERROR: {e}') + return {'answer': f'Could not detect any reactions in the uploaded images. Error: {e}'} + + +@mcp.tool() +def detect_molecules() -> Dict: + """Detects molecular structures in uploaded images and converts them into SMILES using the `molecules_ocr` pipeline. + + Image paths are read from the directory given by the `IMG_STORAGE_PATH` environment variable. + + Returns: + dict: On success: dictionary from `molecules_ocr` (image filenames to SMILES and errors), + with annotated images saved as _annotated.jpg. On failure or no images: + dict with an `"answer"` key and an explanatory message. + """ + logger.info('Running extract_molecules tool...') + try: + directory = Path(os.environ.get('IMG_STORAGE_PATH')) + images = [str(f.resolve()) for f in directory.iterdir() if f.is_file() and f.suffix.lower() in ['.jpg', '.png', '.jpeg']] + if not images: + return {'answer': 'No images provided for molecules recognition.'} + return molecules_ocr(images) + except Exception as e: + logger.error(f'molecules_recognition ERROR: {e}') + return {'answer': 'Could not detect any molecules in the uploaded images.'} + + +@mcp.tool() +def calculate_docking(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 (dict): Dictionary containing docking score for the molecule and the HTML file. + """ + response = calculate_docking_score(smiles, pdb_id) + data = response.get("data", None) + errors = response.get("error") + + if data: + affinity = data['affinity'] + visualization = data['visualization'] + + output_file = None + if visualization: + html_base64 = data.get("visualization", None) + html_content = base64.b64decode(html_base64) + output_file = os.path.join(os.environ.get("PATH_TO_RESULTS"), f"docking_result_{pdb_id}.html") + with open(output_file, "wb") as f: + f.write(html_content) + + result = {'affinity': affinity, "errors": errors} + return { + "answer": result, + "metadata": { + "html_file": output_file + } + } + + +if __name__ == "__main__": + mcp.run(transport="http", host="0.0.0.0", port=7331, path="/mcp") From d5b4cf58723cbc7a9f02925038be511a42e3fb40 Mon Sep 17 00:00:00 2001 From: rodiongolovinsky Date: Thu, 12 Feb 2026 19:07:26 +0300 Subject: [PATCH 3/8] fix docking visualization --- ChemCoScientist/agents/agents.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/ChemCoScientist/agents/agents.py b/ChemCoScientist/agents/agents.py index c39e81d6..a0b2566f 100644 --- a/ChemCoScientist/agents/agents.py +++ b/ChemCoScientist/agents/agents.py @@ -9,8 +9,6 @@ import streamlit as st from langchain_core.language_models import BaseChatModel from langchain_core.messages import ToolMessage -import logging -from langchain_mcp_adapters.client import MultiServerMCPClient from langgraph.types import Command from langgraph.graph import END @@ -246,10 +244,9 @@ def chemist_node(state: dict, config: dict) -> Command: llm = config["configurable"]["llm"] current_prompt = f'{chemist_prompt}\nPass {{"session_id": None}} as a parameter to the detect_molecules and detect_reactions tools' - chemical_tools = _get_chemical_mcp_tools() chem_agent = create_react_agent( - llm, chemical_tools, state_modifier=current_prompt + llm, chem_tools, state_modifier=current_prompt ) task_formatted = f"""For the following plan:\n{str(plan)}\n\nYou are tasked with executing: {task}.""" From bf665325530d1606e086272864f95c821bd42f94 Mon Sep 17 00:00:00 2001 From: rodiongolovinsky Date: Tue, 3 Mar 2026 20:01:22 +0300 Subject: [PATCH 4/8] fix mcp server and dataset processing --- ChemCoScientist/mcp/chemical_server.py | 677 ------------------------- 1 file changed, 677 deletions(-) delete mode 100644 ChemCoScientist/mcp/chemical_server.py diff --git a/ChemCoScientist/mcp/chemical_server.py b/ChemCoScientist/mcp/chemical_server.py deleted file mode 100644 index f62fa80d..00000000 --- a/ChemCoScientist/mcp/chemical_server.py +++ /dev/null @@ -1,677 +0,0 @@ -import base64 -from fastmcp import FastMCP -import logging -from pathlib import Path -import os -from ChemCoScientist.chemical_utils.ocr_pipeline import * -from ChemCoScientist.chemical_utils.chemical_functions import * -import os -from typing import Annotated, Optional, List, Dict -from urllib.parse import quote - -import pubchempy as pcp -import py3Dmol -import rdkit.Chem as Chem -import requests -from langchain_core.runnables.config import RunnableConfig -from langchain_experimental.utilities import PythonREPL -from rdkit.Chem import AllChem -from rdkit.Chem.Descriptors import CalcMolDescriptors -from typing import Dict, List, Optional -from definitions import CONFIG_PATH -from pathlib import Path -from dotenv import load_dotenv -from ChemCoScientist.chemical_utils.ocr_pipeline import molecules_ocr, reactions_ocr -from ChemCoScientist.chemical_utils.chemical_functions import calculate_docking_score - -import aiohttp -import base64 -import asyncio -import json -import re -import pandas as pd -from io import StringIO - -load_dotenv(CONFIG_PATH) - -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - -CHEMBL_BASE = "https://www.ebi.ac.uk/chembl/api/data" -VALID_AFFINITY_TYPES = {"Ki", "Kd", "IC50", "EC50"} -repl = PythonREPL() - - -mcp = FastMCP("ChemTools") - - -def _run_async(coro): - """Run async coroutine from both sync and async contexts.""" - try: - loop = asyncio.get_running_loop() - except RuntimeError: - loop = None - - if loop and loop.is_running(): - # running in async env (e.g. LangChain) - return asyncio.ensure_future(coro) - else: - # safe to call asyncio.run() - return asyncio.run(coro) - - -async def fetch_uniprot_id( - session: aiohttp.ClientSession, - protein_name: str, - organism_id: int = 9606, - max_retries: int = 5, - delay: float = 0.5 -) -> Optional[str]: - """ - Asynchronously fetch UniProt ID for a given protein name. - Retries up to `max_retries` times in case of network or transient API errors. - """ - url = "https://rest.uniprot.org/uniprotkb/search" - params = { - "query": f"{protein_name} AND organism_id:{organism_id}", - "format": "json", - "size": 1, - "fields": "accession", - } - - for attempt in range(max_retries): - try: - async with session.get(url, params=params, timeout=aiohttp.ClientTimeout(total=15)) as resp: - if resp.status != 200: - await asyncio.sleep(delay * (1 + attempt * 0.5)) - continue - data = await resp.json() - results = data.get("results", []) - if results: - return results[0].get("primaryAccession") - return None - except asyncio.CancelledError: - raise - except Exception as e: - logger.error(f"[UniProt] Attempt {attempt+1} failed: {str(e)}") - await asyncio.sleep(delay * (1 + attempt * 0.5)) - return None - - -async def fetch_affinity_bindingdb( - session: aiohttp.ClientSession, - uniprot_id: str, - affinity_type: str, - cutoff: int, - max_retries: int = 5, - delay: float = 0.5 -) -> List[Dict]: - """ - Asynchronously retrieve affinity values from BindingDB for a given UniProt ID. - Retries on network errors or incomplete data. - """ - url = ( - f"http://bindingdb.org/rest/getLigandsByUniprot?" - f"uniprot={uniprot_id};{cutoff}&response=application/json" - ) - - get_smiles = lambda x: re.sub(r'\s*\|.*\|$', '', x) - - for attempt in range(max_retries): - try: - async with session.get(url, timeout=aiohttp.ClientTimeout(total=60)) as resp: - if resp.status != 200: - logger.error(f"[BindingDB] HTTP {resp.status} for {uniprot_id}, retrying...") - await asyncio.sleep(delay * (1 + attempt * 0.5)) - continue - data = json.loads(await resp.text()) - affinities = ( - data.get("getLindsByUniprotResponse", {}).get("bdb.affinities", []) - or data.get("bdb.affinities", []) - or [] - ) - - result = [{'monomerid': a.get('bdb.monomerid'), - 'smiles': get_smiles(a.get('bdb.smile')), - 'affinity_type': a.get('bdb.affinity_type'), - 'affinity': a.get('bdb.affinity')} for a in affinities if a.get("bdb.affinity_type") == affinity_type] - return result - except asyncio.CancelledError: - raise - except Exception as e: - logger.error(f"[BindingDB] Attempt {attempt+1} failed: {str(e)}") - await asyncio.sleep(delay * (1 + attempt * 0.5)) - return [] - - -async def _aio_fetch_json( - session: aiohttp.ClientSession, - url: str, - timeout: int = 30, - max_retries: int = 4, - retry_delay: float = 0.5, - semaphore: Optional[asyncio.Semaphore] = None -) -> dict: - for attempt in range(max_retries): - try: - if semaphore: - async with semaphore: - async with session.get(url, timeout=aiohttp.ClientTimeout(total=timeout)) as resp: - if resp.status == 200: - data = await resp.json() - if isinstance(data, dict): - return data - else: - async with session.get(url, timeout=aiohttp.ClientTimeout(total=timeout)) as resp: - if resp.status == 200: - data = await resp.json() - if isinstance(data, dict): - return data - except asyncio.CancelledError: - raise - except Exception: - pass - await asyncio.sleep(retry_delay * (1 + 0.5 * attempt)) - return {} - - -async def _resolve_chembl_target_id( - session: aiohttp.ClientSession, - target_name: str, - limit: int = 5, - max_retries: int = 3 -) -> str: - """ - Search ChEMBL target endpoint and return first target_chembl_id. - Retries until valid data is obtained. Returns empty string if none found. - """ - if not target_name: - return "" - - for attempt in range(max_retries): - url = f"{CHEMBL_BASE}/target/search?q={quote(target_name)}&format=json&limit={limit}" - data = await _aio_fetch_json(session, url) - targets = data.get("targets", []) - if targets and isinstance(targets, list): - chembl_ids = [(target.get('target_chembl_id'), target.get('organism')) for target in targets] - if chembl_ids: - return chembl_ids - await asyncio.sleep(0.5 * (1 + attempt * 0.5)) - - return "" - - -def _normalize_activities(activities, target_id, affinity_type): - out = [] - for act in activities: - val = act.get("standard_value") - out.append({ - "smiles": act.get("canonical_smiles") or "", - "affinity_type": affinity_type, - "affinity_value": float(val) if val not in (None, "", "NA") else None, - "affinity_units": act.get("standard_units") or "", - "source": "ChEMBL", - "target_id": target_id - }) - return out - - -async def _fetch_chembl_activity_async( - session: aiohttp.ClientSession, - target_id: str, - affinity_type: str = "Ki", - limit_per_page: int = 1000, - max_records: int = 100000, - semaphore: Optional[asyncio.Semaphore] = None, -) -> List[Dict]: - """Fetch all ChEMBL activity pages concurrently for a given target. - Correct handling of total_count vs max_records (no premature return).""" - if affinity_type not in VALID_AFFINITY_TYPES: - return [] - - # First page (offset=0) to determine total_count and collect initial activities - base_url = ( - f"{CHEMBL_BASE}/activity.json?" - f"target_chembl_id={quote(target_id)}&" - f"standard_type={quote(affinity_type)}&" - f"limit={limit_per_page}&offset=0&include=molecule" - ) - first_data = await _aio_fetch_json(session, base_url, semaphore=semaphore) - if not first_data: - return [] - - activities = first_data.get("activities", []) or [] - results = _normalize_activities(activities, target_id, affinity_type) - - page_meta = first_data.get("page_meta", {}) or {} - try: - total_count = int(page_meta.get("total_count", len(results))) - except Exception: - total_count = len(results) - - # Determine how many records we should fetch (cap by max_records) - desired_total = min(total_count, max_records) - already = len(results) - remaining = max(0, desired_total - already) - if remaining <= 0: - return results - - # Build offsets for the remaining pages (start at limit_per_page) - offsets = list(range(limit_per_page, limit_per_page + remaining, limit_per_page)) - # But ensure offsets do not exceed desired_total - offsets = [o for o in offsets if o < desired_total] - - async def fetch_page(offset: int) -> List[Dict]: - url = ( - f"{CHEMBL_BASE}/activity.json?" - f"target_chembl_id={quote(target_id)}&" - f"standard_type={quote(affinity_type)}&" - f"limit={limit_per_page}&offset={offset}&include=molecule" - ) - data = await _aio_fetch_json(session, url, semaphore=semaphore) - acts = data.get("activities", []) if data else [] - return _normalize_activities(acts, target_id, affinity_type) - - # Limit concurrency across all pages + other targets using provided semaphore - tasks = [fetch_page(off) for off in offsets] - # run and collect (exceptions are returned) - page_results = await asyncio.gather(*tasks, return_exceptions=True) - - for pr in page_results: - if isinstance(pr, list): - results.extend(pr) - # if pr is Exception or unexpected, skip (we keep previous retry logic in _aio_fetch_json) - - # Trim results to desired_total in case last page(s) overshot - if len(results) > desired_total: - results = results[:desired_total] - - return results - - -async def fetch_chembl_data( - target_name: str, - target_id: Optional[str] = None, - affinity_type: str = "Ki", - max_records: int = 10000, - concurrency_limit: int = 10, -) -> List[Dict]: - """ - High-performance concurrent ChEMBL data fetcher. - Fetches multiple targets and multiple pages concurrently with controlled concurrency. - """ - semaphore = asyncio.Semaphore(concurrency_limit) - results: List[Dict] = [] - - async with aiohttp.ClientSession() as session: - # Resolve targets - if not target_id: - chembl_targets = await _resolve_chembl_target_id(session, target_name) - if not chembl_targets: - return [] - else: - chembl_targets = [(target_id, "unknown")] - - # Concurrently fetch all targets - tasks = [ - _fetch_chembl_activity_async( - session=session, - target_id=tid, - affinity_type=affinity_type, - max_records=max_records, - semaphore=semaphore, - ) - for tid, _ in chembl_targets - ] - - all_data = await asyncio.gather(*tasks, return_exceptions=True) - - for (tid, organism), data in zip(chembl_targets, all_data): - if isinstance(data, Exception): - continue - for rec in data: - rec["target_id"] = tid - rec["organism"] = organism - results.extend(data) - - return results - - -@mcp.tool() -def fetch_activity_data( - source: str, - protein_name: str, - dir_to_save: str, - protein_id: Optional[str] = None, - affinity_type: str = "IC50", - cutoff: int = 10000, -) -> str: - """ - Unified data retrieval tool for biochemical databases. - - This function fetches protein-ligand interaction or activity data from supported sources - such as BindingDB and ChEMBL. It automatically handles protein ID resolution and - standardized affinity type filtering. - - Args: - source (str): Name of data source ("bindingdb" or "chembl"). - protein_name (str): Target protein name. - dir_to_save (str): directory to save parsed data in csv format - protein_id (str, optional): Target protein id. If passed, protein_name is ignored - affinity_type (str, optional): Type of affinity (Ki, Kd, IC50). Defaults to "Ki". - cutoff (int, optional): Optional threshold (nM) for BindingDB. Defaults to 10000. - - Returns: - str: Summary of results with path to file and some statistics - Returns error string if data not found or error occurs. - """ - source = source.lower().strip() - if affinity_type not in VALID_AFFINITY_TYPES: - return f"Invalid affinity type '{affinity_type}'. Must be one of {VALID_AFFINITY_TYPES}" - - async def _main(): - async with aiohttp.ClientSession() as session: - if source == "bindingdb": - target_id = protein_id # avoid shadowing outer var - if not target_id: - resolved_id = await fetch_uniprot_id(session, protein_name) - if not resolved_id: - return f"[BindingDB] Could not find UniProt ID for '{protein_name}'" - target_id = resolved_id - - entries = await fetch_affinity_bindingdb( - session, target_id, affinity_type, cutoff - ) - return entries - - elif source == "chembl": - entries = await fetch_chembl_data( - target_name=protein_name, - target_id=protein_id, - affinity_type=affinity_type - ) - return entries - - else: - return f"Unsupported data source '{source}'. Use 'bindingdb' or 'chembl'." - - try: - results = _run_async(_main()) - file_name = os.path.join(dir_to_save, f'{protein_name}_{affinity_type}_{source}.csv') - if isinstance(results, list): - os.makedirs(dir_to_save, exist_ok=True) - df = pd.DataFrame(results) - if len(df)>0: - df.to_csv(file_name, index=False) - buffer = StringIO() - df.info(buf=buffer) - info_str = buffer.getvalue() - return_str = f"The data was saved to {file_name}. Here is info about dataset: {info_str}" - del df - else: - return_str = f"The data was not saved because it is empty" - return return_str - else: - return results - except Exception as e: - return f"[fetch_activity_data] Error: {str(e)}" - - -@mcp.tool() -def python_repl_tool( - code: Annotated[str, "The python code to execute"], -): - """ - Use this tool to perform calculations or execute Python code. It provides a safe environment for code execution without access to external resources like files, networks, or external libraries. - - Args: - code (str): The Python code to execute. - - Returns: - str: The result of the execution, including the code and its standard output. If an error occurs during execution, the error message is returned instead. - """ - try: - result = repl.run(code) - except BaseException as e: - # logger.exception(f"'python_repl_tool' failed with error: {e}") - return f"Failed to execute. Error: {repr(e)}" - result_str = ( - f"Successfully executed:\n```python\n{code}\n```\nStdout: {result}" - ) - return result_str - - -@mcp.tool() -def name2smiles( - mol: Annotated[str, "Name of a molecule"], -): - """ - Convert a molecule name to its SMILES representation. - - This method attempts to retrieve the SMILES string for a given molecule name using a chemical database. It handles potential errors during the retrieval process and provides informative messages if the conversion fails. - - Args: - mol (str): The name of the molecule to convert. - - Returns: - str: The SMILES string representation of the molecule if successful, - an error message if the conversion fails after multiple attempts, - or a "couldn't obtain smiles" message if the name is invalid. - """ - max_attempts = 3 - for attempts in range(max_attempts): - try: - compound = pcp.get_compounds(mol, "name") - smiles = compound[0].canonical_smiles - return smiles - except BaseException as e: - # logger.exception(f"'name2smiles' failed with error: {e}") - return f"Failed to execute. Error: {repr(e)}" - return "I've couldn't obtain smiles, the name is wrong" - - -@mcp.tool() -def smiles2name(smiles: Annotated[str, "SMILES of a molecule"]): - """ - Converts a SMILES string representing a molecule into its IUPAC name. - - Args: - smiles (str): The SMILES string of the molecule. - - Returns: - str: The IUPAC name of the molecule, or an error message if the conversion fails. - """ - - url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/smiles/{smiles}/property/IUPACName/JSON" - max_attempts = 3 - for attempts in range(max_attempts): - try: - response = requests.get(url) - if response.status_code == 200: - data = response.json() - iupac_name = data["PropertyTable"]["Properties"][0]["IUPACName"] - return iupac_name - else: - return "I've couldn't get iupac name" - - except BaseException as e: - # logger.exception(f"'smiles2name' failed with error: {e}") - return f"Failed to execute. Error: {repr(e)}" - return "I've couldn't get iupac name" - - -@mcp.tool() -def smiles2prop( - smiles: Annotated[str, "SMILES of a molecule"], iupac: Optional[str] = None -): - """ - Calculate molecular properties from a SMILES string or IUPAC name. - - Args: - smiles (str): The SMILES string of the molecule. - iupac (str, optional): The IUPAC name of the molecule. If provided, the SMILES string will be derived from it. Defaults to None. - - Returns: - CalcMolDescriptors: An object containing calculated molecular properties. - Returns an error message as a string if the calculation fails. - """ - - try: - if iupac: - compound = pcp.get_compounds(iupac, "name") - if len(compound): - smiles = compound[0].canonical_smiles - - res = CalcMolDescriptors(Chem.MolFromSmiles(smiles)) - return res - except BaseException as e: - # logger.exception(f"'smiles2prop' failed with error: {e}") - return f"Failed to execute. Error: {repr(e)}" - - -@mcp.tool() -def visualize_molecule( - smiles: Annotated[str, "SMILES of a molecule"], - config: RunnableConfig, -): - """ - Visualizes a molecule from its SMILES representation and saves the 3D structure as an HTML file. - - Args: - smiles (str): The SMILES string representing the molecule to visualize. - config (RunnableConfig): Configuration object containing necessary settings, - including the path to save the visualization. - - Returns: - str: A message indicating success or failure of the visualization process. - On success, it confirms the molecule was visualized and saved. - On failure, it provides an error message. - """ - try: - mol = Chem.MolFromSmiles(smiles) - if mol: - mol = Chem.Mol(mol) - mol = AllChem.AddHs(mol, addCoords=True) - AllChem.EmbedMolecule(mol) - AllChem.MMFFOptimizeMolecule(mol) - - view = py3Dmol.view( - data=Chem.MolToMolBlock(mol), # Convert the RDKit molecule for py3Dmol - style={ - "stick": {}, - "sphere": {"scale": 0.3}, - }, - width=600, - height=400, - ) - view.setBackgroundColor("#b8bfcc") - view.zoomTo() - html_content = view.write_html() - - state = config["configurable"].get("state") - # tool_call_id: Annotated[str, InjectedToolCallId] = state['messages'][-1]["tool_calls"][0]['id'] - - path_to_results = os.path.join( - os.environ.get("PATH_TO_RESULTS"), "vis_mols" - ) - if not os.path.exists(path_to_results): - os.makedirs(path_to_results) - - with open( - os.path.join(path_to_results, "vis.html"), "w", encoding="utf-8" - ) as f: - f.write(html_content) - - answer = f"I've successfully generated images of {smiles} molecule" - return answer - else: - return f"I've couldn't visualize this molecule. Perhaps SMILES is invalid" - - except BaseException as e: - return f"Failed to execute. Error: {repr(e)}" - - -@mcp.tool() -def extract_reactions() -> Dict: - """Detects chemical reactions in uploaded images and converts them into structured reaction elements using the `reactions_ocr` pipeline. - - Image paths are read from the directory given by the `IMG_STORAGE_PATH` environment variable. - - Returns: - dict: On success: dictionary from `reactions_ocr` (image filenames to reactants, - conditions, products) and annotated images saved as _annotated.jpg. - On failure or no images: dict with an `"answer"` key and an explanatory message. - """ - logger.info('Running extract_reactions tool...') - try: - directory = Path(os.environ.get('IMG_STORAGE_PATH')) - images = [str(f.resolve()) for f in directory.iterdir() if f.is_file() and f.suffix.lower() in ['.jpg', '.png', '.jpeg']] - if not images: - return {'answer': 'No images provided for reactions recognition.'} - return reactions_ocr(images) - except Exception as e: - logger.error(f'reactions_recognition ERROR: {e}') - return {'answer': f'Could not detect any reactions in the uploaded images. Error: {e}'} - - -@mcp.tool() -def detect_molecules() -> Dict: - """Detects molecular structures in uploaded images and converts them into SMILES using the `molecules_ocr` pipeline. - - Image paths are read from the directory given by the `IMG_STORAGE_PATH` environment variable. - - Returns: - dict: On success: dictionary from `molecules_ocr` (image filenames to SMILES and errors), - with annotated images saved as _annotated.jpg. On failure or no images: - dict with an `"answer"` key and an explanatory message. - """ - logger.info('Running extract_molecules tool...') - try: - directory = Path(os.environ.get('IMG_STORAGE_PATH')) - images = [str(f.resolve()) for f in directory.iterdir() if f.is_file() and f.suffix.lower() in ['.jpg', '.png', '.jpeg']] - if not images: - return {'answer': 'No images provided for molecules recognition.'} - return molecules_ocr(images) - except Exception as e: - logger.error(f'molecules_recognition ERROR: {e}') - return {'answer': 'Could not detect any molecules in the uploaded images.'} - - -@mcp.tool() -def calculate_docking(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 (dict): Dictionary containing docking score for the molecule and the HTML file. - """ - response = calculate_docking_score(smiles, pdb_id) - data = response.get("data", None) - errors = response.get("error") - - if data: - affinity = data['affinity'] - visualization = data['visualization'] - - output_file = None - if visualization: - html_base64 = data.get("visualization", None) - html_content = base64.b64decode(html_base64) - output_file = os.path.join(os.environ.get("PATH_TO_RESULTS"), f"docking_result_{pdb_id}.html") - with open(output_file, "wb") as f: - f.write(html_content) - - result = {'affinity': affinity, "errors": errors} - return { - "answer": result, - "metadata": { - "html_file": output_file - } - } - - -if __name__ == "__main__": - mcp.run(transport="http", host="0.0.0.0", port=7331, path="/mcp") From 4ffec885d38203050613a80381e66d56b7139a40 Mon Sep 17 00:00:00 2001 From: AerDragon Date: Mon, 9 Feb 2026 11:12:20 +0300 Subject: [PATCH 5/8] chore(chemical-utils): add retrosynthesis tools and forward pred --- ChemCoScientist/agents/agents.py | 21 ++ .../chemical_utils/chemical_functions.py | 19 +- .../chemical_utils/retrosynthesis.py | 191 ++++++++++ ChemCoScientist/frontend/chat.py | 342 ++++++++++++++++++ ChemCoScientist/tools/chemist_tools.py | 97 ++++- 5 files changed, 659 insertions(+), 11 deletions(-) create mode 100644 ChemCoScientist/chemical_utils/retrosynthesis.py diff --git a/ChemCoScientist/agents/agents.py b/ChemCoScientist/agents/agents.py index a0b2566f..8ff363f1 100644 --- a/ChemCoScientist/agents/agents.py +++ b/ChemCoScientist/agents/agents.py @@ -275,6 +275,27 @@ def chemist_node(state: dict, config: dict) -> Command: updated_metadata["docking"].update(docking_metadata["docking"]) else: 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: + continue + updated_metadata.update({"retrosynthesis": result}) + + elif isinstance(message, ToolMessage) and message.name in ["classify_reaction"]: + try: + result = ast.literal_eval(message.content) + except Exception: + 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: + continue + updated_metadata.update({"forward_prediction": result}) return Command(update={ "past_steps": Annotated[set, operator.or_](set([ diff --git a/ChemCoScientist/chemical_utils/chemical_functions.py b/ChemCoScientist/chemical_utils/chemical_functions.py index ae8b589b..738afc40 100644 --- a/ChemCoScientist/chemical_utils/chemical_functions.py +++ b/ChemCoScientist/chemical_utils/chemical_functions.py @@ -18,7 +18,7 @@ REQUEST_TIMEOUT = 60 -def handle_api_request(endpoint: str, file_param_name: str = None, ): +def handle_api_request(host: str, endpoint: str, file_param_name: str = None, ): """ Decorator for handling requests to Chemical ToolsService API. @@ -45,7 +45,7 @@ def wrapper(*args, **kwargs) -> Any: Data from the "data" field of the API response """ try: - api_url = f"{CHEM_SERVICES_URL}{endpoint}" + api_url = f"{host}{endpoint}" logger.info(f"Calling ChemService API: {api_url}") if file_param_name: @@ -102,7 +102,7 @@ def wrapper(*args, **kwargs) -> Any: return decorator -@handle_api_request(endpoint="/extract_reactions_from_pdf/", file_param_name="pdf_file") +@handle_api_request(host=CHEM_SERVICES_URL, 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. @@ -121,7 +121,7 @@ def extract_reactions_from_pdf(file: bytes) -> List[Dict]: pass -@handle_api_request(endpoint="/extract_reactions_from_figure/", file_param_name="image") +@handle_api_request(host=CHEM_SERVICES_URL, endpoint="/extract_reactions_from_figure/", file_param_name="image") def extract_reactions_from_figure(image: bytes) -> List[Dict]: """ Extract reactions information from an image. @@ -136,7 +136,7 @@ def extract_reactions_from_figure(image: bytes) -> List[Dict]: pass -@handle_api_request(endpoint="/extract_molecules_from_pdf/", file_param_name="pdf_file") +@handle_api_request(host=CHEM_SERVICES_URL, 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. @@ -151,7 +151,7 @@ def extract_molecules_from_pdf(file: bytes) -> List[Dict]: pass -@handle_api_request(endpoint="/extract_molecules_from_figure/", file_param_name="image") +@handle_api_request(host=CHEM_SERVICES_URL, endpoint="/extract_molecules_from_figure/", file_param_name="image") def extract_molecules_from_figure(image: bytes) -> List[Dict]: """ Extract molecules information from an image. @@ -166,7 +166,7 @@ def extract_molecules_from_figure(image: bytes) -> List[Dict]: pass -@handle_api_request(endpoint="/convert_image_to_smiles/", file_param_name="image") +@handle_api_request(host=CHEM_SERVICES_URL, endpoint="/convert_image_to_smiles/", file_param_name="image") def convert_image_to_smiles(image: bytes) -> str: """ Convert an image to a smiles string. @@ -178,7 +178,7 @@ def convert_image_to_smiles(image: bytes) -> str: """ pass -@handle_api_request(endpoint="/docking/", file_param_name=None) +@handle_api_request(host=CHEM_SERVICES_URL, endpoint="/docking/", file_param_name=None) def calculate_docking_score(smiles: str, pdb_id: str) -> str: """ Calculate docking score for a molecule. @@ -191,7 +191,6 @@ def calculate_docking_score(smiles: str, pdb_id: str) -> str: """ 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): @@ -204,7 +203,7 @@ def remove_keys(obj: Any, keys_to_remove: set[str] = {"bbox", "score"}) -> Any: remove_keys(item, keys_to_remove) return obj - if __name__ == "__main__": result = calculate_docking_score(smiles="C1CCCCC1", pdb_id="5vfi") print(result) + diff --git a/ChemCoScientist/chemical_utils/retrosynthesis.py b/ChemCoScientist/chemical_utils/retrosynthesis.py new file mode 100644 index 00000000..20f3c950 --- /dev/null +++ b/ChemCoScientist/chemical_utils/retrosynthesis.py @@ -0,0 +1,191 @@ +from typing import List, Dict, Any +import requests +from dotenv import load_dotenv +import os +from definitions import CONFIG_PATH +import logging + +load_dotenv(CONFIG_PATH) + +RETROSYNTHESIS_SERVICES_HOST = os.environ.get("RETROSYNTHESIS_SERVICES_HOST") +RETROSYNTHESIS_SERVICES_PORT = os.environ.get("RETROSYNTHESIS_SERVICES_PORT") +RETROSYNTHESIS_SERVICES_URL = f"http://{RETROSYNTHESIS_SERVICES_HOST}:{RETROSYNTHESIS_SERVICES_PORT}" +REQUEST_TIMEOUT = 60 + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +def retrosynthesis_result(smiles: str, mode: str = "fast", max_routes: int = 5) -> Dict[str, Any]: + """ + Proxy request to the Retrosynthesis service tree-search endpoint + + Args: + smiles (str): Target molecule SMILES. + mode (str): One of "fast", "balanced", "deep". + max_routes (int): Maximum number of routes to return. + Returns: + response (Dict[str, Any]): Retrosynthesis result payload with: + - target (str | None): input target SMILES returned by ASKCOS. + - routes (List[Dict[str, Any]]): list of retrosynthesis routes: + - id (str): unique route identifier. + - depth (int | None): longest path length in the route. + - precursor_cost (float | None): summed precursor cost metric. + - score (float | None): overall route score. + - min_step_plausibility (float | None): lowest step plausibility. + - avg_step_plausibility (float | None): average step plausibility. + - steps (List[Dict[str, Any]]): ordered reaction steps: + - reaction_smiles (str): step reaction SMILES. + - mapped_smiles (str | None): atom-mapped reaction SMILES. + - plausibility (float | None): step plausibility score. + - precursor_rank (int | None): ranking of precursor set. + - precursor_score (float | None): model score for precursors. + - model_score (float | None): model score for the step. + - template (Dict[str, Any] | None): template metadata: + reaction_smarts (str): reaction SMARTS pattern. + template_rank (int | None): rank among templates. + num_examples (int | None): template training examples count. + - reactants (List[Dict[str, Any]]): precursor molecules: + smiles (str): molecule SMILES. + terminal (bool | None): True if purchasable/terminal. + buy_link (str | None): vendor link if available. + stoichiometry (int): reagent count (default 1). + - products (List[Dict[str, Any]]): products, same schema + """ + api_url = f"{RETROSYNTHESIS_SERVICES_URL}/api/v1/retrosynthesis/result" + logger.info(f"Calling Retrosynthesis API: {api_url}") + try: + response = requests.post( + api_url, + json={"smiles": smiles}, + params={"mode": mode}, + timeout=REQUEST_TIMEOUT, + ) + if response.status_code != 200: + error_msg = f"Retrosynthesis API returned status {response.status_code}: {response.text[:500]}" + logger.error(error_msg) + raise ValueError(error_msg) + json_response = response.json() + if json_response is None: + error_msg = "Retrosynthesis API returned None JSON response" + logger.error(error_msg) + raise ValueError(error_msg) + if isinstance(json_response, dict) and isinstance(json_response.get("routes"), list): + json_response["routes"] = json_response["routes"][:max(0, int(max_routes))] + return json_response + except requests.exceptions.RequestException as e: + error_msg = f"Failed to connect to Retrosynthesis API at {RETROSYNTHESIS_SERVICES_URL}: {str(e)}" + logger.error(error_msg) + raise ConnectionError(error_msg) + +def classify_reaction_smiles(smiles: List[str], num_results: int = 10) -> Dict[str, Any]: + """ + Proxy request to the ASKCOS reaction-classification endpoint. + + Args: + smiles (List[str]): List of reaction SMILES, e.g. ["A.B>>C"]. + num_results (int): Max number of classes per reaction (1..50). + Returns: + response (Dict[str, Any]): ASKCOS classification payload with: + - status_code (int): upstream status code. + - message (str): upstream message. + - result (List[Dict[str, Any]]): list of hits with: + - rank (int): hit rank. + - reaction_num (str): reaction identifier. + - reaction_name (str): reaction name. + - reaction_classnum (str): class number. + - reaction_classname (str): class name. + - reaction_superclassnum (str): superclass number. + - reaction_superclassname (str): superclass name. + - prediction_certainty (float): confidence score. + """ + api_url = f"{RETROSYNTHESIS_SERVICES_URL}/api/v1/reaction-classification/classify" + logger.info(f"Calling Reaction Classification API: {api_url}") + try: + response = requests.post( + api_url, + json={"smiles": smiles, "num_results": num_results}, + timeout=REQUEST_TIMEOUT, + ) + if response.status_code != 200: + error_msg = ( + f"Reaction Classification API returned status {response.status_code}: " + f"{response.text[:500]}" + ) + logger.error(error_msg) + raise ValueError(error_msg) + json_response = response.json() + if json_response is None: + error_msg = "Reaction Classification API returned None JSON response" + logger.error(error_msg) + raise ValueError(error_msg) + return json_response + except requests.exceptions.RequestException as e: + error_msg = ( + f"Failed to connect to Reaction Classification API at " + f"{RETROSYNTHESIS_SERVICES_URL}: {str(e)}" + ) + logger.error(error_msg) + raise ConnectionError(error_msg) + +def forward_predict_products( + smiles: List[str], + backend: str, + model_name: str = "wldn5", + reagents: str = "", + solvent: str = "", +) -> Dict[str, Any]: + """ + Proxy request to the ASKCOS forward prediction endpoint. + + Args: + smiles (List[str]): Batch of reaction inputs (reactants). + backend (str): One of "wldn5", "graph2smiles", "augmented_transformer". + model_name (str): Model name for the backend (default "pistachio"). + reagents (str): Reagents string as in ASKCOS controller. + solvent (str): Solvent string as in ASKCOS controller. + Returns: + response (Dict[str, Any]): ASKCOS forward payload with: + - inputs (List[str]): normalized inputs (reactants+reagents+solvent). + - backend (str): backend identifier used. + - model_name (str): model name used. + - predictions (List[Dict[str, Any]]): predicted products: + - smiles (str): product SMILES. + - score (float): model probability/score. + """ + api_url = f"{RETROSYNTHESIS_SERVICES_URL}/api/v1/forward/predict" + logger.info(f"Calling Forward Prediction API: {api_url}") + try: + response = requests.post( + api_url, + json={ + "smiles": smiles, + "backend": backend, + "model_name": model_name, + "reagents": reagents, + "solvent": solvent, + }, + timeout=REQUEST_TIMEOUT, + ) + if response.status_code != 200: + error_msg = ( + f"Forward Prediction API returned status {response.status_code}: " + f"{response.text[:500]}" + ) + logger.error(error_msg) + raise ValueError(error_msg) + json_response = response.json() + if json_response is None: + error_msg = "Forward Prediction API returned None JSON response" + logger.error(error_msg) + raise ValueError(error_msg) + return json_response + except requests.exceptions.RequestException as e: + error_msg = ( + f"Failed to connect to Forward Prediction API at " + f"{RETROSYNTHESIS_SERVICES_URL}: {str(e)}" + ) + logger.error(error_msg) + raise ConnectionError(error_msg) + +if __name__ == "__main__": + print(retrosynthesis_result(smiles="C1CCCCC1")) diff --git a/ChemCoScientist/frontend/chat.py b/ChemCoScientist/frontend/chat.py index 60eb6547..22d39699 100644 --- a/ChemCoScientist/frontend/chat.py +++ b/ChemCoScientist/frontend/chat.py @@ -1,14 +1,19 @@ import asyncio +import ast import glob import logging import os import streamlit as st import threading +from typing import Optional from io import BytesIO from langgraph.errors import GraphRecursionError from pathlib import Path from PIL import Image + +from rdkit import Chem +from rdkit.Chem import Draw, rdChemReactions from queue import Queue, Empty from urllib.parse import urlparse @@ -115,6 +120,15 @@ def chat(): if message.get("chem_ocr") and message["role"] == "assistant": display_chem_ocr_metadata(message) + if message.get("retrosynthesis") and message["role"] == "assistant": + display_retrosynthesis_metadata(message) + + if message.get("forward_prediction") and message["role"] == "assistant": + display_forward_prediction_metadata(message) + + if message.get("reaction_classification") and message["role"] == "assistant": + display_reaction_classification_metadata(message) + if message.get("docking") and message["role"] == "assistant": display_docking_metadata(message) @@ -410,6 +424,18 @@ def message_handler(user_query: str, placeholder: st.delta_generator.DeltaGenera st.session_state.messages[-1]["docking"] = result["metadata"]["docking"] display_docking_metadata(st.session_state.messages[-1]) + if "retrosynthesis" in result["metadata"].keys(): + st.session_state.messages[-1]["retrosynthesis"] = result["metadata"]["retrosynthesis"] + display_retrosynthesis_metadata(st.session_state.messages[-1]) + + if "forward_prediction" in result["metadata"].keys(): + st.session_state.messages[-1]["forward_prediction"] = result["metadata"]["forward_prediction"] + display_forward_prediction_metadata(st.session_state.messages[-1]) + + if "reaction_classification" in result["metadata"].keys(): + st.session_state.messages[-1]["reaction_classification"] = result["metadata"]["reaction_classification"] + display_reaction_classification_metadata(st.session_state.messages[-1]) + if mols := msg.get("molecules_vis"): for mol in mols: st.components.v1.html(mol, height=400) @@ -495,6 +521,322 @@ def pdf_viewer(folder: str): st.divider() +def _reaction_smiles_to_image(reaction_smiles: str): + if not rdChemReactions or not Draw: + return None + try: + rxn = rdChemReactions.ReactionFromSmarts(reaction_smiles, useSmiles=True) + if rxn is None: + return None + return Draw.ReactionToImage(rxn, subImgSize=(300, 200)) + except Exception: + return None + + +def _render_reaction_smiles(reaction_smiles: str, caption: Optional[str] = None): + normalized = reaction_smiles + if isinstance(normalized, str): + normalized = normalized.replace(" -> ", ">>").replace(" → ", ">>") + normalized = normalized.replace(" + ", ".") + img = _reaction_smiles_to_image(normalized) + if img is not None: + st.image(img, caption=caption) + else: + st.code(reaction_smiles) + + +def display_retrosynthesis_metadata(message): + data = message.get("retrosynthesis") or {} + if not data: + return + if isinstance(data, str): + try: + data = ast.literal_eval(data) + except Exception: + st.markdown("### Retrosynthesis") + st.code(data) + return + payload = data + if isinstance(data, dict): + if isinstance(data.get("data"), dict): + payload = data["data"] + elif isinstance(data.get("result"), dict): + payload = data["result"] + st.markdown("### Retrosynthesis") + target = payload.get("target") if isinstance(payload, dict) else None + if target: + st.markdown(f"**Target:** `{target}`") + routes = payload.get("routes") if isinstance(payload, dict) else None + if not isinstance(routes, list): + st.write("No routes returned.") + return + routes = [r for r in routes if isinstance(r, dict)] + if not routes: + st.write("No routes returned.") + return + # pick the best route by score (fallback: lowest precursor_cost) + def _route_sort_key(r): + score = r.get("score") + precursor_cost = r.get("precursor_cost") + return ( + score is not None, + score if score is not None else float("-inf"), + precursor_cost is not None, + -precursor_cost if precursor_cost is not None else 0, + ) + + try: + best_route = max(routes, key=_route_sort_key) + except Exception: + best_route = routes[0] + for idx, route in enumerate([best_route], start=1): + title_parts = ["Best route"] + meta_parts = [] + for key in ["score", "depth", "precursor_cost"]: + if route.get(key) is not None: + meta_parts.append(f"{key}={route.get(key)}") + if meta_parts: + title_parts.append(f"({', '.join(meta_parts)})") + with st.expander(" ".join(title_parts), expanded=idx == 1): + steps = route.get("steps") or [] + if not steps: + st.write("No steps returned.") + continue + for step_idx, step in enumerate(steps, start=1): + reaction_smiles = step.get("reaction_smiles") or step.get("mapped_smiles") + caption = f"Step {step_idx}" + if step.get("plausibility") is not None: + caption += f" | plausibility={step.get('plausibility')}" + if reaction_smiles: + _render_reaction_smiles(reaction_smiles, caption=caption) + else: + st.markdown(f"**Step {step_idx}**") + st.write(step) + reactants = [r.get("smiles") for r in (step.get("reactants") or []) if r.get("smiles")] + products = [p.get("smiles") for p in (step.get("products") or []) if p.get("smiles")] + if reactants: + st.caption(f"Reactants: {', '.join(reactants)}") + if products: + st.caption(f"Products: {', '.join(products)}") + + +def display_forward_prediction_metadata(message): + data = message.get("forward_prediction") or {} + if not data: + return + st.markdown("### Forward Prediction") + backend = data.get("backend") + model_name = data.get("model_name") + 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( + predictions, + key=lambda p: p.get("score") if isinstance(p, dict) else None, + reverse=True, + )[:2] + if not predictions: + st.write("No predictions returned.") + return + if len(inputs) == 1: + base = inputs[0] + for idx, pred in enumerate(predictions, start=1): + 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) + else: + rows = [{"smiles": p.get("smiles"), "score": p.get("score")} for p in predictions] + st.dataframe(rows) + + +def display_reaction_classification_metadata(message): + data = message.get("reaction_classification") or {} + if not data: + return + st.markdown("### Reaction Classification") + status_code = data.get("status_code") + if status_code is not None: + st.markdown(f"**Status:** {status_code}") + msg = data.get("message") + if msg: + st.markdown(f"**Message:** {msg}") + results = data.get("result") or [] + if results: + st.dataframe(results) + else: + st.write("No classification results.") + + +def _reaction_smiles_to_image(reaction_smiles: str): + if not rdChemReactions or not Draw: + return None + try: + rxn = rdChemReactions.ReactionFromSmarts(reaction_smiles, useSmiles=True) + if rxn is None: + return None + return Draw.ReactionToImage(rxn, subImgSize=(300, 200)) + except Exception: + return None + + +def _render_reaction_smiles(reaction_smiles: str, caption: Optional[str] = None): + normalized = reaction_smiles + if isinstance(normalized, str): + normalized = normalized.replace(" -> ", ">>").replace(" → ", ">>") + normalized = normalized.replace(" + ", ".") + img = _reaction_smiles_to_image(normalized) + if img is not None: + st.image(img, caption=caption) + else: + st.code(reaction_smiles) + + +def display_retrosynthesis_metadata(message): + data = message.get("retrosynthesis") or {} + if not data: + return + if isinstance(data, str): + try: + data = ast.literal_eval(data) + except Exception: + st.markdown("### Retrosynthesis") + st.code(data) + return + payload = data + if isinstance(data, dict): + if isinstance(data.get("data"), dict): + payload = data["data"] + elif isinstance(data.get("result"), dict): + payload = data["result"] + st.markdown("### Retrosynthesis") + target = payload.get("target") if isinstance(payload, dict) else None + if target: + st.markdown(f"**Target:** `{target}`") + routes = payload.get("routes") if isinstance(payload, dict) else None + if not isinstance(routes, list): + st.write("No routes returned.") + return + routes = [r for r in routes if isinstance(r, dict)] + if not routes: + st.write("No routes returned.") + return + # pick the best route by score (fallback: lowest precursor_cost) + def _route_sort_key(r): + score = r.get("score") + precursor_cost = r.get("precursor_cost") + return ( + score is not None, + score if score is not None else float("-inf"), + precursor_cost is not None, + -precursor_cost if precursor_cost is not None else 0, + ) + + try: + best_route = max(routes, key=_route_sort_key) + except Exception: + best_route = routes[0] + for idx, route in enumerate([best_route], start=1): + title_parts = ["Best route"] + meta_parts = [] + for key in ["score", "depth", "precursor_cost"]: + if route.get(key) is not None: + meta_parts.append(f"{key}={route.get(key)}") + if meta_parts: + title_parts.append(f"({', '.join(meta_parts)})") + with st.expander(" ".join(title_parts), expanded=idx == 1): + steps = route.get("steps") or [] + if not steps: + st.write("No steps returned.") + continue + for step_idx, step in enumerate(steps, start=1): + reaction_smiles = step.get("reaction_smiles") or step.get("mapped_smiles") + caption = f"Step {step_idx}" + if step.get("plausibility") is not None: + caption += f" | plausibility={step.get('plausibility')}" + if reaction_smiles: + _render_reaction_smiles(reaction_smiles, caption=caption) + else: + st.markdown(f"**Step {step_idx}**") + st.write(step) + reactants = [r.get("smiles") for r in (step.get("reactants") or []) if r.get("smiles")] + products = [p.get("smiles") for p in (step.get("products") or []) if p.get("smiles")] + if reactants: + st.caption(f"Reactants: {', '.join(reactants)}") + if products: + st.caption(f"Products: {', '.join(products)}") + + +def display_forward_prediction_metadata(message): + data = message.get("forward_prediction") or {} + if not data: + return + st.markdown("### Forward Prediction") + backend = data.get("backend") + model_name = data.get("model_name") + 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( + predictions, + key=lambda p: p.get("score") if isinstance(p, dict) else None, + reverse=True, + )[:2] + if not predictions: + st.write("No predictions returned.") + return + if len(inputs) == 1: + base = inputs[0] + for idx, pred in enumerate(predictions, start=1): + 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) + else: + rows = [{"smiles": p.get("smiles"), "score": p.get("score")} for p in predictions] + st.dataframe(rows) + + +def display_reaction_classification_metadata(message): + data = message.get("reaction_classification") or {} + if not data: + return + st.markdown("### Reaction Classification") + status_code = data.get("status_code") + if status_code is not None: + st.markdown(f"**Status:** {status_code}") + msg = data.get("message") + if msg: + st.markdown(f"**Message:** {msg}") + results = data.get("result") or [] + if results: + st.dataframe(results) + else: + st.write("No classification results.") + + def display_paper_analysis_metadata(message, message_index): """ Display analysis details extracted from scientific papers, allowing users to selectively view text, images, and metadata. diff --git a/ChemCoScientist/tools/chemist_tools.py b/ChemCoScientist/tools/chemist_tools.py index 04422887..ab0b14bb 100644 --- a/ChemCoScientist/tools/chemist_tools.py +++ b/ChemCoScientist/tools/chemist_tools.py @@ -18,7 +18,10 @@ from pathlib import Path from dotenv import load_dotenv from ChemCoScientist.chemical_utils.ocr_pipeline import molecules_ocr, reactions_ocr -from ChemCoScientist.chemical_utils.chemical_functions import calculate_docking_score +from ChemCoScientist.chemical_utils.chemical_functions import ( + calculate_docking_score, +) +from ChemCoScientist.chemical_utils.retrosynthesis import retrosynthesis_result, classify_reaction_smiles, forward_predict_products import aiohttp import base64 @@ -655,6 +658,95 @@ def detect_reactions() -> Dict: logger.error(f'reactions_recognition ERROR: {e}') return {'answer': 'Could not detect any reactions in the uploaded images.'} +@tool +def retrosynthesis_tree_search( + smiles: Annotated[str, "Target molecule SMILES"], + mode: Annotated[str, "One of: fast, balanced, deep"] = "fast", +) -> Dict: + """ + Plan a retrosynthesis route for a target molecule. + + Use this when the user asks for possible synthetic routes or precursors + for a target SMILES. This calls the retrosynthesis service and returns + ASKCOS-like routes with steps, reactants, and scores. + + Args: + smiles (str): Target molecule SMILES. + mode (str): Search depth/quality preset ("fast", "balanced", "deep"). + + Returns: + dict: Retrosynthesis result with target and routes. + On failure returns a dict with an "answer" message. + """ + try: + return retrosynthesis_result(smiles=smiles, mode=mode) + except Exception as e: + logger.error(f"retrosynthesis_tree_search ERROR: {e}") + return {"answer": "Could not run retrosynthesis tree search."} + +@tool +def classify_reaction( + smiles: Annotated[List[str], "List of reaction SMILES, e.g. ['A.B>>C']"], + num_results: Annotated[int, "Max classes per reaction (1..50)"] = 10, +) -> Dict: + """ + Classify reaction SMILES into reaction classes. + + Use this when the user provides reaction SMILES and wants the reaction + type/class (e.g., named reactions or class labels). Returns ASKCOS-like + classification hits with ranks and confidence. + + Args: + smiles (List[str]): List of reaction SMILES to classify. + num_results (int): Max classes per reaction (1..50). + + Returns: + dict: status_code/message/result list with classification hits. + On failure returns a dict with an "answer" message. + """ + try: + return classify_reaction_smiles(smiles=smiles, num_results=num_results) + except Exception as e: + logger.error(f"classify_reaction ERROR: {e}") + return {"answer": "Could not classify reaction SMILES."} + +@tool +def forward_predict( + smiles: Annotated[List[str], "Batch of reaction inputs (reactants)"], + backend: Annotated[str, "One of: wldn5, graph2smiles, augmented_transformer"], + model_name: Annotated[str, "Model name for backend"] = "pistachio", + reagents: Annotated[str, "Reagents string"] = "", + solvent: Annotated[str, "Solvent string"] = "", +) -> Dict: + """ + Predict reaction products from reactants (forward synthesis). + + Use this when the user provides reactants and wants predicted products. + You can specify backend/model_name and optional reagents/solvent strings. + + Args: + smiles (List[str]): Batch of reaction inputs (reactants). + backend (str): Model backend ("wldn5", "graph2smiles", "augmented_transformer"). + model_name (str): Model name for backend (default "pistachio"). + reagents (str): Reagents string. + solvent (str): Solvent string. + + Returns: + dict: inputs/backend/model_name/predictions with product SMILES and scores. + On failure returns a dict with an "answer" message. + """ + try: + return forward_predict_products( + smiles=smiles, + backend=backend, + model_name=model_name, + reagents=reagents, + solvent=solvent, + ) + except Exception as e: + logger.error(f"forward_predict ERROR: {e}") + return {"answer": "Could not run forward prediction."} + @tool def calculate_docking(smiles: str, pdb_id: str) -> str: @@ -699,6 +791,9 @@ def calculate_docking(smiles: str, pdb_id: str) -> str: visualize_molecule, detect_molecules, detect_reactions, + retrosynthesis_tree_search, + classify_reaction, + forward_predict, calculate_docking, ] From 811531000f2a4e08e45713a509a8fd6610f36ffb Mon Sep 17 00:00:00 2001 From: AerDragon Date: Mon, 9 Feb 2026 18:11:23 +0300 Subject: [PATCH 6/8] chore(retrosynthesis): modification tools --- ChemCoScientist/agents/agents.py | 28 +++++++++++------ ChemCoScientist/agents/agents_prompts.py | 2 +- .../chemical_utils/retrosynthesis.py | 5 +-- ChemCoScientist/frontend/chat.py | 31 ++++++++++++------- 4 files changed, 42 insertions(+), 24 deletions(-) diff --git a/ChemCoScientist/agents/agents.py b/ChemCoScientist/agents/agents.py index 8ff363f1..79a90577 100644 --- a/ChemCoScientist/agents/agents.py +++ b/ChemCoScientist/agents/agents.py @@ -256,8 +256,19 @@ 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", "extract_reactions"]: result = json.loads(message.content) ocr_metadata = {"chem_ocr": result.get("metadata", None)} @@ -277,23 +288,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}) diff --git a/ChemCoScientist/agents/agents_prompts.py b/ChemCoScientist/agents/agents_prompts.py index 88ea7670..060a6f36 100644 --- a/ChemCoScientist/agents/agents_prompts.py +++ b/ChemCoScientist/agents/agents_prompts.py @@ -131,4 +131,4 @@ papers_search_prompt = """ You are a helpful assistant. You search for papers in OpenAlex based on a user query and download papers' PDFs. -""" \ No newline at end of file +""" diff --git a/ChemCoScientist/chemical_utils/retrosynthesis.py b/ChemCoScientist/chemical_utils/retrosynthesis.py index 20f3c950..064eb15f 100644 --- a/ChemCoScientist/chemical_utils/retrosynthesis.py +++ b/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 = ( diff --git a/ChemCoScientist/frontend/chat.py b/ChemCoScientist/frontend/chat.py index 22d39699..6ba0a165 100644 --- a/ChemCoScientist/frontend/chat.py +++ b/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 @@ -409,6 +411,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 @@ -692,10 +696,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) @@ -704,6 +711,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 @@ -762,10 +770,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}**") @@ -788,10 +799,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( @@ -804,16 +811,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) From c2aa0c48bebb6e3c02133c148268f81ff36a2def Mon Sep 17 00:00:00 2001 From: AerDragon Date: Wed, 25 Feb 2026 16:31:54 +0300 Subject: [PATCH 7/8] add retrosynthesis tool to mcp server --- ChemCoScientist/tools/chemist_tools.py | 4 +- example_config.env | 4 +- .../server/chemical_server.py | 94 ++++++++++++++++++- poetry.lock | 13 ++- 4 files changed, 110 insertions(+), 5 deletions(-) diff --git a/ChemCoScientist/tools/chemist_tools.py b/ChemCoScientist/tools/chemist_tools.py index ab0b14bb..ceaa24af 100644 --- a/ChemCoScientist/tools/chemist_tools.py +++ b/ChemCoScientist/tools/chemist_tools.py @@ -714,7 +714,7 @@ def classify_reaction( def forward_predict( smiles: Annotated[List[str], "Batch of reaction inputs (reactants)"], backend: Annotated[str, "One of: wldn5, graph2smiles, augmented_transformer"], - model_name: Annotated[str, "Model name for backend"] = "pistachio", + retrosynthesis_model_name: Annotated[str, "Model name for backend"] = "pistachio", reagents: Annotated[str, "Reagents string"] = "", solvent: Annotated[str, "Solvent string"] = "", ) -> Dict: @@ -739,7 +739,7 @@ def forward_predict( return forward_predict_products( smiles=smiles, backend=backend, - model_name=model_name, + model_name=retrosynthesis_model_name, reagents=reagents, solvent=solvent, ) diff --git a/example_config.env b/example_config.env index c68e66cb..6eaf254a 100644 --- a/example_config.env +++ b/example_config.env @@ -42,6 +42,8 @@ CHROMA_PORT=9999 EMBEDDING_PORT=9999 RERANKER_PORT=9999 CHEM_SERVICES_PORT=8005 +RETROSYNTHESIS_SERVICES_HOST="0.0.0.0" +RETROSYNTHESIS_SERVICES_PORT=8001 SUMMARIES_COLLECTION_NAME="your-name-here" TEXTS_COLLECTION_NAME="your-name-here" @@ -67,4 +69,4 @@ BACKEND_TYPE='ChemCoScientist' OPIK_API_KEY="your-api-key" OPIK_URL_OVERRIDE="https://www.comet.com/opik/api" OPIK_PROJECT_NAME="your-project-name" -OPIK_WORKSPACE="your-workspace-name" \ No newline at end of file +OPIK_WORKSPACE="your-workspace-name" diff --git a/mcp-servers/chemical-mcp-server/server/chemical_server.py b/mcp-servers/chemical-mcp-server/server/chemical_server.py index a482ade9..fc4f6186 100644 --- a/mcp-servers/chemical-mcp-server/server/chemical_server.py +++ b/mcp-servers/chemical-mcp-server/server/chemical_server.py @@ -21,7 +21,7 @@ from dotenv import load_dotenv from ChemCoScientist.chemical_utils.ocr_pipeline import molecules_ocr, reactions_ocr from ChemCoScientist.chemical_utils.chemical_functions import calculate_docking_score - +from ChemCoScientist.chemical_utils.retrosynthesis import retrosynthesis_result, classify_reaction_smiles, forward_predict_products import aiohttp import base64 import asyncio @@ -652,6 +652,98 @@ def calculate_docking( "answer": {"affinity": affinity, "errors": errors}, "metadata": {"html_file": output_file} if output_file else {}, } + +@mcp.tool() +def retrosynthesis_tree_search( + smiles: Annotated[str, "Target molecule SMILES"], + mode: Annotated[str, "One of: fast, balanced, deep"] = "fast", +) -> Dict: + """ + Plan a retrosynthesis route for a target molecule. + + Use this when the user asks for possible synthetic routes or precursors + for a target SMILES. This calls the retrosynthesis service and returns + ASKCOS-like routes with steps, reactants, and scores. + + Args: + smiles (str): Target molecule SMILES. + mode (str): Search depth/quality preset ("fast", "balanced", "deep"). + + Returns: + dict: Retrosynthesis result with target and routes. + On failure returns a dict with an "answer" message. + """ + try: + return retrosynthesis_result(smiles=smiles, mode=mode) + except Exception as e: + logger.error(f"retrosynthesis_tree_search ERROR: {e}") + return {"answer": "Could not run retrosynthesis tree search."} + +@mcp.tool() +def classify_reaction( + smiles: Annotated[List[str], "List of reaction SMILES, e.g. ['A.B>>C']"], + num_results: Annotated[int, "Max classes per reaction (1..50)"] = 10, +) -> Dict: + """ + Classify reaction SMILES into reaction classes. + + Use this when the user provides reaction SMILES and wants the reaction + type/class (e.g., named reactions or class labels). Returns ASKCOS-like + classification hits with ranks and confidence. + + Args: + smiles (List[str]): List of reaction SMILES to classify. + num_results (int): Max classes per reaction (1..50). + + Returns: + dict: status_code/message/result list with classification hits. + On failure returns a dict with an "answer" message. + """ + try: + return classify_reaction_smiles(smiles=smiles, num_results=num_results) + except Exception as e: + logger.error(f"classify_reaction ERROR: {e}") + return {"answer": "Could not classify reaction SMILES."} + +@mcp.tool() +def forward_predict( + smiles: Annotated[List[str], "Batch of reaction inputs (reactants)"], + backend: Annotated[str, "One of: wldn5, graph2smiles, augmented_transformer"], + retrosynthesis_model_name: Annotated[str, "Model name for backend"] = "pistachio", + reagents: Annotated[str, "Reagents string"] = "", + solvent: Annotated[str, "Solvent string"] = "", +) -> Dict: + """ + Predict reaction products from reactants (forward synthesis). + + Use this when the user provides reactants and wants predicted products. + You can specify backend/model_name and optional reagents/solvent strings. + + Args: + smiles (List[str]): Batch of reaction inputs (reactants). + backend (str): Model backend ("wldn5", "graph2smiles", "augmented_transformer"). + model_name (str): Model name for backend (default "pistachio"). + reagents (str): Reagents string. + solvent (str): Solvent string. + + Returns: + dict: inputs/backend/model_name/predictions with product SMILES and scores. + On failure returns a dict with an "answer" message. + """ + try: + return forward_predict_products( + smiles=smiles, + backend=backend, + model_name=retrosynthesis_model_name, + reagents=reagents, + solvent=solvent, + ) + except Exception as e: + logger.error(f"forward_predict ERROR: {e}") + return {"answer": "Could not run forward prediction."} + "answer": {"affinity": affinity, "errors": errors}, + "metadata": {"html_file": output_file} if output_file else {}, + } if __name__ == "__main__": diff --git a/poetry.lock b/poetry.lock index c39027bf..a140ce94 100644 --- a/poetry.lock +++ b/poetry.lock @@ -490,7 +490,6 @@ files = [ pyparsing = ">=2.0.3" [[package]] - name = "biopython" version = "1.86" description = "Freely available tools for computational molecular biology." @@ -2564,6 +2563,8 @@ files = [ {file = "greenlet-3.2.4-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c2ca18a03a8cfb5b25bc1cbe20f3d9a4c80d8c3b13ba3df49ac3961af0b1018d"}, {file = "greenlet-3.2.4-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:9fe0a28a7b952a21e2c062cd5756d34354117796c6d9215a87f55e38d15402c5"}, {file = "greenlet-3.2.4-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:8854167e06950ca75b898b104b63cc646573aa5fef1353d4508ecdd1ee76254f"}, + {file = "greenlet-3.2.4-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:f47617f698838ba98f4ff4189aef02e7343952df3a615f847bb575c3feb177a7"}, + {file = "greenlet-3.2.4-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:af41be48a4f60429d5cad9d22175217805098a9ef7c40bfef44f7669fb9d74d8"}, {file = "greenlet-3.2.4-cp310-cp310-win_amd64.whl", hash = "sha256:73f49b5368b5359d04e18d15828eecc1806033db5233397748f4ca813ff1056c"}, {file = "greenlet-3.2.4-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:96378df1de302bc38e99c3a9aa311967b7dc80ced1dcc6f171e99842987882a2"}, {file = "greenlet-3.2.4-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1ee8fae0519a337f2329cb78bd7a8e128ec0f881073d43f023c7b8d4831d5246"}, @@ -2573,6 +2574,8 @@ files = [ {file = "greenlet-3.2.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2523e5246274f54fdadbce8494458a2ebdcdbc7b802318466ac5606d3cded1f8"}, {file = "greenlet-3.2.4-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:1987de92fec508535687fb807a5cea1560f6196285a4cde35c100b8cd632cc52"}, {file = "greenlet-3.2.4-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:55e9c5affaa6775e2c6b67659f3a71684de4c549b3dd9afca3bc773533d284fa"}, + {file = "greenlet-3.2.4-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:c9c6de1940a7d828635fbd254d69db79e54619f165ee7ce32fda763a9cb6a58c"}, + {file = "greenlet-3.2.4-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:03c5136e7be905045160b1b9fdca93dd6727b180feeafda6818e6496434ed8c5"}, {file = "greenlet-3.2.4-cp311-cp311-win_amd64.whl", hash = "sha256:9c40adce87eaa9ddb593ccb0fa6a07caf34015a29bf8d344811665b573138db9"}, {file = "greenlet-3.2.4-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:3b67ca49f54cede0186854a008109d6ee71f66bd57bb36abd6d0a0267b540cdd"}, {file = "greenlet-3.2.4-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ddf9164e7a5b08e9d22511526865780a576f19ddd00d62f8a665949327fde8bb"}, @@ -2582,6 +2585,8 @@ files = [ {file = "greenlet-3.2.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3b3812d8d0c9579967815af437d96623f45c0f2ae5f04e366de62a12d83a8fb0"}, {file = "greenlet-3.2.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:abbf57b5a870d30c4675928c37278493044d7c14378350b3aa5d484fa65575f0"}, {file = "greenlet-3.2.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:20fb936b4652b6e307b8f347665e2c615540d4b42b3b4c8a321d8286da7e520f"}, + {file = "greenlet-3.2.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ee7a6ec486883397d70eec05059353b8e83eca9168b9f3f9a361971e77e0bcd0"}, + {file = "greenlet-3.2.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:326d234cbf337c9c3def0676412eb7040a35a768efc92504b947b3e9cfc7543d"}, {file = "greenlet-3.2.4-cp312-cp312-win_amd64.whl", hash = "sha256:a7d4e128405eea3814a12cc2605e0e6aedb4035bf32697f72deca74de4105e02"}, {file = "greenlet-3.2.4-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:1a921e542453fe531144e91e1feedf12e07351b1cf6c9e8a3325ea600a715a31"}, {file = "greenlet-3.2.4-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cd3c8e693bff0fff6ba55f140bf390fa92c994083f838fece0f63be121334945"}, @@ -2591,6 +2596,8 @@ files = [ {file = "greenlet-3.2.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:23768528f2911bcd7e475210822ffb5254ed10d71f4028387e5a99b4c6699671"}, {file = "greenlet-3.2.4-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:00fadb3fedccc447f517ee0d3fd8fe49eae949e1cd0f6a611818f4f6fb7dc83b"}, {file = "greenlet-3.2.4-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:d25c5091190f2dc0eaa3f950252122edbbadbb682aa7b1ef2f8af0f8c0afefae"}, + {file = "greenlet-3.2.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6e343822feb58ac4d0a1211bd9399de2b3a04963ddeec21530fc426cc121f19b"}, + {file = "greenlet-3.2.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ca7f6f1f2649b89ce02f6f229d7c19f680a6238af656f61e0115b24857917929"}, {file = "greenlet-3.2.4-cp313-cp313-win_amd64.whl", hash = "sha256:554b03b6e73aaabec3745364d6239e9e012d64c68ccd0b8430c64ccc14939a8b"}, {file = "greenlet-3.2.4-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:49a30d5fda2507ae77be16479bdb62a660fa51b1eb4928b524975b3bde77b3c0"}, {file = "greenlet-3.2.4-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:299fd615cd8fc86267b47597123e3f43ad79c9d8a22bebdce535e53550763e2f"}, @@ -2598,6 +2605,8 @@ files = [ {file = "greenlet-3.2.4-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b4a1870c51720687af7fa3e7cda6d08d801dae660f75a76f3845b642b4da6ee1"}, {file = "greenlet-3.2.4-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:061dc4cf2c34852b052a8620d40f36324554bc192be474b9e9770e8c042fd735"}, {file = "greenlet-3.2.4-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:44358b9bf66c8576a9f57a590d5f5d6e72fa4228b763d0e43fee6d3b06d3a337"}, + {file = "greenlet-3.2.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:2917bdf657f5859fbf3386b12d68ede4cf1f04c90c3a6bc1f013dd68a22e2269"}, + {file = "greenlet-3.2.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:015d48959d4add5d6c9f6c5210ee3803a830dce46356e3bc326d6776bde54681"}, {file = "greenlet-3.2.4-cp314-cp314-win_amd64.whl", hash = "sha256:e37ab26028f12dbb0ff65f29a8d3d44a765c61e729647bf2ddfbbed621726f01"}, {file = "greenlet-3.2.4-cp39-cp39-macosx_11_0_universal2.whl", hash = "sha256:b6a7c19cf0d2742d0809a4c05975db036fdff50cd294a93632d6a310bf9ac02c"}, {file = "greenlet-3.2.4-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:27890167f55d2387576d1f41d9487ef171849ea0359ce1510ca6e06c8bece11d"}, @@ -2607,6 +2616,8 @@ files = [ {file = "greenlet-3.2.4-cp39-cp39-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c9913f1a30e4526f432991f89ae263459b1c64d1608c0d22a5c79c287b3c70df"}, {file = "greenlet-3.2.4-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:b90654e092f928f110e0007f572007c9727b5265f7632c2fa7415b4689351594"}, {file = "greenlet-3.2.4-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:81701fd84f26330f0d5f4944d4e92e61afe6319dcd9775e39396e39d7c3e5f98"}, + {file = "greenlet-3.2.4-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:28a3c6b7cd72a96f61b0e4b2a36f681025b60ae4779cc73c1535eb5f29560b10"}, + {file = "greenlet-3.2.4-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:52206cd642670b0b320a1fd1cbfd95bca0e043179c1d8a045f2c6109dfe973be"}, {file = "greenlet-3.2.4-cp39-cp39-win32.whl", hash = "sha256:65458b409c1ed459ea899e939f0e1cdb14f58dbc803f2f93c5eab5694d32671b"}, {file = "greenlet-3.2.4-cp39-cp39-win_amd64.whl", hash = "sha256:d2e685ade4dafd447ede19c31277a224a239a0a1a4eca4e6390efedf20260cfb"}, {file = "greenlet-3.2.4.tar.gz", hash = "sha256:0dca0d95ff849f9a364385f36ab49f50065d76964944638be9691e1832e9f86d"}, From e3c41989f10656377853d28ad2955d2e260009cf Mon Sep 17 00:00:00 2001 From: AerDragon Date: Thu, 5 Mar 2026 15:15:13 +0300 Subject: [PATCH 8/8] move retrosynthesis tools to mcp --- ChemCoScientist/agents/agents.py | 4 ---- .../chemical_utils/chemical_functions.py | 19 ++++++++++--------- ChemCoScientist/conf/create_conf.py | 2 -- .../server/chemical_server.py | 4 ---- 4 files changed, 10 insertions(+), 19 deletions(-) diff --git a/ChemCoScientist/agents/agents.py b/ChemCoScientist/agents/agents.py index 79a90577..a6a02567 100644 --- a/ChemCoScientist/agents/agents.py +++ b/ChemCoScientist/agents/agents.py @@ -267,8 +267,6 @@ def _parse_tool_content(content): 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", "extract_reactions"]: result = json.loads(message.content) ocr_metadata = {"chem_ocr": result.get("metadata", None)} @@ -291,8 +289,6 @@ def _parse_tool_content(content): 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"]: diff --git a/ChemCoScientist/chemical_utils/chemical_functions.py b/ChemCoScientist/chemical_utils/chemical_functions.py index 738afc40..ae8b589b 100644 --- a/ChemCoScientist/chemical_utils/chemical_functions.py +++ b/ChemCoScientist/chemical_utils/chemical_functions.py @@ -18,7 +18,7 @@ REQUEST_TIMEOUT = 60 -def handle_api_request(host: str, endpoint: str, file_param_name: str = None, ): +def handle_api_request(endpoint: str, file_param_name: str = None, ): """ Decorator for handling requests to Chemical ToolsService API. @@ -45,7 +45,7 @@ def wrapper(*args, **kwargs) -> Any: Data from the "data" field of the API response """ try: - api_url = f"{host}{endpoint}" + api_url = f"{CHEM_SERVICES_URL}{endpoint}" logger.info(f"Calling ChemService API: {api_url}") if file_param_name: @@ -102,7 +102,7 @@ def wrapper(*args, **kwargs) -> Any: return decorator -@handle_api_request(host=CHEM_SERVICES_URL, endpoint="/extract_reactions_from_pdf/", file_param_name="pdf_file") +@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. @@ -121,7 +121,7 @@ def extract_reactions_from_pdf(file: bytes) -> List[Dict]: pass -@handle_api_request(host=CHEM_SERVICES_URL, endpoint="/extract_reactions_from_figure/", file_param_name="image") +@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. @@ -136,7 +136,7 @@ def extract_reactions_from_figure(image: bytes) -> List[Dict]: pass -@handle_api_request(host=CHEM_SERVICES_URL, endpoint="/extract_molecules_from_pdf/", file_param_name="pdf_file") +@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. @@ -151,7 +151,7 @@ def extract_molecules_from_pdf(file: bytes) -> List[Dict]: pass -@handle_api_request(host=CHEM_SERVICES_URL, endpoint="/extract_molecules_from_figure/", file_param_name="image") +@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. @@ -166,7 +166,7 @@ def extract_molecules_from_figure(image: bytes) -> List[Dict]: pass -@handle_api_request(host=CHEM_SERVICES_URL, endpoint="/convert_image_to_smiles/", file_param_name="image") +@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. @@ -178,7 +178,7 @@ def convert_image_to_smiles(image: bytes) -> str: """ pass -@handle_api_request(host=CHEM_SERVICES_URL, endpoint="/docking/", file_param_name=None) +@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. @@ -191,6 +191,7 @@ def calculate_docking_score(smiles: str, pdb_id: str) -> str: """ 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): @@ -203,7 +204,7 @@ def remove_keys(obj: Any, keys_to_remove: set[str] = {"bbox", "score"}) -> Any: remove_keys(item, keys_to_remove) return obj + if __name__ == "__main__": result = calculate_docking_score(smiles="C1CCCCC1", pdb_id="5vfi") print(result) - diff --git a/ChemCoScientist/conf/create_conf.py b/ChemCoScientist/conf/create_conf.py index 32a5ea54..2ee137b7 100644 --- a/ChemCoScientist/conf/create_conf.py +++ b/ChemCoScientist/conf/create_conf.py @@ -159,8 +159,6 @@ and an empty or absent `metadata.papers`. """ -======= ->>>>>>> beee5da (add docking score calculation in chemist agent) additional_agents_description = ( automl_agent_description diff --git a/mcp-servers/chemical-mcp-server/server/chemical_server.py b/mcp-servers/chemical-mcp-server/server/chemical_server.py index fc4f6186..98ae2c8d 100644 --- a/mcp-servers/chemical-mcp-server/server/chemical_server.py +++ b/mcp-servers/chemical-mcp-server/server/chemical_server.py @@ -741,10 +741,6 @@ def forward_predict( except Exception as e: logger.error(f"forward_predict ERROR: {e}") return {"answer": "Could not run forward prediction."} - "answer": {"affinity": affinity, "errors": errors}, - "metadata": {"html_file": output_file} if output_file else {}, - } - if __name__ == "__main__": mcp.run(transport="http", host="0.0.0.0", port=7331, path="/mcp")