Add num_tokens_prompt_total and num_tokens_generated to tests

Aleph-Alpha · Dec 5, 2023 · ac2a9bd · ac2a9bd
1 parent e7d6129
commit ac2a9bd
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git 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
@@ -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
diff --git 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: