move from AutonomousUseCase to SimpleStrategy

Author: andreashappe

src/hackingBuddyGPT/strategies.py

@@ -7,6 +7,7 @@
 from hackingBuddyGPT.usecases.base import UseCase
 from hackingBuddyGPT.utils import llm_util
 from hackingBuddyGPT.utils.histories import HistoryCmdOnly, HistoryFull, HistoryNone
+from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib
 from hackingBuddyGPT.utils.openai.openai_llm import OpenAIConnection
 from hackingBuddyGPT.utils.logging import log_conversation, Logger, log_param, log_section
 from hackingBuddyGPT.utils.capability_manager import CapabilityManager
@@ -140,3 +141,59 @@ def check_success(self, cmd:str, result:str) -> bool:
 
     def postprocess_commands(self, cmd:str) -> List[str]:
         return [cmd]
+
+@dataclass
+class SimpleStrategy(UseCase, abc.ABC):
+    max_turns: int = 10
+
+    llm: OpenAILib = None
+
+    log: Logger = log_param
+
+    _got_root: bool = False
+
+    _capabilities: CapabilityManager = None
+
+    def init(self):
+        super().init()
+        self._capabilities = CapabilityManager(self.log)
+
+    @abc.abstractmethod
+    def perform_round(self, turn: int):
+        pass
+
+    def before_run(self):
+        pass
+
+    def after_run(self):
+        pass
+
+    def run(self, configuration):
+        self.configuration = configuration
+        self.log.start_run(self.get_name(), self.serialize_configuration(configuration))
+
+        self.before_run()
+
+        turn = 1
+        try:
+            while turn <= self.max_turns and not self._got_root:
+                with self.log.section(f"round {turn}"):
+                    self.log.console.log(f"[yellow]Starting turn {turn} of {self.max_turns}")
+
+                    self._got_root = self.perform_round(turn)
+
+                    turn += 1
+
+            self.after_run()
+
+            # write the final result to the database and console
+            if self._got_root:
+                self.log.run_was_success()
+            else:
+                self.log.run_was_failure("maximum turn number reached")
+
+            return self._got_root
+        except Exception:
+            import traceback
+            self.log.run_was_failure("exception occurred", details=f":\n\n{traceback.format_exc()}")
+            raise

src/hackingBuddyGPT/usecases/agents.py

@@ -1,12 +1,11 @@
 import abc
 import json
-import argparse
 from dataclasses import dataclass
 
 from hackingBuddyGPT.utils.logging import Logger, log_param
-from typing import Dict, Type, TypeVar, Generic
+from typing import Dict, Type
 
-from hackingBuddyGPT.utils.configurable import Transparent, configurable
+from hackingBuddyGPT.utils.configurable import configurable
 
 @dataclass
 class UseCase(abc.ABC):
