diff --git a/tests/local_testing/test_pass_through_endpoints.py b/tests/local_testing/test_pass_through_endpoints.py
index cc2c5b9b28..282c264883 100644
--- a/tests/local_testing/test_pass_through_endpoints.py
+++ b/tests/local_testing/test_pass_through_endpoints.py
@@ -14,7 +14,7 @@ import asyncio
 from unittest.mock import Mock
 
 import httpx
-
+import litellm_proxy
 from litellm_proxy.proxy_server import initialize_pass_through_endpoints
 
 
diff --git a/tests/logging_callback_tests/test_opentelemetry_unit_tests.py b/tests/logging_callback_tests/test_opentelemetry_unit_tests.py
index a9a1652353..7e0881968a 100644
--- a/tests/logging_callback_tests/test_opentelemetry_unit_tests.py
+++ b/tests/logging_callback_tests/test_opentelemetry_unit_tests.py
@@ -26,7 +26,7 @@ class TestOpentelemetryUnitTests(BaseLoggingCallbackTest):
     def test_parallel_tool_calls(self, mock_response_obj: ModelResponse):
         tool_calls = mock_response_obj.choices[0].message.tool_calls
         from litellm.integrations.opentelemetry import OpenTelemetry
-        from litellm_proxy._types import SpanAttributes
+        from litellm.types.integrations.opentelemetry import SpanAttributes
 
         kv_pair_dict = OpenTelemetry._tool_calls_kv_pair(tool_calls)
 
diff --git a/tests/logging_callback_tests/test_otel_logging.py b/tests/logging_callback_tests/test_otel_logging.py
index b8f0aba1d9..ce7d328d5e 100644
--- a/tests/logging_callback_tests/test_otel_logging.py
+++ b/tests/logging_callback_tests/test_otel_logging.py
@@ -15,7 +15,7 @@ import asyncio
 import logging
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from litellm._logging import verbose_logger
-from litellm_proxy._types import SpanAttributes
+from litellm.types.integrations.opentelemetry import SpanAttributes
 
 verbose_logger.setLevel(logging.DEBUG)