fix missing 1 required positional argument bug & sqlite3.operationalerror: database is locked issue

src/hackingBuddyGPT/usecases/agents.py

@@ -55,9 +55,8 @@ def set_initial_state(self, initial_state:AgentWorldview):
     def set_template(self, template:str):
         self._template = Template(filename=template)
         self._template_size = self.llm.count_tokens(self._template.source)
-
-    def perform_round(self, turn:int) -> bool:
-        got_root : bool = False
+    def perform_round(self, turn: int) -> bool:##fix code
+        got_root: bool = False
 
         with self.console.status("[bold green]Asking LLM for a new command..."):
             # TODO output/log state
@@ -71,15 +70,35 @@ def perform_round(self, turn:int) -> bool:
             cmd = llm_util.cmd_output_fixer(answer.result)
 
         with self.console.status("[bold green]Executing that command..."):
-                self.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_db.add_log_query(self._run_id, turn, cmd, result, answer)
-        self._state.update(capability, cmd, result)
-        # TODO output/log new state
-        self.console.print(Panel(result, title=f"[bold cyan]{cmd}"))
-
-        # if we got root, we can stop the loop
-        return got_root
+            self.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:"))
+
+            # Assuming command is of the form "capability_name arg1 arg2"
+            parts = cmd.split(" ", 1)
+            if len(parts) == 2:
+                capability_name, args = parts
+                capability = self.get_capability(capability_name)
+
+                if capability:
+                    # Assuming capability requires multiple arguments
+                    # Adjust the argument unpacking based on capability's requirements
+                    args_list = args.split()  # Split arguments into a list
+
+                    try:
+                        result, got_root = capability(*args_list)
+                    except TypeError as e:
+                        result = f"Error executing command: {e}"
+                        got_root = False
+                else:
+                    result = f"Unknown capability: {capability_name}"
+                    got_root = False
+            else:
+                result = "Command format error. Expected 'capability_name arg1 arg2'."
+                got_root = False
+
+            # Log and output the command and its result
+            self.log_db.add_log_query(self._run_id, turn, cmd, result, answer)
+            self._state.update(capability, cmd, result)  # Assuming capability is available
+            self.console.print(Panel(result, title=f"[bold cyan]{cmd}"))
+
+            # If we got root, we can stop the loop
+            return got_root

src/hackingBuddyGPT/usecases/minimal/agent.py

@@ -26,7 +26,6 @@ def init(self):
         self.add_capability(SSHRunCommand(conn=self.conn), default=True)
         self.add_capability(SSHTestCredential(conn=self.conn))
         self._template_size = self.llm.count_tokens(template_next_cmd.source)
-
     def perform_round(self, turn):
         got_root : bool = False
 
@@ -39,13 +38,27 @@ def perform_round(self, turn):
             cmd = llm_util.cmd_output_fixer(answer.result)
 
         with self.console.status("[bold green]Executing that command..."):
-                self.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:"))
-                result, got_root = self.get_capability(cmd.split(" ", 1)[0])(cmd)
+            self.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:"))
+
+            # Assuming cmd is of the form "username password"
+            parts = cmd.split(" ", 1)
+            if len(parts) == 2:
+                username, password = parts
+                ##here fix!
+                result, got_root = self.get_capability("test_credential")(username, password)
+            else:
+                # Handle other cases or log error
+                result = "Command format error. Expected 'username password'."
+                got_root = False
 
-        # log and output the command and its result
-        self.log_db.add_log_query(self._run_id, turn, cmd, result, answer)
-        self._sliding_history.add_command(cmd, result)
-        self.console.print(Panel(result, title=f"[bold cyan]{cmd}"))
+            #self.log_db.add_log_query(self._run_id, cmd, result, answer)
+            self.log_db.add_log_query(self._run_id, turn, cmd, result, answer)
+            self._sliding_history.add_command(cmd, result)
+            self.console.print(Panel(result, title=f"[bold cyan]{cmd}"))
 
-        # if we got root, we can stop the loop
         return got_root
+
+
+
+
+

src/hackingBuddyGPT/utils/db_storage/db_storage.py

@@ -11,10 +11,12 @@ def __init__(self, connection_string: str = parameter(desc="sqlite3 database con
     def init(self):
         self.connect()
         self.setup_db()
-
+    
     def connect(self):
-        self.db = sqlite3.connect(self.connection_string)
+#        self.db = sqlite3.connect(self.connection_string, timeout=10)  # Set timeout to 10 seconds
+        self.db = sqlite3.connect(self.connection_string, check_same_thread=False, timeout=10)
         self.cursor = self.db.cursor()
+
 
     def insert_or_select_cmd(self, name: str) -> int:
         results = self.cursor.execute("SELECT id, name FROM commands WHERE name = ?", (name,)).fetchall()
@@ -80,20 +82,21 @@ def setup_db(self):
         self.query_cmd_id = self.insert_or_select_cmd('query_cmd')
         self.analyze_response_id = self.insert_or_select_cmd('analyze_response')
         self.state_update_id = self.insert_or_select_cmd('update_state')
-
+   
     def create_new_run(self, model, context_size, tag):
-        self.cursor.execute(
-            "INSERT INTO runs (model, context_size, state, tag, started_at) VALUES (?, ?, ?, ?, datetime('now'))",
-            (model, context_size, "in progress", tag))
-        return self.cursor.lastrowid
+        with self.db:
+            self.cursor.execute(
+                "INSERT INTO runs (model, context_size, state, tag, started_at) VALUES (?, ?, ?, ?, datetime('now'))",
+                (model, context_size, "in progress", tag))
+            return self.cursor.lastrowid
 
     def add_log_query(self, run_id, round, cmd, result, answer):
-        self.cursor.execute(
-            "INSERT INTO queries (run_id, round, cmd_id, query, response, duration, tokens_query, tokens_response, prompt, answer) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
-            (
-            run_id, round, self.query_cmd_id, cmd, result, answer.duration, answer.tokens_query, answer.tokens_response,
-            answer.prompt, answer.answer))
+        with self.db:
+            self.cursor.execute(
+                "INSERT INTO queries (run_id, round, cmd_id, query, response, duration, tokens_query, tokens_response, prompt, answer) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+                (run_id, round, self.query_cmd_id, cmd, result, answer.duration, answer.tokens_query, answer.tokens_response, answer.prompt, answer.answer))
 
+
     def add_log_analyze_response(self, run_id, round, cmd, result, answer):
         self.cursor.execute(
             "INSERT INTO queries (run_id, round, cmd_id, query, response, duration, tokens_query, tokens_response, prompt, answer) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
@@ -206,4 +209,4 @@ def run_was_failure(self, run_id, round):
         self.db.commit()
 
     def commit(self):
-        self.db.commit()
+        self.db.commit()