From c08344f98613057e2649d6cb8127a82c8e959e6c Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Wed, 5 Jul 2023 18:31:38 -0700
Subject: [PATCH] llm logs now decodes JSON for prompt_json etc

---
 llm/cli.py        |  8 ++++++--
 tests/test_llm.py | 26 ++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/llm/cli.py b/llm/cli.py
index 2164367..b9909b8 100644
--- a/llm/cli.py
+++ b/llm/cli.py
@@ -376,10 +376,14 @@ def logs_list(count, path, truncate):
     db = sqlite_utils.Database(path)
     migrate(db)
     rows = list(db["logs"].rows_where(order_by="-id", limit=count or None))
-    if truncate:
-        for row in rows:
+    for row in rows:
+        if truncate:
             row["prompt"] = _truncate_string(row["prompt"])
             row["response"] = _truncate_string(row["response"])
+        # Decode all JSON keys
+        for key in row:
+            if key.endswith("_json") and row[key] is not None:
+                row[key] = json.loads(row[key])
     click.echo(json.dumps(list(rows), indent=2))
 
 
diff --git a/tests/test_llm.py b/tests/test_llm.py
index 9033fed..b300839 100644
--- a/tests/test_llm.py
+++ b/tests/test_llm.py
@@ -111,3 +111,29 @@ def test_llm_default_prompt(mocked_openai, use_stdin, user_path):
         "usage": {},
         "choices": [{"message": {"content": "Bob, Alice, Eve"}}],
     }
+
+    # Test "llm logs"
+    log_result = runner.invoke(cli, ["logs", "-n", "1"], catch_exceptions=False)
+    log_json = json.loads(log_result.output)
+    assert (
+        log_json[0].items()
+        >= {
+            "model": "gpt-3.5-turbo",
+            "prompt": "three names for a pet pelican",
+            "system": None,
+            "prompt_json": {
+                "messages": [
+                    {"role": "user", "content": "three names for a pet pelican"}
+                ]
+            },
+            "options_json": {},
+            "response": "Bob, Alice, Eve",
+            "response_json": {
+                "model": "gpt-3.5-turbo",
+                "usage": {},
+                "choices": [{"message": {"content": "Bob, Alice, Eve"}}],
+            },
+            "reply_to_id": None,
+            "chat_id": None,
+        }.items()
+    )