Fix litellm max callbacks override (#180)

Author: nabinchha

src/data_designer/engine/models/litellm_overrides.py

@@ -10,6 +10,7 @@
 import litellm
 from litellm import RetryPolicy
 from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
 from litellm.router import Router
 from pydantic import BaseModel, Field
 from typing_extensions import override
@@ -154,7 +155,7 @@ def apply_litellm_patches():
     litellm.in_memory_llm_clients_cache = ThreadSafeCache()
 
     # Workaround for the litellm issue described in https://github.com/BerriAI/litellm/issues/9792
-    litellm.litellm_core_utils.logging_callback_manager.LoggingCallbackManager.MAX_CALLBACKS = DEFAULT_MAX_CALLBACKS
+    LoggingCallbackManager.MAX_CALLBACKS = DEFAULT_MAX_CALLBACKS
 
     quiet_noisy_logger("httpx")
     quiet_noisy_logger("LiteLLM")

tests/engine/models/test_litellm_overrides.py

@@ -47,6 +47,13 @@ def test_custom_router_calculate_exponential_backoff(retry_count: int, jitter: f
         )
 
 
+def test_apply_litellm_patches_no_exceptions():
+    try:
+        apply_litellm_patches()
+    except Exception as e:
+        pytest.fail(f"apply_litellm_patches() raised an unexpected exception: {e}")
+
+
 @patch("data_designer.engine.models.litellm_overrides.quiet_noisy_logger", autospec=True)
 def test_apply_litellm_patches(mock_quiet_noisy_logger):
     apply_litellm_patches()