From fdaee84b827d69e79aec71f85ffdd208beae8d54 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 21 Nov 2024 09:40:26 -0800 Subject: [PATCH] fix TEXT_COMPLETION_CODESTRAL --- litellm/llms/text_completion_codestral.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/litellm/llms/text_completion_codestral.py b/litellm/llms/text_completion_codestral.py index 21582d26c..d3c1ae3cb 100644 --- a/litellm/llms/text_completion_codestral.py +++ b/litellm/llms/text_completion_codestral.py @@ -18,7 +18,10 @@ import litellm from litellm import verbose_logger from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, +) from litellm.types.llms.databricks import GenericStreamingChunk from litellm.utils import ( Choices, @@ -479,8 +482,9 @@ class CodestralTextCompletion(BaseLLM): headers={}, ) -> TextCompletionResponse: - async_handler = AsyncHTTPHandler( - timeout=httpx.Timeout(timeout=timeout), concurrent_limit=1 + async_handler = get_async_httpx_client( + llm_provider=litellm.LlmProviders.TEXT_COMPLETION_CODESTRAL, + params={"timeout": timeout}, ) try: