diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 879914e5ac6..515a6e1011b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1892,10 +1892,9 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False): def cost_tracking(): global prisma_client if prisma_client is not None: - litellm.logging_callback_manager.add_litellm_callback(_ProxyDBLogger()) - litellm.logging_callback_manager.add_litellm_async_success_callback( - _ProxyDBLogger() - ) + db_logger = _ProxyDBLogger() + litellm.logging_callback_manager.add_litellm_callback(db_logger) + litellm.logging_callback_manager.add_litellm_async_success_callback(db_logger) async def get_current_spend(counter_key: str, fallback_spend: float) -> float: @@ -2710,11 +2709,9 @@ def run_ollama_serve(): with open(os.devnull, "w") as devnull: subprocess.Popen(command, stdout=devnull, stderr=devnull) except Exception as e: - verbose_proxy_logger.debug( - f""" + verbose_proxy_logger.debug(f""" LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve` - """ - ) + """) def _get_process_rss_mb() -> Optional[float]: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 6d10d2a6353..4d9a7e690f5 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7487,3 +7487,36 @@ class TestSortModelsByDisplayName: all_models=models, sort_by="model_name", sort_order="asc" ) assert [m["model_name"] for m in sorted_models] == ["alpha", "beta"] + + +class TestCostTrackingSingleDBLogger: + def test_cost_tracking_creates_single_db_logger_instance(self): + from unittest.mock import MagicMock, patch + + from litellm.proxy.proxy_server import cost_tracking + + mock_callback_manager = MagicMock() + + with patch("litellm.proxy.proxy_server.prisma_client", new="fake-client"): + with patch("litellm.logging_callback_manager", mock_callback_manager): + cost_tracking() + + add_callback_call = mock_callback_manager.add_litellm_callback.call_args[0][0] + add_async_call = ( + mock_callback_manager.add_litellm_async_success_callback.call_args[0][0] + ) + assert add_callback_call is add_async_call + + def test_cost_tracking_noop_when_prisma_client_is_none(self): + from unittest.mock import MagicMock, patch + + from litellm.proxy.proxy_server import cost_tracking + + mock_callback_manager = MagicMock() + + with patch("litellm.proxy.proxy_server.prisma_client", new=None): + with patch("litellm.logging_callback_manager", mock_callback_manager): + cost_tracking() + + mock_callback_manager.add_litellm_callback.assert_not_called() + mock_callback_manager.add_litellm_async_success_callback.assert_not_called()