Check num_prompt_tokens

h2oai · Mar 29, 2024 · 51fe7d8 · 51fe7d8
1 parent 51fbc95
commit 51fe7d8
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tests/test_client_calls.py b/tests/test_client_calls.py
@@ -4882,11 +4882,14 @@ def test_client1_image_qa(langchain_mode, base_model):
             raise
 
     # string of dict for output
-    response = ast.literal_eval(res)['response']
+    res_dict = ast.literal_eval(res)
+    response = res_dict['response']
     print('base_model: %s langchain_mode: %s response: %s' % (base_model, langchain_mode, response), file=sys.stderr)
     print(response)
     assert 'license' in response.lower()
 
+    assert res_dict['save_dict']['extra_dict']['num_prompt_tokens'] > 1000
+
 
 @wrap_test_forked
 @pytest.mark.parametrize("base_model", vision_models)
@@ -4936,6 +4939,8 @@ def test_client1_images_qa(langchain_mode, base_model):
     print(response)
     assert 'REINFORCE'.lower() in response.lower()
 
+    assert res_dict['save_dict']['extra_dict']['num_prompt_tokens'] > 1000
+
 
 @wrap_test_forked
 def test_pdf_to_base_64_images():