From 729b5b065be7432a9afa7831ccc1edc60c481cac Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Fri, 26 Jul 2024 11:05:09 -0400 Subject: [PATCH] update --- libs/partners/ai21/tests/integration_tests/test_chat_models.py | 3 ++- libs/partners/ai21/tests/integration_tests/test_standard.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/libs/partners/ai21/tests/integration_tests/test_chat_models.py b/libs/partners/ai21/tests/integration_tests/test_chat_models.py index 36181bdfb4373..d0e6a7e09bbfb 100644 --- a/libs/partners/ai21/tests/integration_tests/test_chat_models.py +++ b/libs/partners/ai21/tests/integration_tests/test_chat_models.py @@ -1,10 +1,11 @@ """Test ChatAI21 chat model.""" import pytest -from langchain_ai21.chat_models import ChatAI21 from langchain_core.messages import AIMessageChunk, HumanMessage from langchain_core.outputs import ChatGeneration from langchain_core.rate_limiters import InMemoryRateLimiter + +from langchain_ai21.chat_models import ChatAI21 from tests.unit_tests.conftest import J2_CHAT_MODEL_NAME, JAMBA_CHAT_MODEL_NAME rate_limiter = InMemoryRateLimiter(requests_per_second=0.5) diff --git a/libs/partners/ai21/tests/integration_tests/test_standard.py b/libs/partners/ai21/tests/integration_tests/test_standard.py index 20f73daae3265..5896573102dca 100644 --- a/libs/partners/ai21/tests/integration_tests/test_standard.py +++ b/libs/partners/ai21/tests/integration_tests/test_standard.py @@ -5,10 +5,10 @@ import pytest from langchain_core.language_models import BaseChatModel +from langchain_core.rate_limiters import InMemoryRateLimiter from langchain_standard_tests.integration_tests import ChatModelIntegrationTests from langchain_ai21 import ChatAI21 -from langchain_core.rate_limiters import InMemoryRateLimiter rate_limiter = InMemoryRateLimiter(requests_per_second=0.5)