From 1a840cd7acb0cbeb13007a65991b7189985dacc0 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 1 Apr 2023 15:28:46 -0700
Subject: [PATCH] Initial prompt/response log, closes #2

---
 README.md  | 16 +++++++++++++++-
 llm/cli.py | 53 +++++++++++++++++++++++++++++++++++++++++++++++++----
 setup.py   |  2 +-
 3 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index a5523077..4210abeb 100644
--- a/README.md
+++ b/README.md
@@ -17,7 +17,7 @@ You need an OpenAI API key, which should either be set in the `OPENAI_API_KEY` e
 
 ## Usage
 
-So far this tool only has one command - `llm chatgpt`. You can just use `llm` as this is the default command.
+The default command for this is `llm chatgpt` - you can use `llm` instead if you prefer.
 
 To run a prompt:
 
@@ -33,6 +33,20 @@ To switch from ChatGPT 3.5 (the default) to GPT-4 if you have access:
 
 Pass `--model <model name>` to use a different model.
 
+## Logging to SQLite
+
+If a SQLite database file exists in `~/.llm/log.db` then the tool will log all prompts and responses to it.
+
+You can create that file by running the `init-db` command:
+
+    llm init-db
+
+Now any prompts you run will be logged to that database.
+
+To avoid logging a prompt, pass `--no-log` or `-n` to the command:
+
+    llm 'Ten names for cheesecakes' -n
+
 ## Help
 
 For help, run:
diff --git a/llm/cli.py b/llm/cli.py
index 546009c0..acb9c21a 100644
--- a/llm/cli.py
+++ b/llm/cli.py
@@ -1,7 +1,9 @@
 import click
 from click_default_group import DefaultGroup
+import datetime
 import openai
 import os
+import sqlite_utils
 import sys
 
 
@@ -21,7 +23,8 @@ def cli():
 @click.option("-4", "--gpt4", is_flag=True, help="Use GPT-4")
 @click.option("-m", "--model", help="Model to use")
 @click.option("-s", "--stream", is_flag=True, help="Stream output")
-def chatgpt(prompt, system, gpt4, model, stream):
+@click.option("-n", "--no-log", is_flag=True, help="Don't log to database")
+def chatgpt(prompt, system, gpt4, model, stream, no_log):
     "Execute prompt against ChatGPT"
     openai.api_key = get_openai_api_key()
     if gpt4:
@@ -33,6 +36,7 @@ def chatgpt(prompt, system, gpt4, model, stream):
         messages.append({"role": "system", "content": system})
     messages.append({"role": "user", "content": prompt})
     if stream:
+        response = []
         for chunk in openai.ChatCompletion.create(
             model=model,
             messages=messages,
@@ -40,23 +44,64 @@ def chatgpt(prompt, system, gpt4, model, stream):
         ):
             content = chunk["choices"][0].get("delta", {}).get("content")
             if content is not None:
+                response.append(content)
                 print(content, end="")
                 sys.stdout.flush()
         print("")
+        log(no_log, "chatgpt", system, prompt, "".join(response), model)
     else:
         response = openai.ChatCompletion.create(
             model=model,
             messages=messages,
         )
-        print(response.choices[0].message.content)
+        content = response.choices[0].message.content
+        print(content)
+        log(no_log, "chatgpt", system, prompt, content, model)
+
+
+@cli.command()
+def init_db():
+    "Ensure ~/.llm/log.db SQLite database exists"
+    path = get_log_db_path()
+    if os.path.exists(path):
+        return
+    # Ensure directory exists
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    db = sqlite_utils.Database(path)
+    db.vacuum()
 
 
 def get_openai_api_key():
     # Expand this to home directory / ~.openai-api-key.txt
     if "OPENAI_API_KEY" in os.environ:
         return os.environ["OPENAI_API_KEY"]
-    path = os.path.expanduser('~/.openai-api-key.txt')
+    path = os.path.expanduser("~/.openai-api-key.txt")
     # If the file exists, read it
     if os.path.exists(path):
         return open(path).read().strip()
-    raise click.ClickException("No OpenAI API key found. Set OPENAI_API_KEY environment variable or create ~/.openai-api-key.txt")
+    raise click.ClickException(
+        "No OpenAI API key found. Set OPENAI_API_KEY environment variable or create ~/.openai-api-key.txt"
+    )
+
+
+def get_log_db_path():
+    return os.path.expanduser("~/.llm/log.db")
+
+
+def log(no_log, provider, system, prompt, response, model):
+    if no_log:
+        return
+    log_path = get_log_db_path()
+    if not os.path.exists(log_path):
+        return
+    db = sqlite_utils.Database(log_path)
+    db["log"].insert(
+        {
+            "provider": provider,
+            "system": system,
+            "prompt": prompt,
+            "response": response,
+            "model": model,
+            "timestamp": str(datetime.datetime.utcnow()),
+        }
+    )
diff --git a/setup.py b/setup.py
index c9313873..ee500ef6 100644
--- a/setup.py
+++ b/setup.py
@@ -31,7 +31,7 @@ def get_long_description():
         [console_scripts]
         llm=llm.cli:cli
     """,
-    install_requires=["click", "openai", "click-default-group-wheel"],
+    install_requires=["click", "openai", "click-default-group-wheel", "sqlite-utils"],
     extras_require={"test": ["pytest"]},
     python_requires=">=3.7",
 )