merge main into rdkit_pdb

README.md

@@ -23,6 +23,13 @@ Other tools require API keys, such as paper-qa for literature searches. We recom
 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.
+
+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).
+
 
 ## Contributing
 

mdagent/mainagent/agent.py

@@ -47,9 +47,14 @@ def __init__(
         resume=False,
         top_k_tools=20,  # set "all" if you want to use all tools (& skills if resume)
         use_human_tool=False,
+        uploaded_files=[],  # user input files to add to path registry
     ):
         if path_registry is None:
             path_registry = PathRegistry.get_instance()
+        self.uploaded_files = uploaded_files
+        for file in uploaded_files:  # todo -> allow users to add descriptions?
+            path_registry.map_path(file, file, description="User uploaded file")
+
         self.agent_type = agent_type
         self.user_tools = tools
         self.tools_llm = _make_llm(tools_model, temp, verbose)

mdagent/subagents/subagent_fxns.py

@@ -2,6 +2,8 @@
 import os
 from typing import Optional
 
+import streamlit as st
+
 from .subagent_setup import SubAgentInitializer, SubAgentSettings
 
 
@@ -76,6 +78,7 @@ def _run_loop(self, task, full_history, skills):
         """
         critique = None
         print("\n\033[46m action agent is running, writing code\033[0m")
+        st.markdown("action agent is running, writing code", unsafe_allow_html=True)
         success, code, fxn_name, code_output = self.action._run_code(
             full_history, task, skills
         )
@@ -126,12 +129,20 @@ def _run_iterations(self, run, task):
 
                 # give successful code to tool/skill manager
                 print("\n\033[46mThe new code is complete, running skill agent\033[0m")
+                st.markdown(
+                    "The new code is complete, running skill agent",
+                    unsafe_allow_html=True,
+                )
                 tool_name = self.skill.add_new_tool(fxn_name, code)
                 return success, tool_name
             iter += 1
 
         # if max iterations reached without success, save failures to file
         print("\n\033[46m Max iterations reached, saving failed history to file\033[0m")
+        st.markdown(
+            "Max iterations reached, saving failed history to file",
+            unsafe_allow_html=True,
+        )
         tool_name = None
         full_failed = self._add_to_history(
             full_history,

mdagent/tools/base_tools/analysis_tools/rmsd_tools.py

@@ -4,6 +4,7 @@
 import matplotlib.pyplot as plt
 import MDAnalysis as mda
 import numpy as np
+import streamlit as st
 from langchain.tools import BaseTool
 from MDAnalysis.analysis import align, diffusionmap, rms
 from pydantic import BaseModel, Field
@@ -44,15 +45,27 @@ def calculate_rmsd(
         if rmsd_type == "rmsd":
             if self.ref_file:
                 print("Calculating 1-D RMSD between two sets of coordinates...")
+                st.markdown(
+                    "Calculating 1-D RMSD between two sets of coordinates...",
+                    unsafe_allow_html=True,
+                )
                 return self.compute_rmsd_2sets(selection=selection)
             else:
                 print("Calculating time-dependent RMSD...")
+                st.markdown(
+                    "Calculating time-dependent RMSD...", unsafe_allow_html=True
+                )
                 return self.compute_rmsd(selection=selection, plot=plot)
         elif rmsd_type == "pairwise_rmsd":
             print("Calculating pairwise RMSD...")
+            st.markdown("Calculating pairwise RMSD...", unsafe_allow_html=True)
             return self.compute_2d_rmsd(selection=selection, plot_heatmap=plot)
         elif rmsd_type == "rmsf":
             print("Calculating root mean square fluctuation (RMSF)...")
+            st.markdown(
+                "Calculating root mean square fluctuation (RMSF)...",
+                unsafe_allow_html=True,
+            )
             return self.compute_rmsf(selection=selection, plot=plot)
         else:
             raise ValueError(

mdagent/tools/base_tools/preprocess_tools/elements.py

@@ -0,0 +1,224 @@
+list_of_elements = [
+    " H",
+    "He",
+    "Li",
+    "Be",
+    " B",
+    " C",
+    " N",
+    " O",
+    " F",
+    "Ne",
+    "Na",
+    "Mg",
+    "Al",
+    "Si",
+    " P",
+    " S",
+    "Cl",
+    "Ar",
+    " K",
+    "Ca",
+    "Sc",
+    "Ti",
+    " V",
+    "Cr",
+    "Mn",
+    "Fe",
+    "Co",
+    "Ni",
+    "Cu",
+    "Zn",
+    "Ga",
+    "Ge",
+    "As",
+    "Se",
+    "Br",
+    "Kr",
+    "Rb",
+    "Sr",
+    " Y",
+    "Zr",
+    "Nb",
+    "Mo",
+    "Tc",
+    "Ru",
+    "Rh",
+    "Pd",
+    "Ag",
+    "Cd",
+    "In",
+    "Sn",
+    "Sb",
+    "Te",
+    " I",
+    "Xe",
+    "Cs",
+    "Ba",
+    "La",
+    "Ce",
+    "Pr",
+    "Nd",
+    "Pm",
+    "Sm",
+    "Eu",
+    "Gd",
+    "Tb",
+    "Dy",
+    "Ho",
+    "Er",
+    "Tm",
+    "Yb",
+    "Lu",
+    "Hf",
+    "Ta",
+    " W",
+    "Re",
+    "Os",
+    "Ir",
+    "Pt",
+    "Au",
+    "Hg",
+    "Tl",
+    "Pb",
+    "Bi",
+    "Po",
+    "At",
+    "Rn",
+    "Fr",
+    "Ra",
+    "Ac",
+    "Th",
+    "Pa",
+    " U",
+    "Np",
+    "Pu",
+    "Am",
+    "Cm",
+    "Bk",
+    "Cf",
+    "Es",
+    "Fm",
+    "Md",
+    "No",
+    "Lr",
+    "Rf",
+    "Db",
+    "Sg",
+    "Bh",
+    "Hs",
+    "Mt",
+    "Ds",
+    "Rg",
+    "Cn",
+    "Nh",
+    "Fl",
+    "Mc",
+    "Lv",
+    "Ts",
+    "Og",
+    "HE",
+    "LI",
+    "BE",
+    "NE",
+    "NA",
+    "MG",
+    "AL",
+    "SI",
+    "CL",
+    "AR",
+    "CA",
+    "SC",
+    "TI",
+    "CR",
+    "MN",
+    "FE",
+    "CO",
+    "NI",
+    "CU",
+    "ZN",
+    "GA",
+    "GE",
+    "AS",
+    "SE",
+    "BR",
+    "KR",
+    "RB",
+    "SR",
+    " Y",
+    "ZR",
+    "NB",
+    "MO",
+    "TC",
+    "RU",
+    "RH",
+    "PD",
+    "AG",
+    "CD",
+    "IN",
+    "SN",
+    "SB",
+    "TE",
+    "XE",
+    "CS",
+    "BA",
+    "LA",
+    "CE",
+    "PR",
+    "ND",
+    "PM",
+    "SM",
+    "EU",
+    "GD",
+    "TB",
+    "DY",
+    "HO",
+    "ER",
+    "TM",
+    "YB",
+    "LU",
+    "HF",
+    "TA",
+    "RE",
+    "OS",
+    "IR",
+    "PT",
+    "AU",
+    "HG",
+    "TL",
+    "PB",
+    "BI",
+    "PO",
+    "AT",
+    "RN",
+    "FR",
+    "RA",
+    "AC",
+    "TH",
+    "PA",
+    "NP",
+    "PU",
+    "AM",
+    "CM",
+    "BK",
+    "CF",
+    "ES",
+    "FM",
+    "MD",
+    "NO",
+    "LR",
+    "RF",
+    "DB",
+    "SG",
+    "BH",
+    "HS",
+    "MT",
+    "DS",
+    "RG",
+    "CN",
+    "NH",
+    "FL",
+    "MC",
+    "LV",
+    "TS",
+]