Adds a tips example

mephisto/abstractions/_subcomponents/agent_state.py

@@ -179,7 +179,7 @@ def set_init_state(self, data: Any) -> bool:
         return True
 
     @abstractmethod
-    def get_init_state(self, get_all_state: Optional[bool] = False) -> Optional[Any]:
+    def get_init_state(self) -> Optional[Any]:
         """
         Return the initial state for this agent,
         None if no such state exists

mephisto/abstractions/blueprints/abstract/static_task/static_agent_state.py

@@ -35,17 +35,13 @@ def _set_init_state(self, data: Any):
         """Set the initial state for this agent"""
         self.state["inputs"] = data
 
-    def get_init_state(
-        self, get_all_state: Optional[bool] = False
-    ) -> Optional[Dict[str, Any]]:
+    def get_init_state(self) -> Optional[Dict[str, Any]]:
 
         """
         Return the initial state for this agent,
         None if no such state exists
         """
-        if get_all_state == True:
-            return self.state
-        elif self.state["inputs"] is None:
+        if self.state["inputs"] is None:
             return None
         return self.state["inputs"].copy()
 

mephisto/abstractions/blueprints/mock/mock_agent_state.py

@@ -29,9 +29,7 @@ def _set_init_state(self, data: Any):
         """Set the initial state for this agent"""
         self.init_state = data
 
-    def get_init_state(
-        self, get_all_state: Optional[bool] = False
-    ) -> Optional[Dict[str, Any]]:
+    def get_init_state(self) -> Optional[Dict[str, Any]]:
         """
         Return the initial state for this agent,
         None if no such state exists

mephisto/abstractions/blueprints/parlai_chat/parlai_chat_agent_state.py

@@ -24,9 +24,7 @@ def _set_init_state(self, data: Any):
         """Set the initial state for this agent"""
         self.init_data: Optional[Any] = data
 
-    def get_init_state(
-        self, get_all_state: Optional[bool] = False
-    ) -> Optional[Dict[str, Any]]:
+    def get_init_state(self) -> Optional[Dict[str, Any]]:
         """
         Return the initial state for this agent,
         None if no such state exists

mephisto/abstractions/blueprints/remote_procedure/remote_procedure_agent_state.py

@@ -33,9 +33,7 @@ def _set_init_state(self, data: Any):
         """Set the initial state for this agent"""
         self.init_data: Optional[Dict[str, Any]] = data
 
-    def get_init_state(
-        self, get_all_state: Optional[bool] = False
-    ) -> Optional[Dict[str, Any]]:
+    def get_init_state(self) -> Optional[Dict[str, Any]]:
         """
         Return the initial state for this agent,
         None if no such state exists

mephisto/tools/data_browser.py

@@ -4,17 +4,14 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from mephisto.abstractions.database import MephistoDB
 from mephisto.data_model.unit import Unit
 from mephisto.data_model.task_run import TaskRun
-from mephisto.abstractions.blueprint import AgentState
-from mephisto.data_model.agent import Agent
 from mephisto.data_model.worker import Worker
 
 from mephisto.abstractions.databases.local_database import LocalMephistoDB
 from mephisto.data_model.constants.assignment_state import AssignmentState
 
-from typing import List, Optional, Any, Dict
+from typing import List, Any, Dict
 
 
 class DataBrowser:
@@ -28,45 +25,47 @@ def __init__(self, db=None):
             db = LocalMephistoDB()
         self.db = db
 
-    def _get_units_for_task_runs(self, task_runs: List[TaskRun]) -> List[Unit]:
+    def collect_matching_units_from_task_runs(
+        self, task_runs: List[TaskRun], statuses: List[str]
+    ) -> List[Unit]:
         """
-        Return a list of all Units in a terminal completed state from all
-        the provided TaskRuns.
+        Loops through task_runs to collect all units in the provided statuses list
         """
         units = []
         for task_run in task_runs:
             assignments = task_run.get_assignments()
             for assignment in assignments:
                 found_units = assignment.get_units()
                 for unit in found_units:
-                    if unit.get_status() in [
-                        AssignmentState.COMPLETED,
-                        AssignmentState.ACCEPTED,
-                        AssignmentState.REJECTED,
-                        AssignmentState.SOFT_REJECTED,
-                    ]:
+                    if unit.get_status() in statuses:
                         units.append(unit)
         return units
 
+    def _get_units_for_task_runs(self, task_runs: List[TaskRun]) -> List[Unit]:
+        """
+        Return a list of all Units in a terminal completed state from all
+        the provided TaskRuns.
+        """
+        statuses = [
+            AssignmentState.COMPLETED,
+            AssignmentState.ACCEPTED,
+            AssignmentState.REJECTED,
+            AssignmentState.SOFT_REJECTED,
+        ]
+        return self.collect_matching_units_from_task_runs(task_runs, statuses)
+
     def _get_all_units_for_task_runs(self, task_runs: List[TaskRun]) -> List[Unit]:
         """
         Does the same as _get_units_for_task_runs except that it includes the EXPIRED state
         """
-        units = []
-        for task_run in task_runs:
-            assignments = task_run.get_assignments()
-            for assignment in assignments:
-                found_units = assignment.get_units()
-                for unit in found_units:
-                    if unit.get_status() in [
-                        AssignmentState.COMPLETED,
-                        AssignmentState.ACCEPTED,
-                        AssignmentState.REJECTED,
-                        AssignmentState.SOFT_REJECTED,
-                        AssignmentState.EXPIRED,
-                    ]:
-                        units.append(unit)
-        return units
+        statuses = [
+            AssignmentState.COMPLETED,
+            AssignmentState.ACCEPTED,
+            AssignmentState.REJECTED,
+            AssignmentState.SOFT_REJECTED,
+            AssignmentState.EXPIRED,
+        ]
+        return self.collect_matching_units_from_task_runs(task_runs, statuses)
 
     def get_task_name_list(self) -> List[str]:
         return [task.task_name for task in self.db.find_tasks()]