From fa67b3fdaf413f7c51397f012a45cc0bea5d03a0 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 11 Jul 2023 20:18:16 -0700
Subject: [PATCH] --log option, closes #68

---
 docs/help.md      |  1 +
 docs/logging.md   | 13 ++++++++-----
 llm/cli.py        |  7 ++++++-
 tests/test_llm.py | 19 ++++++++++++++++---
 4 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/docs/help.md b/docs/help.md
index 9f089bb..a11a6ba 100644
--- a/docs/help.md
+++ b/docs/help.md
@@ -82,6 +82,7 @@ Options:
   -p, --param <TEXT TEXT>...   Parameters for template
   --no-stream                  Do not stream output
   -n, --no-log                 Don't log to database
+  --log                        Log prompt and response to the database
   -c, --continue               Continue the most recent conversation.
   --cid, --conversation TEXT   Continue the conversation with the given ID.
   --key TEXT                   API key to use
diff --git a/docs/logging.md b/docs/logging.md
index afc7c9f..28c1182 100644
--- a/docs/logging.md
+++ b/docs/logging.md
@@ -24,13 +24,16 @@ To turn logging by default off:
 ```bash
 llm logs off
 ```
-To turn it back on again:
+If you've turned off logging you can still log an individual prompt and response by adding `--log`:
+```bash
+llm 'Five ambitious names for a pet pterodactyl' --log
+```
+To turn logging by default back on again:
 
 ```bash
 llm logs on
 ```
-
-To see the status of that database, run this:
+To see the status of the logs database, run this:
 ```bash
 llm logs status
 ```
@@ -38,8 +41,8 @@ Example output:
 ```
 Logging is ON for all prompts
 Found log database at /Users/simon/Library/Application Support/io.datasette.llm/logs.db
-Number of conversations logged: 32
-Number of responses logged:     47
+Number of conversations logged: 33
+Number of responses logged:     48
 Database file size:             19.96MB
 ```
 
diff --git a/llm/cli.py b/llm/cli.py
index d30e025..e0de029 100644
--- a/llm/cli.py
+++ b/llm/cli.py
@@ -80,6 +80,7 @@ def cli():
 )
 @click.option("--no-stream", is_flag=True, help="Do not stream output")
 @click.option("-n", "--no-log", is_flag=True, help="Don't log to database")
+@click.option("--log", is_flag=True, help="Log prompt and response to the database")
 @click.option(
     "_continue",
     "-c",
@@ -105,6 +106,7 @@ def prompt(
     param,
     no_stream,
     no_log,
+    log,
     _continue,
     conversation_id,
     key,
@@ -115,6 +117,9 @@ def prompt(
 
     Documentation: https://llm.datasette.io/en/stable/usage.html
     """
+    if log and no_log:
+        raise click.ClickException("--log and --no-log are mutually exclusive")
+
     model_aliases = get_model_aliases()
 
     def read_prompt():
@@ -247,7 +252,7 @@ def prompt(
         raise click.ClickException(str(ex))
 
     # Log to the database
-    if logs_on():
+    if (logs_on() or log) and not no_log:
         log_path = logs_db_path()
         db = sqlite_utils.Database(log_path)
         migrate(db)
diff --git a/tests/test_llm.py b/tests/test_llm.py
index e93b583..1dec322 100644
--- a/tests/test_llm.py
+++ b/tests/test_llm.py
@@ -65,8 +65,20 @@ def test_logs_path(monkeypatch, env, user_path):
 
 @mock.patch.dict(os.environ, {"OPENAI_API_KEY": "X"})
 @pytest.mark.parametrize("use_stdin", (True, False))
-@pytest.mark.parametrize("logs_off", (True, False))
-def test_llm_default_prompt(mocked_openai, use_stdin, user_path, logs_off):
+@pytest.mark.parametrize(
+    "logs_off,logs_args,should_log",
+    (
+        (True, [], False),
+        (False, [], True),
+        (False, ["--no-log"], False),
+        (False, ["--log"], True),
+        (True, ["-n"], False),  # Short for --no-log
+        (True, ["--log"], True),
+    ),
+)
+def test_llm_default_prompt(
+    mocked_openai, use_stdin, user_path, logs_off, logs_args, should_log
+):
     # Reset the log_path database
     log_path = user_path / "logs.db"
     log_db = sqlite_utils.Database(str(log_path))
@@ -92,6 +104,7 @@ def test_llm_default_prompt(mocked_openai, use_stdin, user_path, logs_off):
         input = prompt
     else:
         args.append(prompt)
+    args += logs_args
     result = runner.invoke(cli, args, input=input, catch_exceptions=False)
     assert result.exit_code == 0
     assert result.output == "Bob, Alice, Eve\n"
@@ -100,7 +113,7 @@ def test_llm_default_prompt(mocked_openai, use_stdin, user_path, logs_off):
     # Was it logged?
     rows = list(log_db["responses"].rows)
 
-    if logs_off:
+    if not should_log:
         assert len(rows) == 0
         return