Fix init order

airtai · Nov 21, 2024 · c23e4ee · c23e4ee
1 parent d8e73e7
commit c23e4ee
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/context_leakage_team/workflow/scenarios/context_leak/context_leak_scenario.py b/context_leakage_team/workflow/scenarios/context_leak/context_leak_scenario.py
@@ -64,7 +64,10 @@ def __init__(self, ui: UI, params: dict[str, Any]) -> None:
         super().__init__(ui, params)
         self.context_leak_log_save_path = self.DEFAULT_LOG_PATH
         self.counter = 0
+        self.model_level = "low"
+        self.max_round = 1
 
+    def setup_environment(self) -> None:
         self.model_level = self.ui.multiple_choice(
             sender="Context leakage team",
             prompt="What model would you like to test?",
@@ -79,9 +82,6 @@ def __init__(self, ui: UI, params: dict[str, Any]) -> None:
             )
         )
 
-    def setup_environment(self) -> None:
-        pass
-
     def setup_agents(self) -> Iterable[Agent]:
         """Create agents specific to context leakage."""
         confidential_text = self.TESTED_MODEL_CONFIDENTIAL_PATH.read_text()