s1 changes update kb paths

gui_agents/s1/core/AgentS.py

@@ -2,6 +2,7 @@ import json
 import logging
 import os
 from typing import Dict, List, Optional, Tuple
+import platform
 
 from gui_agents.s1.aci.ACI import ACI
 from gui_agents.s1.core.Manager import Manager
@@ -19,7 +20,7 @@ class UIAgent:
         self,
         engine_params: Dict,
         grounding_agent: ACI,
-        platform: str = "macos",
+        platform: str = platform.system().lower(),
         action_space: str = "pyautogui",
         observation_type: str = "a11y_tree",
         search_engine: str = "perplexica",
@@ -85,7 +86,7 @@ class GraphSearchAgent(UIAgent):
         self,
         engine_params: Dict,
         grounding_agent: ACI,
-        platform: str = "macos",
+        platform: str = platform.system().lower(),
         action_space: str = "pyatuogui",
         observation_type: str = "mixed",
         search_engine: Optional[str] = None,
@@ -146,7 +147,6 @@ class GraphSearchAgent(UIAgent):
                 "Note, the knowledge is continually updated during inference. Deleting the knowledge base will wipe out all experience gained since the last knowledge base download."
             )
 
-
         self.reset()
 
     def reset(self) -> None:
@@ -160,8 +160,8 @@ class GraphSearchAgent(UIAgent):
             local_kb_path=self.local_kb_path,
         )
         self.executor = Worker(
-            self.engine_params, 
-            self.grounding_agent, 
+            self.engine_params,
+            self.grounding_agent,
             platform=self.platform,
             local_kb_path=self.local_kb_path,
         )
