neo4j-graphrag-python
neo4j-graphrag-python copied to clipboard
Migrate VertexAIEmbeddings to use google-genai SDK
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:
- Update the dependency to use
google-genaiinstead of or alongsidevertexai. - Modify the
VertexAIEmbeddingsclass to use the corresponding methods and classes fromgoogle-genai(e.g.,client.models.generate_embeddings). - 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
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.
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