diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index 2a36e5cc785..40503edbb9e 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -109,8 +109,6 @@ class VertexAIPartnerModels(VertexBase): client=None, ): try: - import vertexai - from litellm.llms.anthropic.chat import AnthropicChatCompletion from litellm.llms.codestral.completion.handler import ( CodestralTextCompletion, @@ -119,14 +117,9 @@ class VertexAIPartnerModels(VertexBase): except Exception as e: raise VertexAIError( status_code=400, - message=f"""vertexai import failed please run `pip install -U "google-cloud-aiplatform>=1.38"`. Got error: {e}""", + message=f"Failed to import a partner model handler. Got error: {e}", ) - if not (hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models")): - raise VertexAIError( - status_code=400, - message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", - ) try: access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py index b20442a032e..8e6270e41a1 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py @@ -127,6 +127,44 @@ class TestPartnerModelsCredentialReuse: assert mock_load.call_count == 1 + def test_completion_works_without_the_vertexai_sdk(self): + """completion() reaches the HTTP handler when `import vertexai` raises ImportError.""" + partner = VertexAIPartnerModels() + + with ( + patch.dict(sys.modules, {"vertexai": None}), + patch.object( + partner, + "_ensure_access_token", + return_value=("cached-token", "test-project"), + ), + patch( + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.base_llm_http_handler" + ) as mock_handler, + ): + mock_handler.completion.return_value = "response" + + result = partner.completion( + model="meta/llama-3.1-405b-instruct-maas", + 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=None, + ) + + assert result == "response" + mock_handler.completion.assert_called_once() + class TestGemmaModelsCredentialReuse: def test_completion_uses_self_ensure_access_token(self):