vertexai: Add embeddings_task_type parameter to embed_query and embed_documents

[mitraan-deshaw]

PR Description

This PR introduces an optional embeddings_task_type parameter to embed_documents and embed_query in VertexAIEmbeddings, mirroring the task_type argument in GoogleGenerativeAIEmbeddings. While working on this, I also addressed a couple of additional improvements:

• Created EmbeddingTaskTypes to consolidate embedding task types, reducing redundancy. I've currently added it as a Literal type but I'm open to suggestions for a more suitable location.
• Fixed an issue where the text-embedding-005 model, which supports embedding tasks, was incorrectly classified as an older model, preventing proper argument passing. I've added a new GoogleEmbeddingModelVersion to address this.
• Added (new) CODE_RETRIEVAL_QUERY as an embedding task.

Sample code:

from langchain_google_vertexai import VertexAIEmbeddings

# Initialize the a specific Embeddings Model version
embeddings = VertexAIEmbeddings(model_name="text-embedding-005")
single_vector = embeddings.embed_query(text, embeddings_task_type="CODE_RETRIEVAL_QUERY")

Relevant issues

Type

🆕 New Feature

Changes(optional)

• Added optional embeddings_task_type parameter to embed_documents and embed_query in VertexAIEmbeddings.
• Consolidated embeddings task types into EmbeddingTaskTypes to reduce redundancy.
• Added support for text-embedding-005 embedding tasks.
• Added CODE_RETRIEVAL_QUERY as an embedding task.

Testing(optional)

Note(optional)

[lkuligin]

please, add a *, so that embeddings_task_type can be provided by name only

[mitraan-deshaw]

fixed

[lkuligin]

do we need this integration test, or would a unit test be enough (we've tested in a test above that embeddings_task is passed to the Google API and it returns a valid output)

we've got too many integration tests now and the execution time gets longer and longer :), it might be a good idea to keep the total amount of them reasonable

[mitraan-deshaw]

Totally makes sense! However, I'm not sure how we would write a clean unit test for the new argument in embed_documents() and embed_query(), as these methods in turn call embed() and return its response. It almost seems like we would be validating the embed() method instead.

If we only want to check if the arguments are passed correctly, one approach could be to mock the embed() method, capture the arguments, and then call the VertexAIEmbedding's embed() method. Happy to remove the tests if you think it's not required.

[lkuligin]

just by mocking the API:

langchain-google/libs/vertexai/tests/unit_tests/test_embeddings.py

Line 20 in c3d4061

def test_langchain_google_vertexai_no_dups_dynamic_batch_size() -> None:

[mitraan-deshaw]

Moved the tests into unit tests.

[lkuligin]

@mitraan-deshaw could you fix the linter, please?