@@ -49,98 +48,8 @@ def get_name(self) -> str:
         """
         pass
 
-
-# this runs the main loop for a bounded amount of turns or until root was achieved
-@dataclass
-class AutonomousUseCase(UseCase, abc.ABC):
-    max_turns: int = 10
-
-    _got_root: bool = False
-
-    @abc.abstractmethod
-    def perform_round(self, turn: int):
-        pass
-
-    def before_run(self):
-        pass
-
-    def after_run(self):
-        pass
-
-    def run(self, configuration):
-        self.configuration = configuration
-        self.log.start_run(self.get_name(), self.serialize_configuration(configuration))
-
-        self.before_run()
-
-        turn = 1
-        try:
-            while turn <= self.max_turns and not self._got_root:
-                with self.log.section(f"round {turn}"):
-                    self.log.console.log(f"[yellow]Starting turn {turn} of {self.max_turns}")
-
-                    self._got_root = self.perform_round(turn)
-
-                    turn += 1
-
-            self.after_run()
-
-            # write the final result to the database and console
-            if self._got_root:
-                self.log.run_was_success()
-            else:
-                self.log.run_was_failure("maximum turn number reached")
-
-            return self._got_root
-        except Exception:
-            import traceback
-            self.log.run_was_failure("exception occurred", details=f":\n\n{traceback.format_exc()}")
-            raise
-
-
 use_cases: Dict[str, configurable] = dict()
 
-
-T = TypeVar("T", bound=type)
-
-
-class AutonomousAgentUseCase(AutonomousUseCase, Generic[T]):
-    agent: T = None
-
-    def perform_round(self, turn: int):
-        raise ValueError("Do not use AutonomousAgentUseCase without supplying an agent type as generic")
-
-    def get_name(self) -> str:
-        raise ValueError("Do not use AutonomousAgentUseCase without supplying an agent type as generic")
-
-    @classmethod
-    def __class_getitem__(cls, item):
-        item = dataclass(item)
-
-        class AutonomousAgentUseCase(AutonomousUseCase):
-            agent: Transparent(item) = None
-
-            def init(self):
-                super().init()
-                self.agent.init()
-
-            def get_name(self) -> str:
-                return self.__class__.__name__
-
-            def before_run(self):
-                return self.agent.before_run()
-
-            def after_run(self):
-                return self.agent.after_run()
-
-            def perform_round(self, turn: int):
-                return self.agent.perform_round(turn)
-
-        constructed_class = dataclass(AutonomousAgentUseCase)
-
-        return constructed_class
-
-
 def use_case(description):
     def inner(cls):
         cls = dataclass(cls)

src/hackingBuddyGPT/usecases/base.py

@@ -1,24 +1,20 @@
 from dataclasses import field
-from typing import List, Any, Union, Dict, Iterable, Optional
-
+from typing import List, Any, Union,  Iterable, Optional
 from openai.types.chat import ChatCompletionMessageParam, ChatCompletionMessage
 from openai.types.chat.chat_completion_chunk import ChoiceDelta
 
-from hackingBuddyGPT.capabilities import Capability
 from hackingBuddyGPT.capabilities.http_request import HTTPRequest
 from hackingBuddyGPT.capabilities.submit_flag import SubmitFlag
-from hackingBuddyGPT.usecases.agents import Agent
-from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case
+from hackingBuddyGPT.strategies import SimpleStrategy
+from hackingBuddyGPT.usecases.base import use_case
 from hackingBuddyGPT.utils import LLMResult, tool_message
 from hackingBuddyGPT.utils.configurable import parameter
-from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib
 
 Prompt = List[Union[ChatCompletionMessage, ChatCompletionMessageParam]]
 Context = Any
 
-
-class WebTestingWithExplanation(Agent):
-    llm: OpenAILib
+@use_case("Minimal implementation of a web testing use case while allowing the llm to 'talk'")
+class WebTestingWithExplanation(SimpleStrategy):
     host: str = parameter(desc="The host to test", default="http://localhost")
     flag_format_description: str = parameter(
         desc="Description of the flag provided to the LLM",
@@ -35,14 +31,16 @@ class WebTestingWithExplanation(Agent):
 
     _prompt_history: Prompt = field(default_factory=list)
     _context: Context = field(default_factory=lambda: {"notes": list()})
-    _capabilities: Dict[str, Capability] = field(default_factory=dict)
     _all_flags_found: bool = False
 
     def init(self):
         super().init()
         self._context["host"] = self.host
-        self.add_capability(SubmitFlag(self.flag_format_description, set(self.flag_template.format(flag=flag) for flag in self.flags.split(",")), success_function=self.all_flags_found))
-        self.add_capability(HTTPRequest(self.host))
+        self._capabilities.add_capability(SubmitFlag(self.flag_format_description, set(self.flag_template.format(flag=flag) for flag in self.flags.split(",")), success_function=self.all_flags_found))
+        self._capabilities.add_capability(HTTPRequest(self.host))
+
+    def get_name(self) -> str:
+        return self.__class__.__name__
 
     def before_run(self):
         system_message = (
@@ -64,7 +62,7 @@ def all_flags_found(self):
     def perform_round(self, turn: int):
         prompt = self._prompt_history  # TODO: in the future, this should do some context truncation
 
-        result_stream: Iterable[Union[ChoiceDelta, LLMResult]] = self.llm.stream_response(prompt, self.log.console, capabilities=self._capabilities, get_individual_updates=True)
+        result_stream: Iterable[Union[ChoiceDelta, LLMResult]] = self.llm.stream_response(prompt, self.log.console, capabilities=self._capabilities._capabilities, get_individual_updates=True)
         result: Optional[LLMResult] = None
         stream_output = self.log.stream_message("assistant")  # TODO: do not hardcode the role
         for delta in result_stream:
@@ -83,12 +81,7 @@ def perform_round(self, turn: int):
 
         if message.tool_calls is not None:
             for tool_call in message.tool_calls:
-                tool_result = self.run_capability_json(message_id, tool_call.id, tool_call.function.name, tool_call.function.arguments)
+                tool_result = self._capabilities.run_capability_json(message_id, tool_call.id, tool_call.function.name, tool_call.function.arguments)
                 self._prompt_history.append(tool_message(tool_result, tool_call.id))
 
         return self._all_flags_found
-
-
-@use_case("Minimal implementation of a web testing use case while allowing the llm to 'talk'")
-class WebTestingWithExplanationUseCase(AutonomousAgentUseCase[WebTestingWithExplanation]):
-    pass