support gemini embeddings

Author: alckasoc

README.md

@@ -307,6 +307,7 @@ agent = AgentS2(
   action_space="pyautogui",
   observation_type="screenshot",
   search_engine="Perplexica"  # Assuming you have set up Perplexica.
+  embedding_engine_type="openai"  # Supports "gemini", "openai"
 )
 ```
 

gui_agents/s2/agents/agent_s.py

@@ -9,6 +9,7 @@
 from gui_agents.s2.agents.manager import Manager
 from gui_agents.s2.utils.common_utils import Node
 from gui_agents.utils import download_kb_data
+from gui_agents.s2.core.engine import OpenAIEmbeddingEngine, GeminiEmbeddingEngine
 
 logger = logging.getLogger("desktopenv.agent")
 
@@ -93,6 +94,7 @@ def __init__(
         memory_root_path: str = os.getcwd(),
         memory_folder_name: str = "kb_s2",
         kb_release_tag: str = "v0.2.2",
+        embedding_engine_type: str = "openai",
     ):
         """Initialize AgentS2
 
@@ -106,6 +108,7 @@ def __init__(
             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".
+            embedding_engine_type: Embedding engine to use for knowledge base. Defaults to "openai". Supports "openai" and "gemini".
         """
         super().__init__(
             engine_params,
@@ -147,23 +150,30 @@ def __init__(
                 "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()
+        elif embedding_engine_type == "gemini":
+            self.embedding_engine = GeminiEmbeddingEngine()
+
         self.reset()
 
     def reset(self) -> None:
         """Reset agent state and initialize components"""
         # Initialize core components
         self.planner = Manager(
-            self.engine_params,
-            self.grounding_agent,
-            platform=self.platform,
-            search_engine=self.engine,
+            engine_params=self.engine_params,
+            grounding_agent=self.grounding_agent,
             local_kb_path=self.local_kb_path,
+            embedding_engine=self.embedding_engine,
+            search_engine=self.engine,
+            platform=self.platform,
         )
         self.executor = Worker(
-            self.engine_params,
-            self.grounding_agent,
-            platform=self.platform,
+            engine_params=self.engine_params,
+            grounding_agent=self.grounding_agent,
             local_kb_path=self.local_kb_path,
+            embedding_engine=self.embedding_engine,
+            platform=self.platform,
         )
 
         # Reset state variables

gui_agents/s2/agents/manager.py

@@ -8,6 +8,7 @@
 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 (
     Dag,
     Node,
@@ -27,6 +28,7 @@ def __init__(
         engine_params: Dict,
         grounding_agent: ACI,
         local_kb_path: str,
+        embedding_engine=OpenAIEmbeddingEngine(),
         search_engine: Optional[str] = None,
         multi_round: bool = False,
         platform: str = platform.system().lower(),
@@ -55,10 +57,12 @@ def __init__(
 
         self.local_kb_path = local_kb_path
 
+        self.embedding_engine = embedding_engine
         self.knowledge_base = KnowledgeBase(
-            self.local_kb_path,
-            platform,
-            engine_params,
+            embedding_engine=self.embedding_engine,
+            local_kb_path=self.local_kb_path,
+            platform=platform,
+            engine_params=engine_params,
         )
 
         self.planner_history = []

gui_agents/s2/agents/worker.py

@@ -8,6 +8,7 @@
 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,
@@ -26,8 +27,8 @@ def __init__(
         engine_params: Dict,
         grounding_agent: ACI,
         local_kb_path: str,
+        embedding_engine=OpenAIEmbeddingEngine(),
         platform: str = platform.system().lower(),
-        search_engine: str = "perplexica",
         enable_reflection: bool = True,
         use_subtask_experience: bool = True,
     ):
@@ -42,8 +43,6 @@ def __init__(
                 Path to knowledge base
             platform: str
                 OS platform the agent runs on (darwin, linux, windows)
-            search_engine: str
-                The search engine to use
             enable_reflection: bool
                 Whether to enable reflection
             use_subtask_experience: bool
@@ -53,7 +52,7 @@ def __init__(
 
         self.grounding_agent = grounding_agent
         self.local_kb_path = local_kb_path
-        self.search_engine = search_engine
+        self.embedding_engine = embedding_engine
         self.enable_reflection = enable_reflection
         self.use_subtask_experience = use_subtask_experience
         self.reset()
@@ -74,6 +73,7 @@ def reset(self):
         )
 
         self.knowledge_base = KnowledgeBase(
+            embedding_engine=self.embedding_engine,
             local_kb_path=self.local_kb_path,
             platform=self.platform,
             engine_params=self.engine_params,

gui_agents/s2/cli_app.py

@@ -206,6 +206,13 @@ def main():
         help="The API key of the grounding model.",
     )
 
+    parser.add_argument(
+        "--embedding_engine_type",
+        type=str,
+        default="openai",
+        help="Specify the embedding engine type (supports openai, gemini)",
+    )
+
     args = parser.parse_args()
     assert (
         args.grounding_model_provider and args.grounding_model
@@ -257,6 +264,7 @@ def main():
         action_space="pyautogui",
         observation_type="mixed",
         search_engine=None,
+        embedding_engine_type=args.embedding_engine_type,
     )
 
     while True:

gui_agents/s2/core/knowledge.py

@@ -7,7 +7,6 @@
 
 from gui_agents.s2.core.module import BaseModule
 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 (
     call_llm_safe,
     load_embeddings,
@@ -20,6 +19,7 @@
 class KnowledgeBase(BaseModule):
     def __init__(
         self,
+        embedding_engine,
         local_kb_path: str,
         platform: str,
         engine_params: Dict,
@@ -30,8 +30,7 @@ def __init__(
         self.local_kb_path = local_kb_path
 
         # initialize embedding engine
-        # TODO: Support other embedding engines
-        self.embedding_engine = OpenAIEmbeddingEngine()
+        self.embedding_engine = embedding_engine
 
         # Initialize paths for different memory types
         self.episodic_memory_path = os.path.join(

osworld_setup/s2/run.py

@@ -219,6 +219,7 @@ def test(args: argparse.Namespace, test_all_meta: dict) -> None:
         memory_root_path=os.getcwd(),
         memory_folder_name=args.kb_name,
         kb_release_tag="v0.2.2",
+        embedding_engine_type="openai",
     )
 
     env = DesktopEnv(