diff --git a/docs/components/embedders/models/google_AI.mdx b/docs/components/embedders/models/google_AI.mdx index d0fb07b1c..7c9db18c1 100644 --- a/docs/components/embedders/models/google_AI.mdx +++ b/docs/components/embedders/models/google_AI.mdx @@ -19,7 +19,7 @@ config = { "embedder": { "provider": "gemini", "config": { - "model": "models/text-embedding-004", + "model": "models/gemini-embedding-001", } } } @@ -66,7 +66,7 @@ Here are the parameters available for configuring Gemini embedder: | Parameter | Description | Default Value | | ---------------- | ------------------------------------ | ----------------------- | -| `model` | The name of the embedding model to use| `models/text-embedding-004` | +| `model` | The name of the embedding model to use| `models/gemini-embedding-001` | | `embedding_dims` | Dimensions of the embedding model | `1536` | | `api_key` | The Google API key | `None` | diff --git a/docs/components/embedders/models/vertexai.mdx b/docs/components/embedders/models/vertexai.mdx index 311ad65a7..6c67ae85d 100644 --- a/docs/components/embedders/models/vertexai.mdx +++ b/docs/components/embedders/models/vertexai.mdx @@ -20,7 +20,7 @@ config = { "embedder": { "provider": "vertexai", "config": { - "model": "text-embedding-004", + "model": "gemini-embedding-001", "memory_add_embedding_type": "RETRIEVAL_DOCUMENT", "memory_update_embedding_type": "RETRIEVAL_DOCUMENT", "memory_search_embedding_type": "RETRIEVAL_QUERY" @@ -51,7 +51,7 @@ Here are the parameters available for configuring the Vertex AI embedder: | Parameter | Description | Default Value | | ------------------------- | ------------------------------------------------ | -------------------- | -| `model` | The name of the Vertex AI embedding model to use | `text-embedding-004` | +| `model` | The name of the Vertex AI embedding model to use | `gemini-embedding-001` | | `vertex_credentials_json` | Path to the Google Cloud credentials JSON file | `None` | | `embedding_dims` | Dimensions of the embedding model | `256` | | `memory_add_embedding_type` | The type of embedding to use for the add memory action | `RETRIEVAL_DOCUMENT` | diff --git a/mem0/embeddings/gemini.py b/mem0/embeddings/gemini.py index 203b311cd..207b2c8aa 100644 --- a/mem0/embeddings/gemini.py +++ b/mem0/embeddings/gemini.py @@ -12,7 +12,7 @@ class GoogleGenAIEmbedding(EmbeddingBase): def __init__(self, config: Optional[BaseEmbedderConfig] = None): super().__init__(config) - self.config.model = self.config.model or "models/text-embedding-004" + self.config.model = self.config.model or "models/gemini-embedding-001" 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") diff --git a/mem0/embeddings/vertexai.py b/mem0/embeddings/vertexai.py index 979edef70..5003b62db 100644 --- a/mem0/embeddings/vertexai.py +++ b/mem0/embeddings/vertexai.py @@ -12,7 +12,7 @@ class VertexAIEmbedding(EmbeddingBase): def __init__(self, config: Optional[BaseEmbedderConfig] = None): super().__init__(config) - self.config.model = self.config.model or "text-embedding-004" + self.config.model = self.config.model or "gemini-embedding-001" self.config.embedding_dims = self.config.embedding_dims or 256 self.embedding_types = { diff --git a/tests/embeddings/test_vertexai_embeddings.py b/tests/embeddings/test_vertexai_embeddings.py index 9f5415267..45644262c 100644 --- a/tests/embeddings/test_vertexai_embeddings.py +++ b/tests/embeddings/test_vertexai_embeddings.py @@ -48,7 +48,7 @@ def mock_text_embedding_input(): @patch("mem0.embeddings.vertexai.TextEmbeddingModel") def test_embed_default_model(mock_text_embedding_model, mock_os_environ, mock_config, mock_text_embedding_input): - mock_config.return_value.model = "text-embedding-004" + mock_config.return_value.model = "gemini-embedding-001" mock_config.return_value.embedding_dims = 256 config = mock_config() @@ -59,7 +59,7 @@ def test_embed_default_model(mock_text_embedding_model, mock_os_environ, mock_co embedder.embed("Hello world") mock_text_embedding_input.assert_called_once_with(text="Hello world", task_type="SEMANTIC_SIMILARITY") - mock_text_embedding_model.from_pretrained.assert_called_once_with("text-embedding-004") + mock_text_embedding_model.from_pretrained.assert_called_once_with("gemini-embedding-001") mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_once_with( texts=[mock_text_embedding_input("Hello world")], output_dimensionality=256 @@ -92,7 +92,7 @@ def test_embed_custom_model(mock_text_embedding_model, mock_os_environ, mock_con def test_embed_with_memory_action( mock_text_embedding_model, mock_os_environ, mock_config, mock_embedding_types, mock_text_embedding_input ): - mock_config.return_value.model = "text-embedding-004" + mock_config.return_value.model = "gemini-embedding-001" mock_config.return_value.embedding_dims = 256 for embedding_type in mock_embedding_types: @@ -103,7 +103,7 @@ def test_embed_with_memory_action( config = mock_config() embedder = VertexAIEmbedding(config) - mock_text_embedding_model.from_pretrained.assert_called_with("text-embedding-004") + mock_text_embedding_model.from_pretrained.assert_called_with("gemini-embedding-001") for memory_action in ["add", "update", "search"]: embedder.embed("Hello world", memory_action=memory_action) @@ -151,7 +151,7 @@ def test_embed_with_different_dimensions(mock_text_embedding_model, mock_os_envi @patch("mem0.embeddings.vertexai.TextEmbeddingModel") def test_invalid_memory_action(mock_text_embedding_model, mock_config): - mock_config.return_value.model = "text-embedding-004" + mock_config.return_value.model = "gemini-embedding-001" mock_config.return_value.embedding_dims = 256 config = mock_config()