From 17a1d57301f55b9635414deeed1a24b433674119 Mon Sep 17 00:00:00 2001 From: Neverbolt Date: Tue, 6 Aug 2024 12:50:54 +0200 Subject: [PATCH] Moves all `_log`s to global `log` --- README.md | 10 ++-- src/hackingBuddyGPT/usecases/agents.py | 13 ++--- src/hackingBuddyGPT/usecases/base.py | 37 +++++++------- src/hackingBuddyGPT/usecases/minimal/agent.py | 10 ++-- .../usecases/privesc/common.py | 26 +++++----- src/hackingBuddyGPT/usecases/privesc/linux.py | 28 ++++++----- src/hackingBuddyGPT/usecases/web/simple.py | 12 ++--- .../usecases/web/with_explanation.py | 8 +-- .../simple_openapi_documentation.py | 6 +-- .../web_api_testing/simple_web_api_testing.py | 8 +-- tests/integration_minimal_test.py | 50 ++++++++++++------- tests/test_web_api_documentation.py | 11 ++-- tests/test_web_api_testing.py | 11 ++-- 13 files changed, 126 insertions(+), 104 deletions(-) diff --git a/README.md b/README.md index b7a64c1..a1eb6bc 100644 --- a/README.md +++ b/README.md @@ -96,7 +96,7 @@ class MinimalLinuxPrivesc(Agent): def perform_round(self, turn: int) -> bool: got_root: bool = False - with self._log.console.status("[bold green]Asking LLM for a new command..."): + with self.log.console.status("[bold green]Asking LLM for a new command..."): # get as much history as fits into the target context size history = self._sliding_history.get_history(self.llm.context_size - llm_util.SAFETY_MARGIN - self._template_size) @@ -104,14 +104,14 @@ class MinimalLinuxPrivesc(Agent): answer = self.llm.get_response(template_next_cmd, capabilities=self.get_capability_block(), history=history, conn=self.conn) cmd = llm_util.cmd_output_fixer(answer.result) - with self._log.console.status("[bold green]Executing that command..."): - self._log.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) + with self.log.console.status("[bold green]Executing that command..."): + self.log.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) result, got_root = self.get_capability(cmd.split(" ", 1)[0])(cmd) # log and output the command and its result - self._log.log_db.add_log_query(self._log.run_id, turn, cmd, result, answer) + self.log.add_log_query(turn, cmd, result, answer) self._sliding_history.add_command(cmd, result) - self._log.console.print(Panel(result, title=f"[bold cyan]{cmd}")) + self.log.console.print(Panel(result, title=f"[bold cyan]{cmd}")) # if we got root, we can stop the loop return got_root diff --git a/src/hackingBuddyGPT/usecases/agents.py b/src/hackingBuddyGPT/usecases/agents.py index 0985e64..05ecfba 100644 --- a/src/hackingBuddyGPT/usecases/agents.py +++ b/src/hackingBuddyGPT/usecases/agents.py @@ -12,9 +12,10 @@ @dataclass class Agent(ABC): + log: Logger = None + _capabilities: Dict[str, Capability] = field(default_factory=dict) _default_capability: Capability = None - _log: Logger = None llm: OpenAIConnection = None @@ -76,7 +77,7 @@ def set_template(self, template:str): def perform_round(self, turn:int) -> bool: got_root : bool = False - with self._log.console.status("[bold green]Asking LLM for a new command..."): + with self.log.console.status("[bold green]Asking LLM for a new command..."): # TODO output/log state options = self._state.to_template() options.update({ @@ -87,16 +88,16 @@ def perform_round(self, turn:int) -> bool: answer = self.llm.get_response(self._template, **options) cmd = llm_util.cmd_output_fixer(answer.result) - with self._log.console.status("[bold green]Executing that command..."): - self._log.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) + with self.log.console.status("[bold green]Executing that command..."): + self.log.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) capability = self.get_capability(cmd.split(" ", 1)[0]) result, got_root = capability(cmd) # log and output the command and its result - self._log.add_log_query(turn, cmd, result, answer) + self.log.add_log_query(turn, cmd, result, answer) self._state.update(capability, cmd, result) # TODO output/log new state - self._log.console.print(Panel(result, title=f"[bold cyan]{cmd}")) + self.log.console.print(Panel(result, title=f"[bold cyan]{cmd}")) # if we got root, we can stop the loop return got_root diff --git a/src/hackingBuddyGPT/usecases/base.py b/src/hackingBuddyGPT/usecases/base.py index 4c24903..e601c19 100644 --- a/src/hackingBuddyGPT/usecases/base.py +++ b/src/hackingBuddyGPT/usecases/base.py @@ -13,18 +13,19 @@ from hackingBuddyGPT.utils.db_storage.db_storage import DbStorage +@configurable("logger", "Logger") @dataclass -class Logger: +class RawLogger: log_db: DbStorage console: Console - model: str = "" tag: str = "" - configuration: str = "" run_id: int = field(init=False, default=None) - def __post_init__(self): - self.run_id = self.log_db.create_new_run(self.model, self.tag, self.configuration) + def start_run(self, name: str, configuration: str): + if self.run_id is not None: + raise ValueError("Run already started") + self.run_id = self.log_db.create_new_run(name, self.tag, configuration) def add_log_query(self, turn: int, command: str, result: str, answer: LLMResult): self.log_db.add_log_query(self.run_id, turn, command, result, answer) @@ -53,6 +54,9 @@ def status_message(self, message: str): self.log_db.add_log_message(self.run_id, "status", message, 0, 0, 0) +Logger = Global(Transparent(RawLogger)) + + @dataclass class UseCase(abc.ABC): """ @@ -65,11 +69,7 @@ class UseCase(abc.ABC): so that they can be automatically discovered and run from the command line. """ - log_db: DbStorage - console: Console - tag: str = "" - - _log: Logger = None + log: Logger def init(self, configuration): """ @@ -77,7 +77,7 @@ def init(self, configuration): perform any dynamic setup that is needed before the run method is called. One of the most common use cases is setting up the llm capabilities from the tools that were injected. """ - self._log = Logger(self.log_db, self.console, self.get_name(), self.tag, self.serialize_configuration(configuration)) + self.log.start_run(self.get_name(), self.serialize_configuration(configuration)) def serialize_configuration(self, configuration) -> str: return json.dumps(configuration) @@ -123,27 +123,27 @@ def run(self): turn = 1 try: while turn <= self.max_turns and not self._got_root: - self._log.console.log(f"[yellow]Starting turn {turn} of {self.max_turns}") + self.log.console.log(f"[yellow]Starting turn {turn} of {self.max_turns}") self._got_root = self.perform_round(turn) # finish turn and commit logs to storage - self._log.log_db.commit() + self.log.log_db.commit() turn += 1 self.after_run() # write the final result to the database and console if self._got_root: - self._log.run_was_success() - self._log.console.print(Panel("[bold green]Got Root!", title="Run finished")) + self.log.run_was_success() + self.log.console.print(Panel("[bold green]Got Root!", title="Run finished")) else: - self._log.run_was_failure("maximum turn number reached") - self._log.console.print(Panel("[green]maximum turn number reached", title="Run finished")) + self.log.run_was_failure("maximum turn number reached") + self.log.console.print(Panel("[green]maximum turn number reached", title="Run finished")) return self._got_root except Exception as e: - self._log.run_was_failure(f"exception occurred: {e}") + self.log.run_was_failure(f"exception occurred: {e}") raise @@ -192,7 +192,6 @@ class AutonomousAgentUseCase(AutonomousUseCase): def init(self, configuration): super().init(configuration) - self.agent._log = self._log self.agent.init() def get_name(self) -> str: diff --git a/src/hackingBuddyGPT/usecases/minimal/agent.py b/src/hackingBuddyGPT/usecases/minimal/agent.py index 0b2e1f0..b419b42 100644 --- a/src/hackingBuddyGPT/usecases/minimal/agent.py +++ b/src/hackingBuddyGPT/usecases/minimal/agent.py @@ -27,7 +27,7 @@ def init(self): def perform_round(self, turn: int) -> bool: got_root: bool = False - with self._log.console.status("[bold green]Asking LLM for a new command..."): + with self.log.console.status("[bold green]Asking LLM for a new command..."): # get as much history as fits into the target context size history = self._sliding_history.get_history(self.llm.context_size - llm_util.SAFETY_MARGIN - self._template_size) @@ -35,14 +35,14 @@ def perform_round(self, turn: int) -> bool: answer = self.llm.get_response(template_next_cmd, capabilities=self.get_capability_block(), history=history, conn=self.conn) cmd = llm_util.cmd_output_fixer(answer.result) - with self._log.console.status("[bold green]Executing that command..."): - self._log.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) + with self.log.console.status("[bold green]Executing that command..."): + self.log.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) result, got_root = self.get_capability(cmd.split(" ", 1)[0])(cmd) # log and output the command and its result - self._log.add_log_query(turn, cmd, result, answer) + self.log.add_log_query(turn, cmd, result, answer) self._sliding_history.add_command(cmd, result) - self._log.console.print(Panel(result, title=f"[bold cyan]{cmd}")) + self.log.console.print(Panel(result, title=f"[bold cyan]{cmd}")) # if we got root, we can stop the loop return got_root diff --git a/src/hackingBuddyGPT/usecases/privesc/common.py b/src/hackingBuddyGPT/usecases/privesc/common.py index 4dd20d7..99ce1ad 100644 --- a/src/hackingBuddyGPT/usecases/privesc/common.py +++ b/src/hackingBuddyGPT/usecases/privesc/common.py @@ -37,7 +37,7 @@ def init(self): def before_run(self): if self.hint != "": - self._log.console.print(f"[bold green]Using the following hint: '{self.hint}'") + self.log.console.print(f"[bold green]Using the following hint: '{self.hint}'") if self.disable_history is False: self._sliding_history = SlidingCliHistory(self.llm) @@ -57,48 +57,48 @@ def before_run(self): def perform_round(self, turn: int) -> bool: got_root: bool = False - with self._log.console.status("[bold green]Asking LLM for a new command..."): + with self.log.console.status("[bold green]Asking LLM for a new command..."): answer = self.get_next_command() cmd = answer.result - with self._log.console.status("[bold green]Executing that command..."): - self._log.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) + with self.log.console.status("[bold green]Executing that command..."): + self.log.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) _capability_descriptions, parser = capabilities_to_simple_text_handler(self._capabilities, default_capability=self._default_capability) success, *output = parser(cmd) if not success: - self._log.console.print(Panel(output[0], title="[bold red]Error parsing command:")) + self.log.console.print(Panel(output[0], title="[bold red]Error parsing command:")) return False assert(len(output) == 1) capability, cmd, (result, got_root) = output[0] # log and output the command and its result - self._log.add_log_query(turn, cmd, result, answer) + self.log.add_log_query(turn, cmd, result, answer) if self._sliding_history: self._sliding_history.add_command(cmd, result) - self._log.console.print(Panel(result, title=f"[bold cyan]{cmd}")) + self.log.console.print(Panel(result, title=f"[bold cyan]{cmd}")) # analyze the result.. if self.enable_explanation: - with self._log.console.status("[bold green]Analyze its result..."): + with self.log.console.status("[bold green]Analyze its result..."): answer = self.analyze_result(cmd, result) - self._log.add_log_analyze_response(turn, cmd, answer.result, answer) + self.log.add_log_analyze_response(turn, cmd, answer.result, answer) # .. and let our local model update its state if self.enable_update_state: # this must happen before the table output as we might include the # status processing time in the table.. - with self._log.console.status("[bold green]Updating fact list.."): + with self.log.console.status("[bold green]Updating fact list.."): state = self.update_state(cmd, result) - self._log.add_log_update_state(turn, "", state.result, state) + self.log.add_log_update_state(turn, "", state.result, state) # Output Round Data.. - self._log.console.print(ui.get_history_table(self.enable_explanation, self.enable_update_state, self._log.run_id, self._log.log_db, turn)) + self.log.console.print(ui.get_history_table(self.enable_explanation, self.enable_update_state, self.log.run_id, self.log.log_db, turn)) # .. and output the updated state if self.enable_update_state: - self._log.console.print(Panel(self._state, title="What does the LLM Know about the system?")) + self.log.console.print(Panel(self._state, title="What does the LLM Know about the system?")) # if we got root, we can stop the loop return got_root diff --git a/src/hackingBuddyGPT/usecases/privesc/linux.py b/src/hackingBuddyGPT/usecases/privesc/linux.py index 37a2070..a5fad40 100644 --- a/src/hackingBuddyGPT/usecases/privesc/linux.py +++ b/src/hackingBuddyGPT/usecases/privesc/linux.py @@ -44,9 +44,9 @@ def read_hint(self): if self.agent.conn.hostname in hints: return hints[self.agent.conn.hostname] except FileNotFoundError: - self._log.console.print("[yellow]Hint file not found") + self.log.console.print("[yellow]Hint file not found") except Exception as e: - self._log.console.print("[yellow]Hint file could not loaded:", str(e)) + self.log.console.print("[yellow]Hint file could not loaded:", str(e)) return "" @@ -61,23 +61,25 @@ class LinuxPrivescWithLSEUseCase(UseCase): _got_root: bool = False - # use either an use-case or an agent to perform the privesc + # use either a use-case or an agent to perform the privesc use_use_case: bool = False + _configuration: any = None def init(self, configuration=None): super().init(configuration) + self._configuration = configuration # simple helper that uses lse.sh to get hints from the system def call_lse_against_host(self): - self._log.console.print("[green]performing initial enumeration with lse.sh") + self.log.console.print("[green]performing initial enumeration with lse.sh") run_cmd = "wget -q 'https://github.com/diego-treitos/linux-smart-enumeration/releases/latest/download/lse.sh' -O lse.sh;chmod 700 lse.sh; ./lse.sh -c -i -l 0 | grep -v 'nope$' | grep -v 'skip$'" result, _ = SSHRunCommand(conn=self.conn, timeout=120)(run_cmd) - self.console.print("[yellow]got the output: " + result) + self.log.console.print("[yellow]got the output: " + result) cmd = self.llm.get_response(template_lse, lse_output=result, number=3) - self.console.print("[yellow]got the cmd: " + cmd.result) + self.log.console.print("[yellow]got the cmd: " + cmd.result) return [x for x in cmd.result.splitlines() if x.strip()] @@ -98,7 +100,7 @@ def run(self): result = self.run_using_agent(hint, turns_per_hint) if result is True: - self.console.print("[green]Got root!") + self.log.console.print("[green]Got root!") return True def run_using_usecases(self, hint, turns_per_hint): @@ -110,13 +112,13 @@ def run_using_usecases(self, hint, turns_per_hint): enable_update_state = self.enable_update_state, disable_history = self.disable_history, llm = self.llm, - hint = hint + hint = hint, + log = self.log, ), max_turns = turns_per_hint, - log_db = self.log_db, - console = self.console + log = self.log, ) - linux_privesc.init() + linux_privesc.init(self._configuration) return linux_privesc.run() def run_using_agent(self, hint, turns_per_hint): @@ -129,7 +131,7 @@ def run_using_agent(self, hint, turns_per_hint): enable_update_state = self.enable_update_state, disable_history = self.disable_history ) - agent._log = self._log + agent.log = self.log agent.init() # perform the privilege escalation @@ -137,7 +139,7 @@ def run_using_agent(self, hint, turns_per_hint): turn = 1 got_root = False while turn <= turns_per_hint and not got_root: - self._log.console.log(f"[yellow]Starting turn {turn} of {turns_per_hint}") + self.log.console.log(f"[yellow]Starting turn {turn} of {turns_per_hint}") if agent.perform_round(turn) is True: got_root = True diff --git a/src/hackingBuddyGPT/usecases/web/simple.py b/src/hackingBuddyGPT/usecases/web/simple.py index 16f0bd4..5989aea 100644 --- a/src/hackingBuddyGPT/usecases/web/simple.py +++ b/src/hackingBuddyGPT/usecases/web/simple.py @@ -52,11 +52,11 @@ def init(self): } def all_flags_found(self): - self._log.console.print(Panel("All flags found! Congratulations!", title="system")) + self.log.console.print(Panel("All flags found! Congratulations!", title="system")) self._all_flags_found = True def perform_round(self, turn: int): - with self._log.console.status("[bold green]Asking LLM for a new command..."): + with self.log.console.status("[bold green]Asking LLM for a new command..."): prompt = self._prompt_history # TODO: in the future, this should do some context truncation tic = time.perf_counter() @@ -66,17 +66,17 @@ def perform_round(self, turn: int): message = completion.choices[0].message tool_call_id = message.tool_calls[0].id command = pydantic_core.to_json(response).decode() - self._log.console.print(Panel(command, title="assistant")) + self.log.console.print(Panel(command, title="assistant")) self._prompt_history.append(message) answer = LLMResult(completion.choices[0].message.content, str(prompt), completion.choices[0].message.content, toc-tic, completion.usage.prompt_tokens, completion.usage.completion_tokens) - with self._log.console.status("[bold green]Executing that command..."): + with self.log.console.status("[bold green]Executing that command..."): result = response.execute() - self._log.console.print(Panel(result, title="tool")) + self.log.console.print(Panel(result, title="tool")) self._prompt_history.append(tool_message(result, tool_call_id)) - self._log.add_log_query(turn, command, result, answer) + self.log.add_log_query(turn, command, result, answer) return self._all_flags_found diff --git a/src/hackingBuddyGPT/usecases/web/with_explanation.py b/src/hackingBuddyGPT/usecases/web/with_explanation.py index f2eba71..967bef8 100644 --- a/src/hackingBuddyGPT/usecases/web/with_explanation.py +++ b/src/hackingBuddyGPT/usecases/web/with_explanation.py @@ -48,15 +48,15 @@ def init(self): } def all_flags_found(self): - self._log.status_message("All flags found! Congratulations!") + self.log.status_message("All flags found! Congratulations!") self._all_flags_found = True def perform_round(self, turn: int): prompt = self._prompt_history # TODO: in the future, this should do some context truncation - result: LLMResult = self.llm.stream_response(prompt, self._log.console, capabilities=self._capabilities) + result: LLMResult = self.llm.stream_response(prompt, self.log.console, capabilities=self._capabilities) message: ChatCompletionMessage = result.result - message_id = self._log.add_log_message(message.role, message.content, result.tokens_query, result.tokens_response, result.duration) + message_id = self.log.add_log_message(message.role, message.content, result.tokens_query, result.tokens_response, result.duration) self._prompt_history.append(result.result) if message.tool_calls is not None: @@ -65,7 +65,7 @@ def perform_round(self, turn: int): tool_call_result = self._capabilities[tool_call.function.name].to_model().model_validate_json(tool_call.function.arguments).execute() toc = time.perf_counter() self._prompt_history.append(tool_message(tool_call_result, tool_call.id)) - self._log.add_log_tool_call(message_id, tool_call.id, tool_call.function.name, tool_call.function.arguments, tool_call_result, toc - tic) + self.log.add_log_tool_call(message_id, tool_call.id, tool_call.function.name, tool_call.function.arguments, tool_call_result, toc - tic) return self._all_flags_found diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py b/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py index 1d44d2e..cd81d6e 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py @@ -100,12 +100,12 @@ def _handle_response(self, completion, response, turn): message = completion.choices[0].message tool_call_id = message.tool_calls[0].id command = pydantic_core.to_json(response).decode() - self._log.console.print(Panel(command, title="assistant")) + self.log.console.print(Panel(command, title="assistant")) self._prompt_history.append(message) - with self._log.console.status("[bold green]Executing that command..."): + with self.log.console.status("[bold green]Executing that command..."): result = response.execute() - self._log.console.print(Panel(result[:30], title="tool")) + self.log.console.print(Panel(result[:30], title="tool")) result_str = self.response_handler.parse_http_status_line(result) self._prompt_history.append(tool_message(result_str, tool_call_id)) invalid_flags = ["recorded","Not a valid HTTP method", "404" ,"Client Error: Not Found"] diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py b/src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py index 20b3af6..84d0ae6 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py @@ -83,7 +83,7 @@ def all_http_methods_found(self): Handles the event when all HTTP methods are found. Displays a congratulatory message and sets the _all_http_methods_found flag to True. """ - self._log.console.print(Panel("All HTTP methods found! Congratulations!", title="system")) + self.log.console.print(Panel("All HTTP methods found! Congratulations!", title="system")) self._all_http_methods_found = True def _setup_capabilities(self): @@ -125,12 +125,12 @@ def _handle_response(self, completion, response): message = completion.choices[0].message tool_call_id = message.tool_calls[0].id command = pydantic_core.to_json(response).decode() - self._log.console.print(Panel(command, title="assistant")) + self.log.console.print(Panel(command, title="assistant")) self._prompt_history.append(message) - with self._log.console.status("[bold green]Executing that command..."): + with self.log.console.status("[bold green]Executing that command..."): result = response.execute() - self._log.console.print(Panel(result[:30], title="tool")) + self.log.console.print(Panel(result[:30], title="tool")) result_str = self.response_handler.parse_http_status_line(result) self._prompt_history.append(tool_message(result_str, tool_call_id)) diff --git a/tests/integration_minimal_test.py b/tests/integration_minimal_test.py index 851dd4f..29e61c0 100644 --- a/tests/integration_minimal_test.py +++ b/tests/integration_minimal_test.py @@ -1,5 +1,7 @@ from typing import Tuple + +from hackingBuddyGPT.usecases.base import RawLogger from hackingBuddyGPT.usecases.minimal.agent import MinimalLinuxPrivesc, MinimalLinuxPrivescUseCase from hackingBuddyGPT.usecases.minimal.agent_with_state import MinimalLinuxTemplatedPrivesc, MinimalLinuxTemplatedPrivescUseCase from hackingBuddyGPT.usecases.privesc.linux import LinuxPrivesc, LinuxPrivescUseCase @@ -74,18 +76,22 @@ def test_linuxprivesc(): log_db.init() + log = RawLogger( + log_db=log_db, + console=console, + tag='integration_test_linuxprivesc', + ) priv_esc = LinuxPrivescUseCase( - agent = LinuxPrivesc( + agent=LinuxPrivesc( conn=conn, enable_explanation=False, disable_history=False, hint='', - llm = llm, + llm=llm, + log=log, ), - log_db = log_db, - console = console, - tag = 'integration_test_linuxprivesc', - max_turns = len(llm.responses) + log=log, + max_turns=len(llm.responses) ) priv_esc.init({}) @@ -101,15 +107,19 @@ def test_minimal_agent(): log_db.init() + log = RawLogger( + log_db=log_db, + console=console, + tag='integration_test_minimallinuxprivesc', + ) priv_esc = MinimalLinuxPrivescUseCase( - agent = MinimalLinuxPrivesc( + agent=MinimalLinuxPrivesc( conn=conn, - llm=llm + llm=llm, + log=log, ), - log_db = log_db, - console = console, - tag = 'integration_test_minimallinuxprivesc', - max_turns = len(llm.responses) + log=log, + max_turns=len(llm.responses) ) priv_esc.init({}) @@ -125,15 +135,19 @@ def test_minimal_agent_state(): log_db.init() + log = RawLogger( + log_db=log_db, + console=console, + tag='integration_test_linuxprivesc', + ) priv_esc = MinimalLinuxTemplatedPrivescUseCase( - agent = MinimalLinuxTemplatedPrivesc( + agent=MinimalLinuxTemplatedPrivesc( conn=conn, - llm = llm, + llm=llm, + log=log, ), - log_db = log_db, - console = console, - tag = 'integration_test_linuxprivesc', - max_turns = len(llm.responses) + log=log, + max_turns=len(llm.responses) ) priv_esc.init({}) diff --git a/tests/test_web_api_documentation.py b/tests/test_web_api_documentation.py index 33ca4a6..0a0be8f 100644 --- a/tests/test_web_api_documentation.py +++ b/tests/test_web_api_documentation.py @@ -15,13 +15,16 @@ def setUp(self, MockOpenAILib): console = Console() log_db.init() - self.agent = SimpleWebAPIDocumentation(llm=self.mock_llm) - self.agent.init() - self.simple_api_testing = SimpleWebAPIDocumentationUseCase( - agent=self.agent, + log = RawLogger( log_db=log_db, console=console, tag='webApiDocumentation', + ) + self.agent = SimpleWebAPIDocumentation(llm=self.mock_llm, log=log) + self.agent.init() + self.simple_api_testing = SimpleWebAPIDocumentationUseCase( + agent=self.agent, + log=log, max_turns=len(self.mock_llm.responses) ) self.simple_api_testing.init({}) diff --git a/tests/test_web_api_testing.py b/tests/test_web_api_testing.py index f84516c..f844fbe 100644 --- a/tests/test_web_api_testing.py +++ b/tests/test_web_api_testing.py @@ -15,13 +15,16 @@ def setUp(self, MockOpenAILib): console = Console() log_db.init() - self.agent = SimpleWebAPITesting(llm=self.mock_llm) - self.agent.init() - self.simple_api_testing = SimpleWebAPITestingUseCase( - agent=self.agent, + log = RawLogger( log_db=log_db, console=console, tag='integration_test_linuxprivesc', + ) + self.agent = SimpleWebAPITesting(llm=self.mock_llm, log=log) + self.agent.init() + self.simple_api_testing = SimpleWebAPITestingUseCase( + agent=self.agent, + log=log, max_turns=len(self.mock_llm.responses) ) self.simple_api_testing.init({})