From 3ee4768c14c42b34c18bbf04c67e02e4bb0c0098 Mon Sep 17 00:00:00 2001
From: Dev Khant <devkhant24@gmail.com>
Date: Sat, 19 Apr 2025 10:53:01 +0530
Subject: [PATCH] Init embedding_model_dims in all vectordbs (#2572)

---
 mem0/memory/setup.py                | 2 +-
 mem0/vector_stores/elasticsearch.py | 4 ++--
 mem0/vector_stores/opensearch.py    | 4 ++--
 mem0/vector_stores/pgvector.py      | 1 +
 mem0/vector_stores/qdrant.py        | 1 +
 mem0/vector_stores/redis.py         | 1 +
 mem0/vector_stores/weaviate.py      | 1 +
 7 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/mem0/memory/setup.py b/mem0/memory/setup.py
index b4fa99ae..90aab55b 100644
--- a/mem0/memory/setup.py
+++ b/mem0/memory/setup.py
@@ -46,7 +46,7 @@ def get_or_create_user_id(vector_store):
 
     # If we get here, we need to insert the user_id
     try:
-        dims = getattr(vector_store, "embedding_model_dims", 1)
+        dims = getattr(vector_store, "embedding_model_dims", 1536)
         vector_store.insert(
             vectors=[[0.0] * dims], payloads=[{"user_id": user_id, "type": "user_identity"}], ids=[VECTOR_ID]
         )
diff --git a/mem0/vector_stores/elasticsearch.py b/mem0/vector_stores/elasticsearch.py
index 8a535429..7b7cc2b4 100644
--- a/mem0/vector_stores/elasticsearch.py
+++ b/mem0/vector_stores/elasticsearch.py
@@ -40,7 +40,7 @@ class ElasticsearchDB(VectorStoreBase):
             )
 
         self.collection_name = config.collection_name
-        self.vector_dim = config.embedding_model_dims
+        self.embedding_model_dims = config.embedding_model_dims
 
         # Create index only if auto_create_index is True
         if config.auto_create_index:
@@ -58,7 +58,7 @@ class ElasticsearchDB(VectorStoreBase):
             "mappings": {
                 "properties": {
                     "text": {"type": "text"},
-                    "vector": {"type": "dense_vector", "dims": self.vector_dim, "index": True, "similarity": "cosine"},
+                    "vector": {"type": "dense_vector", "dims": self.embedding_model_dims, "index": True, "similarity": "cosine"},
                     "metadata": {"type": "object", "properties": {"user_id": {"type": "keyword"}}},
                 }
             },
diff --git a/mem0/vector_stores/opensearch.py b/mem0/vector_stores/opensearch.py
index dca5b287..18b0063c 100644
--- a/mem0/vector_stores/opensearch.py
+++ b/mem0/vector_stores/opensearch.py
@@ -37,7 +37,7 @@ class OpenSearchDB(VectorStoreBase):
         )
 
         self.collection_name = config.collection_name
-        self.vector_dim = config.embedding_model_dims
+        self.embedding_model_dims = config.embedding_model_dims
 
         # Create index only if auto_create_index is True
         if config.auto_create_index:
@@ -54,7 +54,7 @@ class OpenSearchDB(VectorStoreBase):
                     "text": {"type": "text"},
                     "vector": {
                         "type": "knn_vector",
-                        "dimension": self.vector_dim,
+                        "dimension": self.embedding_model_dims,
                         "method": {"engine": "lucene", "name": "hnsw", "space_type": "cosinesimil"},
                     },
                     "metadata": {"type": "object", "properties": {"user_id": {"type": "keyword"}}},
diff --git a/mem0/vector_stores/pgvector.py b/mem0/vector_stores/pgvector.py
index 3997f0af..85ff3e14 100644
--- a/mem0/vector_stores/pgvector.py
+++ b/mem0/vector_stores/pgvector.py
@@ -51,6 +51,7 @@ class PGVector(VectorStoreBase):
         self.collection_name = collection_name
         self.use_diskann = diskann
         self.use_hnsw = hnsw
+        self.embedding_model_dims = embedding_model_dims
 
         self.conn = psycopg2.connect(dbname=dbname, user=user, password=password, host=host, port=port)
         self.cur = self.conn.cursor()
diff --git a/mem0/vector_stores/qdrant.py b/mem0/vector_stores/qdrant.py
index 7e16089f..d50054ca 100644
--- a/mem0/vector_stores/qdrant.py
+++ b/mem0/vector_stores/qdrant.py
@@ -66,6 +66,7 @@ class Qdrant(VectorStoreBase):
             self.client = QdrantClient(**params)
 
         self.collection_name = collection_name
+        self.embedding_model_dims = embedding_model_dims
         self.create_col(embedding_model_dims, on_disk)
 
     def create_col(self, vector_size: int, on_disk: bool, distance: Distance = Distance.COSINE):
diff --git a/mem0/vector_stores/redis.py b/mem0/vector_stores/redis.py
index 8947beeb..25cdcbb5 100644
--- a/mem0/vector_stores/redis.py
+++ b/mem0/vector_stores/redis.py
@@ -59,6 +59,7 @@ class RedisDB(VectorStoreBase):
             collection_name (str): Collection name.
             embedding_model_dims (int): Embedding model dimensions.
         """
+        self.embedding_model_dims = embedding_model_dims
         index_schema = {
             "name": collection_name,
             "prefix": f"mem0:{collection_name}",
diff --git a/mem0/vector_stores/weaviate.py b/mem0/vector_stores/weaviate.py
index 0988b103..cdf20dee 100644
--- a/mem0/vector_stores/weaviate.py
+++ b/mem0/vector_stores/weaviate.py
@@ -57,6 +57,7 @@ class Weaviate(VectorStoreBase):
             )
 
         self.collection_name = collection_name
+        self.embedding_model_dims = embedding_model_dims
         self.create_col(embedding_model_dims)
 
     def _parse_output(self, data: Dict) -> List[OutputData]: