From 488fec5718f44a3ec321cf0e4f8c38a4471850a9 Mon Sep 17 00:00:00 2001
From: gtarpenning <gtarpenning@gmail.com>
Date: Wed, 19 Jun 2024 16:04:12 -0400
Subject: [PATCH] feat: add sqlite logging to cmd plugin if configured

---
 llm_cmd.py | 55 ++++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 45 insertions(+), 10 deletions(-)

diff --git a/llm_cmd.py b/llm_cmd.py
index 7d749cd..70485fe 100644
--- a/llm_cmd.py
+++ b/llm_cmd.py
@@ -27,7 +27,7 @@ def register_commands(cli):
     @click.option("--key", help="API key to use")
     def cmd(args, model, system, key):
         """Generate and execute commands in your shell"""
-        from llm.cli import get_default_model
+        from llm.cli import get_default_model, logs_on, logs_db_path
 
         prompt = " ".join(args)
 
@@ -39,19 +39,54 @@ def cmd(args, model, system, key):
 
         result = model_obj.prompt(prompt, system=system or SYSTEM_PROMPT)
 
-        interactive_exec(str(result))
+        output = interactive_exec(str(result))
+        print(output)
 
+        if logs_on():
+            import sqlite_utils
+            from llm.migrations import migrate
+
+            log_path = logs_db_path()
+            (log_path.parent).mkdir(parents=True, exist_ok=True)
+            db = sqlite_utils.Database(log_path)
+            migrate(db)
+            _log_cmd(db, model_id, result, output, prompt)
+
+
+def interactive_exec(command) -> str:
+    try:
+        if "\n" in command:
+            print("Multiline command - Meta-Enter or Esc Enter to execute")
+            edited_command = prompt(
+                "> ", default=command, lexer=PygmentsLexer(BashLexer), multiline=True
+            )
+        else:
+            edited_command = prompt(
+                "> ", default=command, lexer=PygmentsLexer(BashLexer)
+            )
+    except KeyboardInterrupt:
+        return "Aborted!"
 
-def interactive_exec(command):
-    if '\n' in command:
-        print("Multiline command - Meta-Enter or Esc Enter to execute")
-        edited_command = prompt("> ", default=command, lexer=PygmentsLexer(BashLexer), multiline=True)
-    else:
-        edited_command = prompt("> ", default=command, lexer=PygmentsLexer(BashLexer))
     try:
         output = subprocess.check_output(
             edited_command, shell=True, stderr=subprocess.STDOUT
         )
-        print(output.decode())
+        return output.decode()
     except subprocess.CalledProcessError as e:
-        print(f"Command failed with error (exit status {e.returncode}): {e.output.decode()}")
+        return f"Command failed with error (exit status {e.returncode}): {e.output.decode()}"
+
+
+def _log_cmd(db, model_id, result, output, prompt):
+    failed = "Command failed with error" in output
+    canceled = output == "Aborted!"
+    db["cmds"].insert(
+        {
+            "model": model_id,
+            "prompt": prompt,
+            "cmd": str(result),
+            "output": output,
+            "failed": failed,
+            "canceled": canceled,
+        },
+        ignore=True,
+    )