From ac2a9bd375f0908e3a9d4d4a03b292313a9ef112 Mon Sep 17 00:00:00 2001 From: Martin Reinhardt Date: Tue, 5 Dec 2023 15:28:34 +0100 Subject: [PATCH] Add num_tokens_prompt_total and num_tokens_generated to tests --- tests/test_clients.py | 4 ++-- tests/test_complete.py | 1 + tests/test_error_handling.py | 3 +-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_clients.py b/tests/test_clients.py index 333c55b..10da424 100644 --- a/tests/test_clients.py +++ b/tests/test_clients.py @@ -51,12 +51,12 @@ async def test_nice_flag_on_async_client(httpserver: HTTPServer): httpserver.expect_request( "/complete", query_string={"nice": "true"}, - num_tokens_prompt_total=2, - num_tokens_generated=1, ).respond_with_json( CompletionResponse( "model_version", [CompletionResult(log_probs=[], completion="foo")], + num_tokens_prompt_total=2, + num_tokens_generated=1, ).to_json() ) diff --git a/tests/test_complete.py b/tests/test_complete.py index 88a0dc3..87f726a 100644 --- a/tests/test_complete.py +++ b/tests/test_complete.py @@ -177,6 +177,7 @@ def test_num_tokens_generated_with_best_of(sync_client: Client, model_name: str) response = sync_client.complete(request, model=model_name) completion_result = response.completions[0] + assert completion_result.completion_tokens is not None number_tokens_completion = len(completion_result.completion_tokens) assert response.num_tokens_generated == best_of * number_tokens_completion \ No newline at end of file diff --git a/tests/test_error_handling.py b/tests/test_error_handling.py index 834d204..485f27e 100644 --- a/tests/test_error_handling.py +++ b/tests/test_error_handling.py @@ -111,8 +111,7 @@ def expect_retryable_error( def expect_valid_completion(httpserver: HTTPServer) -> None: httpserver.expect_ordered_request("/complete").respond_with_json( - {"model_version": "1", "completions": []} - ) + {"model_version": "1", "completions": [], "num_tokens_prompt_total": 0, "num_tokens_generated": 0}) def expect_valid_version(httpserver: HTTPServer) -> None: