From ae23c02b2f6dcd4ae20b5fd53a3b3046b9666b0b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 10:34:08 -0800 Subject: [PATCH] fix merge updates --- litellm/proxy/hooks/proxy_track_cost_callback.py | 9 +++------ tests/test_callbacks_on_proxy.py | 1 - 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 2de2b0673..28518a68a 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -13,15 +13,12 @@ from typing import Optional import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy.utils import ( - _get_parent_otel_span_from_kwargs, - get_litellm_metadata_from_kwargs, - log_to_opentelemetry, -) +from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs +from litellm.proxy.utils import log_db_metrics from litellm.types.utils import StandardLoggingPayload -@log_to_opentelemetry +@log_db_metrics async def _PROXY_track_cost_callback( kwargs, # kwargs to completion completion_response: litellm.ModelResponse, # response from completion diff --git a/tests/test_callbacks_on_proxy.py b/tests/test_callbacks_on_proxy.py index 3677a99ad..42665c35b 100644 --- a/tests/test_callbacks_on_proxy.py +++ b/tests/test_callbacks_on_proxy.py @@ -147,7 +147,6 @@ async def test_check_num_callbacks(): @pytest.mark.asyncio @pytest.mark.order2 -@pytest.mark.skip(reason="skipping this test for now") async def test_check_num_callbacks_on_lowest_latency(): """ Test 1: num callbacks should NOT increase over time