refactor code (run-llama#46)

1_🏠_Home.py

@@ -1,14 +1,7 @@
 import streamlit as st
 from streamlit_pills import pills
 
-from agent_utils import (
-    load_meta_agent_and_tools,
-    load_agent_ids_from_directory,
-)
-from st_utils import add_sidebar
-from constants import (
-    AGENT_CACHE_DIR,
-)
+from st_utils import add_sidebar, get_current_state
 
 
 ####################
@@ -35,29 +28,11 @@
     st.info("**NOTE**: The ability to add web search is enabled.")
 
 
+current_state = get_current_state()
 add_sidebar()
 
 
-if (
-    "selected_cache" in st.session_state.keys()
-    and st.session_state.selected_cache is not None
-):
-    # create builder agent / tools from selected cache
-    builder_agent, agent_builder = load_meta_agent_and_tools(
-        cache=st.session_state.selected_cache
-    )
-else:
-    # create builder agent / tools from new cache
-    builder_agent, agent_builder = load_meta_agent_and_tools()
-
-
-st.info(f"Currently building/editing agent: {agent_builder.cache.agent_id}", icon="ℹ️")
-
-
-if "builder_agent" not in st.session_state.keys():
-    st.session_state.builder_agent = builder_agent
-if "agent_builder" not in st.session_state.keys():
-    st.session_state.agent_builder = agent_builder
+st.info(f"Currently building/editing agent: {current_state.cache.agent_id}", icon="ℹ️")
 
 # add pills
 selected = pills(
@@ -85,32 +60,75 @@ def add_to_message_history(role: str, content: str) -> None:
     with st.chat_message(message["role"]):
         st.write(message["content"])
 
+
+# def handle_user_input() -> None:
+#     """Handle user input."""
+#     prompt = st.session_state.user_question_st
+#     print(f"USER PROMPT: {prompt}")
+#     add_to_message_history("user", prompt)
+#     with st.chat_message("user"):
+#         st.write(prompt)
+#     # If last message is not from assistant, generate a new response
+#     if st.session_state.messages[-1]["role"] != "assistant":
+#         with st.chat_message("assistant"):
+#             with st.spinner("Thinking..."):
+#                 response = current_state.builder_agent.chat(prompt)
+#                 st.write(str(response))
+#                 add_to_message_history("assistant", str(response))
+
+#     else:
+#         pass
+
+#     # check agent_ids again
+#     # if it doesn't match, add to directory and refresh
+#     agent_ids = current_state.agent_registry.get_agent_ids()
+#     # check diff between agent_ids and cur agent ids
+#     diff_ids = list(set(agent_ids) - set(st.session_state.cur_agent_ids))
+#     if len(diff_ids) > 0:
+#         # # clear streamlit cache, to allow you to generate a new agent
+#         # st.cache_resource.clear()
+#         st.rerun()
+
 # handle user input
+# st.chat_input(
+#     "Your question", key="user_question_st", on_submit=handle_user_input
+# )  # Prompt for user input and save to chat history
+
+# TODO: this is really hacky, only because st.rerun is jank
 if prompt := st.chat_input(
-    "Your question"
+    "Your question",
 ):  # Prompt for user input and save to chat history
-    add_to_message_history("user", prompt)
-    with st.chat_message("user"):
-        st.write(prompt)
-
-# If last message is not from assistant, generate a new response
-if st.session_state.messages[-1]["role"] != "assistant":
-    with st.chat_message("assistant"):
-        with st.spinner("Thinking..."):
-            response = st.session_state.builder_agent.chat(prompt)
-            st.write(str(response))
-            add_to_message_history("assistant", str(response))
-
-            # check agent_ids again, if it doesn't match, add to directory and refresh
-            agent_ids = load_agent_ids_from_directory(str(AGENT_CACHE_DIR))
-            # check diff between agent_ids and cur agent ids
-            diff_ids = list(set(agent_ids) - set(st.session_state.cur_agent_ids))
-            if len(diff_ids) > 0:
-                # clear streamlit cache, to allow you to generate a new agent
-                st.cache_resource.clear()
-
-                # trigger refresh
-                st.rerun()
+    # TODO: hacky
+    if "has_rerun" in st.session_state.keys() and st.session_state.has_rerun:
+        # if this is true, skip the user input
+        st.session_state.has_rerun = False
+    else:
+        add_to_message_history("user", prompt)
+        with st.chat_message("user"):
+            st.write(prompt)
+
+        # If last message is not from assistant, generate a new response
+        if st.session_state.messages[-1]["role"] != "assistant":
+            with st.chat_message("assistant"):
+                with st.spinner("Thinking..."):
+                    response = current_state.builder_agent.chat(prompt)
+                    st.write(str(response))
+                    add_to_message_history("assistant", str(response))
+
+        else:
+            pass
+
+        # check agent_ids again
+        # if it doesn't match, add to directory and refresh
+        agent_ids = current_state.agent_registry.get_agent_ids()
+        # check diff between agent_ids and cur agent ids
+        diff_ids = list(set(agent_ids) - set(st.session_state.cur_agent_ids))
+        if len(diff_ids) > 0:
+            # # clear streamlit cache, to allow you to generate a new agent
+            # st.cache_resource.clear()
+            st.session_state.has_rerun = True
+            st.rerun()
 
 else:
-    pass
+    # TODO: set has_rerun to False
+    st.session_state.has_rerun = False

README.md

@@ -83,7 +83,7 @@ It will be able to pick the right RAG tools (either top-k vector search or optio
 
 ### Builder Agent
 
-By default the builder agent uses OpenAI. This is defined in the `builder_config.py` file.
+By default the builder agent uses OpenAI. This is defined in the `core/builder_config.py` file.
 
 You can customize this to whatever LLM you want (an example is provided for Anthropic).