klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix: Migrate Gemini Embeddings (#3002)

Browse Source 
Co-authored-by: Dev-Khant <devkhant24@gmail.com>
This commit is contained in:
Akshat Jain
2025-06-23 13:16:10 +05:30
committed by GitHub
parent c173ec32d0
commit 386d8b87ae
5 changed files with 124 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
mem0/embeddings/gemini.py
View File

@@ -1,7 +1,7 @@
import os
from typing import Literal, Optional
import google.generativeai as genai
import google.genai as genai
from mem0.configs.embeddings.base import BaseEmbedderConfig
from mem0.embeddings.base import EmbeddingBase
@@ -12,23 +12,28 @@ class GoogleGenAIEmbedding(EmbeddingBase):
super().__init__(config)
self.config.model = self.config.model or "models/text-embedding-004"
self.config.embedding_dims = self.config.embedding_dims or 768
self.config.embedding_dims = self.config.embedding_dims or self.config.output_dimensionality or 768
api_key = self.config.api_key or os.getenv("GOOGLE_API_KEY")
genai.configure(api_key=api_key)
if api_key:
self.client = genai.Client(api_key="api_key")
else:
self.client = genai.Client()
def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
"""
Get the embedding for the given text using Google Generative AI.
Args:
text (str): The text to embed.
memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
memory_action (optional): The type of embedding to use. (Currently not used by Gemini for task_type)
Returns:
list: The embedding vector.
"""
text = text.replace("\n", " ")
response = genai.embed_content(
response = self.client.models.embed_content(
model=self.config.model, content=text, output_dimensionality=self.config.embedding_dims
)
return response["embedding"]