Merge pull request #119 from simular-ai/vincent/kb_fix

Bug fixes

gui_agents/s2/agents/agent_s.py

@@ -96,6 +96,7 @@ class AgentS2(UIAgent):
         observation_type: str = "mixed",
         search_engine: Optional[str] = None,
         memory_root_path: str = os.getcwd(),
+        use_default_kb: bool = False,
         memory_folder_name: str = "kb_s2",
         kb_release_tag: str = "v0.2.2",
         embedding_engine_type: str = "openai",
@@ -110,6 +111,7 @@ class AgentS2(UIAgent):
             action_space: Type of action space to use (pyautogui, other)
             observation_type: Type of observations to use (a11y_tree, screenshot, mixed)
             search_engine: Search engine to use (LLM, perplexica)
+            use_default_kb: True to use the default OpenAI kb.
             memory_root_path: Path to memory directory. Defaults to current working directory.
             memory_folder_name: Name of memory folder. Defaults to "kb_s2".
             kb_release_tag: Release tag for knowledge base. Defaults to "v0.2.2".
@@ -130,31 +132,32 @@ class AgentS2(UIAgent):
         self.kb_release_tag = kb_release_tag
 
         # Initialize agent's knowledge base on user's current working directory.
-        print("Downloading knowledge base initial Agent-S knowledge...")
         self.local_kb_path = os.path.join(
             self.memory_root_path, self.memory_folder_name
         )
 
-        if not os.path.exists(os.path.join(self.local_kb_path, self.platform)):
-            download_kb_data(
-                version="s2",
-                release_tag=kb_release_tag,
-                download_dir=self.local_kb_path,
-                platform=self.platform,
-            )
-            print(
-                f"Successfully completed download of knowledge base for version s2, tag {self.kb_release_tag}, platform {self.platform}."
-            )
-        else:
-            print(
-                f"Path local_kb_path {self.local_kb_path} already exists. Skipping download."
-            )
-            print(
-                f"If you'd like to re-download the initial knowledge base, please delete the existing knowledge base at {self.local_kb_path}."
-            )
-            print(
-                "Note, the knowledge is continually updated during inference. Deleting the knowledge base will wipe out all experience gained since the last knowledge base download."
-            )
+        if use_default_kb:
+            if not os.path.exists(os.path.join(self.local_kb_path, self.platform)):
+                print("Downloading Agent S2's default knowledge base...")
+                download_kb_data(
+                    version="s2",
+                    release_tag=kb_release_tag,
+                    download_dir=self.local_kb_path,
+                    platform=self.platform,
+                )
+                print(
+                    f"Successfully completed download of knowledge base for version s2, tag {self.kb_release_tag}, platform {self.platform}."
+                )
+            else:
+                print(
+                    f"Path local_kb_path {self.local_kb_path} already exists. Skipping download."
+                )
+                print(
+                    f"If you'd like to re-download the initial knowledge base, please delete the existing knowledge base at {self.local_kb_path}."
+                )
+                print(
+                    "Note, the knowledge is continually updated during inference. Deleting the knowledge base will wipe out all experience gained since the last knowledge base download."
+                )
 
         if embedding_engine_type == "openai":
             self.embedding_engine = OpenAIEmbeddingEngine(**embedding_engine_params)

gui_agents/s2/agents/manager.py

@@ -28,7 +28,7 @@ class Manager(BaseModule):
         engine_params: Dict,
         grounding_agent: ACI,
         local_kb_path: str,
-        embedding_engine=OpenAIEmbeddingEngine(),
+        embedding_engine,
         search_engine: Optional[str] = None,
         multi_round: bool = False,
         platform: str = platform.system().lower(),

gui_agents/s2/agents/worker.py

@@ -8,7 +8,6 @@ from gui_agents.s2.agents.grounding import ACI
 from gui_agents.s2.core.module import BaseModule
 from gui_agents.s2.core.knowledge import KnowledgeBase
 from gui_agents.s2.memory.procedural_memory import PROCEDURAL_MEMORY
-from gui_agents.s2.core.engine import OpenAIEmbeddingEngine
 from gui_agents.s2.utils.common_utils import (
     Node,
     calculate_tokens,
@@ -27,7 +26,7 @@ class Worker(BaseModule):
         engine_params: Dict,
         grounding_agent: ACI,
         local_kb_path: str,
-        embedding_engine=OpenAIEmbeddingEngine(),
+        embedding_engine,
         platform: str = platform.system().lower(),
         enable_reflection: bool = True,
         use_subtask_experience: bool = True,

gui_agents/s2/core/knowledge.py

@@ -278,10 +278,9 @@ class KnowledgeBase(BaseModule):
             subtask_summarization = self.summarize_episode(subtask_traj)
             kb[subtask_key] = subtask_summarization
 
-            if self.save_knowledge:
-                os.makedirs(os.path.dirname(self.episodic_memory_path), exist_ok=True)
-                with open(self.episodic_memory_path, "w") as fout:
-                    json.dump(kb, fout, indent=2)
+            os.makedirs(os.path.dirname(self.episodic_memory_path), exist_ok=True)
+            with open(self.episodic_memory_path, "w") as fout:
+                json.dump(kb, fout, indent=2)
 
         return kb.get(subtask_key)
 
@@ -304,10 +303,9 @@ class KnowledgeBase(BaseModule):
             task_summarization = self.summarize_narrative(task_traj)
             kb[task_key] = task_summarization
 
-            if self.save_knowledge:
-                os.makedirs(os.path.dirname(self.narrative_memory_path), exist_ok=True)
-                with open(self.narrative_memory_path, "w") as fout:
-                    json.dump(kb, fout, indent=2)
+            os.makedirs(os.path.dirname(self.narrative_memory_path), exist_ok=True)
+            with open(self.narrative_memory_path, "w") as fout:
+                json.dump(kb, fout, indent=2)
 
         return kb.get(task_key)