diff --git a/litellm/__init__.py b/litellm/__init__.py index 80625e0b189..389ee4d4792 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -26,7 +26,6 @@ from typing import ( ) from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams from litellm.types.integrations.datadog import DatadogInitParams -from litellm.caching.llm_caching_handler import LLMClientCache from litellm.types.llms.bedrock import COHERE_EMBEDDING_INPUT_TYPES from litellm.types.utils import ( ImageObject, @@ -285,7 +284,7 @@ disable_token_counter: bool = False disable_add_transform_inline_image_block: bool = False disable_add_user_agent_to_request_tags: bool = False extra_spend_tag_headers: Optional[List[str]] = None -in_memory_llm_clients_cache: LLMClientCache = LLMClientCache() +in_memory_llm_clients_cache: "LLMClientCache" safe_memory_mode: bool = False enable_azure_ad_token_refresh: Optional[bool] = False ### DEFAULT AZURE API VERSION ### @@ -1515,6 +1514,7 @@ if TYPE_CHECKING: from litellm.types.utils import ModelInfo as _ModelInfoType from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.caching.caching import Cache + from litellm.caching.llm_caching_handler import LLMClientCache # Cost calculator functions cost_per_token: Callable[..., Tuple[float, float]] @@ -1567,6 +1567,7 @@ def __getattr__(name: str) -> Any: LITELLM_LOGGING_NAMES, UTILS_NAMES, TOKEN_COUNTER_NAMES, + LLM_CLIENT_CACHE_NAMES, CACHING_NAMES, HTTP_HANDLER_NAMES, ) @@ -1591,6 +1592,11 @@ def __getattr__(name: str) -> Any: from ._lazy_imports import _lazy_import_token_counter return _lazy_import_token_counter(name) + # Lazy load LLM client cache and its singleton + if name in LLM_CLIENT_CACHE_NAMES: + from ._lazy_imports import _lazy_import_llm_client_cache + return _lazy_import_llm_client_cache(name) + # Lazy load caching classes if name in CACHING_NAMES: from ._lazy_imports import _lazy_import_caching diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 1fbf3f1be0f..a639e5b5b5e 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -39,6 +39,12 @@ TOKEN_COUNTER_NAMES = ( "get_modified_max_tokens", ) +# LLM client cache names that support lazy loading via _lazy_import_llm_client_cache +LLM_CLIENT_CACHE_NAMES = ( + "LLMClientCache", + "in_memory_llm_clients_cache", +) + # Caching / cache classes that support lazy loading via _lazy_import_caching CACHING_NAMES = ( "Cache", @@ -330,6 +336,28 @@ def _lazy_import_caching(name: str) -> Any: raise AttributeError(f"Caching lazy import: unknown attribute {name!r}") +def _lazy_import_llm_client_cache(name: str) -> Any: + """Lazy import for LLM client cache class and singleton.""" + _globals = _get_litellm_globals() + + if name == "LLMClientCache": + from litellm.caching.llm_caching_handler import LLMClientCache as _LLMClientCache + + _globals["LLMClientCache"] = _LLMClientCache + return _LLMClientCache + + if name == "in_memory_llm_clients_cache": + from litellm.caching.llm_caching_handler import LLMClientCache as _LLMClientCache + + instance = _LLMClientCache() + # Only populate the requested singleton name to keep lazy-import + # semantics consistent with other helpers (no extra symbols). + _globals["in_memory_llm_clients_cache"] = instance + return instance + + raise AttributeError(f"LLM client cache lazy import: unknown attribute {name!r}") + + def _lazy_import_litellm_logging(name: str) -> Any: """Lazy import for litellm_logging module.""" _globals = _get_litellm_globals() diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 6f46933cf9f..9352eab6c8e 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,12 +6,11 @@ This module provides fake streaming by converting non-streaming responses into s """ import asyncio -from typing import Any, AsyncIterator, Dict +from typing import Any, AsyncIterator, Dict, cast from uuid import uuid4 -import httpx - from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client class PydanticAITransformation: @@ -78,7 +77,7 @@ class PydanticAITransformation: @staticmethod async def _poll_for_completion( - client: httpx.AsyncClient, + client: AsyncHTTPHandler, endpoint: str, task_id: str, request_id: str, @@ -179,34 +178,37 @@ class PydanticAITransformation: f"Pydantic AI: Sending non-streaming request to {endpoint}" ) - # Send request to Pydantic AI agent - async with httpx.AsyncClient(timeout=timeout) as client: - response = await client.post( - endpoint, - json=a2a_request, - headers={"Content-Type": "application/json"}, - ) - response.raise_for_status() - response_data = response.json() - - # Check if task is already completed - result = response_data.get("result", {}) - status = result.get("status", {}) - state = status.get("state", "") - - if state != "completed": - # Need to poll for completion - task_id = result.get("id") - if task_id: - verbose_logger.info( - f"Pydantic AI: Task {task_id} submitted, polling for completion..." - ) - response_data = await PydanticAITransformation._poll_for_completion( - client=client, - endpoint=endpoint, - task_id=task_id, - request_id=request_id, - ) + # Send request to Pydantic AI agent using shared async HTTP client + client = get_async_httpx_client( + llm_provider=cast(Any, "pydantic_ai_agent"), + params={"timeout": timeout}, + ) + response = await client.post( + endpoint, + json=a2a_request, + headers={"Content-Type": "application/json"}, + ) + response.raise_for_status() + response_data = response.json() + + # Check if task is already completed + result = response_data.get("result", {}) + status = result.get("status", {}) + state = status.get("state", "") + + if state != "completed": + # Need to poll for completion + task_id = result.get("id") + if task_id: + verbose_logger.info( + f"Pydantic AI: Task {task_id} submitted, polling for completion..." + ) + response_data = await PydanticAITransformation._poll_for_completion( + client=client, + endpoint=endpoint, + task_id=task_id, + request_id=request_id, + ) verbose_logger.info(f"Pydantic AI: Received completed response for request_id={request_id}") diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 6488128b215..6771999cd35 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -16,7 +16,6 @@ from typing import ( from pydantic import BaseModel from litellm._logging import verbose_logger -from litellm.caching.caching import DualCache from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER from litellm.types.integrations.argilla import ArgillaItem from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest @@ -33,6 +32,7 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + from litellm.caching.caching import DualCache from opentelemetry.trace import Span as _Span from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -334,7 +334,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, - cache: DualCache, + cache: "DualCache", data: dict, call_type: CallTypesLiteral, ) -> Optional[ diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 5697700b46d..1b4d04af80a 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1153,7 +1153,17 @@ def get_async_httpx_client( pass _cache_key_name = "async_httpx_client" + _params_key_name + llm_provider - _cached_client = litellm.in_memory_llm_clients_cache.get_cache(_cache_key_name) + + # Lazily initialize the global in-memory client cache to avoid relying on + # litellm globals being fully populated during import time. + cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if cache is None: + from litellm.caching.llm_caching_handler import LLMClientCache + + cache = LLMClientCache() + setattr(litellm, "in_memory_llm_clients_cache", cache) + + _cached_client = cache.get_cache(_cache_key_name) if _cached_client: return _cached_client @@ -1166,7 +1176,7 @@ def get_async_httpx_client( shared_session=shared_session, ) - litellm.in_memory_llm_clients_cache.set_cache( + cache.set_cache( key=_cache_key_name, value=_new_client, ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, @@ -1191,7 +1201,16 @@ def _get_httpx_client(params: Optional[dict] = None) -> HTTPHandler: _cache_key_name = "httpx_client" + _params_key_name - _cached_client = litellm.in_memory_llm_clients_cache.get_cache(_cache_key_name) + # Lazily initialize the global in-memory client cache to avoid relying on + # litellm globals being fully populated during import time. + cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if cache is None: + from litellm.caching.llm_caching_handler import LLMClientCache + + cache = LLMClientCache() + setattr(litellm, "in_memory_llm_clients_cache", cache) + + _cached_client = cache.get_cache(_cache_key_name) if _cached_client: return _cached_client @@ -1200,7 +1219,7 @@ def _get_httpx_client(params: Optional[dict] = None) -> HTTPHandler: else: _new_client = HTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0)) - litellm.in_memory_llm_clients_cache.set_cache( + cache.set_cache( key=_cache_key_name, value=_new_client, ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py index 08737df4e4a..15eb7a5be6c 100644 --- a/tests/test_litellm/test_lazy_imports.py +++ b/tests/test_litellm/test_lazy_imports.py @@ -14,12 +14,14 @@ from litellm._lazy_imports import ( UTILS_NAMES, TOKEN_COUNTER_NAMES, CACHING_NAMES, + LLM_CLIENT_CACHE_NAMES, HTTP_HANDLER_NAMES, _lazy_import_cost_calculator, _lazy_import_litellm_logging, _lazy_import_utils, _lazy_import_token_counter, _lazy_import_caching, + _lazy_import_llm_client_cache, _lazy_import_http_handlers, ) @@ -111,6 +113,18 @@ def test_token_counter_lazy_imports(): _verify_only_requested_name_imported(name, TOKEN_COUNTER_NAMES) +def test_llm_client_cache_lazy_imports(): + """Test that LLM client cache class and singleton can be lazy imported.""" + for name in LLM_CLIENT_CACHE_NAMES: + _clear_names_from_globals(LLM_CLIENT_CACHE_NAMES) + + obj = _lazy_import_llm_client_cache(name) + assert obj is not None + assert name in litellm.__dict__ + + _verify_only_requested_name_imported(name, LLM_CLIENT_CACHE_NAMES) + + def test_http_handler_lazy_imports(): """Test that HTTP handler singletons can be lazy imported.""" for name in HTTP_HANDLER_NAMES: @@ -140,3 +154,6 @@ def test_unknown_attribute_raises_error(): with pytest.raises(AttributeError): _lazy_import_token_counter("unknown") + with pytest.raises(AttributeError): + _lazy_import_llm_client_cache("unknown") +