From 7cb69ddcb961beff2d7fbc86a8d00754191ff6ea Mon Sep 17 00:00:00 2001 From: Prithvi Kannan Date: Thu, 12 Dec 2024 00:16:31 -0800 Subject: [PATCH 1/2] Fix test_chat_databricks_invoke with llama 3.3 Signed-off-by: Prithvi Kannan --- .../langchain/tests/integration_tests/test_chat_models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/integrations/langchain/tests/integration_tests/test_chat_models.py b/integrations/langchain/tests/integration_tests/test_chat_models.py index cbce738..06cb3f1 100644 --- a/integrations/langchain/tests/integration_tests/test_chat_models.py +++ b/integrations/langchain/tests/integration_tests/test_chat_models.py @@ -43,9 +43,9 @@ 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["prompt_tokens"] == 25 assert response.response_metadata["completion_tokens"] == 3 - assert response.response_metadata["total_tokens"] == 27 + assert response.response_metadata["total_tokens"] == 28 response = chat.invoke("How to learn Python? Start the response by 'To learn Python,'") assert response.content.startswith("To learn Python,") From d4ccfe7ebf142512125e3fbd159a8cd2e8b22ad6 Mon Sep 17 00:00:00 2001 From: Prithvi Kannan Date: Thu, 12 Dec 2024 00:34:08 -0800 Subject: [PATCH 2/2] fix Signed-off-by: Prithvi Kannan --- .../tests/integration_tests/test_chat_models.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/integrations/langchain/tests/integration_tests/test_chat_models.py b/integrations/langchain/tests/integration_tests/test_chat_models.py index 06cb3f1..ab67e3d 100644 --- 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"] == 25 - assert response.response_metadata["completion_tokens"] == 3 - assert response.response_metadata["total_tokens"] == 28 + 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,")