From eabe323eaa90ad412868474370d694a16f8ceec0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 22 Nov 2024 17:34:08 -0800 Subject: [PATCH] fix ollama --- litellm/llms/ollama.py | 6 +++++- litellm/llms/ollama_chat.py | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index 896b93be5..e9dd2b53f 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -14,6 +14,7 @@ import requests # type: ignore import litellm from litellm import verbose_logger +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.secret_managers.main import get_secret_str from litellm.types.utils import ModelInfo, ProviderField, StreamingChoices @@ -456,7 +457,10 @@ def ollama_completion_stream(url, data, logging_obj): async def ollama_async_streaming(url, data, model_response, encoding, logging_obj): try: - client = httpx.AsyncClient() + _async_http_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders.OLLAMA + ) + client = _async_http_client.client async with client.stream( url=f"{url}", json=data, method="POST", timeout=litellm.request_timeout ) as response: diff --git a/litellm/llms/ollama_chat.py b/litellm/llms/ollama_chat.py index 536f766e0..ce0df139d 100644 --- a/litellm/llms/ollama_chat.py +++ b/litellm/llms/ollama_chat.py @@ -13,6 +13,7 @@ from pydantic import BaseModel import litellm from litellm import verbose_logger +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.ollama import OllamaToolCall, OllamaToolCallFunction from litellm.types.llms.openai import ChatCompletionAssistantToolCall from litellm.types.utils import StreamingChoices @@ -445,7 +446,10 @@ async def ollama_async_streaming( url, api_key, data, model_response, encoding, logging_obj ): try: - client = httpx.AsyncClient() + _async_http_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders.OLLAMA + ) + client = _async_http_client.client _request = { "url": f"{url}", "json": data,