Merge branch 'main' into litellm_reliability_fix_db_txs

.circleci/requirements.txt

@@ -10,6 +10,6 @@ anthropic
 orjson==3.9.15
 pydantic==2.10.2
 google-cloud-aiplatform==1.43.0
-fastapi-sso==0.10.0
+fastapi-sso==0.16.0 
 uvloop==0.21.0
 mcp==1.5.0    # for MCP server  

litellm/integrations/gcs_bucket/gcs_bucket.py

@@ -121,6 +121,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
             gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config(
                 kwargs
             )
+
             headers = await self.construct_request_headers(
                 vertex_instance=gcs_logging_config["vertex_instance"],
                 service_account_json=gcs_logging_config["path_service_account"],

litellm/llms/vertex_ai/vertex_llm_base.py

@@ -6,7 +6,7 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S
 
 import json
 import os
-from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple
 
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.asyncify import asyncify
@@ -28,6 +28,10 @@ class VertexBase(BaseLLM):
         self.access_token: Optional[str] = None
         self.refresh_token: Optional[str] = None
         self._credentials: Optional[GoogleCredentialsObject] = None
+        self._credentials_project_mapping: Dict[
+            Tuple[Optional[VERTEX_CREDENTIALS_TYPES], Optional[str]],
+            GoogleCredentialsObject,
+        ] = {}
         self.project_id: Optional[str] = None
         self.async_handler: Optional[AsyncHTTPHandler] = None
 
@@ -128,32 +132,11 @@ class VertexBase(BaseLLM):
         """
         if custom_llm_provider == "gemini":
             return "", ""
-        if self.access_token is not None:
-            if project_id is not None:
-                return self.access_token, project_id
-            elif self.project_id is not None:
-                return self.access_token, self.project_id
-
-        if not self._credentials:
-            self._credentials, cred_project_id = self.load_auth(
-                credentials=credentials, project_id=project_id
-            )
-            if not self.project_id:
-                self.project_id = project_id or cred_project_id
         else:
-            if self._credentials.expired or not self._credentials.token:
-                self.refresh_auth(self._credentials)
-
-            if not self.project_id:
-                self.project_id = self._credentials.quota_project_id
-
-        if not self.project_id:
-            raise ValueError("Could not resolve project_id")
-
-        if not self._credentials or not self._credentials.token:
-            raise RuntimeError("Could not resolve API token from the environment")
-
-        return self._credentials.token, project_id or self.project_id
+            return self.get_access_token(
+                credentials=credentials,
+                project_id=project_id,
+            )
 
     def is_using_v1beta1_features(self, optional_params: dict) -> bool:
         """
@@ -259,6 +242,101 @@ class VertexBase(BaseLLM):
             url=url,
         )
 
