Merge pull request #136 from ipa-lab/restructuring

Restructuring: AutonomousUseCase -> SimpleStrategy

Author: andreashappe

pyproject.toml

@@ -41,7 +41,6 @@ dependencies = [
     'nltk',
     'fastapi == 0.114.0',
     'fastapi-utils == 0.7.0',
-    'jinja2 == 3.1.6',
     'uvicorn[standard] == 0.30.6',
     'dataclasses_json == 0.6.7',
     'websockets == 13.1',

src/hackingBuddyGPT/cli/wintermute.py

@@ -1,7 +1,7 @@
 import argparse
 import sys
 
-from hackingBuddyGPT.usecases.base import use_cases
+from hackingBuddyGPT.usecases.usecase import use_cases
 from hackingBuddyGPT.utils.configurable import CommandMap, InvalidCommand, Parseable, instantiate
 
 

src/hackingBuddyGPT/strategies.py

@@ -4,9 +4,10 @@
 from dataclasses import dataclass
 from mako.template import Template
 from hackingBuddyGPT.capability import capabilities_to_simple_text_handler
-from hackingBuddyGPT.usecases.base import UseCase
+from hackingBuddyGPT.usecases.usecase 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,7 +1,7 @@
 from mako.template import Template
 
 from hackingBuddyGPT.capabilities import SSHRunCommand
-from hackingBuddyGPT.usecases.base import UseCase, use_case
+from hackingBuddyGPT.usecases.usecase import UseCase, use_case
 from hackingBuddyGPT.utils.connectors.ssh_connection import SSHConnection
 from hackingBuddyGPT.utils.openai.openai_llm import OpenAIConnection
 

src/hackingBuddyGPT/usecases/base.py

@@ -6,7 +6,7 @@
 
 from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential
 from hackingBuddyGPT.strategies import CommandStrategy
-from hackingBuddyGPT.usecases.base import use_case
+from hackingBuddyGPT.usecases.usecase import use_case
 from hackingBuddyGPT.utils import llm_util
 from hackingBuddyGPT.utils.logging import log_conversation
 from hackingBuddyGPT.utils.rag import RagBackground