diff --git a/docs/help.md b/docs/help.md index 34f429b8..ba5d3f0d 100644 --- a/docs/help.md +++ b/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 diff --git a/llm/cli.py b/llm/cli.py index 22699a92..f78840de 100644 --- a/llm/cli.py +++ b/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( diff --git a/llm/models.py b/llm/models.py index 759f9f76..c4a00db6 100644 --- a/llm/models.py +++ b/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 diff --git a/llm/utils.py b/llm/utils.py index 20510006..a47984d2 100644 --- a/llm/utils.py +++ b/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) diff --git a/tests/test_llm.py b/tests/test_llm.py index 0e54cc91..49795ca2 100644 --- a/tests/test_llm.py +++ b/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 "" )