diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py
index baa3570bfd..d0adbb663f 100644
--- a/litellm/types/proxy/management_endpoints/ui_sso.py
+++ b/litellm/types/proxy/management_endpoints/ui_sso.py
@@ -2,7 +2,7 @@ from typing import List, Literal, Optional, TypedDict
 
 from pydantic import Field
 
-from litellm_proxy._types import LiteLLMPydanticObjectBase, LitellmUserRoles
+from litellm.types.utils import LiteLLMPydanticObjectBase
 
 
 class MicrosoftGraphAPIUserGroupDirectoryObject(TypedDict, total=False):
diff --git a/tests/local_testing/test_config.py b/tests/local_testing/test_config.py
index 6aa29d3af2..15cf762543 100644
--- a/tests/local_testing/test_config.py
+++ b/tests/local_testing/test_config.py
@@ -21,6 +21,7 @@ import pytest
 from pydantic import BaseModel, ConfigDict
 
 import litellm
+import litellm_proxy
 from litellm_proxy.common_utils.encrypt_decrypt_utils import encrypt_value
 from litellm_proxy.proxy_server import ProxyConfig
 from litellm_proxy.utils import DualCache, ProxyLogging