Migrate VertexAIEmbeddings to use google-genai SDK

[datasherlock]

The current implementation of the VertexAIEmbeddings class, found in the provided code snippet, uses the deprecated vertexai SDK for Google's Generative AI models.

To ensure compatibility, access to the latest features, and adherence to Google's recommended library usage, this component should be migrated to use the google-genai Python SDK.

The specific parts of the code that use the legacy SDK are the imports and calls related to TextEmbeddingModel and TextEmbeddingInput:

try:
    from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
except (ImportError, AttributeError):
    TextEmbeddingModel = TextEmbeddingInput = None  # type: ignore[misc, assignment]

# ... inside VertexAIEmbeddings __init__
self.model = TextEmbeddingModel.from_pretrained(model)

# ... inside VertexAIEmbeddings _embed_query
inputs: list[str | TextEmbeddingInput] = [
    TextEmbeddingInput(text, task_type)
]
embeddings = self.model.get_embeddings(inputs, **kwargs)

Action Items:

1. Update the dependency to use google-genai instead of or alongside vertexai.
2. Modify the VertexAIEmbeddings class to use the corresponding methods and classes from google-genai (e.g., client.models.generate_embeddings).
3. Ensure the migration follows the official guide provided by Google for the SDK deprecation.

Reference:
The migration guide for the Vertex AI SDK deprecation can be found here:
https://cloud.google.com/vertex-ai/generative-ai/docs/deprecations/genai-vertexai-sdk.md

neo4j-graphrag-python/src/neo4j_graphrag/embeddings/vertexai.py

Line 24 in 7b4e2d3

from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel

[stellasia]

Hey @datasherlock ,

Thanks for the heads-up.

Note that the VertexAILLM also needs update when we move to the new google-genai package.

We've added an item to our backlog, but PRs are welcome on this topic.