neo4j-graphrag-python icon indicating copy to clipboard operation
neo4j-graphrag-python copied to clipboard

Migrate VertexAIEmbeddings to use google-genai SDK

Open datasherlock opened this issue 2 months ago • 2 comments

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

https://github.com/neo4j/neo4j-graphrag-python/blob/7b4e2d39f6931721c0899f37a1c428238990cbe6/src/neo4j_graphrag/embeddings/vertexai.py#L24

datasherlock avatar Oct 10 '25 15:10 datasherlock

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.

stellasia avatar Oct 14 '25 11:10 stellasia

Agreed. VertexAILLM is using a deprecated model too https://github.com/neo4j/neo4j-graphrag-python/blob/291d919ae4a4eed1001bad99199b74fd7896f252/src/neo4j_graphrag/llm/vertexai_llm.py#L70 as the default.

I'll try finding time to address these

datasherlock avatar Oct 16 '25 04:10 datasherlock