llm logs -u/--usage option, refs #610

Refs #642 (comment)

docs/help.md

@@ -293,6 +293,7 @@ Options:
   -m, --model TEXT            Filter by model or model alias
   -q, --query TEXT            Search for logs matching this string
   -t, --truncate              Truncate long strings in output
+  -u, --usage                 Include token usage
   -r, --response              Just output the last response
   -c, --current               Show logs from the current conversation
   --cid, --conversation TEXT  Show logs for this conversation ID

llm/cli.py

@@ -33,7 +33,7 @@
 
 from .migrations import migrate
 from .plugins import pm, load_plugins
-from .utils import mimetype_from_path, mimetype_from_string
+from .utils import mimetype_from_path, mimetype_from_string, token_usage_string
 import base64
 import httpx
 import pathlib
@@ -824,6 +824,7 @@ def logs_turn_off():
 @click.option("-m", "--model", help="Filter by model or model alias")
 @click.option("-q", "--query", help="Search for logs matching this string")
 @click.option("-t", "--truncate", is_flag=True, help="Truncate long strings in output")
+@click.option("-u", "--usage", is_flag=True, help="Include token usage")
 @click.option("-r", "--response", is_flag=True, help="Just output the last response")
 @click.option(
     "current_conversation",
@@ -851,6 +852,7 @@ def logs_list(
     model,
     query,
     truncate,
+    usage,
     response,
     current_conversation,
     conversation_id,
@@ -1013,6 +1015,14 @@ def logs_list(
                         )
 
             click.echo("\n## Response:\n\n{}\n".format(row["response"]))
+            if usage:
+                token_usage = token_usage_string(
+                    row["input_tokens"],
+                    row["output_tokens"],
+                    json.loads(row["token_details"]) if row["token_details"] else None,
+                )
+                if token_usage:
+                    click.echo("## Token usage:\n\n{}".format(token_usage))
 
 
 @cli.group(

llm/models.py

@@ -18,7 +18,7 @@
     Set,
     Union,
 )
-from .utils import mimetype_from_path, mimetype_from_string
+from .utils import mimetype_from_path, mimetype_from_string, token_usage_string
 from abc import ABC, abstractmethod
 import json
 from pydantic import BaseModel
@@ -261,14 +261,9 @@ def from_row(cls, db, row):
         return response
 
     def token_usage(self) -> str:
-        bits = []
-        if self.input_tokens is not None:
-            bits.append(f"{self.input_tokens} input")
-        if self.output_tokens is not None:
-            bits.append(f"{self.output_tokens} output")
-        if self.token_details:
-            bits.append(json.dumps(self.token_details))
-        return ", ".join(bits)
+        return token_usage_string(
+            self.input_tokens, self.output_tokens, self.token_details
+        )
 
     def log_to_db(self, db):
         conversation = self.conversation

llm/utils.py

@@ -142,3 +142,14 @@ def remove_empty_and_zero(obj):
         return obj
 
     return remove_empty_and_zero(d) or {}
+
+
+def token_usage_string(input_tokens, output_tokens, token_details) -> str:
+    bits = []
+    if input_tokens is not None:
+        bits.append(f"{input_tokens} input")
+    if output_tokens is not None:
+        bits.append(f"{output_tokens} output")
+    if token_details:
+        bits.append(json.dumps(token_details))
+    return ", ".join(bits)

tests/test_llm.py

@@ -37,6 +37,8 @@ def log_path(user_path):
             "model": "davinci",
             "datetime_utc": (start + datetime.timedelta(seconds=i)).isoformat(),
             "conversation_id": "abc123",
+            "input_tokens": 2,
+            "output_tokens": 5,
         }
         for i in range(100)
     )
@@ -46,9 +48,12 @@ def log_path(user_path):
 datetime_re = re.compile(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}")
 
 
-def test_logs_text(log_path):
+@pytest.mark.parametrize("usage", (False, True))
+def test_logs_text(log_path, usage):
     runner = CliRunner()
     args = ["logs", "-p", str(log_path)]
+    if usage:
+        args.append("-u")
     result = runner.invoke(cli, args, catch_exceptions=False)
     assert result.exit_code == 0
     output = result.output
@@ -64,18 +69,24 @@ def test_logs_text(log_path):
         "system\n\n"
         "## Response:\n\n"
         "response\n\n"
+    ) + ("## Token usage:\n\n2 input, 5 output\n" if usage else "") + (
         "# YYYY-MM-DDTHH:MM:SS    conversation: abc123\n\n"
         "Model: **davinci**\n\n"
         "## Prompt:\n\n"
         "prompt\n\n"
         "## Response:\n\n"
         "response\n\n"
+    ) + (
+        "## Token usage:\n\n2 input, 5 output\n" if usage else ""
+    ) + (
         "# YYYY-MM-DDTHH:MM:SS    conversation: abc123\n\n"
         "Model: **davinci**\n\n"
         "## Prompt:\n\n"
         "prompt\n\n"
         "## Response:\n\n"
         "response\n\n"
+    ) + (
+        "## Token usage:\n\n2 input, 5 output\n" if usage else ""
     )