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 index f71e88a..ff94923 100644 --- 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()