Fix test_chat_databricks_invoke with llama 3.3 (#46)

* Fix test_chat_databricks_invoke with llama 3.3 Signed-off-by: Prithvi Kannan <[email protected]> * make test less brittle Signed-off-by: Prithvi Kannan <[email protected]> * lint Signed-off-by: Prithvi Kannan <[email protected]> --------- Signed-off-by: Prithvi Kannan <[email protected]>
langchain-ai · Dec 12, 2024 · 85cdc7d · 85cdc7d
1 parent 0aa4d60
commit 85cdc7d
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/libs/databricks/tests/integration_tests/test_chat_models.py b/libs/databricks/tests/integration_tests/test_chat_models.py
@@ -45,9 +45,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,'"