diff --git a/litellm/llms/vertex_ai.py b/litellm/llms/vertex_ai.py index 5b99683ca..0a7980fda 100644 --- a/litellm/llms/vertex_ai.py +++ b/litellm/llms/vertex_ai.py @@ -1047,6 +1047,7 @@ def embedding( vertex_project=None, vertex_location=None, aembedding=False, + print_verbose=None, ): # logic for parsing in - calling - parsing out model embedding calls try: @@ -1062,7 +1063,13 @@ def embedding( ## Load credentials with the correct quota project ref: https://github.com/googleapis/python-aiplatform/issues/2557#issuecomment-1709284744 try: + print_verbose( + f"VERTEX AI: vertex_project={vertex_project}; vertex_location={vertex_location}" + ) creds, _ = google.auth.default(quota_project_id=vertex_project) + print_verbose( + f"VERTEX AI: creds={creds}; google application credentials: {os.getenv('GOOGLE_APPLICATION_CREDENTIALS')}" + ) vertexai.init( project=vertex_project, location=vertex_location, credentials=creds ) diff --git a/litellm/main.py b/litellm/main.py index 60effd96f..63649844a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2617,6 +2617,7 @@ def embedding( vertex_project=vertex_ai_project, vertex_location=vertex_ai_location, aembedding=aembedding, + print_verbose=print_verbose, ) elif custom_llm_provider == "oobabooga": response = oobabooga.embedding( diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 8ed15db65..67db9b61c 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -451,6 +451,34 @@ async def test_gemini_pro_async_function_calling(): # asyncio.run(gemini_pro_async_function_calling()) + +def test_vertexai_embedding(): + try: + load_vertex_ai_credentials() + # litellm.set_verbose=True + response = embedding( + model="textembedding-gecko@001", + input=["good morning from litellm", "this is another item"], + ) + print(f"response:", response) + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + +@pytest.mark.asyncio +async def test_vertexai_aembedding(): + try: + load_vertex_ai_credentials() + # litellm.set_verbose=True + response = await litellm.aembedding( + model="textembedding-gecko@001", + input=["good morning from litellm", "this is another item"], + ) + print(f"response: {response}") + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + # Extra gemini Vision tests for completion + stream, async, async + stream # if we run into issues with gemini, we will also add these to our ci/cd pipeline # def test_gemini_pro_vision_stream(): diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index a2f71eb98..c32a55353 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -231,31 +231,6 @@ def test_cohere_embedding3(): # test_cohere_embedding3() -def test_vertexai_embedding(): - try: - # litellm.set_verbose=True - response = embedding( - model="textembedding-gecko@001", - input=["good morning from litellm", "this is another item"], - ) - print(f"response:", response) - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - -@pytest.mark.asyncio -async def test_vertexai_aembedding(): - try: - # litellm.set_verbose=True - response = await litellm.aembedding( - model="textembedding-gecko@001", - input=["good morning from litellm", "this is another item"], - ) - print(f"response: {response}") - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - def test_bedrock_embedding_titan(): try: # this tests if we support str input for bedrock embedding diff --git a/litellm/tests/test_get_optional_params_embeddings.py b/litellm/tests/test_get_optional_params_embeddings.py new file mode 100644 index 000000000..41396b531 --- /dev/null +++ b/litellm/tests/test_get_optional_params_embeddings.py @@ -0,0 +1,42 @@ +# What is this? +## This tests the `get_optional_params_embeddings` function +import sys, os +import traceback +from dotenv import load_dotenv + +load_dotenv() +import os, io + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import pytest +import litellm +from litellm import embedding +from litellm.utils import get_optional_params_embeddings, get_llm_provider + + +def test_vertex_projects(): + litellm.drop_params = True + model, custom_llm_provider, _, _ = get_llm_provider( + model="vertex_ai/textembedding-gecko" + ) + optional_params = get_optional_params_embeddings( + model=model, + user="test-litellm-user-5", + dimensions=None, + encoding_format="base64", + custom_llm_provider=custom_llm_provider, + **{ + "vertex_ai_project": "my-test-project", + "vertex_ai_location": "us-east-1", + }, + ) + + print(f"received optional_params: {optional_params}") + + assert "vertex_ai_project" in optional_params + assert "vertex_ai_location" in optional_params + + +# test_vertex_projects()