Fix test_chat_databricks_invoke with llama 3.3 (#29)

* Fix test_chat_databricks_invoke with llama 3.3 Signed-off-by: Prithvi Kannan <[email protected]> * fix Signed-off-by: Prithvi Kannan <[email protected]> --------- Signed-off-by: Prithvi Kannan <[email protected]>
databricks · Dec 12, 2024 · c9946ad · c9946ad
1 parent 681d59d
commit c9946ad
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/integrations/langchain/tests/integration_tests/test_chat_models.py b/integrations/langchain/tests/integration_tests/test_chat_models.py
@@ -43,9 +43,13 @@ def test_chat_databricks_invoke():
     response = chat.invoke("How to learn Java? Start the response by 'To learn Java,'")
     assert isinstance(response, AIMessage)
     assert response.content == "To learn "
-    assert response.response_metadata["prompt_tokens"] == 24
-    assert response.response_metadata["completion_tokens"] == 3
-    assert response.response_metadata["total_tokens"] == 27
+    assert 20 <= response.response_metadata["prompt_tokens"] <= 30
+    assert 1 <= response.response_metadata["completion_tokens"] <= 10
+    expected_total = (
+        response.response_metadata["prompt_tokens"]
+        + response.response_metadata["completion_tokens"]
+    )
+    assert response.response_metadata["total_tokens"] == expected_total
 
     response = chat.invoke("How to learn Python? Start the response by 'To learn Python,'")
     assert response.content.startswith("To learn Python,")