+    def get_access_token(
+        self,
+        credentials: Optional[VERTEX_CREDENTIALS_TYPES],
+        project_id: Optional[str],
+    ) -> Tuple[str, str]:
+        """
+        Get access token and project id
+
+        1. Check if credentials are already in self._credentials_project_mapping
+        2. If not, load credentials and add to self._credentials_project_mapping
+        3. Check if loaded credentials have expired
+        4. If expired, refresh credentials
+        5. Return access token and project id
+        """
+
+        # Convert dict credentials to string for caching
+        cache_credentials = (
+            json.dumps(credentials) if isinstance(credentials, dict) else credentials
+        )
+        credential_cache_key = (cache_credentials, project_id)
+        _credentials: Optional[GoogleCredentialsObject] = None
+
+        verbose_logger.debug(
+            f"Checking cached credentials for project_id: {project_id}"
+        )
+
+        if credential_cache_key in self._credentials_project_mapping:
+            verbose_logger.debug(
+                f"Cached credentials found for project_id: {project_id}."
+            )
+            _credentials = self._credentials_project_mapping[credential_cache_key]
+            verbose_logger.debug("Using cached credentials")
+            credential_project_id = _credentials.quota_project_id or getattr(
+                _credentials, "project_id", None
+            )
+
+        else:
+            verbose_logger.debug(
+                f"Credential cache key not found for project_id: {project_id}, loading new credentials"
+            )
+
+            try:
+                _credentials, credential_project_id = self.load_auth(
+                    credentials=credentials, project_id=project_id
+                )
+            except Exception as e:
+                verbose_logger.exception(
+                    "Failed to load vertex credentials. Check to see if credentials containing partial/invalid information."
+                )
+                raise e
+
+            if _credentials is None:
+                raise ValueError(
+                    "Could not resolve credentials - either dynamically or from environment, for project_id: {}".format(
+                        project_id
+                    )
+                )
+
+            self._credentials_project_mapping[credential_cache_key] = _credentials
+
+        ## VALIDATE CREDENTIALS
+        verbose_logger.debug(f"Validating credentials for project_id: {project_id}")
+        if (
+            project_id is not None
+            and credential_project_id
+            and credential_project_id != project_id
+        ):
+            raise ValueError(
+                "Could not resolve project_id. Credential project_id: {} does not match requested project_id: {}".format(
+                    _credentials.quota_project_id, project_id
+                )
+            )
+        elif (
+            project_id is None
+            and credential_project_id is not None
+            and isinstance(credential_project_id, str)
+        ):
+            project_id = credential_project_id
+
+        if _credentials.expired:
+            self.refresh_auth(_credentials)
+
+        ## VALIDATION STEP
+        if _credentials.token is None or not isinstance(_credentials.token, str):
+            raise ValueError(
+                "Could not resolve credentials token. Got None or non-string token - {}".format(
+                    _credentials.token
+                )
+            )
+
+        if project_id is None:
+            raise ValueError("Could not resolve project_id")
+
+        return _credentials.token, project_id
+
     async def _ensure_access_token_async(
         self,
         credentials: Optional[VERTEX_CREDENTIALS_TYPES],
@@ -272,38 +350,14 @@ class VertexBase(BaseLLM):
         """
         if custom_llm_provider == "gemini":
             return "", ""
-        if self.access_token is not None:
-            if project_id is not None:
-                return self.access_token, project_id
-            elif self.project_id is not None:
-                return self.access_token, self.project_id
-
-        if not self._credentials:
-            try:
-                self._credentials, cred_project_id = await asyncify(self.load_auth)(
-                    credentials=credentials, project_id=project_id
-                )
-            except Exception:
-                verbose_logger.exception(
-                    "Failed to load vertex credentials. Check to see if credentials containing partial/invalid information."
-                )
-                raise
-            if not self.project_id:
-                self.project_id = project_id or cred_project_id
         else:
-            if self._credentials.expired or not self._credentials.token:
-                await asyncify(self.refresh_auth)(self._credentials)
-
-            if not self.project_id:
-                self.project_id = self._credentials.quota_project_id
-
-        if not self.project_id:
-            raise ValueError("Could not resolve project_id")
-
-        if not self._credentials or not self._credentials.token:
-            raise RuntimeError("Could not resolve API token from the environment")
-
-        return self._credentials.token, project_id or self.project_id
+            try:
+                return await asyncify(self.get_access_token)(
+                    credentials=credentials,
+                    project_id=project_id,
+                )
+            except Exception as e:
+                raise e
 
     def set_headers(
         self, auth_header: Optional[str], extra_headers: Optional[dict]

litellm/proxy/management_endpoints/ui_sso.py

@@ -485,7 +485,14 @@ async def auth_callback(request: Request):  # noqa: PLR0915
             redirect_uri=redirect_url,
             allow_insecure_http=True,
         )
-        result = await microsoft_sso.verify_and_process(request)
+        original_msft_result = await microsoft_sso.verify_and_process(
+            request=request,
+            convert_response=False,
+        )
+        result = MicrosoftSSOHandler.openid_from_response(
+            response=original_msft_result,
+            jwt_handler=jwt_handler,
+        )
     elif generic_client_id is not None:
         result = await get_generic_sso_response(
             request=request,
@@ -494,6 +501,7 @@ async def auth_callback(request: Request):  # noqa: PLR0915
             redirect_url=redirect_url,
         )
     # User is Authe'd in - generate key for the UI to access Proxy
+    verbose_proxy_logger.debug(f"SSO callback result: {result}")
     user_email: Optional[str] = getattr(result, "email", None)
     user_id: Optional[str] = getattr(result, "id", None) if result is not None else None
 
@@ -779,3 +787,27 @@ async def get_ui_settings(request: Request):
         ),
         "DISABLE_EXPENSIVE_DB_QUERIES": disable_expensive_db_queries,
     }
+
+
+class MicrosoftSSOHandler:
+    """
+    Handles Microsoft SSO callback response and returns a CustomOpenID object
+    """
+
+    @staticmethod
+    def openid_from_response(
+        response: Optional[dict], jwt_handler: JWTHandler
+    ) -> CustomOpenID:
+        response = response or {}
+        verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}")
+        openid_response = CustomOpenID(
+            email=response.get("mail"),
+            display_name=response.get("displayName"),
+            provider="microsoft",
+            id=response.get("id"),
+            first_name=response.get("givenName"),
+            last_name=response.get("surname"),
+            team_ids=jwt_handler.get_team_ids_from_jwt(cast(dict, response)),
+        )
+        verbose_proxy_logger.debug(f"Microsoft SSO OpenID Response: {openid_response}")
+        return openid_response

tests/litellm/llms/vertex_ai/test_vertex_llm_base.py

@@ -0,0 +1,176 @@
+import os
+import sys
+from unittest.mock import MagicMock, call, patch
+
+import pytest
+
+from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
+
+sys.path.insert(
+    0, os.path.abspath("../../..")
+)  # Adds the parent directory to the system path
+
+import litellm
+from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
+
+
+def run_sync(coro):
+    """Helper to run coroutine synchronously for testing"""
+    import asyncio
+
+    return asyncio.run(coro)
+
+
+class TestVertexBase:
+    @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
+    @pytest.mark.asyncio
+    async def test_credential_project_validation(self, is_async):
+        vertex_base = VertexBase()
+
+        # Mock credentials with project_id "project-1"
+        mock_creds = MagicMock()
+        mock_creds.project_id = "project-1"
+        mock_creds.token = "fake-token-1"
+        mock_creds.expired = False
+        mock_creds.quota_project_id = "project-1"
+
+        # Test case 1: Ensure credentials match project
+        with patch.object(
+            vertex_base, "load_auth", return_value=(mock_creds, "project-1")
+        ):
+            if is_async:
+                token, project = await vertex_base._ensure_access_token_async(
+                    credentials={"type": "service_account", "project_id": "project-1"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            else:
+                token, project = vertex_base._ensure_access_token(
+                    credentials={"type": "service_account", "project_id": "project-1"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            assert project == "project-1"
+            assert token == "fake-token-1"
+
+        # Test case 2: Prevent using credentials from different project
+        with patch.object(
+            vertex_base, "load_auth", return_value=(mock_creds, "project-1")
+        ):
+            with pytest.raises(ValueError, match="Could not resolve project_id"):
+                if is_async:
+                    result = await vertex_base._ensure_access_token_async(
+                        credentials={"type": "service_account"},
+                        project_id="different-project",
+                        custom_llm_provider="vertex_ai",
+                    )
+                else:
+                    result = vertex_base._ensure_access_token(
+                        credentials={"type": "service_account"},
+                        project_id="different-project",
+                        custom_llm_provider="vertex_ai",
+                    )
+                print(f"result: {result}")
+
+    @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
+    @pytest.mark.asyncio
+    async def test_cached_credentials(self, is_async):
+        vertex_base = VertexBase()
+
+        # Initial credentials
+        mock_creds = MagicMock()
+        mock_creds.token = "token-1"
+        mock_creds.expired = False
+        mock_creds.project_id = "project-1"
+        mock_creds.quota_project_id = "project-1"
+
+        # Test initial credential load and caching
+        with patch.object(
+            vertex_base, "load_auth", return_value=(mock_creds, "project-1")
+        ):
+            # First call should load credentials
+            if is_async:
+                token, project = await vertex_base._ensure_access_token_async(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            else:
+                token, project = vertex_base._ensure_access_token(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            assert token == "token-1"
+
+            # Second call should use cached credentials
+            if is_async:
+                token2, project2 = await vertex_base._ensure_access_token_async(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            else:
+                token2, project2 = vertex_base._ensure_access_token(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            assert token2 == "token-1"
+            assert project2 == "project-1"
+
+    @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
+    @pytest.mark.asyncio
+    async def test_credential_refresh(self, is_async):
+        vertex_base = VertexBase()
+
+        # Create expired credentials
+        mock_creds = MagicMock()
+        mock_creds.token = "my-token"
+        mock_creds.expired = True
+        mock_creds.project_id = "project-1"
+        mock_creds.quota_project_id = "project-1"
+
+        with patch.object(
+            vertex_base, "load_auth", return_value=(mock_creds, "project-1")
+        ), patch.object(vertex_base, "refresh_auth") as mock_refresh:
+
+            def mock_refresh_impl(creds):
+                creds.token = "refreshed-token"
+                creds.expired = False
+
+            mock_refresh.side_effect = mock_refresh_impl
+
+            if is_async:
+                token, project = await vertex_base._ensure_access_token_async(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            else:
+                token, project = vertex_base._ensure_access_token(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+
+            assert mock_refresh.called
+            assert token == "refreshed-token"
+            assert not mock_creds.expired
+
+    @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
+    @pytest.mark.asyncio
+    async def test_gemini_credentials(self, is_async):
+        vertex_base = VertexBase()
+
+        # Test that Gemini requests bypass credential checks
+        if is_async:
+            token, project = await vertex_base._ensure_access_token_async(
+                credentials=None, project_id=None, custom_llm_provider="gemini"
+            )
+        else:
+            token, project = vertex_base._ensure_access_token(
+                credentials=None, project_id=None, custom_llm_provider="gemini"
+            )
+        assert token == ""
+        assert project == ""

tests/litellm/proxy/management_endpoints/test_ui_sso.py

@@ -0,0 +1,81 @@
+import json
+import os
+import sys
+from typing import Optional, cast
+from unittest.mock import MagicMock, patch
+
+import pytest
+from fastapi.testclient import TestClient
+
+sys.path.insert(
+    0, os.path.abspath("../../..")
+)  # Adds the parent directory to the system path
+
+from litellm.proxy.auth.handle_jwt import JWTHandler
+from litellm.proxy.management_endpoints.types import CustomOpenID
+from litellm.proxy.management_endpoints.ui_sso import MicrosoftSSOHandler
+
+
+def test_microsoft_sso_handler_openid_from_response():
+    # Arrange
+    # Create a mock response similar to what Microsoft SSO would return
+    mock_response = {
+        "mail": "test@example.com",
+        "displayName": "Test User",
+        "id": "user123",
+        "givenName": "Test",
+        "surname": "User",
+        "some_other_field": "value",
+    }
+
+    # Create a mock JWTHandler that returns predetermined team IDs
+    mock_jwt_handler = MagicMock(spec=JWTHandler)
+    expected_team_ids = ["team1", "team2"]
+    mock_jwt_handler.get_team_ids_from_jwt.return_value = expected_team_ids
+
+    # Act
+    # Call the method being tested
+    result = MicrosoftSSOHandler.openid_from_response(
+        response=mock_response, jwt_handler=mock_jwt_handler
+    )
+
+    # Assert
+    # Verify the JWT handler was called with the correct parameters
+    mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with(
+        cast(dict, mock_response)
+    )
+
+    # Check that the result is a CustomOpenID object with the expected values
+    assert isinstance(result, CustomOpenID)
+    assert result.email == "test@example.com"
+    assert result.display_name == "Test User"
+    assert result.provider == "microsoft"
+    assert result.id == "user123"
+    assert result.first_name == "Test"
+    assert result.last_name == "User"
+    assert result.team_ids == expected_team_ids
+
+
+def test_microsoft_sso_handler_with_empty_response():
+    # Arrange
+    # Test with None response
+    mock_jwt_handler = MagicMock(spec=JWTHandler)
+    mock_jwt_handler.get_team_ids_from_jwt.return_value = []
+
+    # Act
+    result = MicrosoftSSOHandler.openid_from_response(
+        response=None, jwt_handler=mock_jwt_handler
+    )
+
+    # Assert
+    assert isinstance(result, CustomOpenID)
+    assert result.email is None
+    assert result.display_name is None
+    assert result.provider == "microsoft"
+    assert result.id is None
+    assert result.first_name is None
+    assert result.last_name is None
+    assert result.team_ids == []
+
+    # Make sure the JWT handler was called with an empty dict
+    mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with({})

tests/llm_translation/test_prompt_factory.py

@@ -795,9 +795,7 @@ def test_hf_chat_template():
     print(chat_template)
     assert (
         chat_template.rstrip()
-        == """You are a helpful assistant.
-What is the weather in Copenhagen?
-"""
+        == "You are a helpful assistant. What is the weather in Copenhagen?"
     )
 
 

tests/local_testing/test_amazing_vertex_completion.py

@@ -59,6 +59,7 @@ VERTEX_MODELS_TO_NOT_TEST = [
     "gemini-pro-experimental",
     "gemini-flash-experimental",
     "gemini-1.5-flash-exp-0827",
+    "gemini-2.0-pro-exp-02-05",
     "gemini-pro-flash",
     "gemini-1.5-flash-exp-0827",
     "gemini-2.0-flash-exp",
@@ -452,6 +453,7 @@ async def test_async_vertexai_response():
             or "ultra" in model
             or "002" in model
             or "gemini-2.0-flash-thinking-exp" in model
+            or "gemini-2.0-pro-exp-02-05" in model
         ):
             # our account does not have access to this model
             continue
@@ -498,6 +500,7 @@ async def test_async_vertexai_streaming_response():
             or "ultra" in model
             or "002" in model
             or "gemini-2.0-flash-thinking-exp" in model
+            or "gemini-2.0-pro-exp-02-05" in model
         ):
             # our account does not have access to this model
             continue