remove chroma & simplify dependencies

.env.example

@@ -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

README.md

@@ -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.

mdagent/agent/agent.py

@@ -4,7 +4,7 @@
 from langchain.agents import AgentExecutor, OpenAIFunctionsAgent
 from langchain.agents.structured_chat.base import StructuredChatAgent
 
-from ..tools import get_tools, make_all_tools
+from ..tools import get_relevant_tools, make_all_tools
 from ..utils import PathRegistry, SetCheckpoint, _make_llm
 from .memory import MemoryManager
 from .prompt import openaifxn_prompt, structured_prompt
@@ -76,7 +76,7 @@ def _initialize_tools_and_agent(self, user_input=None):
         else:
             if self.top_k_tools != "all" and user_input is not None:
                 # retrieve only tools relevant to user input
-                self.tools = get_tools(
+                self.tools = get_relevant_tools(
                     query=user_input,
                     llm=self.tools_llm,
                     top_k_tools=self.top_k_tools,

mdagent/tools/__init__.py

@@ -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"]

mdagent/tools/base_tools/__init__.py

@@ -45,7 +45,6 @@
 )
 from .simulation_tools.create_simulation import ModifyBaseSimulationScriptTool
 from .simulation_tools.setup_and_run import SetUpandRunFunction
-from .util_tools.git_issues_tool import SerpGitTool
 from .util_tools.registry_tools import ListRegistryPaths, MapPath2Name
 from .util_tools.search_tools import Scholar2ResultLLM
 
@@ -87,7 +86,6 @@
     "RDFTool",
     "RMSDCalculator",
     "Scholar2ResultLLM",
-    "SerpGitTool",
     "SetUpandRunFunction",
     "SimulationOutputFigures",
     "SmallMolPDB",

mdagent/tools/base_tools/preprocess_tools/pdb_get.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
 import requests
-import streamlit as st
 from langchain.tools import BaseTool
 from rdkit import Chem
 from rdkit.Chem import AllChem
@@ -36,7 +35,6 @@ def get_pdb(query_string: str, path_registry: PathRegistry):
         results = r.json()["result_set"]
         pdbid = max(results, key=lambda x: x["score"])["identifier"]
         print(f"PDB file found with this ID: {pdbid}")
-        st.markdown(f"PDB file found with this ID: {pdbid}", unsafe_allow_html=True)
         url = f"https://files.rcsb.org/download/{pdbid}.{filetype}"
         pdb = requests.get(url)
         filename = path_registry.write_file_name(

mdagent/tools/base_tools/simulation_tools/setup_and_run.py

@@ -7,7 +7,6 @@
 from typing import Any, Dict, List, Optional, Type
 
 import requests
-import streamlit as st
 from langchain.tools import BaseTool
 from openff.toolkit.topology import Molecule
 from openmm import (
@@ -251,7 +250,6 @@ def __init__(
 
     def setup_system(self):
         print("Building system...")
-        st.markdown("Building system", unsafe_allow_html=True)
         self.pdb_id = self.params["pdb_id"]
         self.pdb_path = self.path_registry.get_mapped_path(self.pdb_id)
         self.pdb = PDBFile(self.pdb_path)
@@ -285,7 +283,6 @@ def setup_system(self):
 
     def setup_integrator(self):
         print("Setting up integrator...")
-        st.markdown("Setting up integrator", unsafe_allow_html=True)
         int_params = self.int_params
         integrator_type = int_params.get("integrator_type", "LangevinMiddle")
 
@@ -310,7 +307,6 @@ def setup_integrator(self):
 
     def create_simulation(self):
         print("Creating simulation...")
-        st.markdown("Creating simulation", unsafe_allow_html=True)
         self.simulation = Simulation(
             self.modeller.topology,
             self.system,
@@ -838,12 +834,10 @@ def remove_leading_spaces(text):
             file.write(script_content)
 
         print(f"Standalone simulation script written to {directory}/{filename}")
-        st.markdown("Standalone simulation script written", unsafe_allow_html=True)
 
     def run(self):
         # Minimize and Equilibrate
         print("Performing energy minimization...")
-        st.markdown("Performing energy minimization", unsafe_allow_html=True)
 
         self.simulation.minimizeEnergy()
         print("Minimization complete!")
@@ -857,19 +851,16 @@ def run(self):
             )
         self.path_registry.map_path(f"top_{self.sim_id}", top_name, top_description)
         print("Initial Positions saved to initial_positions.pdb")
-        st.markdown("Minimization complete! Equilibrating...", unsafe_allow_html=True)
         print("Equilibrating...")
         _temp = self.int_params["Temperature"]
         self.simulation.context.setVelocitiesToTemperature(_temp)
         _eq_steps = self.sim_params.get("equilibrationSteps", 1000)
         self.simulation.step(_eq_steps)
         # Simulate
         print("Simulating...")
-        st.markdown("Simulating...", unsafe_allow_html=True)
         self.simulation.currentStep = 0
         self.simulation.step(self.sim_params["Number of Steps"])
         print("Done!")
-        st.markdown("Done!", unsafe_allow_html=True)
         if not self.save:
             if os.path.exists("temp_trajectory.dcd"):
                 os.remove("temp_trajectory.dcd")
@@ -950,7 +941,6 @@ def _run(self, **input_args):
             openmmsim.create_simulation()
 
             print("simulation set!")
-            st.markdown("simulation set!", unsafe_allow_html=True)
         except ValueError as e:
             msg = str(e) + f"This were the inputs {input_args}"
             if "No template for" in msg:
@@ -1492,11 +1482,9 @@ def check_system_params(cls, values):
         forcefield_files = values.get("forcefield_files")
         if forcefield_files is None or forcefield_files is []:
             print("Setting default forcefields")
-            st.markdown("Setting default forcefields", unsafe_allow_html=True)
             forcefield_files = ["amber14-all.xml", "amber14/tip3pfb.xml"]
         elif len(forcefield_files) == 0:
             print("Setting default forcefields v2")
-            st.markdown("Setting default forcefields", unsafe_allow_html=True)
             forcefield_files = ["amber14-all.xml", "amber14/tip3pfb.xml"]
         else:
             for file in forcefield_files:

mdagent/tools/base_tools/util_tools/__init__.py

@@ -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",
 ]