@@ -301,7 +301,7 @@ class GraphSearchAgent(UIAgent):
         """
         try:
             reflection_path = os.path.join(
-                working_dir, "../kb", self.platform, "narrative_memory.json"
+                self.local_kb_path, self.platform, "narrative_memory.json"
             )
             try:
                 reflections = json.load(open(reflection_path))
@@ -341,7 +341,7 @@ class GraphSearchAgent(UIAgent):
                 )[0]
                 try:
                     subtask_path = os.path.join(
-                        working_dir, "../kb", self.platform, "episodic_memory.json"
+                        self.local_kb_path, self.platform, "episodic_memory.json"
                     )
                     kb = json.load(open(subtask_path))
                 except:

gui_agents/s1/core/Evaluator.py

@@ -1,196 +0,0 @@
-import json
-import logging
-import os
-import pickle
-import re
-import time
-from collections import defaultdict
-from typing import Dict, List, Tuple, Union
-
-import numpy as np
-from desktop_env.desktop_env import DesktopEnv
-from pydantic import BaseModel
-from sklearn.metrics.pairwise import cosine_similarity
-
-from gui_agents import osworld_utils
-from gui_agents.MultimodalAgent import LMMAgent
-from gui_agents.MultimodalEngine import OpenAIEmbeddingEngine
-from gui_agents.osworld.GroundingAgent import GroundingAgent
-from gui_agents.osworld_utils import Dag, Node
-from gui_agents.ProceduralMemory import PROCEDURAL_MEMORY
-from gui_agents.query_perplexica import query_to_perplexica
-
-
-class Evaluator:
-    def __init__(self, instruction, engine_params: Dict, script_check: bool = False):
-        self.instruction = instruction
-        self.engine_params = engine_params
-        self.reset()
-
-    def reset(self):
-        self.state_evaluator = LMMAgent(self.engine_params)
-        self.state_evaluator_system_prompt = (
-            PROCEDURAL_MEMORY.STATE_EVALUATOR_SYSTEM_PROMPT
-        )
-        self.state_evaluator.add_system_prompt(
-            self.state_evaluator_system_prompt.replace(
-                "TASK_DESCRIPTION", self.instruction
-            )
-        )
-
-        self.obs_evaluator = LMMAgent(self.engine_params)
-        self.obs_evaluator_system_prompt = PROCEDURAL_MEMORY.OBS_EVALUATOR_SYSTEM_PROMPT
-        self.obs_evaluator.add_system_prompt(
-            self.obs_evaluator_system_prompt.replace(
-                "TASK_DESCRIPTION", self.instruction
-            )
-        )
-
-    def state_evaluate(self, input, input_img, plan_codes, env: DesktopEnv = None):
-        init_obs, last_obs = input[0], input[-1]
-        init_obs_img, last_obs_img = input_img[0], input_img[-1]
-
-        input_message_1 = f"""
-                        The accessibility tree at the first step:{init_obs}, and the screenshot at the first step: \n
-                        """
-        input_message_2 = f"""
-                        The accessibility tree at the last step:{last_obs}, and the screenshot at the last step: \n
-                        """
-        input_message_3 = (
-            "\nThe whole actions performed by the digital agent:\n"
-            + "\n".join(plan_codes)
-        )
-        self.state_evaluator.add_message(
-            text_content=input_message_1, image_content=init_obs_img, role="user"
-        )
-        self.state_evaluator.add_message(
-            text_content=input_message_2, image_content=last_obs_img, role="user"
-        )
-        self.state_evaluator.add_message(text_content=input_message_3, role="user")
-        script_response = self.state_evaluator.get_response()
-        print(
-            f"The evaluation result of current task: {self.instruction}:\n {script_response}"
-        )
-        self.state_evaluator.add_message(script_response)
-        try:
-            script_response = script_response.split("Judgment:")[1]
-        except:
-            script_response = script_response
-
-        if "Yes" in script_response:
-            eval_result = 1.0
-            with open("script_result.txt", "a", encoding="utf-8") as f:
-                f.write(
-                    f"The task {self.instruction} generate the eval script: \n{script_response}\n"
-                )
-            return eval_result
-        elif "No" in script_response:
-            eval_result = 0.0
-            with open("script_result.txt", "a", encoding="utf-8") as f:
-                f.write(
-                    f"The task {self.instruction} generate the eval script: \n{script_response}\n"
-                )
-            return eval_result
-        else:
-            script = osworld_utils.parse_single_code_from_string(script_response)
-            script_run_output = env.controller.execute_python_command(script)
-            matching_message = f"""
-                                The output after executing the script is: {script_run_output}, now please do the subsequent task to judge the completeness of the task based on the script's result and the task information(Like accessibility trees, screenshots, whole actions).
-                                The Script and the task regarding information is also aforementioned. Provide your analysis and put the judgment at the end of the response in this format: Judgment: Yes/No
-                                """
-            print(matching_message)
-            self.state_evaluator.add_message(matching_message)
-            matching_response = self.state_evaluator.get_response()
-            print(
-                f"The matching result of current task{self.instruction}:\n {matching_response}"
-            )
-            with open("script_result.txt", "a", encoding="utf-8") as f:
-                f.write(
-                    f"The task: {self.instruction} generate the eval script:\n {script_response} \n, and the script_run_output is {script_run_output} \n, the matching response is {matching_response} \n\n"
-                )
-            try:
-                if "Yes" in matching_response.split("Judgment:")[1]:
-                    eval_result = 1.0
-                else:
-                    eval_result = 0.0
-            except:
-                if "Yes" in matching_response:
-                    eval_result = 1.0
-                else:
-                    eval_result = 0.0
-            return eval_result
-
-    def obs_evaluate(
-        self,
-        instruction: str,
-        input: List[str],
-        input_img: List[bytes],
-        plan_codes: List[str],
-    ):
-
-        self.obs_evaluator.add_system_prompt(
-            self.obs_evaluator_system_prompt.replace("TASK_DESCRIPTION", instruction)
-        )
-
-        init_obs, last_obs = input[0], input[-1]
-        init_obs_img, last_obs_img = input_img[0], input_img[-1]
-
-        input_message_1 = f"""
-                        The accessibility tree at the first step:{init_obs}, and the screenshot at the first step: \n
-                        """
-        input_message_2 = f"""
-                        The accessibility tree at the last step:{last_obs}, and the screenshot at the last step: \n
-                        """
-        input_message_3 = (
-            "\nThe whole actions performed by the digital agent:\n"
-            + "\n".join(plan_codes)
-        )
-
-        input_message = input_message_1 + input_message_2 + input_message_3
-
-        self.obs_evaluator.add_message(
-            input_message, image_content=[init_obs_img, last_obs_img]
-        )
-        response = call_llm_safe(self.obs_evaluator)
-        logger.info(
-            f"The evaluation result of current subtask: {instruction}:\n {response}"
-        )
-
-        # TODO: Expand coverage
-        def check_judgment(response):
-            # Improved regex pattern to match "Judgment: Yes" or "Judgment: No" at the end of the response, allowing extra spaces or newlines
-            pattern = r"Judgment:\s*(yes|no)\s*$"
-
-            # Search for the pattern in the response, case insensitive
-            match = re.search(pattern, response.strip(), re.IGNORECASE | re.MULTILINE)
-            eval_result = 0
-            if match:
-                # Normalize the judgment (capitalize the first letter)
-                judgment = match.group(1).capitalize()
-                # Set eval_result based on the judgment
-                eval_result = 1 if judgment == "Yes" else 0
-
-            return eval_result
-
-        try:
-            eval_result = check_judgment(response)
-        except:
-            logger.error(
-                "Failed to extract judgment from the response. Defaulting to 0."
-            )
-            eval_result = 0
-
-        input_tokens, output_tokens = calculate_tokens(self.obs_evaluator.messages)
-
-        # Set Cost based on GPT-4o
-        cost = input_tokens * (0.0050 / 1000) + output_tokens * (0.0150 / 1000)
-        logger.info("EVALUATION COST: %s", cost)
-
-        evaluator_info = {
-            "obs_evaluator_response": response,
-            "num_input_tokens_evaluator": input_tokens,
-            "num_output_tokens_evaluator": output_tokens,
-            "evaluator_cost": cost,
-        }
-
-        return evaluator_info, eval_result, response

gui_agents/s1/core/Knowledge.py

@@ -15,15 +15,19 @@ from gui_agents.s1.utils.common_utils import (
 )
 from gui_agents.s1.utils.query_perplexica import query_to_perplexica
 
-working_dir = os.path.dirname(os.path.abspath(__file__))
-
 
 class KnowledgeBase(BaseModule):
     def __init__(
-        self, platform: str, engine_params: Dict, use_image_for_search: bool = False
+        self,
+        local_kb_path: str,
+        platform: str,
+        engine_params: Dict,
+        use_image_for_search: bool = False,
     ):
         super().__init__(engine_params, platform)
 
+        self.local_kb_path = local_kb_path
+
         # initialize embedding engine
         # TODO: Support other embedding engines
         self.embedding_engine = OpenAIEmbeddingEngine(
@@ -34,6 +38,17 @@ class KnowledgeBase(BaseModule):
             )
         )
 
+        # Initialize paths for different memory types
+        self.episodic_memory_path = os.path.join(
+            self.local_kb_path, self.platform, "episodic_memory.json"
+        )
+        self.narrative_memory_path = os.path.join(
+            self.local_kb_path, self.platform, "narrative_memory.json"
+        )
+        self.embeddings_path = os.path.join(
+            self.local_kb_path, self.platform, "embeddings.pkl"
+        )
+
         self.rag_module_system_prompt = PROCEDURAL_MEMORY.RAG_AGENT.replace(
             "CURRENT_OS", self.platform
         )
@@ -62,7 +77,7 @@ class KnowledgeBase(BaseModule):
     def formulate_query(self, instruction: str, observation: Dict) -> str:
         """Formulate search query based on instruction and current state"""
         query_path = os.path.join(
-            working_dir, "../kb", self.platform, "formulate_query.json"
+            self.local_kb_path, self.platform, "formulate_query.json"
         )
         try:
             with open(query_path, "r") as f:
@@ -102,7 +117,7 @@ class KnowledgeBase(BaseModule):
 
         # Default to perplexica rag knowledge to see if the query exists
         file = os.path.join(
-            working_dir, "../kb", self.platform, f"{search_engine}_rag_knowledge.json"
+            self.local_kb_path, self.platform, f"{search_engine}_rag_knowledge.json"
         )
 
         try:
@@ -126,8 +141,7 @@ class KnowledgeBase(BaseModule):
         exist_search_results[instruction] = search_results.strip()
         with open(
             os.path.join(
-                working_dir,
-                "../kb",
+                self.local_kb_path,
                 self.platform,
                 f"{search_engine}_rag_knowledge.json",
             ),
@@ -139,18 +153,11 @@ class KnowledgeBase(BaseModule):
 
     def retrieve_narrative_experience(self, instruction: str) -> Tuple[str, str]:
         """Retrieve narrative experience using embeddings"""
-        kb_path = os.path.join(
-            working_dir, "../kb", self.platform, "narrative_memory.json"
-        )
-        embeddings_path = os.path.join(
-            working_dir, "../kb", self.platform, "embeddings.pkl"
-        )
-
-        knowledge_base = load_knowledge_base(kb_path)
+        knowledge_base = load_knowledge_base(self.narrative_memory_path)
         if not knowledge_base:
             return "None", "None"
 
-        embeddings = load_embeddings(embeddings_path)
+        embeddings = load_embeddings(self.embeddings_path)
 
         # Get or create instruction embedding
         instruction_embedding = embeddings.get(instruction)
@@ -169,7 +176,7 @@ class KnowledgeBase(BaseModule):
 
             candidate_embeddings.append(candidate_embedding)
 
-        save_embeddings(embeddings_path, embeddings)
+        save_embeddings(self.embeddings_path, embeddings)
 
         similarities = cosine_similarity(
             instruction_embedding, np.vstack(candidate_embeddings)
@@ -182,18 +189,11 @@ class KnowledgeBase(BaseModule):
 
     def retrieve_episodic_experience(self, instruction: str) -> Tuple[str, str]:
         """Retrieve similar task experience using embeddings"""
-        kb_path = os.path.join(
-            working_dir, "../kb", self.platform, "episodic_memory.json"
-        )
-        embeddings_path = os.path.join(
-            working_dir, "../kb", self.platform, "embeddings.pkl"
-        )
-
-        knowledge_base = load_knowledge_base(kb_path)
+        knowledge_base = load_knowledge_base(self.episodic_memory_path)
         if not knowledge_base:
             return "None", "None"
 
-        embeddings = load_embeddings(embeddings_path)
+        embeddings = load_embeddings(self.embeddings_path)
 
         # Get or create instruction embedding
         instruction_embedding = embeddings.get(instruction)
@@ -212,7 +212,7 @@ class KnowledgeBase(BaseModule):
 
             candidate_embeddings.append(candidate_embedding)
 
-        save_embeddings(embeddings_path, embeddings)
+        save_embeddings(self.embeddings_path, embeddings)
 
         similarities = cosine_similarity(
             instruction_embedding, np.vstack(candidate_embeddings)

gui_agents/s1/core/Manager.py

@@ -1,7 +1,7 @@
 import logging
-import os
 from collections import defaultdict
 from typing import Dict, List, Optional, Tuple
+import platform
 
 from gui_agents.s1.aci.ACI import ACI
 from gui_agents.s1.core.BaseModule import BaseModule
@@ -17,9 +17,6 @@ from gui_agents.s1.utils.common_utils import (
 
 logger = logging.getLogger("desktopenv.agent")
 
-# Get the directory of the current script
-working_dir = os.path.dirname(os.path.abspath(__file__))
-
 NUM_IMAGE_TOKEN = 1105  # Value set of screen of size 1920x1080 for openai vision
 
 
@@ -28,9 +25,10 @@ class Manager(BaseModule):
         self,
         engine_params: Dict,
         grounding_agent: ACI,
+        local_kb_path: str,
         search_engine: Optional[str] = None,
         multi_round: bool = False,
-        platform: str = "macos",
+        platform: str = platform.system().lower(),
     ):
         # TODO: move the prompt to Procedural Memory
         super().__init__(engine_params, platform)
@@ -49,9 +47,10 @@ class Manager(BaseModule):
         self.episode_summarization_agent = self._create_agent(
             PROCEDURAL_MEMORY.SUBTASK_SUMMARIZATION_PROMPT
         )
-        self.rag_agent = self._create_agent(PROCEDURAL_MEMORY.RAG_AGENT)
 
-        self.knowldge_base = KnowledgeBase(platform, engine_params)
+        self.local_kb_path = local_kb_path
+
+        self.knowledge_base = KnowledgeBase(self.local_kb_path, platform, engine_params)
 
         self.planner_history = []
 
@@ -97,7 +96,7 @@ class Manager(BaseModule):
         # Perform Retrieval only at the first planning step
         if self.turn_count == 0:
 
-            self.search_query = self.knowldge_base.formulate_query(
+            self.search_query = self.knowledge_base.formulate_query(
                 instruction, observation
             )
 
@@ -105,7 +104,7 @@ class Manager(BaseModule):
             integrated_knowledge = ""
             # Retrieve most similar narrative (task) experience
             most_similar_task, retrieved_experience = (
-                self.knowldge_base.retrieve_narrative_experience(instruction)
+                self.knowledge_base.retrieve_narrative_experience(instruction)
             )
             logger.info(
                 "SIMILAR TASK EXPERIENCE: %s",
@@ -114,7 +113,7 @@ class Manager(BaseModule):
 
             # Retrieve knowledge from the web if search_engine is provided
             if self.search_engine is not None:
-                retrieved_knowledge = self.knowldge_base.retrieve_knowledge(
+                retrieved_knowledge = self.knowledge_base.retrieve_knowledge(
                     instruction=instruction,
                     search_query=self.search_query,
                     search_engine=self.search_engine,
@@ -123,7 +122,7 @@ class Manager(BaseModule):
 
                 if retrieved_knowledge is not None:
                     # Fuse the retrieved knowledge and experience
-                    integrated_knowledge = self.knowldge_base.knowledge_fusion(
+                    integrated_knowledge = self.knowledge_base.knowledge_fusion(
                         observation=observation,
                         instruction=instruction,
                         web_knowledge=retrieved_knowledge,

gui_agents/s1/core/Worker.py

@@ -1,19 +1,17 @@
-import ast
 import logging
 import os
 import re
 from typing import Dict, List, Tuple
+import platform
 
 from gui_agents.s1.aci.ACI import ACI
 from gui_agents.s1.core.BaseModule import BaseModule
 from gui_agents.s1.core.Knowledge import KnowledgeBase
 from gui_agents.s1.core.ProceduralMemory import PROCEDURAL_MEMORY
-from gui_agents.s1.mllm.MultimodalEngine import OpenAIEmbeddingEngine
 from gui_agents.s1.utils import common_utils
 from gui_agents.s1.utils.common_utils import Node, calculate_tokens, call_llm_safe
 
 logger = logging.getLogger("desktopenv.agent")
-working_dir = os.path.dirname(os.path.abspath(__file__))
 
 
 class Worker(BaseModule):
@@ -21,7 +19,8 @@ class Worker(BaseModule):
         self,
         engine_params: Dict,
         grounding_agent: ACI,
-        platform: str = "macos",
+        local_kb_path: str,
+        platform: str = platform.system().lower(),
         search_engine: str = "perplexica",
         enable_reflection: bool = True,
         use_subtask_experience: bool = True,
@@ -33,6 +32,8 @@ class Worker(BaseModule):
                 Parameters for the multimodal engine
             grounding_agent: Agent
                 The grounding agent to use
+            local_kb_path: str
+                Path to knowledge base
             search_engine: str
                 The search engine to use
             enable_reflection: bool
@@ -40,10 +41,11 @@ class Worker(BaseModule):
             use_subtask_experience: bool
                 Whether to use subtask experience
         """
+        super().__init__(engine_params, platform)
+
         self.grounding_agent = grounding_agent
-        self.platform = platform
+        self.local_kb_path = local_kb_path
         self.enable_reflection = enable_reflection
-        self.engine_params = engine_params
         self.search_engine = search_engine
         self.use_subtask_experience = use_subtask_experience
         self.reset()
@@ -67,7 +69,9 @@ class Worker(BaseModule):
         )
 
         self.knowledge_base = KnowledgeBase(
-            platform=self.platform, engine_params=self.engine_params
+            local_kb_path=self.local_kb_path,
+            platform=self.platform,
+            engine_params=self.engine_params,
         )
 
         self.turn_count = 0

gui_agents/s2/agents/agent_s.py

@@ -320,7 +320,9 @@ class GraphSearchAgent(UIAgent):
             trajectory: String containing task execution trajectory
         """
         try:
-            reflection_path = os.path.join(self.local_kb_path, self.platform, "narrative_memory.json")
+            reflection_path = os.path.join(
+                self.local_kb_path, self.platform, "narrative_memory.json"
+            )
             try:
                 reflections = json.load(open(reflection_path))
             except: