Init embedding_model_dims in all vectordbs (#2572)

This commit is contained in:
Dev Khant
2025-04-19 10:53:01 +05:30
committed by GitHub
parent 78912928bc
commit 3ee4768c14
7 changed files with 9 additions and 5 deletions

View File

@@ -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]
)

View File

@@ -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"}}},
}
},

View File

@@ -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"}}},

View File

@@ -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()

View File

@@ -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):

View File

@@ -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}",

View File

@@ -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]: