-
Notifications
You must be signed in to change notification settings - Fork 15
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Loading status checks…
merge remote branch 'main' into experiments
Showing
14 changed files
with
181 additions
and
256 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,8 +1,11 @@ | ||
# Copy this file to a new file named .env and replace the placeholders with your actual keys. | ||
# REMOVE "pragma: allowlist secret" when you replace with actual keys. | ||
# DO NOT fill your keys directly into this file. | ||
|
||
# OpenAI API Key | ||
OPENAI_API_KEY=YOUR_OPENAI_API_KEY_GOES_HERE # pragma: allowlist secret | ||
|
||
# Serp API key | ||
SERP_API_KEY=YOUR_SERP_API_KEY_GOES_HERE # pragma: allowlist secret | ||
# PQA API Key to use LiteratureSearch tool (optional) -- it also requires OpenAI key | ||
PQA_API_KEY=YOUR_PQA_API_KEY_GOES_HERE # pragma: allowlist secret | ||
|
||
# Optional: add TogetherAI, Fireworks, or Anthropic API key here to use their models |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,38 +1,48 @@ | ||
MD-Agent is a LLM-agent based toolset for Molecular Dynamics. | ||
MDAgent is a LLM-agent based toolset for Molecular Dynamics. | ||
It's built using Langchain and uses a collection of tools to set up and execute molecular dynamics simulations, particularly in OpenMM. | ||
|
||
|
||
## Environment Setup | ||
To use the OpenMM features in the agent, please set up a conda environment, following these steps. | ||
- Create conda environment: `conda env create -n mdagent -f environment.yaml` | ||
- Activate your environment: `conda activate mdagent` | ||
``` | ||
conda env create -n mdagent -f environment.yaml | ||
conda activate mdagent | ||
``` | ||
|
||
If you already have a conda environment, you can install dependencies before you activate it with the following step. | ||
- Install the necessary conda dependencies: `conda env update -n <YOUR_CONDA_ENV_HERE> -f environment.yaml` | ||
|
||
If you already have a conda environment, you can install dependencies with the following step. | ||
- Install the necessary conda dependencies: `conda install -c conda-forge openmm pdbfixer mdtraj` | ||
|
||
|
||
## Installation | ||
``` | ||
pip install git+https://github.com/ur-whitelab/md-agent.git | ||
``` | ||
|
||
|
||
## Usage | ||
The first step is to set up your API keys in your environment. An OpenAI key is necessary for this project. | ||
The next step is to set up your API keys in your environment. An API key for LLM provider is necessary for this project. Supported LLM providers are OpenAI, TogetherAI, Fireworks, and Anthropic. | ||
Other tools require API keys, such as paper-qa for literature searches. We recommend setting up the keys in a .env file. You can use the provided .env.example file as a template. | ||
1. Copy the `.env.example` file and rename it to `.env`: `cp .env.example .env` | ||
2. Replace the placeholder values in `.env` with your actual keys | ||
|
||
<!-- ## Using Streamlit Interface | ||
If you'd like to use MDAgent via the streamlit app, make sure you have completed the steps above. Then, in your terminal, run `streamlit run st_app.py` in the project root directory. | ||
You can ask MDAgent to conduct molecular dynamics tasks using OpenAI's GPT model | ||
``` | ||
from mdagent import MDAgent | ||
agent = MDAgent(model="gpt-3.5-turbo") | ||
agent.run("Simulate protein 1ZNI at 300 K for 0.1 ps and calculate the RMSD over time.") | ||
``` | ||
Note: to distinguish Together models from the rest, you'll need to add "together\" prefix in model flag, such as `agent = MDAgent(model="together/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo")` | ||
|
||
From there you may upload files to use during the run. Note: the app is currently limited to uploading .pdb and .cif files, and the max size is defaulted at 200MB. | ||
- To upload larger files, instead run `streamlit run st_app.py --server.maxUploadSize=some_large_number` | ||
- To add different file types, you can add your desired file type to the list in the [streamlit app file](https://github.com/ur-whitelab/md-agent/blob/main/st_app.py). --> | ||
## LLM Providers | ||
By default, we support LLMs through OpenAI API. However, feel free to use other LLM providers. Make sure to install the necessary package for it. Here's list of packages required for alternative LLM providers we support: | ||
- `pip install langchain-together` to use models from TogetherAI | ||
- `pip install langchain-anthropic` to use models from Anthropic | ||
- `pip install langchain-fireworks` to use models from Fireworks | ||
|
||
|
||
## Contributing | ||
|
||
We welcome contributions to MD-Agent! If you're interested in contributing to the project, please check out our [Contributor's Guide](CONTRIBUTING.md) for detailed instructions on getting started, feature development, and the pull request process. | ||
We welcome contributions to MDAgent! If you're interested in contributing to the project, please check out our [Contributor's Guide](CONTRIBUTING.md) for detailed instructions on getting started, feature development, and the pull request process. | ||
|
||
We value and appreciate all contributions to MD-Agent. | ||
We value and appreciate all contributions to MDAgent. |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,3 +1,3 @@ | ||
from .maketools import get_tools, make_all_tools | ||
from .maketools import get_relevant_tools, make_all_tools | ||
|
||
__all__ = ["get_tools", "make_all_tools"] | ||
__all__ = ["get_relevant_tools", "make_all_tools"] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,10 +1,8 @@ | ||
from .git_issues_tool import SerpGitTool | ||
from .registry_tools import ListRegistryPaths, MapPath2Name | ||
from .search_tools import Scholar2ResultLLM | ||
|
||
__all__ = [ | ||
"ListRegistryPaths", | ||
"MapPath2Name", | ||
"Scholar2ResultLLM", | ||
"SerpGitTool", | ||
] |
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,96 @@ | ||
import os | ||
from unittest.mock import MagicMock, patch | ||
|
||
import numpy as np | ||
import pytest | ||
|
||
from mdagent.tools.maketools import get_relevant_tools | ||
|
||
|
||
@pytest.fixture | ||
def mock_llm(): | ||
return MagicMock() | ||
|
||
|
||
@pytest.fixture | ||
def mock_tools(): | ||
Tool = MagicMock() | ||
tool1 = Tool(name="Tool1", description="This is the first tool") | ||
tool2 = Tool(name="Tool2", description="This is the second tool") | ||
tool3 = Tool(name="Tool3", description="This is the third tool") | ||
return [tool1, tool2, tool3] | ||
|
||
|
||
@patch("mdagent.tools.maketools.make_all_tools") | ||
@patch("mdagent.tools.maketools.OpenAIEmbeddings") | ||
def test_get_relevant_tools_with_openai_embeddings( | ||
mock_openai_embeddings, mock_make_all_tools, mock_llm, mock_tools | ||
): | ||
mock_make_all_tools.return_value = mock_tools | ||
mock_embed_documents = mock_openai_embeddings.return_value.embed_documents | ||
mock_embed_query = mock_openai_embeddings.return_value.embed_query | ||
mock_embed_documents.return_value = np.random.rand(3, 512) | ||
mock_embed_query.return_value = np.random.rand(512) | ||
|
||
with patch.dict( | ||
os.environ, {"OPENAI_API_KEY": "test_key"} # pragma: allowlist secret | ||
): | ||
relevant_tools = get_relevant_tools("test query", mock_llm, top_k_tools=2) | ||
assert len(relevant_tools) == 2 | ||
assert relevant_tools[0] in mock_tools | ||
assert relevant_tools[1] in mock_tools | ||
|
||
|
||
@patch("mdagent.tools.maketools.make_all_tools") | ||
@patch("mdagent.tools.maketools.TfidfVectorizer") | ||
def test_get_relevant_tools_with_tfidf( | ||
mock_tfidf_vectorizer, mock_make_all_tools, mock_llm, mock_tools | ||
): | ||
mock_make_all_tools.return_value = mock_tools | ||
mock_vectorizer = mock_tfidf_vectorizer.return_value | ||
mock_vectorizer.fit_transform.return_value = np.random.rand(3, 10) | ||
mock_vectorizer.transform.return_value = np.random.rand(1, 10) | ||
|
||
with patch.dict(os.environ, {}, clear=True): # ensure OPENAI_API_KEY is not set | ||
relevant_tools = get_relevant_tools("test query", mock_llm, top_k_tools=2) | ||
assert len(relevant_tools) == 2 | ||
assert relevant_tools[0] in mock_tools | ||
assert relevant_tools[1] in mock_tools | ||
|
||
|
||
@patch("mdagent.tools.maketools.make_all_tools") | ||
def test_get_relevant_tools_with_no_tools(mock_make_all_tools, mock_llm): | ||
mock_make_all_tools.return_value = [] | ||
|
||
with patch.dict(os.environ, {}, clear=True): | ||
relevant_tools = get_relevant_tools("test query", mock_llm) | ||
assert relevant_tools is None | ||
|
||
|
||
@patch("mdagent.tools.maketools.make_all_tools") | ||
@patch("mdagent.tools.maketools.OpenAIEmbeddings") | ||
def test_get_relevant_tools_with_openai_exception( | ||
mock_openai_embeddings, mock_make_all_tools, mock_llm, mock_tools | ||
): | ||
mock_make_all_tools.return_value = mock_tools | ||
mock_embed_documents = mock_openai_embeddings.return_value.embed_documents | ||
mock_embed_documents.side_effect = Exception("Embedding error") | ||
|
||
with patch.dict( | ||
os.environ, {"OPENAI_API_KEY": "test_key"} # pragma: allowlist secret | ||
): | ||
relevant_tools = get_relevant_tools("test query", mock_llm) | ||
assert relevant_tools is None | ||
|
||
|
||
@patch("mdagent.tools.maketools.make_all_tools") | ||
def test_get_relevant_tools_top_k(mock_make_all_tools, mock_llm, mock_tools): | ||
mock_make_all_tools.return_value = mock_tools | ||
|
||
with patch.dict(os.environ, {}, clear=True): | ||
relevant_tools = get_relevant_tools("test query", mock_llm, top_k_tools=1) | ||
assert len(relevant_tools) == 1 | ||
assert relevant_tools[0] in mock_tools | ||
|
||
relevant_tools = get_relevant_tools("test query", mock_llm, top_k_tools=5) | ||
assert len(relevant_tools) == len(mock_tools) |