Add gcp auth for better supporton service account

This commit is contained in:
parshvadaftari
2025-09-16 00:13:42 +05:30
parent 2dd4e93af4
commit a93a7ea6cf
2 changed files with 184 additions and 7 deletions
+17 -7
View File
@@ -5,6 +5,7 @@ from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
from mem0.configs.embeddings.base import BaseEmbedderConfig
from mem0.embeddings.base import EmbeddingBase
from mem0.utils.gcp_auth import GCPAuthenticator
class VertexAIEmbedding(EmbeddingBase):
@@ -20,14 +21,23 @@ class VertexAIEmbedding(EmbeddingBase):
"search": self.config.memory_search_embedding_type or "RETRIEVAL_QUERY",
}
credentials_path = self.config.vertex_credentials_json
if credentials_path:
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credentials_path
elif not os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
raise ValueError(
"Google application credentials JSON is not provided. Please provide a valid JSON path or set the 'GOOGLE_APPLICATION_CREDENTIALS' environment variable."
# Set up authentication using centralized GCP authenticator
# This supports multiple authentication methods while preserving environment variable support
try:
GCPAuthenticator.setup_vertex_ai(
service_account_json=getattr(self.config, 'google_service_account_json', None),
credentials_path=self.config.vertex_credentials_json,
project_id=getattr(self.config, 'google_project_id', None)
)
except Exception as e:
# Fall back to original behavior for backward compatibility
credentials_path = self.config.vertex_credentials_json
if credentials_path:
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credentials_path
elif not os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
raise ValueError(
"Google application credentials JSON is not provided. Please provide a valid JSON path or set the 'GOOGLE_APPLICATION_CREDENTIALS' environment variable."
)
self.model = TextEmbeddingModel.from_pretrained(self.config.model)