VertexAIEmbeddings class failed with AttributeError: module 'vertexai' has no attribute 'language_models'

[yohei1126]

pyproject.toml

[project]
name = "dde-graph-rag-builder"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
    "neo4j-graphrag[google]==1.8.0",
    "google-cloud-aiplatform==1.97.0",
    "openai>=1.98.0",
]

main.py

from neo4j_graphrag.llm import VertexAILLM
from neo4j_graphrag.generation import GraphRAG
from neo4j import GraphDatabase
from neo4j_graphrag.embeddings.vertexai import VertexAIEmbeddings
from neo4j_graphrag.retrievers import VectorRetriever

URI = "neo4j://localhost:7687"
AUTH = ("neo4j", "password123")
INDEX_NAME = "index-name"

if __name__ == "__main__":
    model_params = {"temperature": 0.0}
    llm = VertexAILLM(
        model_name="gemini-2.5-flash",
        model_params=model_params
    )
    result = llm.invoke("say something")
    print(result)

    driver = GraphDatabase.driver(URI, auth=AUTH)
    # https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings
    embedder = VertexAIEmbeddings(model="text-embedding-005")
    retriever = VectorRetriever(driver, INDEX_NAME, embedder)
    rag = GraphRAG(retriever=retriever, llm=llm)

    query_text = "How do I do similarity search in Neo4j?"
    response = rag.search(query_text=query_text, retriever_config={"top_k": 5})
    print(response.answer)

This script failed with the following error

uv sync
uv run main.py
content="Hello! I'm here and ready to chat. What's on your mind?"
Traceback (most recent call last):
  File "/Users/yohei.onishi/work/dde_graph_rag_builder/main.py", line 22, in <module>
    embedder = VertexAIEmbeddings(model="text-embedding-005")
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/yohei.onishi/work/dde_graph_rag_builder/.venv/lib/python3.12/site-packages/neo4j_graphrag/embeddings/vertexai.py", line 43, in __init__
    vertexai.language_models.TextEmbeddingModel.from_pretrained(model)
    ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/yohei.onishi/work/dde_graph_rag_builder/.venv/lib/python3.12/site-packages/vertexai/__init__.py", line 52, in __getattr__
    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
AttributeError: module 'vertexai' has no attribute 'language_models'

[stellasia]

Hi @yohei1126 ,

Thank you for raising the issue, we will investigate and fix as soon as possible.