From f0de3a729ad2eddc19291a88ffbe4df345f538dd Mon Sep 17 00:00:00 2001 From: Shehtab Date: Wed, 19 Nov 2025 00:37:39 -0500 Subject: [PATCH 01/11] Updated SMILES and SMARTS server implementations with new ServerToolkits --- charge/servers/SMARTS_reactions.py | 54 ++++++++++++++-------- charge/servers/SMARTS_reactions_utils.py | 2 +- charge/servers/SMILES.py | 57 ++++++++++++++++-------- 3 files changed, 75 insertions(+), 38 deletions(-) diff --git a/charge/servers/SMARTS_reactions.py b/charge/servers/SMARTS_reactions.py index 87f1443..3ba7aff 100644 --- a/charge/servers/SMARTS_reactions.py +++ b/charge/servers/SMARTS_reactions.py @@ -7,10 +7,12 @@ from mcp.server.fastmcp import FastMCP from loguru import logger +from charge.servers.ServerToolkit import ServerToolkit try: - from rdkit import Chem - from rdkit.Chem import AllChem, rdChemReactions + from charge.servers.SMARTS_reactions_utils import verify_reaction_SMARTS + from charge.servers.SMILES_utils import verify_smiles + from charge.servers.SMARTS_reactions_utils import verify_reaction HAS_SMARTS = True except (ImportError, ModuleNotFoundError) as e: HAS_SMARTS = False @@ -19,26 +21,42 @@ "Install it with: pip install charge[rdkit]", ) -from charge.servers.server_utils import add_server_arguments -import argparse -parser = argparse.ArgumentParser() -add_server_arguments(parser) -args = parser.parse_args() +class SMARTSServer(ServerToolkit): + """ + A ChARGe server that provides tools for Chemistry and reaction verification. + """ + def __init__(self, mcp: FastMCP): + """ + Initialize the SMARTSServer. -SMARTS_mcp = FastMCP( - "[RDKit-SMARTS] Chemistry and reaction verification MCP Server", - port=args.port, - website_url=f"{args.host}", -) + Args: + mcp (FastMCP): The MCP instance to register the function with. + """ + super().__init__(mcp) -logger.info("[RDKit-SMARTS] Starting Chemistry and reaction verification MCP Server") + if HAS_SMARTS: + self.register_function_as_tool(self.mcp, verify_reaction_SMARTS) + self.register_function_as_tool(self.mcp, verify_smiles) + self.register_function_as_tool(self.mcp, verify_reaction) -import charge.servers.SMARTS_reactions_utils as smarts -import charge.servers.SMILES_utils as smiles -SMARTS_mcp.tool()(smarts.verify_reaction_SMARTS) +if __name__ == "__main__": + from charge.servers.server_utils import add_server_arguments + import argparse -SMARTS_mcp.tool()(smiles.verify_smiles) + logger.info("[RDKit-SMARTS] Starting Chemistry and reaction verification MCP Server") -SMARTS_mcp.tool()(smarts.verify_reaction) + parser = argparse.ArgumentParser() + add_server_arguments(parser) + args = parser.parse_args() + + SMARTS_mcp = FastMCP( + "[RDKit-SMARTS] Chemistry and reaction verification MCP Server", + port=args.port, + website_url=f"{args.host}", + ) + + sm = SMARTSServer(SMARTS_mcp) + mcp = sm.return_mcp() + mcp.run() \ No newline at end of file diff --git a/charge/servers/SMARTS_reactions_utils.py b/charge/servers/SMARTS_reactions_utils.py index 43a990e..fcf8974 100644 --- a/charge/servers/SMARTS_reactions_utils.py +++ b/charge/servers/SMARTS_reactions_utils.py @@ -27,7 +27,7 @@ def verify_reaction_SMARTS(smarts: str) -> Tuple[bool, str]: The bool indicates if the SMARTS is valid, and the str is an error message if it is not. Args: - smiles (str): The input SMILES string. + smarts (str): The input SMARTS string. Returns: A tuple containing: bool: True if the SMARTS is valid, False if it is invalid. diff --git a/charge/servers/SMILES.py b/charge/servers/SMILES.py index a02cfcc..36773a0 100644 --- a/charge/servers/SMILES.py +++ b/charge/servers/SMILES.py @@ -7,11 +7,11 @@ from mcp.server.fastmcp import FastMCP from loguru import logger +from charge.servers.ServerToolkit import ServerToolkit + try: - from rdkit import Chem - from rdkit.Chem import AllChem, Descriptors - from rdkit.Contrib.SA_Score import sascorer + import charge.servers.SMILES_utils as smiles HAS_SMILES = True except (ImportError, ModuleNotFoundError) as e: HAS_SMILES = False @@ -20,27 +20,46 @@ "Install it with: pip install charge[rdkit]", ) -from charge.servers.server_utils import add_server_arguments -import argparse -parser = argparse.ArgumentParser() -add_server_arguments(parser) -args = parser.parse_args() -SMILES_mcp = FastMCP( - "[RDKit-SMILES] Chem and BioInformatics MCP Server", - port=args.port, - website_url=f"{args.host}", -) +class SMILESServer(ServerToolkit): + """ + A ChARGe server that provides SMILES based tools. + """ + def __init__(self, mcp: FastMCP): + """ + Initialize the SMILESServer. + + Args: + mcp (FastMCP): The MCP instance to register the function with. + """ + super().__init__(mcp) + + if HAS_SMILES: + self.register_function_as_tool(self.mcp, smiles.canonicalize_smiles) + self.register_function_as_tool(self.mcp, smiles.verify_smiles) + self.register_function_as_tool(self.mcp, smiles.get_synthesizability) + self.register_function_as_tool(self.mcp, smiles.known_smiles) + -logger.info("[RDKit-SMILES] Starting Chem and BioInformatics MCP Server") -import charge.servers.SMILES_utils as smiles -SMILES_mcp.tool()(smiles.canonicalize_smiles) +if __name__ == "__main__": + from charge.servers.server_utils import add_server_arguments + import argparse -SMILES_mcp.tool()(smiles.verify_smiles) + parser = argparse.ArgumentParser() + add_server_arguments(parser) + args = parser.parse_args() + + SMILES_mcp = FastMCP( + "[RDKit-SMILES] Chem and BioInformatics MCP Server", + port=args.port, + website_url=f"{args.host}", + ) -SMILES_mcp.tool()(smiles.get_synthesizability) + logger.info("[RDKit-SMILES] Starting Chem and BioInformatics MCP Server") -SMILES_mcp.tool()(smiles.known_smiles) + sm = SMILESServer(SMILES_mcp) + mcp = sm.return_mcp() + mcp.run() \ No newline at end of file From f3445589127351806696f958876d6c487941f8d4 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Wed, 19 Nov 2025 00:38:11 -0500 Subject: [PATCH 02/11] Add ServerToolkit class --- charge/servers/ServerToolkit.py | 90 +++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 charge/servers/ServerToolkit.py diff --git a/charge/servers/ServerToolkit.py b/charge/servers/ServerToolkit.py new file mode 100644 index 0000000..52233e1 --- /dev/null +++ b/charge/servers/ServerToolkit.py @@ -0,0 +1,90 @@ +################################################################################ +## Copyright 2025 Lawrence Livermore National Security, LLC. and Binghamton University. +## See the top-level LICENSE file for details. +## +## SPDX-License-Identifier: Apache-2.0 +################################################################################ + +from mcp.server.fastmcp import FastMCP +from functools import wraps +from typing import Callable + + +class ServerToolkit: + """ + A class that provides a toolkit for registering methods as MCP tools. + """ + def __init__(self, mcp: FastMCP): + self.mcp = mcp + self._pending_methods = [] + + def _register_methods(self): + """ + Register all methods marked with the @mcp_tool decorator. + """ + for name in dir(self): + # Skip private/magic methods + if name.startswith("_"): + continue + + attr = getattr(self, name) + + # Check if this is a method marked for registration + if callable(attr) and hasattr(attr, "_is_mcp_tool"): + self._register_single_method(attr) + + def _register_single_method(self, method: Callable): + """ + Register a single method as an MCP tool. + + Args: + method (Callable): The method to register as an MCP tool. + + Returns: + None + """ + + # Create a closure that captures the bound method + @wraps(method) + def tool_wrapper(*args, **kwargs): + return method(*args, **kwargs) + + self.mcp.tool()(tool_wrapper) + + @staticmethod + def mcp_tool(func: Callable) -> Callable: + """ + Decorator to mark methods for MCP registration. + + Args: + func (Callable): The function to mark for MCP registration. + + Returns: + Callable: The marked function. + """ + func._is_mcp_tool = True + return func + + @staticmethod + def register_function_as_tool(mcp: FastMCP, func: Callable) -> None: + """ + Register an external function as an MCP tool. + + Args: + mcp (FastMCP): The MCP instance to register the function with. + func (Callable): The function to register as an MCP tool. + + Returns: + None + """ + mcp.tool()(func) + + def return_mcp(self) -> FastMCP: + """ + Return the MCP instance. + + Returns: + FastMCP: The MCP instance. + """ + self._register_methods() + return self.mcp From 1c8777db3eac15208029c32bf785415fa2252ae2 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Wed, 19 Nov 2025 01:10:26 -0500 Subject: [PATCH 03/11] - Convert molecule generation server to use AgentPool rather than deprecated Client - Update documentation to use Google style --- charge/servers/AiZynthTools.py | 24 +- charge/servers/SMARTS_reactions.py | 8 +- charge/servers/SMARTS_reactions_utils.py | 13 +- charge/servers/SMILES.py | 9 +- charge/servers/SMILES_utils.py | 14 +- charge/servers/ServerToolkit.py | 14 +- charge/servers/log_progress.py | 3 +- charge/servers/molecular_generation_server.py | 235 +++++++++--------- charge/servers/molecular_property_utils.py | 125 ++++++---- charge/servers/molecule_pricer.py | 63 ++--- .../servers/retrosynthesis_reaction_server.py | 98 ++++++-- charge/servers/server_utils.py | 13 +- 12 files changed, 373 insertions(+), 246 deletions(-) diff --git a/charge/servers/AiZynthTools.py b/charge/servers/AiZynthTools.py index 952cc82..a7c9b6c 100644 --- a/charge/servers/AiZynthTools.py +++ b/charge/servers/AiZynthTools.py @@ -57,11 +57,15 @@ def __init__(self, route): self.nodes: Dict[int, Node] = {} self.num_nodes = 0 self._build_path() - self.leaf_nodes = [node_id for node_id, node in self.nodes.items() if node.is_leaf] + self.leaf_nodes = [ + node_id for node_id, node in self.nodes.items() if node.is_leaf + ] def _build_path(self): - self.root = Node(node_id=0, smiles=self.route["smiles"], children=[], is_root=True) + self.root = Node( + node_id=0, smiles=self.route["smiles"], children=[], is_root=True + ) self.nodes[0] = self.root self.num_nodes += 1 reaction_node = self.route["children"][0] @@ -86,7 +90,9 @@ def _add_children(self, parent_node, reaction, children): self.num_nodes += 1 if "children" in child: reaction_node = child["children"][0] - self._add_children(child_node, reaction_node, reaction_node["children"]) + self._add_children( + child_node, reaction_node, reaction_node["children"] + ) else: child_node.is_leaf = True @@ -143,7 +149,9 @@ def is_molecule_synthesizable(smiles: str) -> bool: ValueError: If the molecule is not valid. """ if not HAS_AIZYNTHFINDER: - raise ImportError("Please install the aizynthfinder support packages to use this module.") + raise ImportError( + "Please install the aizynthfinder support packages to use this module." + ) logger.info(f"Checking if molecule {smiles} is synthesizable.") @@ -163,7 +171,9 @@ def is_molecule_synthesizable(smiles: str) -> bool: for route in routes: path = ReactionPath(route=route) # check if all leaf nodes are purchasable - all_purchasable = all(path.nodes[node_id].purchasable is True for node_id in path.leaf_nodes) + all_purchasable = all( + path.nodes[node_id].purchasable is True for node_id in path.leaf_nodes + ) if all_purchasable: return True return False @@ -181,7 +191,9 @@ def find_synthesis_routes(smiles: str) -> list[dict]: ValueError: If the molecule is not valid. """ if not HAS_AIZYNTHFINDER: - raise ImportError("Please install the aizynthfinder support packages to use this module.") + raise ImportError( + "Please install the aizynthfinder support packages to use this module." + ) logger.info(f"Find a synthesis route for molecule {smiles}.") diff --git a/charge/servers/SMARTS_reactions.py b/charge/servers/SMARTS_reactions.py index 3ba7aff..8f09938 100644 --- a/charge/servers/SMARTS_reactions.py +++ b/charge/servers/SMARTS_reactions.py @@ -13,6 +13,7 @@ from charge.servers.SMARTS_reactions_utils import verify_reaction_SMARTS from charge.servers.SMILES_utils import verify_smiles from charge.servers.SMARTS_reactions_utils import verify_reaction + HAS_SMARTS = True except (ImportError, ModuleNotFoundError) as e: HAS_SMARTS = False @@ -26,6 +27,7 @@ class SMARTSServer(ServerToolkit): """ A ChARGe server that provides tools for Chemistry and reaction verification. """ + def __init__(self, mcp: FastMCP): """ Initialize the SMARTSServer. @@ -45,7 +47,9 @@ def __init__(self, mcp: FastMCP): from charge.servers.server_utils import add_server_arguments import argparse - logger.info("[RDKit-SMARTS] Starting Chemistry and reaction verification MCP Server") + logger.info( + "[RDKit-SMARTS] Starting Chemistry and reaction verification MCP Server" + ) parser = argparse.ArgumentParser() add_server_arguments(parser) @@ -59,4 +63,4 @@ def __init__(self, mcp: FastMCP): sm = SMARTSServer(SMARTS_mcp) mcp = sm.return_mcp() - mcp.run() \ No newline at end of file + mcp.run() diff --git a/charge/servers/SMARTS_reactions_utils.py b/charge/servers/SMARTS_reactions_utils.py index fcf8974..ef40fda 100644 --- a/charge/servers/SMARTS_reactions_utils.py +++ b/charge/servers/SMARTS_reactions_utils.py @@ -7,9 +7,11 @@ from loguru import logger from mcp.server.fastmcp import FastMCP + try: from rdkit import Chem from rdkit.Chem import AllChem, rdChemReactions + HAS_SMARTS = True except (ImportError, ModuleNotFoundError) as e: HAS_SMARTS = False @@ -20,6 +22,7 @@ from typing import Tuple + def verify_reaction_SMARTS(smarts: str) -> Tuple[bool, str]: """ Verify if a SMARTS string is valid. @@ -34,7 +37,9 @@ def verify_reaction_SMARTS(smarts: str) -> Tuple[bool, str]: str: Error message if the SMARTS reaction is valid. """ if not HAS_SMARTS: - raise ImportError("Please install the rdkit support packages to use this module.") + raise ImportError( + "Please install the rdkit support packages to use this module." + ) try: logger.info(f"Verifying SMARTS: {smarts}") rxn = AllChem.ReactionFromSmarts(smarts) @@ -59,6 +64,7 @@ def verify_reaction_SMARTS(smarts: str) -> Tuple[bool, str]: logger.error(f"Invalid SMARTS string: {e}") return False, f"Invalid Syntax for SMARTS string. The error is: {e}" + def verify_reaction( smarts: str, reactants: list[str], products: list[str] ) -> Tuple[bool, str]: @@ -79,7 +85,9 @@ def verify_reaction( str: Error message if the SMARTS reaction is valid. """ if not HAS_SMARTS: - raise ImportError("Please install the rdkit support packages to use this module.") + raise ImportError( + "Please install the rdkit support packages to use this module." + ) try: logger.info( f"Verifying reaction with SMARTS: {smarts}, Reactants: {reactants}, Products: {products}" @@ -131,4 +139,3 @@ def verify_reaction( except Exception as e: logger.error(f"Error verifying reaction: {e}") return False, f"Error verifying reaction: {e}" - diff --git a/charge/servers/SMILES.py b/charge/servers/SMILES.py index 36773a0..fa0c84d 100644 --- a/charge/servers/SMILES.py +++ b/charge/servers/SMILES.py @@ -12,6 +12,7 @@ try: import charge.servers.SMILES_utils as smiles + HAS_SMILES = True except (ImportError, ModuleNotFoundError) as e: HAS_SMILES = False @@ -21,11 +22,11 @@ ) - class SMILESServer(ServerToolkit): """ A ChARGe server that provides SMILES based tools. - """ + """ + def __init__(self, mcp: FastMCP): """ Initialize the SMILESServer. @@ -40,8 +41,6 @@ def __init__(self, mcp: FastMCP): self.register_function_as_tool(self.mcp, smiles.verify_smiles) self.register_function_as_tool(self.mcp, smiles.get_synthesizability) self.register_function_as_tool(self.mcp, smiles.known_smiles) - - if __name__ == "__main__": @@ -62,4 +61,4 @@ def __init__(self, mcp: FastMCP): sm = SMILESServer(SMILES_mcp) mcp = sm.return_mcp() - mcp.run() \ No newline at end of file + mcp.run() diff --git a/charge/servers/SMILES_utils.py b/charge/servers/SMILES_utils.py index fd2b0d9..35a60b2 100644 --- a/charge/servers/SMILES_utils.py +++ b/charge/servers/SMILES_utils.py @@ -12,6 +12,7 @@ from rdkit import Chem from rdkit.Chem import AllChem, Descriptors from rdkit.Contrib.SA_Score import sascorer + HAS_SMILES = True except (ImportError, ModuleNotFoundError) as e: HAS_SMILES = False @@ -20,6 +21,7 @@ "Install it with: pip install charge[rdkit]", ) + def canonicalize_smiles(smiles: str) -> str: """ Canonicalize a SMILES string. Returns the canonical SMILES. @@ -55,7 +57,9 @@ def verify_smiles(smiles: str) -> bool: bool: True if the SMILES is valid, False otherwise. """ if not HAS_SMILES: - raise ImportError("Please install the rdkit support packages to use this module.") + raise ImportError( + "Please install the rdkit support packages to use this module." + ) try: global SMILES_VERIFICATION_COUNTER SMILES_VERIFICATION_COUNTER += 1 @@ -85,7 +89,9 @@ def get_synthesizability(smiles: str) -> float: float: The synthesizability score. """ if not HAS_SMILES: - raise ImportError("Please install the rdkit support packages to use this module.") + raise ImportError( + "Please install the rdkit support packages to use this module." + ) try: # logger.info(f"Calculating synthesizability for SMILES: {smiles}") mol = Chem.MolFromSmiles(smiles) @@ -115,7 +121,9 @@ def known_smiles(smiles: str) -> bool: """ if not HAS_SMILES: - raise ImportError("Please install the rdkit support packages to use this module.") + raise ImportError( + "Please install the rdkit support packages to use this module." + ) try: global NUM_HITS logger.info(f"Tool has been call: {NUM_HITS} times") diff --git a/charge/servers/ServerToolkit.py b/charge/servers/ServerToolkit.py index 52233e1..e949ca2 100644 --- a/charge/servers/ServerToolkit.py +++ b/charge/servers/ServerToolkit.py @@ -7,13 +7,14 @@ from mcp.server.fastmcp import FastMCP from functools import wraps -from typing import Callable +from typing import Callable, Literal class ServerToolkit: """ A class that provides a toolkit for registering methods as MCP tools. """ + def __init__(self, mcp: FastMCP): self.mcp = mcp self._pending_methods = [] @@ -67,7 +68,7 @@ def mcp_tool(func: Callable) -> Callable: @staticmethod def register_function_as_tool(mcp: FastMCP, func: Callable) -> None: - """ + """ Register an external function as an MCP tool. Args: @@ -88,3 +89,12 @@ def return_mcp(self) -> FastMCP: """ self._register_methods() return self.mcp + + def run(self, transport: Literal["sse", "stdio"] = "sse") -> None: + """ + Run the MCP server. + + Args: + transport (Literal["sse", "stdio"], optional): The transport to use. Defaults to "sse". + """ + self.mcp.run(transport=transport) diff --git a/charge/servers/log_progress.py b/charge/servers/log_progress.py index e50bc60..a388bdf 100644 --- a/charge/servers/log_progress.py +++ b/charge/servers/log_progress.py @@ -8,7 +8,7 @@ from loguru import logger -LOG_PROGRESS_SYSTEM_PROMPT="At each step of your reasoning use the log_progress tool to report your current prograss, current thinking, and plan." +LOG_PROGRESS_SYSTEM_PROMPT = "At each step of your reasoning use the log_progress tool to report your current prograss, current thinking, and plan." def log_progress(log_msg: str) -> None: @@ -22,4 +22,3 @@ def log_progress(log_msg: str) -> None: """ logger.info(f"[ChARGe Orchestrator Inner Monologue] {log_msg}") - diff --git a/charge/servers/molecular_generation_server.py b/charge/servers/molecular_generation_server.py index 3e5160b..ff7c5a8 100644 --- a/charge/servers/molecular_generation_server.py +++ b/charge/servers/molecular_generation_server.py @@ -6,8 +6,10 @@ ################################################################################ from loguru import logger + try: from rdkit import Chem + HAS_RDKIT = True except (ImportError, ModuleNotFoundError) as e: HAS_RDKIT = False @@ -21,33 +23,20 @@ import os from charge.tasks.Task import Task from charge.servers.server_utils import add_server_arguments -from mcp.server.fastmcp import FastMCP from charge.clients.autogen import AutoGenClient from charge.clients.Client import Client import asyncio from charge.servers import SMILES_utils import charge.utils.helper_funcs as hf import argparse - -parser = argparse.ArgumentParser() -add_server_arguments(parser) -args = parser.parse_args() - -mcp = FastMCP( - "SMILES Diagnosis and retrieval MCP Server", - port=args.port, - website_url=f"{args.host}", -) - -MODEL = "gpt-oss:latest" -BACKEND = "ollama" -API_KEY = None -KWARGS = {} -JSON_FILE_PATH = f"{os.getcwd()}/known_molecules.json" +from charge.servers import ServerToolkit +from charge.servers.SMILES import SMILESServer +from charge.clients.autogen import AutoGenPool class DiagnoseSMILESTask(Task): - def __init__(self): + def __init__(self, smiles: str): + self.smiles = smiles system_prompt = ( "You are a world-class chemist. Your task is to diagnose and evaluate " "the quality of the provided SMILES strings. You will be given invalid" @@ -56,122 +45,136 @@ def __init__(self): ) user_prompt = ( - "Diagnose the followig SMILES string {0}. Give it a short and concise " + f"Diagnose the following SMILES string {self.smiles}. Give it a short and concise " "explanation of what is wrong with it, and if possible, provide a corrected " "version of the SMILES string. If the SMILES string is valid, simply state " "'The SMILES string is valid.'" ) super().__init__(system_prompt=system_prompt, user_prompt=user_prompt) - def update_user_prompt(self, smiles: str) -> None: - assert self.user_prompt is not None - self.user_prompt.format(smiles) - -@mcp.tool() -def diagnose_smiles(smiles: str) -> str: +class MoleculeGenerationServer(SMILESServer): """ - Diagnose a SMILES string. Returns a diagnosis of the SMILES string. - - Args: - smiles (str): The input SMILES string. - Returns: - str: The diagnosis of the SMILES string. + A ChARGe server that provides molecular generation tools. """ - if not HAS_RDKIT: - raise ImportError("Please install the rdkit support packages to use this module.") - logger.info(f"Diagnosing SMILES string: {smiles}") - task = DiagnoseSMILESTask() - task.update_user_prompt(smiles) - diagnose_agent = AutoGenClient( - task=task, - model=MODEL, - backend=BACKEND, - api_key=API_KEY, - model_kwargs=KWARGS, - ) - - try: - response = asyncio.run(diagnose_agent.run()) - assert response is not None - assert len(response.messages) > 0 # type: ignore - assert response.messages[-1] is not None # type: ignore - diagnoses = response.messages[-1].content # type: ignore - logger.info(f"Diagnosis: {diagnoses}") - return f"SMILES diagnoses: {diagnoses}" # type: ignore + def __init__(self, mcp: FastMCP, model: str, backend: str, json_file_path: str): + """ + Initialize the MoleculeGenerationServer. Inherits from SMILESServer. + + Args: + mcp (FastMCP): The MCP instance to register the function with. + model (str): The model to use for generation. + backend (str): The backend to use for generation. + json_file_path (str): The path to the JSON file containing known molecules. + """ + super().__init__(mcp=mcp) + self.model = model + self.backend = backend + self.json_file_path = json_file_path + + self.agent_pool = AutoGenPool( + model=self.model, + backend=self.backend, + ) - except Exception as e: - logger.error(f"An error occurred: {e}") - return "Error: Unable to process the SMILES string at this time." + @mcp_tool + def diagnose_smiles(self, smiles: str) -> str: + """ + Diagnose a SMILES string. Returns a diagnosis of the SMILES string. + Args: + smiles (str): The input SMILES string. -@mcp.tool() -def is_already_known(smiles: str) -> bool: - """ - Check if a SMILES string provided is already known. Only provide - valid SMILES strings. Returns True if the SMILES string is valid, and - already in the database, False otherwise. - Args: - smiles (str): The input SMILES string. - Returns: - bool: True if the SMILES string is valid and known, False otherwise. - - Raises: - ValueError: If the SMILES string is invalid. - """ - if not HAS_RDKIT: - raise ImportError("Please install the rdkit support packages to use this module.") - if not Chem.MolFromSmiles(smiles): - raise ValueError("Invalid SMILES string.") + Returns: + str: The diagnosis of the SMILES string. + """ + if not HAS_RDKIT: + raise ImportError( + "Please install the rdkit support packages to use this module." + ) + logger.info(f"Diagnosing SMILES string: {smiles}") + task = DiagnoseSMILESTask(smiles=smiles) - try: - canonical_smiles = SMILES_utils.canonicalize_smiles(smiles) + try: + diagnose_agent = self.agent_pool.create_agent(task=task) + response = asyncio.run(diagnose_agent.run()) + assert response is not None + assert len(response.messages) > 0 # type: ignore + assert response.messages[-1] is not None # type: ignore + + diagnoses = response.messages[-1].content # type: ignore + logger.info(f"Diagnosis: {diagnoses}") + return f"SMILES diagnoses: {diagnoses}" # type: ignore + + except Exception as e: + logger.error(f"An error occurred: {e}") + return "Error: Unable to process the SMILES string at this time." + + @mcp_tool + def is_already_known(self, smiles: str) -> bool: + """ + Check if a SMILES string provided is already known. Only provide + valid SMILES strings. Returns True if the SMILES string is valid, and + already in the database, False otherwise. + + Args: + smiles (str): The input SMILES string. + + Returns: + bool: True if the SMILES string is valid and known, False otherwise. + + Raises: + ValueError: If the SMILES string is invalid. + """ + if not HAS_RDKIT: + raise ImportError( + "Please install the rdkit support packages to use this module." + ) + if not Chem.MolFromSmiles(smiles): + raise ValueError("Invalid SMILES string.") try: - with open(JSON_FILE_PATH) as f: - known_mols = json.load(f) - known_smiles = [mol["smiles"] for mol in known_mols] + canonical_smiles = SMILES_utils.canonicalize_smiles(smiles) - except FileNotFoundError: - logger.warning(f"{JSON_FILE_PATH} not found. Creating a new one.") - known_mols = [] + try: + with open(self.json_file_path) as f: + known_mols = json.load(f) + known_smiles = [mol["smiles"] for mol in known_mols] - except Exception as e: - raise ValueError("Error in canonicalizing SMILES string.") from e + except FileNotFoundError: + logger.warning(f"{self.json_file_path} not found. Creating a new one.") + known_mols = [] - # Check if the SMILES string is already known (in the database) - # This is a placeholder for the actual database check - return canonical_smiles in known_smiles + except Exception as e: + raise ValueError("Error in canonicalizing SMILES string.") from e + # Check if the SMILES string is already known (in the database) + return canonical_smiles in known_smiles -@mcp.tool() -def get_density(smiles: str) -> float: - """ - Calculate the density of a molecule given its SMILES string. - - Args: - smiles (str): The input SMILES string. - Returns: - float: The density of the molecule. - """ - if not HAS_RDKIT: - raise ImportError("Please install the rdkit support packages to use this module.") - density = hf.get_density(smiles) - logger.info(f"Density for SMILES {smiles}: {density}") - return density + @mcp_tool + def get_density(self, smiles: str) -> float: + """ + Calculate the density of a molecule given its SMILES string. + Args: + smiles (str): The input SMILES string. -# Add the SMILES utility functions as MCP tools -mcp.tool()(SMILES_utils.canonicalize_smiles) -mcp.tool()(SMILES_utils.verify_smiles) -mcp.tool()(SMILES_utils.get_synthesizability) + Returns: + float: The density of the molecule. + """ + if not HAS_RDKIT: + raise ImportError( + "Please install the rdkit support packages to use this module." + ) + density = hf.get_density(smiles) + logger.info(f"Density for SMILES {smiles}: {density}") + return density if __name__ == "__main__": - if not HAS_RDKIT: - raise ImportError("Please install the rdkit support packages to use this module.") parser = argparse.ArgumentParser(description="Molecule Tools Server") + add_server_arguments(parser) Client.add_std_parser_arguments(parser) parser.add_argument( "--json_file", @@ -181,14 +184,14 @@ def get_density(smiles: str) -> float: ) args = parser.parse_args() - # global MODEL, BACKEND, API_KEY, KWARGS, JSON_FILE_PATH - MODEL = args.model if args.model else MODEL - BACKEND = args.backend if args.backend else BACKEND - MODEL, BACKEND, API_KEY, KWARGS = AutoGenClient.configure( - model=MODEL, backend=BACKEND + + mcp = FastMCP( + "Molecule Generation MCP Server", + port=args.port, + website_url=f"{args.host}", ) - logger.info(f"Using model: {MODEL} on backend: {BACKEND}") - JSON_FILE_PATH = args.json_file if args.json_file else JSON_FILE_PATH - logger.info(f"Using known molecules database at: {JSON_FILE_PATH}") - mcp.run(transport="sse") + server = MoleculeGenerationServer( + mcp=mcp, model=args.model, backend=args.backend, json_file_path=args.json_file + ) + server.run(transport="sse") diff --git a/charge/servers/molecular_property_utils.py b/charge/servers/molecular_property_utils.py index c2244c3..c012b2f 100644 --- a/charge/servers/molecular_property_utils.py +++ b/charge/servers/molecular_property_utils.py @@ -6,10 +6,12 @@ ################################################################################ from loguru import logger + try: from rdkit import Chem from rdkit.Chem import AllChem, Descriptors from rdkit.Contrib.SA_Score import sascorer + HAS_RDKIT = True except (ImportError, ModuleNotFoundError) as e: HAS_RDKIT = False @@ -24,6 +26,7 @@ import sys import os + def get_density(smiles: str) -> float: """ Calculate the density of a molecule given its SMILES string. @@ -32,11 +35,14 @@ def get_density(smiles: str) -> float: Args: smiles (str): The input SMILES string. + Returns: float: Density of the molecule, returns 0.0 if there is an error. """ if not HAS_RDKIT: - raise ImportError("Please install the rdkit support packages to use this module.") + raise ImportError( + "Please install the rdkit support packages to use this module." + ) try: # logger.info(f"Calculating density for SMILES: {smiles}") mol = Chem.MolFromSmiles(smiles) @@ -75,6 +81,7 @@ def get_density_and_synthesizability(smiles: str) -> tuple[float, float]: Args: smiles (str): The input SMILES string. + Returns: A tuple containing: float: Density of the molecule, returns 0.0 if there is an error. @@ -82,23 +89,25 @@ def get_density_and_synthesizability(smiles: str) -> tuple[float, float]: """ if not HAS_RDKIT: - raise ImportError("Please install the rdkit support packages to use this module.") + raise ImportError( + "Please install the rdkit support packages to use this module." + ) density = get_density(smiles) synthesizability = get_synthesizability(smiles) return density, synthesizability -def chemprop_preds_server(smiles: str,property:str) -> float: - + +def chemprop_preds_server(smiles: str, property: str) -> float: """ Predict molecular properties using pre-trained Chemprop models. - This function returns property predictions from Chemprop models. It validates the requested property name, + This function returns property predictions from Chemprop models. It validates the requested property name, constructs the appropriate model, and returns predictions for the provided SMILES input. Valid properties ---------------- ChARGe can request any of the following property names: - - density : Predicted density (g/cm³) - - hof : Heat of formation (kcal/mol) + - density : Predicted density (g/cm³) + - hof : Heat of formation (kcal/mol) - alpha : Polarizability (a0³) - cv : Heat capacity at constant volume (cal/mol·K) - gap : HOMO–LUMO energy gap (Hartree) @@ -109,69 +118,77 @@ def chemprop_preds_server(smiles: str,property:str) -> float: - zpve : Zero-point vibrational energy (Hartree) - lipo : Octanol–water partition coefficient (logD) - Parameters - ---------- - smiles : str - A SMILES string representing the molecule to be evaluated. - property : str - The property to predict. Must be one of the valid property names listed above. - - Returns - ------- - float - A float representing the predicted value for the specified property. - - Raises - ------ - SystemExit - If the environment variable `CHEMPROP_BASE_PATH` is not set. - - Examples - -------- - >>> chemprop_preds_server("CCO", "gap") - 6.73 - - >>> chemprop_preds_server("c1ccccc1", "lipo") - 2.94 + Args: + smiles (str): A SMILES string representing the molecule to be evaluated. + property (str): The property to predict. Must be one of the valid property names listed above. + + Returns: + float + A float representing the predicted value for the specified property. + + Raises: + SystemExit + If the environment variable `CHEMPROP_BASE_PATH` is not set. + + Examples: + >>> chemprop_preds_server("CCO", "gap") + 6.73 + + >>> chemprop_preds_server("c1ccccc1", "lipo") + 2.94 """ if not HAS_RDKIT: - raise ImportError("Please install the rdkit support packages to use this module.") - valid_properties = {'density', 'hof', 'alpha','cv','gap','homo','lumo','mu','r2','zpve','lipo'} + raise ImportError( + "Please install the rdkit support packages to use this module." + ) + valid_properties = { + "density", + "hof", + "alpha", + "cv", + "gap", + "homo", + "lumo", + "mu", + "r2", + "zpve", + "lipo", + } if property not in valid_properties: raise ValueError( f"Invalid property '{property}'. Must be one of {valid_properties}." ) - chemprop_base_path=os.environ.get("CHEMPROP_BASE_PATH") - if(chemprop_base_path): - model_path=os.path.join(chemprop_base_path, property) - model_path=os.path.join(model_path, 'model_0/best.pt') - return(predict_with_chemprop(model_path,[smiles])[0][0]) + chemprop_base_path = os.environ.get("CHEMPROP_BASE_PATH") + if chemprop_base_path: + model_path = os.path.join(chemprop_base_path, property) + model_path = os.path.join(model_path, "model_0/best.pt") + return predict_with_chemprop(model_path, [smiles])[0][0] else: - print('CHEMPROP_BASE_PATH environment variable not set!') + print("CHEMPROP_BASE_PATH environment variable not set!") sys.exit(2) + def get_molecule_price(smiles): """ Retrieve vendor pricing from ChemSpace for the molecule specified by the SMILES string, smiles. - Parameters - ---------- - smiles : str - A SMILES string for the molecule of interest. + Args: + smiles : str + A SMILES string for the molecule of interest. - Returns - ------- - float - Returns float representing the lowest price (in USD/g) among all vendors for the specified molecules in SMILES_list. + Returns: + float + Returns float representing the lowest price (in USD/g) among all vendors for the specified molecules in SMILES_list. - Examples - -------- - >>> get_molecule_price("CCO") - 0.1056 + Examples: + >>> get_molecule_price("CCO") + 0.1056 """ if not HAS_RDKIT: - raise ImportError("Please install the rdkit support packages to use this module.") - price=get_chemspace_prices([smiles]) - return(price[0]) + raise ImportError( + "Please install the rdkit support packages to use this module." + ) + price = get_chemspace_prices([smiles]) + return price[0] diff --git a/charge/servers/molecule_pricer.py b/charge/servers/molecule_pricer.py index cb96b42..bf497ca 100644 --- a/charge/servers/molecule_pricer.py +++ b/charge/servers/molecule_pricer.py @@ -1,7 +1,9 @@ from loguru import logger + try: import chemprice from chemprice import PriceCollector + HAS_CHEMPRICE = True except (ImportError, ModuleNotFoundError) as e: HAS_CHEMPRICE = False @@ -11,8 +13,8 @@ ) import os, sys -def get_chemspace_prices(SMILES_list,best_only=True): - + +def get_chemspace_prices(SMILES_list, best_only=True): """ Retrieve vendor pricing from ChemSpace for one or more molecules specified by SMILES. @@ -29,49 +31,47 @@ def get_chemspace_prices(SMILES_list,best_only=True): - USD/g : float # Price of the chemical in U.S. dollars per gram of chemical. - USD/mol : float # Price of the chemical in U.S. dollars per mol of chemical. - Parameters - ---------- - SMILES_list : list[str] - A list of SMILES strings to query prices for. - best_only : bool, default True - If True, return only the cheapest molecule price in USD/g; if False, return - all collected vendor offers and all properties. - - Returns - ------- - list[float] or pandas.DataFrame + Args: + SMILES_list (list[str]): A list of SMILES strings to query prices for. + best_only (bool, optional): If True, return only the cheapest molecule price in USD/g; if False, return + all collected vendor offers and all properties. Defaults to True. + + Returns: + list[float] or pandas.DataFrame If `best_only=True`, returns a list of floats representing the lowest price (in USD/g) among all vendors for the specified molecules in SMILES_list. Otherwise, returns a pandas DataFrame containing detailed vendor information for the molecule, including columns as listed above. - Examples - -------- - >>> get_chemspace_prices(["CCO","CO","CCC"], best_only=True) - [0.1056, 9.57, nan] + Examples: + >>> get_chemspace_prices(["CCO","CO","CCC"], best_only=True) + [0.1056, 9.57, nan] - """ + """ if not HAS_CHEMPRICE: - raise ImportError("Please install the chemprice support packages to use this module.") + raise ImportError( + "Please install the chemprice support packages to use this module." + ) pc = PriceCollector() chemspace_api_key = os.getenv("CHEMSPACE_API_KEY") - if(chemspace_api_key): + if chemspace_api_key: pc.setChemSpaceApiKey(chemspace_api_key) else: - print('CHEMPROP_API_KEY environment variable not set!') + print("CHEMPROP_API_KEY environment variable not set!") sys.exit(2) print(pc.check()) all_prices = pc.collect(SMILES_list) - if(best_only): - best_price=pc.selectBest(all_prices) - return(best_price["USD/g"].astype(float).tolist()) + if best_only: + best_price = pc.selectBest(all_prices) + return best_price["USD/g"].astype(float).tolist() else: - return(all_prices) + return all_prices -def main(smiles_list,price_source='Chemspace'): + +def main(smiles_list, price_source="Chemspace"): """ Main function that retrieves prices for a list of SMILES strings. Default to Chemspace because it doesn't have API limits. Keep this function to add future price_sources (Molport). @@ -80,12 +80,15 @@ def main(smiles_list,price_source='Chemspace'): """ if not HAS_CHEMPRICE: - raise ImportError("Please install the chemprice support packages to use this module.") - if(price_source=='Chemspace'): - prices=get_chemspace_prices(smiles_list) - print("Retrieved Prices from "+price_source+":") + raise ImportError( + "Please install the chemprice support packages to use this module." + ) + if price_source == "Chemspace": + prices = get_chemspace_prices(smiles_list) + print("Retrieved Prices from " + price_source + ":") print(prices) + if __name__ == "__main__": # Example usage example_smiles = ["CCO"] diff --git a/charge/servers/retrosynthesis_reaction_server.py b/charge/servers/retrosynthesis_reaction_server.py index e71ade4..8cebb2c 100644 --- a/charge/servers/retrosynthesis_reaction_server.py +++ b/charge/servers/retrosynthesis_reaction_server.py @@ -8,14 +8,77 @@ from mcp.server.fastmcp import FastMCP from charge.servers.SMILES_utils import verify_smiles, canonicalize_smiles from charge.servers.log_progress import log_progress +from charge.servers.ServerToolkit import ServerToolkit import argparse import os -template_free_mcp = FastMCP("template_free_reaction_server") +try: + from charge.servers.AiZynthTools import is_molecule_synthesizable, RetroPlanner + + HAS_AIZYNTH = True +except (ImportError, ModuleNotFoundError) as e: + HAS_AIZYNTH = False + logger.warning( + "Please install the aiZynthFinder support packages to use this module." + "Install it with: pip install charge[aiZynthFinder]", + ) + + +class RetroSynthesisServer(ServerToolkit): + """ + A ChARGe server that provides common retrosynthesis reaction tools. + """ + + def __init__(self, mcp: FastMCP): + """ + Initialize the RetroSynthesisServer. + + Args: + mcp (FastMCP): The MCP instance to register the function with. + """ + super().__init__(mcp) + + self.register_function_as_tool(self.mcp, verify_smiles) + self.register_function_as_tool(self.mcp, canonicalize_smiles) + self.register_function_as_tool(self.mcp, log_progress) + + +class TemplateFreeRetroSynthesisServer(RetroSynthesisServer): + """ + A ChARGe server that provides common retrosynthesis reaction tools for template free retrosynthesis. + """ + + def __init__(self, mcp: FastMCP): + """ + Initialize the TemplateFreeRetroSynthesisServer. + + Args: + mcp (FastMCP): The MCP instance to register the function with. + """ + super().__init__(mcp) -template_free_mcp.tool()(verify_smiles) -template_free_mcp.tool()(canonicalize_smiles) -template_free_mcp.tool()(log_progress) + +class TemplateRetroSynthesisServer(RetroSynthesisServer): + """ + A ChARGe server that provides common retrosynthesis reaction tools for template retrosynthesis. + """ + + def __init__(self, mcp: FastMCP, configfile: str): + """ + Initialize the TemplateRetroSynthesisServer. + + Args: + mcp (FastMCP): The MCP instance to register the function with. + configfile (str): The path to the configuration file for the AiZynthFinder. + """ + super().__init__(mcp) + if HAS_AIZYNTH: + RetroPlanner.initialize(configfile=configfile) + self.register_function_as_tool(self.mcp, is_molecule_synthesizable) + else: + raise ImportError( + "Please install the aiZynthFinder support packages to use this module." + ) if __name__ == "__main__": @@ -34,22 +97,19 @@ args = parser.parse_args() exp_type = args.exp_type - if exp_type == "template": - from charge.servers.SMARTS_reactions import SMARTS_mcp - SMARTS_mcp.tool()(log_progress) + mcp = FastMCP( + "RetroSynthesis Reaction Server", port=args.port, website_url=f"{args.host}" + ) - SMARTS_mcp.run( - transport="sse", - ) + if exp_type == "template": + logger.info("Starting Template RetroSynthesis Server") + server = TemplateRetroSynthesisServer(mcp, args.config) elif exp_type == "template-free": - - from charge.servers.AiZynthTools import is_molecule_synthesizable, RetroPlanner - - RetroPlanner.initialize(configfile=args.config) - - template_free_mcp.tool()(is_molecule_synthesizable) - template_free_mcp.run( - transport="sse", - ) + logger.info("Starting Template Free RetroSynthesis Server") + server = TemplateFreeRetroSynthesisServer(mcp) else: raise ValueError(f"Unknown task type: {exp_type}") + + server.run( + transport="sse", + ) diff --git a/charge/servers/server_utils.py b/charge/servers/server_utils.py index 393101d..cafae12 100644 --- a/charge/servers/server_utils.py +++ b/charge/servers/server_utils.py @@ -17,10 +17,11 @@ def add_server_arguments(parser: argparse.ArgumentParser) -> None: "--host", type=str, default=None, help="Host to run the server on" ) parser.add_argument( - '--transport', type=str, - help='MCP transport type', - choices=['stdio', 'streamable-http', 'sse'], - default='sse' + "--transport", + type=str, + help="MCP transport type", + choices=["stdio", "streamable-http", "sse"], + default="sse", ) @@ -28,8 +29,10 @@ def update_mcp_network(mcp: FastMCP, host: str, port: str): mcp.settings.host = host mcp.settings.port = port + def get_hostname(): import socket + hostname = socket.gethostname() try: host = socket.gethostbyname(hostname) @@ -37,8 +40,10 @@ def get_hostname(): host = "127.0.0.1" return hostname, host + def try_get_public_hostname(): import socket + hostname = socket.gethostname() try: public_hostname = hostname + "-pub" From 665212a9910c9d5adfc2379c56e8eefb5e8d4563 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Wed, 19 Nov 2025 01:13:23 -0500 Subject: [PATCH 04/11] Fix formatting on AiZynthTools --- charge/servers/AiZynthTools.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/charge/servers/AiZynthTools.py b/charge/servers/AiZynthTools.py index a7c9b6c..6ec6f1c 100644 --- a/charge/servers/AiZynthTools.py +++ b/charge/servers/AiZynthTools.py @@ -139,6 +139,7 @@ def is_molecule_synthesizable(smiles: str) -> bool: """Checks if a given molecule is synthesizable. First checks if it is available in a stock database, otherwise runs a retrosynthesis to see if a synthesis route can be found. + Args: smiles (str): The SMILES string of the molecule to check. @@ -185,8 +186,10 @@ def find_synthesis_routes(smiles: str) -> list[dict]: Args: smiles (str): the target molecule in SMILES representation. + Returns: list[dict]: a list of synthesis routes, each of which is a reaction tree in json/dict format. + Raises: ValueError: If the molecule is not valid. """ From 3c1b7ce316862f16e44babc7a5c9f57ae65db415 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Wed, 19 Nov 2025 01:18:32 -0500 Subject: [PATCH 05/11] Clean up local running infrastructure --- charge/servers/SMARTS_reactions.py | 3 +-- charge/servers/SMARTS_reactions_utils.py | 2 ++ charge/servers/SMILES.py | 3 +-- charge/servers/retrosynthesis_reaction_server.py | 3 ++- 4 files changed, 6 insertions(+), 5 deletions(-) diff --git a/charge/servers/SMARTS_reactions.py b/charge/servers/SMARTS_reactions.py index 8f09938..fa6d6fa 100644 --- a/charge/servers/SMARTS_reactions.py +++ b/charge/servers/SMARTS_reactions.py @@ -62,5 +62,4 @@ def __init__(self, mcp: FastMCP): ) sm = SMARTSServer(SMARTS_mcp) - mcp = sm.return_mcp() - mcp.run() + sm.run(transport=args.transport) diff --git a/charge/servers/SMARTS_reactions_utils.py b/charge/servers/SMARTS_reactions_utils.py index ef40fda..adb6324 100644 --- a/charge/servers/SMARTS_reactions_utils.py +++ b/charge/servers/SMARTS_reactions_utils.py @@ -31,6 +31,7 @@ def verify_reaction_SMARTS(smarts: str) -> Tuple[bool, str]: Args: smarts (str): The input SMARTS string. + Returns: A tuple containing: bool: True if the SMARTS is valid, False if it is invalid. @@ -78,6 +79,7 @@ def verify_reaction( smarts (str): The input SMARTS string. reactants (list[str]): The input list of reactants in SMILES strings products (list[str]): The input list of products created by the reaction in SMILES strings + Returns: A tuple containing: bool: True if a reaction can be performed given the SMARTS and reactants. diff --git a/charge/servers/SMILES.py b/charge/servers/SMILES.py index fa0c84d..de91a33 100644 --- a/charge/servers/SMILES.py +++ b/charge/servers/SMILES.py @@ -60,5 +60,4 @@ def __init__(self, mcp: FastMCP): logger.info("[RDKit-SMILES] Starting Chem and BioInformatics MCP Server") sm = SMILESServer(SMILES_mcp) - mcp = sm.return_mcp() - mcp.run() + sm.run(transport=args.transport) diff --git a/charge/servers/retrosynthesis_reaction_server.py b/charge/servers/retrosynthesis_reaction_server.py index 8cebb2c..6b86848 100644 --- a/charge/servers/retrosynthesis_reaction_server.py +++ b/charge/servers/retrosynthesis_reaction_server.py @@ -94,6 +94,7 @@ def __init__(self, mcp: FastMCP, configfile: str): default=os.path.join(os.getcwd(), "config.yml"), help="Path to the configuration file for the AiZynthFinder", ) + add_server_arguments(parser) args = parser.parse_args() exp_type = args.exp_type @@ -111,5 +112,5 @@ def __init__(self, mcp: FastMCP, configfile: str): raise ValueError(f"Unknown task type: {exp_type}") server.run( - transport="sse", + transport=args.transport, ) From 3a14e6875907528ea8188f0ce4837e305ae1bd7f Mon Sep 17 00:00:00 2001 From: Shehtab Date: Wed, 19 Nov 2025 01:45:44 -0500 Subject: [PATCH 06/11] Convert FLASKv2 implementation --- charge/servers/FLASKv2_reactions.py | 277 +++++++++++++++++----------- charge/servers/ServerToolkit.py | 10 +- 2 files changed, 171 insertions(+), 116 deletions(-) diff --git a/charge/servers/FLASKv2_reactions.py b/charge/servers/FLASKv2_reactions.py index 26af5b5..0f8a6b9 100644 --- a/charge/servers/FLASKv2_reactions.py +++ b/charge/servers/FLASKv2_reactions.py @@ -5,10 +5,16 @@ from typing import Optional try: - from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaForCausalLM, PreTrainedTokenizer + from transformers import ( + AutoTokenizer, + AutoModelForCausalLM, + LlamaForCausalLM, + PreTrainedTokenizer, + ) from peft import PeftModel from trl import apply_chat_template import torch + HAS_FLASKV2 = True except (ImportError, ModuleNotFoundError) as e: HAS_FLASKV2 = False @@ -18,58 +24,174 @@ ) from charge.servers.server_utils import update_mcp_network, get_hostname +from charge.servers.ServerToolkit import ServerToolkit + def format_rxn_prompt(data: dict, forward: bool) -> dict: - required_keys = ['reactants', 'products', 'agents', 'solvents', 'catalysts', 'atmospheres'] - non_product_keys = [k for k in required_keys if k != 'products'] + required_keys = [ + "reactants", + "products", + "agents", + "solvents", + "catalysts", + "atmospheres", + ] + non_product_keys = [k for k in required_keys if k != "products"] if forward: d = {k: data[k] for k in non_product_keys if data.get(k, None)} prompt = json.dumps(d) else: - d = {'products': data['products']} + d = {"products": data["products"]} prompt = json.dumps(d) - data['prompt'] = [{'role': 'user', 'content': prompt}] + data["prompt"] = [{"role": "user", "content": prompt}] return data -def predict_reaction_internal(molecules: list[str], retrosynthesis: bool) -> list[str]: - if not HAS_FLASKV2: - raise ImportError( - "Please install the [flask] optional packages to use this module." +class FlaskV2ReactionServer(ServerToolkit): + def __init__( + self, + mcp: FastMCP, + model_dir_fwd: Optional[str], + model_dir_retro: Optional[str], + adapter_weights_fwd: Optional[str], + adapter_weights_retro: Optional[str], + ): + super().__init__(mcp) + + if not HAS_FLASKV2: + raise ImportError( + "Please install the [flask] optional packages to use this module." + ) + assert ( + model_dir_fwd or model_dir_retro + ), "At least one model has to be given to the MCP server" + + # Load tokenizer and models + self.tokenizer = AutoTokenizer.from_pretrained( + model_dir_fwd or model_dir_retro, padding_side="left" ) - model = retro_model if retrosynthesis else fwd_model - data = {'products': molecules} if retrosynthesis else {'reactants': molecules} - with torch.inference_mode(): - prompt = format_rxn_prompt(data, forward=(not retrosynthesis)) - prompt = apply_chat_template(prompt, tokenizer=tokenizer) - inputs = tokenizer(prompt["prompt"], return_tensors="pt", padding="longest").to('cuda') - prompt_length = inputs["input_ids"].size(1) - outputs = model.generate( - **inputs, - max_new_tokens=2048, - num_return_sequences=3, - # do_sample=True, - num_beams=3, - pad_token_id=tokenizer.pad_token_id, - eos_token_id=tokenizer.eos_token_id, - use_cache=True, # enable KV cache + self.tokenizer.add_special_tokens({"pad_token": "<|finetune_right_pad_id|>"}) + self.fwd_model = self._load_model(model_dir_fwd, adapter_weights_fwd) + self.retro_model = self._load_model(model_dir_retro, adapter_weights_retro) + + self.available_tools = [] + if self.fwd_model is not None: + self.available_tools.append("Forward Prediction") + self._register_single_method(self.predict_reaction_products) + if self.retro_model is not None: + self.available_tools.append("Single-Step Retrosynthesis") + self._register_single_method(self.predict_reaction_reactants) + + def _load_model(self, model_dir: Optional[str], adapter_weights: Optional[str]): + if model_dir is None: + return None + model = AutoModelForCausalLM.from_pretrained( + model_dir, + device_map="cuda", + torch_dtype=torch.bfloat16, ) - processed_outputs = [tokenizer.decode(out[prompt_length:], skip_special_tokens=True) for out in outputs] - logger.debug(f'Model input: {prompt["prompt"]}') - processed_outs = "\n".join(processed_outputs) - logger.debug(f'Model output: {processed_outs}') - return processed_outputs + if adapter_weights is not None: + model = PeftModel.from_pretrained(model, adapter_weights) + model = model.merge_and_unload() + + model.eval() + # Enable model optimizations + if hasattr(model, "config") and hasattr(model.config, "use_cache"): + model.config.use_cache = True # enable KV caching + return model + + def predict_reaction_products(self, reactants: list[str]) -> list[str]: + """ + Given a set of reactant molecules, predict the likely product molecule(s). + + Args: + reactants (list[str]): a list of reactant molecules in SMILES representation. + + Returns: + list[str]: a list of predictions, each of which is a json string listing the predicted product molecule(s) in SMILES. + """ + + return self._predict_reaction_internal(reactants, False) + + def predict_reaction_reactants(self, products: list[str]) -> list[str]: + """ + Given a set of product molecules, predict the likely reactant molecule(s). + + Args: + products (list[str]): a list of product molecules in SMILES representation. + + Returns: + list[str]: a list of predictions, each of which is a json string listing the predicted reactant molecule(s) in SMILES. + """ + return self._predict_reaction_internal(products, True) + + def _predict_reaction_internal( + self, molecules: list[str], retrosynthesis: bool + ) -> list[str]: + + model = self.retro_model if retrosynthesis else self.fwd_model + data = {"products": molecules} if retrosynthesis else {"reactants": molecules} + with torch.inference_mode(): + prompt = format_rxn_prompt(data, forward=(not retrosynthesis)) + prompt = apply_chat_template(prompt, tokenizer=self.tokenizer) + inputs = self.tokenizer( + prompt["prompt"], return_tensors="pt", padding="longest" + ).to("cuda") + prompt_length = inputs["input_ids"].size(1) + outputs = model.generate( + **inputs, + max_new_tokens=2048, + num_return_sequences=3, + # do_sample=True, + num_beams=3, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + use_cache=True, # enable KV cache + ) + processed_outputs = [ + tokenizer.decode(out[prompt_length:], skip_special_tokens=True) + for out in outputs + ] + logger.debug(f'Model input: {prompt["prompt"]}') + processed_outs = "\n".join(processed_outputs) + logger.debug(f"Model output: {processed_outs}") + return processed_outputs + + def get_available_tools(self) -> list[str]: + return self.available_tools @click.command() -@click.option("--model-dir-fwd", envvar="FLASKV2_MODEL_FWD", help="Path to flaskv2 model") -@click.option("--model-dir-retro", envvar="FLASKV2_MODEL_RETRO", help="Path to flaskv2 model for retrosynthesis") +@click.option( + "--model-dir-fwd", envvar="FLASKV2_MODEL_FWD", help="Path to flaskv2 model" +) +@click.option( + "--model-dir-retro", + envvar="FLASKV2_MODEL_RETRO", + help="Path to flaskv2 model for retrosynthesis", +) @click.option("--adapter-weights-fwd", help="LoRA adapter weights, if used") -@click.option("--adapter-weights-retro", help="LoRA adapter weights for retrosynthesis model, if used") -@click.option("--transport", type=click.Choice(['stdio', 'streamable-http', 'sse']), help="MCP transport type", default="sse") +@click.option( + "--adapter-weights-retro", + help="LoRA adapter weights for retrosynthesis model, if used", +) +@click.option( + "--transport", + type=click.Choice(["stdio", "streamable-http", "sse"]), + help="MCP transport type", + default="sse", +) @click.option("--port", type=int, default=8125, help="Port to run the server on") @click.option("--host", type=str, default=None, help="Host to run the server on") -def main(model_dir_fwd: str, model_dir_retro: str, adapter_weights_fwd: str, adapter_weights_retro: str, transport: str, port: str, host: Optional[str]): +def main( + model_dir_fwd: str, + model_dir_retro: str, + adapter_weights_fwd: str, + adapter_weights_retro: str, + transport: str, + port: str, + host: Optional[str], +): if not HAS_FLASKV2: raise ImportError( "Please install the [flask] optional packages to use this module." @@ -80,87 +202,20 @@ def main(model_dir_fwd: str, model_dir_retro: str, adapter_weights_fwd: str, ada if host is None: _, host = get_hostname() - mcp = FastMCP("FLASKv2 Reaction Predictor", - port=port, - website_url=f"{host}", + mcp = FastMCP( + "FLASKv2 Reaction Predictor", + port=port, + website_url=f"{host}", ) # Init MCP server mcp = FastMCP("FLASKv2 Reaction Predictor", host=host, port=port) - # Make HF models and tokenizer global objects - global fwd_model, retro_model, tokenizer - fwd_model = None - retro_model = None - - # Load tokenizer and models - tokenizer = AutoTokenizer.from_pretrained(model_dir_fwd or model_dir_retro, padding_side="left") - tokenizer.add_special_tokens({"pad_token": "<|finetune_right_pad_id|>"}) - if model_dir_fwd: - fwd_model = AutoModelForCausalLM.from_pretrained( - model_dir_fwd, - device_map='cuda', - torch_dtype=torch.bfloat16, - ) - if adapter_weights_fwd is not None: - fwd_model = PeftModel.from_pretrained(fwd_model, adapter_weights_fwd) - fwd_model = fwd_model.merge_and_unload() - if model_dir_retro: - retro_model = AutoModelForCausalLM.from_pretrained( - model_dir_retro, - device_map='cuda', - torch_dtype=torch.bfloat16, - ) - if adapter_weights_retro is not None: - retro_model = PeftModel.from_pretrained(retro_model, adapter_weights_retro) - retro_model = retro_model.merge_and_unload() - - # Enable model optimizations - if fwd_model is not None: - fwd_model.eval() - if hasattr(fwd_model, "config") and hasattr(fwd_model.config, "use_cache"): - fwd_model.config.use_cache = True # enable KV caching - if retro_model is not None: - retro_model.eval() - if hasattr(retro_model, "config") and hasattr(retro_model.config, "use_cache"): - retro_model.config.use_cache = True # enable KV caching - - # Dynamic tool creation based on input models - available_tools = [] - if fwd_model is not None: - available_tools.append("Forward Prediction") - - @mcp.tool() - def predict_reaction_products(reactants: list[str]) -> list[str]: - """ - Given a set of reactant molecules, predict the likely product molecule(s). - - Args: - reactants (list[str]): a list of reactant molecules in SMILES representation. - Returns: - list[str]: a list of predictions, each of which is a json string listing the predicted product molecule(s) in SMILES. - """ - logger.debug('Calling `predict_reaction_products`') - return predict_reaction_internal(reactants, False) - - if retro_model is not None: - available_tools.append("Single-Step Retrosynthesis") - - @mcp.tool() - def predict_reaction_reactants(products: list[str]) -> list[str]: - """ - Given a product molecule, predict the likely reactants and other chemical species (e.g., agents, solvents). - - Args: - products (list[str]): a list of product molecules in SMILES representation. - Returns: - list[str]: a list of predictions, each of which is a json string listing the predicted reactant molecule(s) in SMILES, - as well as potential (re)agents and solvents used in the reaction. - """ - logger.debug('Calling `predict_reaction_reactants`') - return predict_reaction_internal(products, True) - - logger.info(f"Available tools: {', '.join(available_tools)}") + server = FlaskV2ReactionServer( + mcp, model_dir_fwd, model_dir_retro, adapter_weights_fwd, adapter_weights_retro + ) + + logger.info(f"Available tools: {', '.join(server.get_available_tools())}") # Run MCP server mcp.run(transport=transport) diff --git a/charge/servers/ServerToolkit.py b/charge/servers/ServerToolkit.py index e949ca2..3bedc6c 100644 --- a/charge/servers/ServerToolkit.py +++ b/charge/servers/ServerToolkit.py @@ -16,8 +16,7 @@ class ServerToolkit: """ def __init__(self, mcp: FastMCP): - self.mcp = mcp - self._pending_methods = [] + self._mcp = mcp def _register_methods(self): """ @@ -50,7 +49,7 @@ def _register_single_method(self, method: Callable): def tool_wrapper(*args, **kwargs): return method(*args, **kwargs) - self.mcp.tool()(tool_wrapper) + self._mcp.tool()(tool_wrapper) @staticmethod def mcp_tool(func: Callable) -> Callable: @@ -88,7 +87,7 @@ def return_mcp(self) -> FastMCP: FastMCP: The MCP instance. """ self._register_methods() - return self.mcp + return self._mcp def run(self, transport: Literal["sse", "stdio"] = "sse") -> None: """ @@ -97,4 +96,5 @@ def run(self, transport: Literal["sse", "stdio"] = "sse") -> None: Args: transport (Literal["sse", "stdio"], optional): The transport to use. Defaults to "sse". """ - self.mcp.run(transport=transport) + self._register_methods() + self._mcp.run(transport=transport) From 233976b0e8b4e1743e4fc321f7768cffd218290f Mon Sep 17 00:00:00 2001 From: Shehtab Date: Wed, 19 Nov 2025 01:53:01 -0500 Subject: [PATCH 07/11] Add additional ServerToolKit method for registering external functions --- charge/servers/FLASKv2_reactions.py | 2 +- charge/servers/ServerToolkit.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/charge/servers/FLASKv2_reactions.py b/charge/servers/FLASKv2_reactions.py index 0f8a6b9..9400cb3 100644 --- a/charge/servers/FLASKv2_reactions.py +++ b/charge/servers/FLASKv2_reactions.py @@ -218,7 +218,7 @@ def main( logger.info(f"Available tools: {', '.join(server.get_available_tools())}") # Run MCP server - mcp.run(transport=transport) + server.run(transport=transport) if __name__ == "__main__": diff --git a/charge/servers/ServerToolkit.py b/charge/servers/ServerToolkit.py index 3bedc6c..77618b9 100644 --- a/charge/servers/ServerToolkit.py +++ b/charge/servers/ServerToolkit.py @@ -79,6 +79,18 @@ def register_function_as_tool(mcp: FastMCP, func: Callable) -> None: """ mcp.tool()(func) + def register_function_to_server(func: Callable) -> None: + """ + Register an external function to the server. + + Args: + func (Callable): The function to register to the server. + + Returns: + None + """ + self._mcp.tool()(func) + def return_mcp(self) -> FastMCP: """ Return the MCP instance. From 548385b26c4359b5787337aaa2802cb003dea8f0 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Wed, 19 Nov 2025 02:04:35 -0500 Subject: [PATCH 08/11] Add MultiServerToolKit --- charge/servers/ServerToolkit.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/charge/servers/ServerToolkit.py b/charge/servers/ServerToolkit.py index 77618b9..d874528 100644 --- a/charge/servers/ServerToolkit.py +++ b/charge/servers/ServerToolkit.py @@ -110,3 +110,35 @@ def run(self, transport: Literal["sse", "stdio"] = "sse") -> None: """ self._register_methods() self._mcp.run(transport=transport) + + def update_mcp(self, mcp: FastMCP) -> None: + """ + Update the MCP instance. + + Args: + mcp (FastMCP): The new MCP instance. + """ + self._mcp = mcp + self._register_methods() + + +class MultiServerToolkit(ServerToolkit): + def __init__( + self, + servers: list[ServerToolkit], + description: str, + host: str, + port: int, + ): + self.mcp = FastMCP(description, host=host, port=port) + + self.servers = servers + # Register all tools from all servers + for server in self.servers: + for name in dir(server): + if name.startswith("_"): + continue + + attr = getattr(server, name) + if callable(attr) and hasattr(attr, "_is_mcp_tool"): + self.mcp.tool()(attr) From 59fe5ed379e5b65e589da95b56cfc4d9842815dc Mon Sep 17 00:00:00 2001 From: Shehtab Date: Wed, 19 Nov 2025 02:06:50 -0500 Subject: [PATCH 09/11] Add documentation --- charge/servers/ServerToolkit.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/charge/servers/ServerToolkit.py b/charge/servers/ServerToolkit.py index d874528..bfa5948 100644 --- a/charge/servers/ServerToolkit.py +++ b/charge/servers/ServerToolkit.py @@ -123,6 +123,11 @@ def update_mcp(self, mcp: FastMCP) -> None: class MultiServerToolkit(ServerToolkit): + """ + A class to combine multiple servers into a single object that + can be run as a single MCP server. + """ + def __init__( self, servers: list[ServerToolkit], @@ -130,7 +135,17 @@ def __init__( host: str, port: int, ): - self.mcp = FastMCP(description, host=host, port=port) + """ + Initialize the MultiServerToolkit. + + Args: + servers (list[ServerToolkit]): The list of servers to register. + description (str): The description of the server. + host (str): The host of the server. + port (int): The port of the server. + """ + mcp = FastMCP(description, host=host, port=port) + super().__init__(mcp) self.servers = servers # Register all tools from all servers From b9ea751da62b34c6910f8e90e31325d277645b58 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Wed, 19 Nov 2025 02:16:36 -0500 Subject: [PATCH 10/11] Fix registration wrapper call --- charge/servers/molecular_generation_server.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/charge/servers/molecular_generation_server.py b/charge/servers/molecular_generation_server.py index ff7c5a8..edc26a5 100644 --- a/charge/servers/molecular_generation_server.py +++ b/charge/servers/molecular_generation_server.py @@ -23,7 +23,6 @@ import os from charge.tasks.Task import Task from charge.servers.server_utils import add_server_arguments -from charge.clients.autogen import AutoGenClient from charge.clients.Client import Client import asyncio from charge.servers import SMILES_utils @@ -78,7 +77,7 @@ def __init__(self, mcp: FastMCP, model: str, backend: str, json_file_path: str): backend=self.backend, ) - @mcp_tool + @ServerToolkit.mcp_tool def diagnose_smiles(self, smiles: str) -> str: """ Diagnose a SMILES string. Returns a diagnosis of the SMILES string. @@ -111,7 +110,7 @@ def diagnose_smiles(self, smiles: str) -> str: logger.error(f"An error occurred: {e}") return "Error: Unable to process the SMILES string at this time." - @mcp_tool + @ServerToolkit.mcp_tool def is_already_known(self, smiles: str) -> bool: """ Check if a SMILES string provided is already known. Only provide @@ -152,7 +151,7 @@ def is_already_known(self, smiles: str) -> bool: # Check if the SMILES string is already known (in the database) return canonical_smiles in known_smiles - @mcp_tool + @ServerToolkit.mcp_tool def get_density(self, smiles: str) -> float: """ Calculate the density of a molecule given its SMILES string. From 14810a23dac3be291adb1d773429bac0d4826283 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Mon, 24 Nov 2025 15:09:19 -0500 Subject: [PATCH 11/11] Experimenting with Automatic Toolkit launcher --- charge/servers/ServerToolkit.py | 82 +++++++++++++++++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/charge/servers/ServerToolkit.py b/charge/servers/ServerToolkit.py index bfa5948..9d11033 100644 --- a/charge/servers/ServerToolkit.py +++ b/charge/servers/ServerToolkit.py @@ -8,6 +8,7 @@ from mcp.server.fastmcp import FastMCP from functools import wraps from typing import Callable, Literal +import time class ServerToolkit: @@ -157,3 +158,84 @@ def __init__( attr = getattr(server, name) if callable(attr) and hasattr(attr, "_is_mcp_tool"): self.mcp.tool()(attr) + + +class ToolKitLauncher: + """ + A class to launch on a seperate process a server with a toolkit + """ + + def __init__(self, toolkit: ServerToolkit): + self.toolkit = toolkit + self.process = None + + def start(self, transport: Literal["sse", "streamable-http"] = "sse"): + """ + Run the toolkit in a separate process. + + Args: + transport (Literal["sse", "streamable-http"], optional): The transport to use. Defaults to "sse". + """ + + import multiprocessing + + self.process = multiprocessing.Process( + target=self.toolkit.run, args=(transport,), daemon=True + ) + self.process.start() + self.wait_for_start() + + def wait_for_start(self, timeout: int = 10): + """ + Wait for the process to start. + + Args: + timeout (int, optional): The timeout in seconds. Defaults to 10. + """ + + start_time = time.time() + while not self.is_running(): + if time.time() - start_time > timeout: + raise TimeoutError("Process did not start within the timeout period.") + time.sleep(0.1) + + def stop(self): + """ + Stop the process. + """ + if self.process is None: + return + self.process.terminate() + self.process.join() + self.process = None + + def is_running(self) -> bool: + """ + Check if the process is running. + + Returns: + bool: True if the process is running, False otherwise. + """ + return self.process is not None and self.process.is_alive() + + def get_process(self) -> Optional[multiprocessing.Process]: + """ + Get the process. + + Returns: + Optional[multiprocessing.Process]: The process. + """ + return self.process + + def __del__(self): + """ + Destructor to stop the process. + """ + self.stop() + + def __enter__(self): + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop()