From ffbf4799f83a314e5a316c134a8757abff141df8 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Tue, 18 Nov 2025 21:53:43 -0500 Subject: [PATCH 1/7] Adding readthedocs generator --- docs/LICENSE.md | 30 ++++++++++++++++ docs/about.md | 1 + docs/index.md | 67 ++++++++++++++++++++++++++++++++++ mkdocs.yml | 36 +++++++++++++++++++ pyproject.toml | 10 +++++- scripts/gen_ref_pages.py | 78 ++++++++++++++++++++++++++++++++++++++++ 6 files changed, 221 insertions(+), 1 deletion(-) create mode 100644 docs/LICENSE.md create mode 100644 docs/about.md create mode 100644 docs/index.md create mode 100644 mkdocs.yml create mode 100644 scripts/gen_ref_pages.py diff --git a/docs/LICENSE.md b/docs/LICENSE.md new file mode 100644 index 0000000..640b202 --- /dev/null +++ b/docs/LICENSE.md @@ -0,0 +1,30 @@ +``` +################################################################################ +## Copyright 2025 Lawrence Livermore National Security, LLC. and Binghamton University. +## See the top-level LICENSE file for details. +## +## SPDX-License-Identifier: Apache-2.0 +################################################################################ +Copyright (c) 2025, Lawrence Livermore National Security, LLC. and Binghamton University. +Produced at the Lawrence Livermore National Laboratory and Binghamton University. +Written by the FLASK Research Team (Shehtab Zaman, Kenneth Chiu, Bhavya Kailkhura, +Tal Ben-Nun, Brian Van Essen, et al.) listed in the CONTRIBUTORS file. + +LLNL-CODE-2006345. +All rights reserved. + +This file is part of FLASK Project: Foundation Learning AI for Synthesis Knowledge. +For details, see https://github.com/FLASK-LLNL + +Licensed under the Apache License, Version 2.0 (the "Licensee"); you +may not use this file except in compliance with the License. You may +obtain a copy of the License at: + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or +implied. See the License for the specific language governing +permissions and limitations under the license. +``` diff --git a/docs/about.md b/docs/about.md new file mode 100644 index 0000000..1a406ba --- /dev/null +++ b/docs/about.md @@ -0,0 +1 @@ +Some input \ No newline at end of file diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..81ae727 --- /dev/null +++ b/docs/index.md @@ -0,0 +1,67 @@ +# ChARGe Documentation + +**ChARGe** is a **Ch**emical tool **A**ugmented **R**easoning models for **Ge**nerating molecules and reactions. + +## Table of Contents + +- [Installation](#installation) +- [Quick Start](#quick-start) +- [Modules Overview](#modules-overview) +- [API Reference](#api-reference) +- [Contributing](#contributing) +- [License](#license) + +--- + +### Installation + +```bash +pip install charge +``` + +or install from source with the GitHub repository: +```bash +pip install git+https://github.com/FLASK-LLNL/charge.git +``` + +For development, clone the repository and install with: +```bash +git clone https://github.com/FLASK-LLNL/charge.git +cd charge +pip install -e . +``` + +After installation, in order to install additional dependencies, run: +```bash +charge-install --extras all +``` + +### Quick Start + +```python +from charge import some_function + +result = some_function(...) +print(result) +``` + +### Modules Overview + + +- [`charge.clients`](reference/charge/clients): Client implementations for various backends. +- [`charge.servers`](reference/charge/servers): Server components. +- [`charge.tasks`](reference/charge/tasks): Task orchestration utilities. +- [`charge.inspector`](reference/charge/inspector): Inspection tools for debugging. +- [`charge.install`](reference/charge/install): Installation helpers. + +### API Reference + +Detailed API documentation can be found in the [Reference](reference/charge/index.md) section. + +### Contributing + +We welcome contributions! Please see the `CONTRIBUTORS` file on our [GitHub](https://github.com/FLASK-LLNL/charge) repository for guidelines. + +### License + +Please see the [LICENSE](LICENSE.md) in the documentation or the file in our repository [here](https://github.com/FLASK-LLNL/charge/blob/main/LICENSE) for details. \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 0000000..5980ae8 --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,36 @@ +site_name: ChARGe + +nav: + - Home: index.md + - About: about.md + - Reference: reference/ # The trailing slash tells literate-nav to look inside this directory + - LICENSE: LICENSE.md +plugins: + - search + - gen-files: + scripts: + # 1. This is the script that will generate the .md files and the SUMMARY.md + - scripts/gen_ref_pages.py + - literate-nav: + # 2. This tells mkdocs-literate-nav where to find the generated nav file + nav_file: SUMMARY.md + - section-index # 3. Optional, but highly recommended + - mkdocstrings: + handlers: + python: + # Optional global settings for the Python handler + options: + show_root_heading: true + docstring_style: google + extra: + show_source: false + show_members_order: true + paths: [.] # IMPORTANT: tell mkdocstrings where to find your package source + +theme: + name: readthedocs + language: en + highlightjs: true + hljs_languages: + - python + diff --git a/pyproject.toml b/pyproject.toml index cfba156..4b12efb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,5 +45,13 @@ test = ["pytest", "pytest-asyncio", "pytest-mock", "responses", "httpx", "reques # Define a set of optional packages for use when deploying persistent data services pds = ["charge[aizynthfinder,autogen,rdkit,chemprice]"] +docs = [ + "mkdocs", + "mkdocstrings[python]", + "mkdocs-gen-files", + "mkdocs-literate-nav", + "mkdocs-section-index", + "mkdocs-readthedocs-theme" +] -all = ["charge[ollama,gemini,autogen,rdkit,flask,aizynthfinder,chemprop,chemprice]"] +all = ["charge[ollama,gemini,autogen,rdkit,flask,aizynthfinder,chemprop,chemprice,docs]"] diff --git a/scripts/gen_ref_pages.py b/scripts/gen_ref_pages.py new file mode 100644 index 0000000..63728a3 --- /dev/null +++ b/scripts/gen_ref_pages.py @@ -0,0 +1,78 @@ +from pathlib import Path +import mkdocs_gen_files + +# --- Configuration --- +# 1. The path to the directory *containing* the package folder 'charge'. (Project Root) +SRC_ROOT = Path(".") +# 2. The name of your top-level package. +PACKAGE_NAME = "charge" +# 3. The target directory inside 'docs' for the generated reference. +REF_DIR = Path("reference") + +nav = mkdocs_gen_files.Nav() + +# Walk through all Python files within the 'charge' package directory +# path will be, e.g., 'charge/_tags.py', 'charge/__init__.py', 'charge/clients.py' +paths = list((SRC_ROOT / PACKAGE_NAME).rglob("*.py")) + +def get_sort_key(path): + """ + Sort key function to ensure directories appear before files. + Returns a list of tuples (priority, name), where priority is 0 for directories and 1 for files. + """ + parts = path.relative_to(SRC_ROOT).parts + return [(0, part) if i < len(parts) - 1 else (1, part) for i, part in enumerate(parts)] + +for path in sorted(paths, key=get_sort_key): + + # 1. Get the path relative to the root, without the extension + # e.g., 'charge/_tags' or 'charge/__init__' + module_path = path.relative_to(SRC_ROOT).with_suffix("") + + # Split into parts: e.g., ['charge', '_tags'] + parts = list(module_path.parts) + + # --- Logic to handle __init__.py files --- + if parts[-1] == "__init__": + # Drop the "__init__" part. Remaining parts: ['charge'] + parts = parts[:-1] + + # Set the documentation filename to index.md for the package root. + doc_path = REF_DIR / Path(*parts) / "index.md" + + # --- Logic to handle regular modules (e.g., _tags.py) --- + else: + # Keep all parts: e.g., ['charge', '_tags'] + + # Set the documentation filename: e.g., reference/charge/_tags.md + doc_path = REF_DIR / module_path.with_suffix(".md") + + # The full Python identifier for mkdocstrings (e.g., 'charge' or 'charge._tags') + ident = ".".join(parts) + + # --- FINAL CHECKS --- + if not ident: + # Skip if somehow the identifier is still empty (shouldn't happen now) + continue + + # 2. Add entry to the Nav object. + # This automatically uses the fully qualified name (ident) as the title + # and doc_path as the link, structured by the 'parts' list. + nav[parts] = doc_path.relative_to(REF_DIR).as_posix() + + # 3. Create the Markdown file with the mkdocstrings injection + with mkdocs_gen_files.open(doc_path, "w") as fd: + fd.write(f"# `{ident}`\n\n") + fd.write(f"::: {ident}\n") + + # Set edit link path to the source file + mkdocs_gen_files.set_edit_path(doc_path, path) + + +# 4. Write the final SUMMARY.md navigation file +with mkdocs_gen_files.open(REF_DIR / "SUMMARY.md", "w") as nav_file: + # This will generate the correct nested structure, like: + # * [charge](charge/index.md) + # * [charge._tags](charge/_tags.md) + # ... + nav_file.writelines(nav.build_literate_nav()) \ No newline at end of file From 0fb5bd41eab657575e399f73de727143d703d589 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Tue, 18 Nov 2025 21:55:19 -0500 Subject: [PATCH 2/7] Updating library documentation to conform to Google docstring style --- charge/_tags.py | 4 ++ charge/_to_mcp.py | 7 +++ charge/_utils.py | 6 +-- charge/clients/Client.py | 58 +++++++++++++++++++++++- charge/clients/autogen.py | 2 +- charge/servers/SMARTS_reactions_utils.py | 2 +- charge/servers/SMILES_utils.py | 4 ++ charge/servers/molecule_pricer.py | 7 +++ charge/servers/server_utils.py | 40 ++++++++++++++-- charge/tasks/LMOTask.py | 54 +++++++++++++++++++++- charge/tasks/Task.py | 4 +- 11 files changed, 175 insertions(+), 13 deletions(-) diff --git a/charge/_tags.py b/charge/_tags.py index 5099a2e..dcfd40b 100644 --- a/charge/_tags.py +++ b/charge/_tags.py @@ -1,16 +1,20 @@ def verifier(func): + """Decorator to mark a method as a verifier.""" func.__verifier_tag = True return func def is_verifier(func): + """Check if a method is marked as a verifier.""" return hasattr(func, "__verifier_tag") def hypothesis(func): + """Decorator to mark a method as a hypothesis.""" func.__hypothesis_tag = True return func def is_hypothesis(func): + """Check if a method is marked as a hypothesis.""" return hasattr(func, "__hypothesis_tag") diff --git a/charge/_to_mcp.py b/charge/_to_mcp.py index 1b7d5c4..3bdbde3 100644 --- a/charge/_to_mcp.py +++ b/charge/_to_mcp.py @@ -7,6 +7,13 @@ def task_to_mcp(class_info, methods_list) -> str: """ Convert an Task class to an MCP server definition string. + + Args: + class_info (dict): A dictionary containing information about the class. + methods_list (list): A list of methods to be converted to MCP server definition strings. + + Returns: + str: A string representing the MCP server definition. """ return_str = "" return_str += "from mcp.server.fastmcp import FastMCP\n" diff --git a/charge/_utils.py b/charge/_utils.py index 3b5cc1c..ca0b3e0 100644 --- a/charge/_utils.py +++ b/charge/_utils.py @@ -2,9 +2,9 @@ import readline import inspect import asyncio +from typing import Any - -def enable_cmd_history_and_shell_integration(history: str): +def enable_cmd_history_and_shell_integration(history: str) -> None: """Enable persistent command-line history and integrate with the interactive shell. Attempts to load an existing readline history file, sets the in-memory history length to 1000, and registers an atexit handler to persist history on exit. @@ -28,7 +28,7 @@ def enable_cmd_history_and_shell_integration(history: str): atexit.register(readline.write_history_file, history) -async def maybe_await_async(var, *args, **kwargs): +async def maybe_await_async(var: Any, *args: Any, **kwargs: Any) -> Any: """Utility function to handle both synchronous and asynchronous callables or values. Args: var: A value, callable, or awaitable. diff --git a/charge/clients/Client.py b/charge/clients/Client.py index 713c7e7..4a7a8b3 100644 --- a/charge/clients/Client.py +++ b/charge/clients/Client.py @@ -1,3 +1,9 @@ +################################################################################ +## 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 typing import Type, Dict, Optional from abc import ABC, abstractmethod from charge.tasks.Task import Task @@ -10,11 +16,25 @@ import argparse import atexit import readline +from charge.agents.Agent import Agent + + class Client: + """Base client class for orchestrating tasks and interacting with MCP servers. + + Subclasses must implement configuration, execution, and interaction methods. + """ def __init__( self, task: Task, path: str = ".", max_retries: int = 3 ): + """Initialize the client with a task instance. + + Args: + task: The Task object this client will manage. + path: Directory path for generated files. + max_retries: Maximum number of retry attempts for server communication. + """ self.task = task self.path = path self.max_retries = max_retries @@ -24,10 +44,15 @@ def __init__( self._setup() def reset(self): + """Reset internal message and reasoning traces to start a fresh run.""" self.messages = [] self.reasoning_trace = [] def _setup(self): + """Inspect the task class and collect verifier methods. + + Populates ``self.verifier_methods`` with methods marked as verifiers. + """ cls_info = inspect_class(self.task) methods = inspect.getmembers(self.task, predicate=inspect.ismethod) name = cls_info["name"] @@ -46,6 +71,10 @@ def _setup(self): self.verifier_methods = verifier_methods def setup_mcp_servers(self): + """Generate MCP server files for hypothesis and verifier methods. + + Creates Python files containing MCP representations of task methods. + """ class_info = inspect_class(self.task) name = class_info["name"] @@ -77,26 +106,53 @@ def setup_mcp_servers(self): @abstractmethod def configure(model: str, backend: str) -> (str, str, str, Dict[str, str]): + """Configure the client with model and backend details. + + Returns a tuple of (model, backend, additional_info, config_dict). + """ raise NotImplementedError("Subclasses must implement this method.") @abstractmethod async def run(self): + """Execute the full task workflow. + + Subclasses should implement the orchestration logic here. + """ raise NotImplementedError("Subclasses must implement this method.") @abstractmethod - async def step(self, agent, task: str): + async def step(self, agent: Agent, task: str): + """Perform a single step of the task using the given agent. + + Args: + agent: The agent performing the step. + task: Description of the task step. + """ raise NotImplementedError("Subclasses must implement this method.") @abstractmethod async def chat(self): + """Interactively chat with the orchestrator. + + Subclasses should handle chat I/O. + """ raise NotImplementedError("Subclasses must implement this method.") @abstractmethod async def refine(self, feedback: str): + """Refine the task based on feedback. + + Args: + feedback: Feedback string to adjust the task execution. + """ raise NotImplementedError("Subclasses must implement this method.") @staticmethod def add_std_parser_arguments(parser: argparse.ArgumentParser): + """Utility method to add standard command‑line arguments for the client. + + Populates an ``argparse.ArgumentParser`` with common options. + """ parser.add_argument( "--model", type=str, diff --git a/charge/clients/autogen.py b/charge/clients/autogen.py index e80752a..fb75151 100644 --- a/charge/clients/autogen.py +++ b/charge/clients/autogen.py @@ -167,7 +167,7 @@ def __init__( timeout: int = 60, backend: Optional[str] = None, model_kwargs: Optional[dict] = None, - **kwargs, + **kwargs: Any, ) -> None: super().__init__(task=task, **kwargs) self.max_retries = max_retries 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_utils.py b/charge/servers/SMILES_utils.py index fd2b0d9..cec5881 100644 --- a/charge/servers/SMILES_utils.py +++ b/charge/servers/SMILES_utils.py @@ -27,6 +27,7 @@ def canonicalize_smiles(smiles: str) -> str: Args: smiles (str): The input SMILES string. + Returns: str: The canonicalized SMILES string. """ @@ -51,6 +52,7 @@ def verify_smiles(smiles: str) -> bool: Args: smiles (str): The input SMILES string. + Returns: bool: True if the SMILES is valid, False otherwise. """ @@ -81,6 +83,7 @@ def get_synthesizability(smiles: str) -> float: Args: smiles (str): The input SMILES string. + Returns: float: The synthesizability score. """ @@ -110,6 +113,7 @@ def known_smiles(smiles: str) -> bool: Args: smiles (str): The input SMILES string. + Returns: bool: True if the SMILES is known to this MCP server, False otherwise. """ diff --git a/charge/servers/molecule_pricer.py b/charge/servers/molecule_pricer.py index cb96b42..532d8f0 100644 --- a/charge/servers/molecule_pricer.py +++ b/charge/servers/molecule_pricer.py @@ -1,3 +1,10 @@ +################################################################################ +## 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 loguru import logger try: import chemprice diff --git a/charge/servers/server_utils.py b/charge/servers/server_utils.py index 393101d..4cfd180 100644 --- a/charge/servers/server_utils.py +++ b/charge/servers/server_utils.py @@ -1,3 +1,10 @@ +################################################################################ +## Copyright 2025 Lawrence Livermore National Security, LLC. and Binghamton University. +## See the top-level LICENSE file for details. +## +## SPDX-License-Identifier: Apache-2.0 +################################################################################ + import argparse from mcp.server.fastmcp import FastMCP @@ -5,8 +12,12 @@ def add_server_arguments(parser: argparse.ArgumentParser) -> None: """ Add standard server arguments to an argparse parser. + Args: parser (argparse.ArgumentParser): The parser to add arguments to. + + Returns: + None """ parser.add_argument( "--port", type=int, default=8000, help="Port to run the server on" @@ -24,11 +35,28 @@ def add_server_arguments(parser: argparse.ArgumentParser) -> None: ) -def update_mcp_network(mcp: FastMCP, host: str, port: str): +def update_mcp_network(mcp: FastMCP, host: str, port: str) -> None: + """ + Update the MCP network settings. + + Args: + mcp (FastMCP): The MCP server to update. + host (str): The host to run the server on. + port (str): The port to run the server on. + + Returns: + None + """ mcp.settings.host = host mcp.settings.port = port -def get_hostname(): +def get_hostname() -> Tuple[str, str]: + """ + Get the hostname and IP address of the host. + + Returns: + Tuple[str, str]: The hostname and IP address of the host. + """ import socket hostname = socket.gethostname() try: @@ -37,7 +65,13 @@ def get_hostname(): host = "127.0.0.1" return hostname, host -def try_get_public_hostname(): +def try_get_public_hostname() -> Tuple[str, str]: + """ + Try to get the public hostname and IP address of the host. + + Returns: + Tuple[str, str]: The public hostname and IP address of the host. + """ import socket hostname = socket.gethostname() try: diff --git a/charge/tasks/LMOTask.py b/charge/tasks/LMOTask.py index 71d3429..3a4a2ea 100644 --- a/charge/tasks/LMOTask.py +++ b/charge/tasks/LMOTask.py @@ -1,3 +1,16 @@ +################################################################################ +## Copyright 2025 Lawrence Livermore National Security, LLC. and Binghamton University. +## See the top-level LICENSE file for details. +## +## SPDX-License-Identifier: Apache-2.0 +################################################################################ + +""" +This module defines the LMOTask class, which is responsible for generating and validating +novel small molecules based on a lead molecule. It uses a Large Language Model (LLM) +to propose molecules that meet specific criteria regarding binding affinity, +synthesizability, and density. +""" import charge from charge.tasks.Task import Task from charge.servers import SMILES_utils @@ -43,7 +56,19 @@ class MoleculeOutputSchema(BaseModel): @field_validator("smiles_list") @classmethod - def validate_smiles_list(cls, smiles_list): + def validate_smiles_list(cls, smiles_list: List[str]) -> List[str]: + """ + Validate that the input is a list of valid SMILES strings. + + Args: + smiles_list (List[str]): The list of SMILES strings to validate. + + Returns: + List[str]: The validated list of SMILES strings. + + Raises: + ValueError: If the input is not a list, contains non-strings, or contains invalid SMILES. + """ if not isinstance(smiles_list, list): raise ValueError("smiles_list must be a list.") for smiles in smiles_list: @@ -54,9 +79,21 @@ def validate_smiles_list(cls, smiles_list): return smiles_list def as_list(self) -> List[str]: + """ + Return the list of SMILES strings. + + Returns: + List[str]: The list of SMILES strings. + """ return self.smiles_list def as_dict(self) -> dict: + """ + Return the object as a dictionary. + + Returns: + dict: A dictionary representation of the object. + """ return { "reasoning_summary": self.reasoning_summary, "smiles_list": self.smiles_list, @@ -70,6 +107,12 @@ def as_dict(self) -> dict: class LMOTask(Task): + """ + This class defines the LMOTask, which is responsible for generating and validating + novel small molecules based on a lead molecule. It uses a Large Language Model (LLM) + to propose molecules that meet specific criteria regarding binding affinity, + synthesizability, and density. + """ def __init__( self, lead_molecule: str, @@ -108,10 +151,13 @@ def check_proposal(self, smiles: str) -> bool: Check if the proposed SMILES string is valid. If it is valid, checks if its synthesizability score is less than or equal to the lead molecule and if its density is greater than or equal to the lead molecule. + Args: smiles (str): The proposed SMILES string. + Returns: bool: True if the proposal is valid and meets the criteria, False otherwise. + Raises: ValueError: If the SMILES string is invalid or does not meet the criteria. """ @@ -137,12 +183,16 @@ def check_final_proposal(self, smiles_list_as_string: str) -> bool: """ Check if the proposed SMILES strings are valid and meet the criteria. The criteria are: + 1. The SMILES must be valid. + 2. The synthesizability score must be less than or equal to the lead molecule. + 3. The density must be greater than or equal to the lead molecule. Args: - smiles (str): The proposed list of SMILES strings. + smiles_list_as_string (str): The proposed list of SMILES strings as a string. + Returns: bool: True if the proposal is valid and meets the criteria, False otherwise. diff --git a/charge/tasks/Task.py b/charge/tasks/Task.py index e3e323a..fc9209c 100644 --- a/charge/tasks/Task.py +++ b/charge/tasks/Task.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod from pydantic import BaseModel -from typing import Type, Union, Optional +from typing import Type, Union, Optional, Any import os import os.path as osp import warnings @@ -18,7 +18,7 @@ def __init__( server_urls: Optional[Union[str, list]] = None, server_files: Optional[Union[str, list]] = None, structured_output_schema: Optional[Type[BaseModel]] = None, - **kwargs, + **kwargs: Any, ): """ Base class for defining an task, which is composed of a set of steps: From 50594d4afa6166b1e2afe34c3798eef75bf39ffd Mon Sep 17 00:00:00 2001 From: Shehtab Date: Tue, 18 Nov 2025 21:58:58 -0500 Subject: [PATCH 3/7] Complete documentation for function tagging in _tags --- charge/_tags.py | 46 ++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 38 insertions(+), 8 deletions(-) diff --git a/charge/_tags.py b/charge/_tags.py index dcfd40b..32822e1 100644 --- a/charge/_tags.py +++ b/charge/_tags.py @@ -1,20 +1,50 @@ -def verifier(func): - """Decorator to mark a method as a verifier.""" +from typing import Callable + +def verifier(func: Callable) -> Callable: + """Decorator to mark a method as a verifier. + + Args: + func (Callable): The function to mark as a verifier. + + Returns: + Callable: The function marked as a verifier. + """ func.__verifier_tag = True return func -def is_verifier(func): - """Check if a method is marked as a verifier.""" +def is_verifier(func: Callable) -> bool: + """Check if a method is marked as a verifier. + + Args: + func (Callable): The function to check. + + Returns: + bool: True if the function is marked as a verifier, False otherwise. + """ return hasattr(func, "__verifier_tag") -def hypothesis(func): - """Decorator to mark a method as a hypothesis.""" +def hypothesis(func: Callable) -> Callable: + """Decorator to mark a method as a hypothesis. + + Args: + func (Callable): The function to mark as a hypothesis. + + Returns: + Callable: The function marked as a hypothesis. + """ func.__hypothesis_tag = True return func -def is_hypothesis(func): - """Check if a method is marked as a hypothesis.""" +def is_hypothesis(func: Callable) -> bool: + """Check if a method is marked as a hypothesis. + + Args: + func (Callable): The function to check. + + Returns: + bool: True if the function is marked as a hypothesis, False otherwise. + """ return hasattr(func, "__hypothesis_tag") From a0864c80c9f474da87f9c84cb430f1f76a989874 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Tue, 18 Nov 2025 22:08:26 -0500 Subject: [PATCH 4/7] Updated documentation for utilites to conform to doc standards --- charge/_tags.py | 5 +++ charge/_utils.py | 2 + charge/inspector.py | 13 +++++-- charge/utils/helper_funcs.py | 9 +++++ charge/utils/system_utils.py | 73 ++++++++++++++++++++++++++++++++++-- 5 files changed, 96 insertions(+), 6 deletions(-) diff --git a/charge/_tags.py b/charge/_tags.py index 32822e1..e89b556 100644 --- a/charge/_tags.py +++ b/charge/_tags.py @@ -33,6 +33,11 @@ def hypothesis(func: Callable) -> Callable: Returns: Callable: The function marked as a hypothesis. + + + Side Effects: + Marks the function as a hypothesis. + `func` is modified in place to include the `__hypothesis_tag` attribute. """ func.__hypothesis_tag = True return func diff --git a/charge/_utils.py b/charge/_utils.py index ca0b3e0..76848d7 100644 --- a/charge/_utils.py +++ b/charge/_utils.py @@ -30,10 +30,12 @@ def enable_cmd_history_and_shell_integration(history: str) -> None: async def maybe_await_async(var: Any, *args: Any, **kwargs: Any) -> Any: """Utility function to handle both synchronous and asynchronous callables or values. + Args: var: A value, callable, or awaitable. *args: Positional arguments to pass if var is callable. **kwargs: Keyword arguments to pass if var is callable. + Returns: The result of the callable or awaitable, or the value itself. """ diff --git a/charge/inspector.py b/charge/inspector.py index 2432c6c..a2363cc 100644 --- a/charge/inspector.py +++ b/charge/inspector.py @@ -1,8 +1,15 @@ import inspect +from typing import Type - -def inspect_class(cls): - +def inspect_class(cls: Type) -> dict: + """Inspect a class and return its type, name, and file path. + + Args: + cls: The class to inspect. + + Returns: + dict: A dictionary containing the type, name, and file path of the class. + """ type_ = type(cls) module = inspect.getmodule(cls.__class__) file = module.__file__ if module else "Unknown" diff --git a/charge/utils/helper_funcs.py b/charge/utils/helper_funcs.py index 45bee2f..aa86957 100644 --- a/charge/utils/helper_funcs.py +++ b/charge/utils/helper_funcs.py @@ -8,8 +8,10 @@ def get_list_from_json_file(file_path: str) -> list: """ Load a list of molecules from a JSON file. + Args: file_path (str): The path to the JSON file. + Returns: list: The list of molecules. """ @@ -28,9 +30,13 @@ def get_list_from_json_file(file_path: str) -> list: def save_list_to_json_file(data: list, file_path: str) -> None: """ Save a list of molecules to a JSON file. + Args: data (list): The list of molecules. file_path (str): The path to the JSON file. + + Returns: + None """ try: with open(file_path, "w") as f: @@ -46,6 +52,9 @@ def post_process_smiles(smiles: str, parent_id: int, node_id: int) -> dict: Args: smiles (str): The input SMILES string. + parent_id (int): The parent ID of the molecule. + node_id (int): The node ID of the molecule. + Returns: dict: The post-processed dictionary. """ diff --git a/charge/utils/system_utils.py b/charge/utils/system_utils.py index 19f3169..34d85d1 100644 --- a/charge/utils/system_utils.py +++ b/charge/utils/system_utils.py @@ -7,6 +7,16 @@ import requests def normalize_string(s: str) -> str: + """ + Normalize a string by converting it to lowercase and removing any whitespace or + hyphens. + + Args: + s (str): The string to normalize. + + Returns: + str: The normalized string. + """ s = s.lower() s = re.sub(r"[\s\-]+", "_", s) s = re.sub(r"_+", "_", s) @@ -15,12 +25,31 @@ def normalize_string(s: str) -> str: def _load_json(file_path: str) -> dict: + """ + Load a JSON file and return its contents. + + Args: + file_path (str): The path to the JSON file. + + Returns: + dict: The contents of the JSON file. + """ with open(file_path, "r") as f: data = json.load(f) return data def _prompt_from_json_file(file_path: str, key: str) -> str: + """ + Load a JSON file and return the value associated with the specified key. + + Args: + file_path (str): The path to the JSON file. + key (str): The key to look for in the JSON file. + + Returns: + str: The value associated with the specified key. + """ data = _load_json(file_path) for k in data.keys(): k = normalize_string(k) @@ -30,17 +59,44 @@ def _prompt_from_json_file(file_path: str, key: str) -> str: def _prompt_from_txt_file(file_path: str) -> str: - + """ + Load a text file and return its contents. + + Args: + file_path (str): The path to the text file. + + Returns: + str: The contents of the text file. + """ with open(file_path, "r") as f: prompt = f.read() return prompt def _check_file_exists(file_path: str) -> bool: + """ + Check if a file exists. + + Args: + file_path (str): The path to the file. + + Returns: + bool: True if the file exists, False otherwise. + """ return osp.isfile(file_path) def read_from_file(self, file_path: str, key: str) -> str: + """ + Read a file and return its contents. + + Args: + file_path (str): The path to the file. + key (str): The key to look for in the file. + + Returns: + str: The contents of the file. + """ assert osp.isfile(file_path), f"File {file_path} does not exist" if file_path.endswith(".txt"): return _prompt_from_txt_file(file_path) @@ -50,6 +106,15 @@ def read_from_file(self, file_path: str, key: str) -> str: raise ValueError("Only .txt and .json files are supported") def check_url_exists(url: str) -> bool: + """ + Check if a URL exists. + + Args: + url (str): The URL to check. + + Returns: + bool: True if the URL exists, False otherwise. + """ if not url.startswith("http://") and not url.startswith("https://"): return False @@ -70,13 +135,15 @@ def check_url_exists(url: str) -> bool: def check_server_paths(server_paths: Optional[Union[str, list]]) -> list: """ Gracefully handle errors in server paths provided by user. + Args: server_paths (Optional[Union[str, list]]): The server paths to check. + Returns: list: A list of valid server paths. + Raises: - FileNotFoundError: If any of the server paths do not exist and - CHARGE_ERROR_ON_MISSING_SERVER is set to 1. + ValueError: If any of the server paths do not exist and `CHARGE_ERROR_ON_MISSING_SERVER` is set to 1. """ if server_paths is None: From 61ca0bb84e5ec3ccef54791593922a5a0789a498 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Tue, 18 Nov 2025 22:12:15 -0500 Subject: [PATCH 5/7] Add Task abstract class documentation --- charge/tasks/Task.py | 65 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 64 insertions(+), 1 deletion(-) diff --git a/charge/tasks/Task.py b/charge/tasks/Task.py index fc9209c..76bb72d 100644 --- a/charge/tasks/Task.py +++ b/charge/tasks/Task.py @@ -82,21 +82,58 @@ def __init__( self.constructor_args = {} def get_system_prompt(self) -> str: + """ + Get the system prompt for the task. + + Returns: + str: The system prompt for the task. + """ return self.system_prompt or "" def get_user_prompt(self) -> str: + """ + Get the user prompt for the task. + + Returns: + str: The user prompt for the task. + """ return self.user_prompt or "" def get_verification_prompt(self) -> str: + """ + Get the verification prompt for the task. + + Returns: + str: The verification prompt for the task. + """ return self.verification_prompt or "" def get_refinement_prompt(self) -> str: + """ + Get the refinement prompt for the task. + + Returns: + str: The refinement prompt for the task. + """ return self.refinement_prompt or "" def register_buffer(self, name: str, value: str): + """ + Register a buffer for the task. + + Args: + name (str): The name of the buffer. + value (str): The value of the buffer. + """ self.constructor_args[name] = value - def get_structured_output_schema(self): + def get_structured_output_schema(self) -> Type[BaseModel]: + """ + Get the structured output schema for the task. + + Returns: + Type[BaseModel]: The structured output schema for the task. + """ assert ( self.has_structured_output_schema() ), "structured_output_schema not implemented" @@ -104,9 +141,21 @@ def get_structured_output_schema(self): return self.structured_output_schema # type: ignore def set_structured_output_schema(self, schema: Type[BaseModel]): + """ + Set the structured output schema for the task. + + Args: + schema (Type[BaseModel]): The structured output schema to set. + """ self.structured_output_schema = schema def has_structured_output_schema(self) -> bool: + """ + Check if the task has a structured output schema. + + Returns: + bool: True if the task has a structured output schema, False otherwise. + """ return ( hasattr(self, "structured_output_schema") and self.structured_output_schema is not None @@ -115,8 +164,10 @@ def has_structured_output_schema(self) -> bool: def set_system_prompt_from_file(self, file_path: str): """ Set the system prompt from a file. + Args: file_path (str): Path to the file containing the system prompt. + Raises: ValueError: If the file is not a .txt or .json file. """ @@ -126,8 +177,10 @@ def set_system_prompt_from_file(self, file_path: str): def set_user_prompt_from_file(self, file_path: str): """ Set the user prompt from a file. + Args: file_path (str): Path to the file containing the user prompt. + Raises: ValueError: If the file is not a .txt or .json file. @@ -138,8 +191,10 @@ def set_user_prompt_from_file(self, file_path: str): def set_verification_prompt_from_file(self, file_path: str): """ Set the verification prompt from a file. + Args: file_path (str): Path to the file containing the verification prompt. + Raises: ValueError: If the file is not a .txt or .json file. """ @@ -149,8 +204,10 @@ def set_verification_prompt_from_file(self, file_path: str): def set_refinement_prompt_from_file(self, file_path: str): """ Set the refinement prompt from a file. + Args: file_path (str): Path to the file containing the refinement prompt. + Raises: ValueError: If the file is not a .txt or .json file. """ @@ -160,6 +217,7 @@ def set_refinement_prompt_from_file(self, file_path: str): def has_verification_prompt(self) -> bool: """ Check if the task has a verification prompt. + Returns: bool: True if the task has a verification prompt, False otherwise. """ @@ -168,6 +226,7 @@ def has_verification_prompt(self) -> bool: def has_refinement_prompt(self) -> bool: """ Check if the task has a refinement prompt. + Returns: bool: True if the task has a refinement prompt, False otherwise. """ @@ -176,6 +235,10 @@ def has_refinement_prompt(self) -> bool: def check_output_formatting(self, content: str | bytes | bytearray) -> bool: """ Check if the task has output formatting requirements. + + Args: + content (str | bytes | bytearray): The content to check. + Returns: bool: True if the task has output formatting requirements, False otherwise. """ From 0095637282365422fbc2dbf098414607696d5780 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Tue, 18 Nov 2025 22:18:22 -0500 Subject: [PATCH 6/7] Add license header --- charge/experiments/Experiment.py | 7 +++++++ charge/tasks/RetrosynthesisTask.py | 7 +++++++ charge/tasks/Task.py | 29 ++++++++++++++--------------- 3 files changed, 28 insertions(+), 15 deletions(-) diff --git a/charge/experiments/Experiment.py b/charge/experiments/Experiment.py index 504b9cf..6f591c5 100644 --- a/charge/experiments/Experiment.py +++ b/charge/experiments/Experiment.py @@ -1,3 +1,10 @@ +################################################################################ +## 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 abc import abstractmethod from typing import Any, List, Union, Optional from charge.tasks.Task import Task diff --git a/charge/tasks/RetrosynthesisTask.py b/charge/tasks/RetrosynthesisTask.py index b30b776..0caa327 100644 --- a/charge/tasks/RetrosynthesisTask.py +++ b/charge/tasks/RetrosynthesisTask.py @@ -1,3 +1,10 @@ +################################################################################ +## 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 charge.tasks.Task import Task from charge.servers.log_progress import LOG_PROGRESS_SYSTEM_PROMPT from typing import List, Optional diff --git a/charge/tasks/Task.py b/charge/tasks/Task.py index 76bb72d..5c715e8 100644 --- a/charge/tasks/Task.py +++ b/charge/tasks/Task.py @@ -35,8 +35,19 @@ def __init__( reasoning engine for self verification. The refinement prompt is used to guide the reasoning engine to refine its response if the verification fails. - **Note**: Automatic verification is an experimental feature and may not work as - expected. + + Args: + system_prompt (str, optional): The system prompt for the task. + user_prompt (str, optional): The user prompt for the task. + verification_prompt (str, optional): The verification prompt for the task. + refinement_prompt (str, optional): The refinement prompt for the task. + server_urls (Union[str, list], optional): The MCP server URLs to use with over SSE protocol + for the task. + server_files (Union[str, list], optional): The MCP server files to use with over STDIO protocl + for the task. + **kwargs: Additional keyword arguments to be stored in the task. + + **Note:** Automatic verification is an experimental feature and may not work as expected. The task class can also be extended to include hypothesis methods (decorated with @hypothesis) and verifier methods (decorated with @verifier). @@ -45,24 +56,12 @@ def __init__( (with type annotations and docstrings) can also be added to the Task via the register__tool functions. - **Note**: Automatic MCP generation is an experimental feature and may not work as + **Note:** Automatic MCP generation is an experimental feature and may not work as expected. All decorated methods must have proper type annotations and be static. The docstring of the methods is used as the docstring in the MCP server. Long running MCPs with high starting costs should be provided separately to the client or the method should call out to an external service / process. - - Args: - system_prompt (str, optional): The system prompt for the task. - user_prompt (str, optional): The user prompt for the task. - verification_prompt (str, optional): The verification prompt for the task. - refinement_prompt (str, optional): The refinement prompt for the task. - server_urls (Union[str, list], optional): The MCP server URLs to use with over SSE protocol - for the task. - server_files (Union[str, list], optional): The MCP server files to use with over STDIO protocl - for the task. - **kwargs: Additional keyword arguments to be stored in the task. - """ self.system_prompt = system_prompt self.user_prompt = user_prompt From 1e0f9a2c9bbf0b89d1ef81fd7d3215477c53254b Mon Sep 17 00:00:00 2001 From: Shehtab Date: Tue, 18 Nov 2025 22:23:44 -0500 Subject: [PATCH 7/7] Add proper admonitions --- charge/tasks/Task.py | 29 +++++++++++++++-------------- mkdocs.yml | 5 +++++ 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/charge/tasks/Task.py b/charge/tasks/Task.py index 5c715e8..2bd1aaa 100644 --- a/charge/tasks/Task.py +++ b/charge/tasks/Task.py @@ -47,20 +47,21 @@ def __init__( for the task. **kwargs: Additional keyword arguments to be stored in the task. - **Note:** Automatic verification is an experimental feature and may not work as expected. - - The task class can also be extended to include hypothesis methods - (decorated with @hypothesis) and verifier methods (decorated with @verifier). - Appropriate MCPs are automatically generated for these methods and used by the - Client class to call these methods in the HVR process. Prewritten functions - (with type annotations and docstrings) can also be added to the Task - via the register__tool functions. - - **Note:** Automatic MCP generation is an experimental feature and may not work as - expected. All decorated methods must have proper type annotations and be static. - The docstring of the methods is used as the docstring in the MCP server. - Long running MCPs with high starting costs should be provided separately to the - client or the method should call out to an external service / process. + !!! note "Automatic verification is an experimental feature and may not work as expected." + + The task class can also be extended to include hypothesis methods + (decorated with @hypothesis) and verifier methods (decorated with @verifier). + Appropriate MCPs are automatically generated for these methods and used by the + Client class to call these methods in the HVR process. Prewritten functions + (with type annotations and docstrings) can also be added to the Task + via the register__tool functions. + + !!! note "Automatic MCP generation is an experimental feature and may not work as expected." + + All decorated methods must have proper type annotations and be static. + The docstring of the methods is used as the docstring in the MCP server. + Long running MCPs with high starting costs should be provided separately to the + client or the method should call out to an external service / process." """ self.system_prompt = system_prompt diff --git a/mkdocs.yml b/mkdocs.yml index 5980ae8..3596c78 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -27,6 +27,11 @@ plugins: show_members_order: true paths: [.] # IMPORTANT: tell mkdocstrings where to find your package source +markdown_extensions: + - admonition + - pymdownx.details + - pymdownx.superfences + theme: name: readthedocs language: en