my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2	from typing import Literal, Optional
#3
#4	from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
#5
#6	from mem0.configs.embeddings.base import BaseEmbedderConfig
#7	from mem0.embeddings.base import EmbeddingBase
#8	from mem0.utils.gcp_auth import GCPAuthenticator
#9
#10
#11	class VertexAIEmbedding(EmbeddingBase):
#12	def __init__(self, config: Optional[BaseEmbedderConfig] = None):
#13	super().__init__(config)
#14
#15	self.config.model = self.config.model or "text-embedding-004"
#16	self.config.embedding_dims = self.config.embedding_dims or 256
#17
#18	self.embedding_types = {
#19	"add": self.config.memory_add_embedding_type or "RETRIEVAL_DOCUMENT",
#20	"update": self.config.memory_update_embedding_type or "RETRIEVAL_DOCUMENT",
#21	"search": self.config.memory_search_embedding_type or "RETRIEVAL_QUERY",
#22	}
#23
#24	# Set up authentication using centralized GCP authenticator
#25	# This supports multiple authentication methods while preserving environment variable support
#26	try:
#27	GCPAuthenticator.setup_vertex_ai(
#28	service_account_json=getattr(self.config, 'google_service_account_json', None),
#29	credentials_path=self.config.vertex_credentials_json,
#30	project_id=getattr(self.config, 'google_project_id', None)
#31	)
#32	except Exception:
#33	# Fall back to original behavior for backward compatibility
#34	credentials_path = self.config.vertex_credentials_json
#35	if credentials_path:
#36	os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credentials_path
#37	elif not os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
#38	raise ValueError(
#39	"Google application credentials JSON is not provided. Please provide a valid JSON path or set the 'GOOGLE_APPLICATION_CREDENTIALS' environment variable."
#40	)
#41
#42	self.model = TextEmbeddingModel.from_pretrained(self.config.model)
#43
#44	def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
#45	"""
#46	Get the embedding for the given text using Vertex AI.
#47
#48	Args:
#49	text (str): The text to embed.
#50	memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
#51	Returns:
#52	list: The embedding vector.
#53	"""
#54	embedding_type = "SEMANTIC_SIMILARITY"
#55	if memory_action is not None:
#56	if memory_action not in self.embedding_types:
#57	raise ValueError(f"Invalid memory action: {memory_action}")
#58
#59	embedding_type = self.embedding_types[memory_action]
#60
#61	text_input = TextEmbeddingInput(text=text, task_type=embedding_type)
#62	embeddings = self.model.get_embeddings(texts=[text_input], output_dimensionality=self.config.embedding_dims)
#63
#64	return embeddings[0].values
#65

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public