From 8308fe5cbffc8d703cd75f59934214274adfa5dd Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Fri, 16 Jun 2023 08:47:15 +0100 Subject: [PATCH] Store debug info, closes #34 --- llm/cli.py | 19 ++++++++++++++++--- llm/migrations.py | 6 ++++++ tests/conftest.py | 13 +++++++++++++ tests/test_keys.py | 11 ++++------- tests/test_llm.py | 9 ++------- tests/test_migrate.py | 44 +++++++++++++++++++++---------------------- 6 files changed, 63 insertions(+), 39 deletions(-) diff --git a/llm/cli.py b/llm/cli.py index 320be7b..2786237 100644 --- a/llm/cli.py +++ b/llm/cli.py @@ -8,6 +8,7 @@ import os import pathlib import sqlite_utils import sys +import time import warnings warnings.simplefilter("ignore", ResourceWarning) @@ -79,28 +80,34 @@ def prompt(prompt, system, model, no_stream, no_log, _continue, chat_id, key): # Resolve model aliases model = MODEL_ALIASES.get(model, model) try: + debug = {} if no_stream: + start = time.time() response = openai.ChatCompletion.create( model=model, messages=messages, ) + debug["model"] = response.model + debug["usage"] = response.usage content = response.choices[0].message.content - log(no_log, system, prompt, content, model, chat_id) + log(no_log, system, prompt, content, model, chat_id, debug, start) print(content) else: + start = time.time() response = [] for chunk in openai.ChatCompletion.create( model=model, messages=messages, stream=True, ): + debug["model"] = chunk.model 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, system, prompt, "".join(response), model, chat_id) + log(no_log, system, prompt, "".join(response), model, chat_id, debug, start) except openai.error.AuthenticationError as ex: raise click.ClickException("{}: {}".format(ex.error.type, ex.error.code)) except openai.error.OpenAIError as ex: @@ -251,7 +258,11 @@ def log_db_path(): return user_dir() / "log.db" -def log(no_log, system, prompt, response, model, chat_id=None): +def log(no_log, system, prompt, response, model, chat_id=None, debug=None, start=None): + duration_ms = None + if start is not None: + end = time.time() + duration_ms = int((end - start) * 1000) if no_log: return log_path = log_db_path() @@ -267,6 +278,8 @@ def log(no_log, system, prompt, response, model, chat_id=None): "response": response, "model": model, "timestamp": str(datetime.datetime.utcnow()), + "debug": debug, + "duration_ms": duration_ms, }, ) diff --git a/llm/migrations.py b/llm/migrations.py index dafb18d..4d76439 100644 --- a/llm/migrations.py +++ b/llm/migrations.py @@ -78,3 +78,9 @@ def m004_column_order(db): @migration def m004_drop_provider(db): db["log"].transform(drop=("provider",)) + + +@migration +def m005_debug(db): + db["log"].add_column("debug", str) + db["log"].add_column("duration_ms", int) diff --git a/tests/conftest.py b/tests/conftest.py index 44bf4ce..c0ea817 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,3 +15,16 @@ def keys_path(tmpdir): def env_setup(monkeypatch, log_path, keys_path): monkeypatch.setenv("LLM_KEYS_PATH", str(keys_path)) monkeypatch.setenv("LLM_LOG_PATH", str(log_path)) + + +@pytest.fixture +def mocked_openai(requests_mock): + return requests_mock.post( + "https://api.openai.com/v1/chat/completions", + json={ + "model": "gpt-3.5-turbo", + "usage": {}, + "choices": [{"message": {"content": "Bob, Alice, Eve"}}], + }, + headers={"Content-Type": "application/json"}, + ) diff --git a/tests/test_keys.py b/tests/test_keys.py index a990b07..0d84de9 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -31,7 +31,7 @@ def test_keys_set(monkeypatch, tmpdir): } -def test_uses_correct_key(requests_mock, monkeypatch, tmpdir): +def test_uses_correct_key(mocked_openai, monkeypatch, tmpdir): keys_path = tmpdir / "keys.json" keys_path.write_text( json.dumps( @@ -44,14 +44,11 @@ def test_uses_correct_key(requests_mock, monkeypatch, tmpdir): ) monkeypatch.setenv("LLM_KEYS_PATH", str(keys_path)) monkeypatch.setenv("OPENAI_API_KEY", "from-env") - mocked = requests_mock.post( - "https://api.openai.com/v1/chat/completions", - json={"choices": [{"message": {"content": "Bob, Alice, Eve"}}]}, - headers={"Content-Type": "application/json"}, - ) def assert_key(key): - assert mocked.last_request.headers["Authorization"] == "Bearer {}".format(key) + assert mocked_openai.last_request.headers[ + "Authorization" + ] == "Bearer {}".format(key) runner = CliRunner() # Called without --key uses environment variable diff --git a/tests/test_llm.py b/tests/test_llm.py index 444e647..33093dc 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -61,15 +61,10 @@ def test_logs_path(monkeypatch, env, log_path): @mock.patch.dict(os.environ, {"OPENAI_API_KEY": "X"}) @pytest.mark.parametrize("use_stdin", (True, False)) -def test_llm_default_prompt(requests_mock, use_stdin, log_path): +def test_llm_default_prompt(mocked_openai, use_stdin, log_path): # Reset the log_path database log_db = sqlite_utils.Database(str(log_path)) log_db["log"].delete_where() - mocked = requests_mock.post( - "https://api.openai.com/v1/chat/completions", - json={"choices": [{"message": {"content": "Bob, Alice, Eve"}}]}, - headers={"Content-Type": "application/json"}, - ) runner = CliRunner() prompt = "three names for a pet pelican" input = None @@ -81,7 +76,7 @@ def test_llm_default_prompt(requests_mock, use_stdin, log_path): result = runner.invoke(cli, args, input=input, catch_exceptions=False) assert result.exit_code == 0 assert result.output == "Bob, Alice, Eve\n" - assert mocked.last_request.headers["Authorization"] == "Bearer X" + assert mocked_openai.last_request.headers["Authorization"] == "Bearer X" # Was it logged? rows = list(log_db["log"].rows) assert len(rows) == 1 diff --git a/tests/test_migrate.py b/tests/test_migrate.py index 1cfd3fc..670eb8f 100644 --- a/tests/test_migrate.py +++ b/tests/test_migrate.py @@ -6,17 +6,17 @@ def test_migrate_blank(): db = sqlite_utils.Database(memory=True) migrate(db) assert set(db.table_names()) == {"_llm_migrations", "log"} - assert db["log"].schema == ( - 'CREATE TABLE "log" (\n' - " [id] INTEGER PRIMARY KEY,\n" - " [model] TEXT,\n" - " [timestamp] TEXT,\n" - " [prompt] TEXT,\n" - " [system] TEXT,\n" - " [response] TEXT,\n" - " [chat_id] INTEGER REFERENCES [log]([id])\n" - ")" - ) + assert db["log"].columns_dict == { + "id": int, + "model": str, + "timestamp": str, + "prompt": str, + "system": str, + "response": str, + "chat_id": int, + "debug": str, + "duration_ms": int, + } def test_migrate_from_original_schema(): @@ -35,14 +35,14 @@ def test_migrate_from_original_schema(): migrate(db) assert set(db.table_names()) == {"_llm_migrations", "log"} schema = db["log"].schema - assert schema == ( - 'CREATE TABLE "log" (\n' - " [id] INTEGER PRIMARY KEY,\n" - " [model] TEXT,\n" - " [timestamp] TEXT,\n" - " [prompt] TEXT,\n" - " [system] TEXT,\n" - " [response] TEXT,\n" - " [chat_id] INTEGER REFERENCES [log]([id])\n" - ")" - ) + assert db["log"].columns_dict == { + "id": int, + "model": str, + "timestamp": str, + "prompt": str, + "system": str, + "response": str, + "chat_id": int, + "debug": str, + "duration_ms": int, + }