From 0863145fbf46fe031562c752bba9d3d65d664c19 Mon Sep 17 00:00:00 2001
From: Jerry Liu <jerryjliu98@gmail.com>
Date: Sat, 2 Dec 2023 19:46:44 -0800
Subject: [PATCH] cr

---
 core/agent_builder.py |  5 +++-
 core/constants.py     |  4 +--
 st_utils.py           | 69 +++++++++++++++++++++++++------------------
 3 files changed, 47 insertions(+), 31 deletions(-)

diff --git a/core/agent_builder.py b/core/agent_builder.py
index 95f02dc..e44a385 100644
--- a/core/agent_builder.py
+++ b/core/agent_builder.py
@@ -91,6 +91,9 @@ def delete_agent_cache(self, agent_id: str) -> None:
             # recursive delete
             shutil.rmtree(full_path)
 
+        print("TMP DIR: ", self._dir)
+        print(f"Deleted agent {agent_id}.")
+
 
 # System prompt tool
 GEN_SYS_PROMPT_STR = """\
@@ -270,7 +273,7 @@ def update_agent(
         NOTE: Currently is manually called, not meant for agent use.
 
         """
-        self._agent_registry.delete_agent_cache(agent_id)
+        self._agent_registry.delete_agent_cache(self.cache.agent_id)
 
         # set agent id
         self.cache.agent_id = agent_id
diff --git a/core/constants.py b/core/constants.py
index a6d4599..9cf91bd 100644
--- a/core/constants.py
+++ b/core/constants.py
@@ -1,4 +1,4 @@
 from pathlib import Path
 
-AGENT_CACHE_DIR = Path(__file__).parent / "cache" / "agents"
-MESSAGES_CACHE_DIR = Path(__file__).parent / "cache" / "messages"
+AGENT_CACHE_DIR = Path(__file__).parent.parent / "cache" / "agents"
+MESSAGES_CACHE_DIR = Path(__file__).parent.parent / "cache" / "messages"
diff --git a/st_utils.py b/st_utils.py
index f71a0bb..6d553e4 100644
--- a/st_utils.py
+++ b/st_utils.py
@@ -21,13 +21,13 @@ def update_selected_agent_with_id(selected_id: Optional[str] = None) -> None:
     st.session_state.selected_id = (
         selected_id if selected_id != "Create a new agent" else None
     )
-    if st.session_state.selected_id is None:
-        st.session_state.selected_cache = None
-    else:
-        # load agent from directory
-        agent_registry = cast(AgentCacheRegistry, st.session_state.agent_registry)
-        agent_cache = agent_registry.get_agent_cache(st.session_state.selected_id)
-        st.session_state.selected_cache = agent_cache
+
+    # clear agent builder and builder agent
+    st.session_state.builder_agent = None
+    st.session_state.agent_builder = None
+
+    # clear selected cache
+    st.session_state.selected_cache = None
 
 
 ## handler for sidebar specifically
@@ -42,6 +42,9 @@ def add_sidebar() -> None:
     """Add sidebar."""
     with st.sidebar:
         agent_registry = cast(AgentCacheRegistry, st.session_state.agent_registry)
+        print("AGENT IDS")
+        print(agent_registry._dir)
+        print(agent_registry.get_agent_ids())
         st.session_state.cur_agent_ids = agent_registry.get_agent_ids()
         choices = ["Create a new agent"] + st.session_state.cur_agent_ids
 
@@ -97,35 +100,45 @@ def get_current_state() -> CurrentSessionState:
     if "selected_id" not in st.session_state.keys():
         st.session_state.selected_id = None
 
-    if "selected_cache" not in st.session_state.keys():
-        st.session_state.selected_cache = None
+    # set selected cache if doesn't exist
+    if "selected_cache" not in st.session_state.keys() or st.session_state.selected_cache is None:
+        # update selected cache
+        if st.session_state.selected_id is None:
+            st.session_state.selected_cache = None
+        else:
+            # load agent from directory
+            agent_registry = cast(AgentCacheRegistry, st.session_state.agent_registry)
+            agent_cache = agent_registry.get_agent_cache(st.session_state.selected_id)
+            st.session_state.selected_cache = agent_cache
 
     # set builder agent / agent builder
-    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,
-            agent_registry=st.session_state.agent_registry,
-        )
-    else:
-        # create builder agent / tools from new cache
-        builder_agent, agent_builder = load_meta_agent_and_tools(
-            agent_registry=st.session_state.agent_registry
-        )
+    if ("builder_agent" not in st.session_state.keys() 
+        or st.session_state.builder_agent is None
+        or "agent_builder" not in st.session_state.keys()
+        or st.session_state.agent_builder is None):
+        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,
+                agent_registry=st.session_state.agent_registry,
+            )
+        else:
+            # create builder agent / tools from new cache
+            builder_agent, agent_builder = load_meta_agent_and_tools(
+                agent_registry=st.session_state.agent_registry
+            )
 
-    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
 
     return CurrentSessionState(
         agent_registry=st.session_state.agent_registry,
         selected_id=st.session_state.selected_id,
         selected_cache=st.session_state.selected_cache,
-        agent_builder=agent_builder,
-        cache=agent_builder.cache,
-        builder_agent=builder_agent,
+        agent_builder=st.session_state.agent_builder,
+        cache=st.session_state.agent_builder.cache,
+        builder_agent=st.session_state.builder_agent,
     )