diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/main.py b/litellm/llms/vertex_ai/vertex_gemma_models/main.py index 82cfe6de984..b6bf2f73b72 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/main.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/main.py @@ -31,7 +31,7 @@ from ..vertex_llm_base import VertexBase class VertexAIGemmaModels(VertexBase): def __init__(self) -> None: - pass + super().__init__() def completion( self, @@ -62,9 +62,6 @@ class VertexAIGemmaModels(VertexBase): try: import vertexai - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexLLM, - ) from litellm.llms.vertex_ai.vertex_gemma_models.transformation import ( VertexGemmaConfig, ) @@ -83,9 +80,8 @@ class VertexAIGemmaModels(VertexBase): ) try: model = get_vertex_base_model_name(model=model) - vertex_httpx_logic = VertexLLM() - access_token, project_id = vertex_httpx_logic._ensure_access_token( + access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", diff --git a/litellm/llms/vertex_ai/vertex_model_garden/main.py b/litellm/llms/vertex_ai/vertex_model_garden/main.py index c37bb449ecf..6234676e184 100644 --- a/litellm/llms/vertex_ai/vertex_model_garden/main.py +++ b/litellm/llms/vertex_ai/vertex_model_garden/main.py @@ -41,7 +41,7 @@ def create_vertex_url( class VertexAIModelGardenModels(VertexBase): def __init__(self) -> None: - pass + super().__init__() def completion( self, @@ -73,9 +73,6 @@ class VertexAIModelGardenModels(VertexBase): import vertexai from litellm.llms.openai_like.chat.handler import OpenAILikeChatHandler - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexLLM, - ) except Exception as e: raise VertexAIError( status_code=400, @@ -91,9 +88,8 @@ class VertexAIModelGardenModels(VertexBase): ) try: model = get_vertex_base_model_name(model=model) - vertex_httpx_logic = VertexLLM() - access_token, project_id = vertex_httpx_logic._ensure_access_token( + access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py index 4f9e784b76d..b20442a032e 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py @@ -1,40 +1,53 @@ """ -Test that VertexAIPartnerModels reuses cached credentials from VertexBase -instead of creating a new VertexLLM instance on every request. +Test that VertexBase subclasses (PartnerModels, Gemma, ModelGarden) reuse +cached credentials instead of creating a new VertexLLM instance on every request. """ import sys from unittest.mock import MagicMock, patch +import pytest + from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( VertexAIPartnerModels, ) +from litellm.llms.vertex_ai.vertex_gemma_models.main import VertexAIGemmaModels +from litellm.llms.vertex_ai.vertex_model_garden.main import VertexAIModelGardenModels + + +def _mock_vertexai(): + """Return a MagicMock that satisfies the vertexai import guards.""" + m = MagicMock() + m.preview = MagicMock() + m.preview.language_models = MagicMock() + return m + + +class TestVertexBaseSubclassInit: + """All VertexBase subclasses must call super().__init__() so that + the credential cache is initialized.""" + + @pytest.mark.parametrize( + "cls", + [VertexAIPartnerModels, VertexAIGemmaModels, VertexAIModelGardenModels], + ids=["PartnerModels", "Gemma", "ModelGarden"], + ) + def test_init_calls_super(self, cls): + instance = cls() + assert hasattr(instance, "_credentials_project_mapping") + assert isinstance(instance._credentials_project_mapping, dict) + assert hasattr(instance, "access_token") + assert hasattr(instance, "project_id") class TestPartnerModelsCredentialReuse: - def test_init_calls_super(self): - """VertexAIPartnerModels.__init__ must call super().__init__() so that - the VertexBase credential cache is initialized.""" - partner = VertexAIPartnerModels() - # These attributes are set by VertexBase.__init__ - assert hasattr(partner, "_credentials_project_mapping") - assert isinstance(partner._credentials_project_mapping, dict) - assert hasattr(partner, "access_token") - assert hasattr(partner, "project_id") - def test_completion_uses_self_ensure_access_token(self): """completion() should call self._ensure_access_token, not create a - throwaway VertexLLM instance. This ensures the credential cache on the - singleton is reused across calls.""" + throwaway VertexLLM instance.""" partner = VertexAIPartnerModels() - # Mock vertexai import and the completion handler - mock_vertexai = MagicMock() - mock_vertexai.preview = MagicMock() - mock_vertexai.preview.language_models = MagicMock() - with ( - patch.dict(sys.modules, {"vertexai": mock_vertexai}), + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), patch.object( partner, "_ensure_access_token", @@ -64,7 +77,6 @@ class TestPartnerModelsCredentialReuse: vertex_credentials='{"type": "service_account"}', ) - # _ensure_access_token should have been called on self mock_ensure.assert_called_once_with( credentials='{"type": "service_account"}', project_id="test-project", @@ -72,8 +84,7 @@ class TestPartnerModelsCredentialReuse: ) def test_credential_cache_shared_across_calls(self): - """Two successive completion() calls should hit load_auth only once, - proving the credential cache on the VertexAIPartnerModels instance works.""" + """Two successive completion() calls should hit load_auth only once.""" partner = VertexAIPartnerModels() mock_creds = MagicMock() @@ -82,12 +93,8 @@ class TestPartnerModelsCredentialReuse: mock_creds.project_id = "proj" mock_creds.quota_project_id = "proj" - mock_vertexai = MagicMock() - mock_vertexai.preview = MagicMock() - mock_vertexai.preview.language_models = MagicMock() - with ( - patch.dict(sys.modules, {"vertexai": mock_vertexai}), + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), patch.object( partner, "load_auth", return_value=(mock_creds, "proj") ) as mock_load, @@ -118,5 +125,96 @@ class TestPartnerModelsCredentialReuse: partner.completion(**common_kwargs) partner.completion(**common_kwargs) - # load_auth should only be called once — second call uses cache assert mock_load.call_count == 1 + + +class TestGemmaModelsCredentialReuse: + def test_completion_uses_self_ensure_access_token(self): + """completion() should call self._ensure_access_token, not create a + throwaway VertexLLM instance.""" + gemma = VertexAIGemmaModels() + + mock_gemma_config = MagicMock() + mock_gemma_config.return_value.completion.return_value = "response" + + with ( + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), + patch.object( + gemma, + "_ensure_access_token", + return_value=("cached-token", "test-project"), + ) as mock_ensure, + patch( + "litellm.llms.vertex_ai.vertex_gemma_models.transformation.VertexGemmaConfig", + mock_gemma_config, + ), + ): + gemma.completion( + model="gemma/gemma-3-12b-it-1234567890", + messages=[{"role": "user", "content": "hello"}], + model_response=MagicMock(), + print_verbose=lambda *a, **kw: None, + encoding=MagicMock(), + logging_obj=MagicMock(), + api_base="https://123.us-central1-1.prediction.vertexai.goog/v1/projects/proj/locations/us-central1/endpoints/456:predict", + optional_params={}, + custom_prompt_dict={}, + headers=None, + timeout=30.0, + litellm_params={}, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials='{"type": "service_account"}', + ) + + mock_ensure.assert_called_once_with( + credentials='{"type": "service_account"}', + project_id="test-project", + custom_llm_provider="vertex_ai", + ) + + +class TestModelGardenCredentialReuse: + def test_completion_uses_self_ensure_access_token(self): + """completion() should call self._ensure_access_token, not create a + throwaway VertexLLM instance.""" + garden = VertexAIModelGardenModels() + + mock_handler = MagicMock() + mock_handler.return_value.completion.return_value = "response" + + with ( + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), + patch.object( + garden, + "_ensure_access_token", + return_value=("cached-token", "test-project"), + ) as mock_ensure, + patch( + "litellm.llms.openai_like.chat.handler.OpenAILikeChatHandler", + mock_handler, + ), + ): + garden.completion( + model="openai/5464397967697903616", + messages=[{"role": "user", "content": "hello"}], + model_response=MagicMock(), + print_verbose=lambda *a, **kw: None, + encoding=MagicMock(), + logging_obj=MagicMock(), + api_base=None, + optional_params={}, + custom_prompt_dict={}, + headers=None, + timeout=30.0, + litellm_params={}, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials='{"type": "service_account"}', + ) + + mock_ensure.assert_called_once_with( + credentials='{"type": "service_account"}', + project_id="test-project", + custom_llm_provider="vertex_ai", + )