added streamlit app with ability to upload files

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)

setup.py

@@ -31,6 +31,7 @@
         "requests",
         "rmrkl",
         "tiktoken",
+        "streamlit",
     ],
     test_suite="tests",
     long_description=long_description,

st_app.py

@@ -0,0 +1,80 @@
+import os
+from typing import List
+
+import streamlit as st
+from dotenv import load_dotenv
+from langchain.callbacks import StreamlitCallbackHandler
+from langchain.callbacks.base import BaseCallbackHandler
+
+from mdagent import MDAgent
+
+load_dotenv()
+
+
+st_callback = StreamlitCallbackHandler(st.container())
+
+
+# Streamlit app
+st.title("MDAgent")
+
+# option = st.selectbox("Choose an option:", ["Explore & Learn", "Use Learned Skills"])
+# if option == "Explore & Learn":
+#    explore = True
+# else:
+#    explore = False
+
+resume_op = st.selectbox("Resume:", ["False", "True"])
+if resume_op == "True":
+    resume = True
+else:
+    resume = False
+
+# for now I'm just going to allow  pdb and csv files - we can add more later
+uploaded_files = st.file_uploader(
+    "Upload a .pdb or .csv file", type=["pdb", "csv"], accept_multiple_files=True
+)
+files: List[str] = []
+# write file to disk
+if uploaded_files:
+    for file in uploaded_files:
+        with open(file.name, "wb") as f:
+            f.write(file.getbuffer())
+
+    st.write("Files successfully uploaded!")
+    uploaded_file = [os.path.join(os.getcwd(), file.name) for file in uploaded_files]
+else:
+    uploaded_file = []
+
+mdagent = MDAgent(resume=resume, uploaded_files=uploaded_file)
+
+
+def generate_response(prompt):
+    result = mdagent.run(prompt)
+    return result
+
+
+# make new container to store scratch
+scratch = st.empty()
+scratch.write(
+    """Hi! I am MDAgent, your MD automation assistant.
+              How can I help you today?"""
+)
+
+
+# This allows streaming of llm tokens
+class TokenStreamlitCallbackHandler(BaseCallbackHandler):
+    def __init__(self, container):
+        self.container = container
+
+    def on_llm_new_token(self, token, **kwargs):
+        self.container.write("".join(token))
+
+
+token_st_callback = TokenStreamlitCallbackHandler(scratch)
+
+if prompt := st.chat_input():
+    st.chat_message("user").write(prompt)
+    with st.chat_message("assistant"):
+        st_callback = StreamlitCallbackHandler(st.container())
+        response = mdagent.run(prompt, callbacks=[st_callback, token_st_callback])
+        st.write(response)