GoogleGenerativeAIEmbeddings - embed_documents throws "Error embedding content: 504 Deadline Exceeded"

[LostInCode404]

The issue

The embed_documents function throws Error: Error embedding content: 504 Deadline Exceeded

Version of langchain-google-genai: 1.0.6

Background

I am using embed_documents function from GoogleGenerativeAIEmbeddings. Sample code:

from langchain_google_genai import GoogleGenerativeAIEmbeddings
embedding = GoogleGenerativeAIEmbeddings(google_api_key=GOOGLE_API_KEY, model="models/embedding-001")
embedding.embed_documents(string_chunks)

Error snippet:

GoogleGenerativeAIError: Error embedding content: 504 Deadline Exceeded

Workaround that works

If I stop passing google_api_key to the constructor and set it as an env variable, the error resolves.

Code that works:

from langchain_google_genai import GoogleGenerativeAIEmbeddings
os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY
embedding = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
embedding.embed_documents(string_chunks)