From 70c62ea54fe948271c7b0233eb42247a62e9c7dd Mon Sep 17 00:00:00 2001 From: Soumil Rathi Date: Thu, 2 Apr 2026 10:34:03 -0700 Subject: [PATCH] fix: add text_lemmatized to create/update memory + fix keyword search across vector stores - Add text_lemmatized to _create_memory() and _update_memory() (sync + async) so BM25/keyword search works for all code paths, not just the batch pipeline - pgvector: add GIN index on text_lemmatized for fast full-text search - pgvector: add error logging to keyword_search - milvus: add sparse BM25 field + function to create_col() so keyword_search works - mongodb: auto-create Atlas Search text index in create_col() - azure_mysql: add generated column + FULLTEXT index for keyword_search --- mem0/memory/main.py | 3 +++ mem0/vector_stores/azure_mysql.py | 17 +++++++++++++- mem0/vector_stores/milvus.py | 39 +++++++++++++++++++++++++++---- mem0/vector_stores/mongodb.py | 35 +++++++++++++++++++++++++++ mem0/vector_stores/pgvector.py | 10 +++++++- 5 files changed, 97 insertions(+), 7 deletions(-) diff --git a/mem0/memory/main.py b/mem0/memory/main.py index e5481783c..23a4f5a1a 100644 --- a/mem0/memory/main.py +++ b/mem0/memory/main.py @@ -1320,6 +1320,7 @@ class Memory(MemoryBase): metadata["data"] = data metadata["hash"] = hashlib.md5(data.encode()).hexdigest() metadata["created_at"] = datetime.now(pytz.timezone("US/Pacific")).isoformat() + metadata["text_lemmatized"] = lemmatize_for_bm25(data) self.vector_store.insert( vectors=[embeddings], @@ -1391,6 +1392,7 @@ class Memory(MemoryBase): new_metadata["data"] = data new_metadata["hash"] = hashlib.md5(data.encode()).hexdigest() + new_metadata["text_lemmatized"] = lemmatize_for_bm25(data) new_metadata["created_at"] = existing_memory.payload.get("created_at") new_metadata["updated_at"] = datetime.now(pytz.timezone("US/Pacific")).isoformat() @@ -2610,6 +2612,7 @@ class AsyncMemory(MemoryBase): new_metadata["data"] = data new_metadata["hash"] = hashlib.md5(data.encode()).hexdigest() + new_metadata["text_lemmatized"] = lemmatize_for_bm25(data) new_metadata["created_at"] = existing_memory.payload.get("created_at") new_metadata["updated_at"] = datetime.now(pytz.timezone("US/Pacific")).isoformat() diff --git a/mem0/vector_stores/azure_mysql.py b/mem0/vector_stores/azure_mysql.py index d95bcd50f..77f7db9e0 100644 --- a/mem0/vector_stores/azure_mysql.py +++ b/mem0/vector_stores/azure_mysql.py @@ -178,17 +178,32 @@ class AzureMySQL(VectorStoreBase): dims = vector_size or self.embedding_model_dims with self._get_cursor(commit=True) as cur: - # Create table with vector column + # Create table with vector column and a generated column for fulltext keyword search cur.execute(f""" CREATE TABLE IF NOT EXISTS `{table_name}` ( id VARCHAR(255) PRIMARY KEY, vector JSON, payload JSON, + text_lemmatized VARCHAR(1000) GENERATED ALWAYS AS + (CAST(payload->>'$.text_lemmatized' AS CHAR(1000))) STORED, INDEX idx_payload_keys ((CAST(payload AS CHAR(255)) ARRAY)) ) """) logger.info(f"Created collection '{table_name}' with vector dimension {dims}") + # Add FULLTEXT index on text_lemmatized for keyword_search() + try: + cur.execute(f""" + CREATE FULLTEXT INDEX ft_text_lemmatized + ON `{table_name}` (text_lemmatized) + """) + logger.info(f"Created FULLTEXT index on '{table_name}.text_lemmatized'") + except Exception as e: + logger.debug( + f"Could not create FULLTEXT index on '{table_name}.text_lemmatized': {e}. " + "It may already exist or FULLTEXT may not be supported." + ) + def insert(self, vectors: List[List[float]], payloads: Optional[List[Dict]] = None, ids: Optional[List[str]] = None): """ Insert vectors into the collection. diff --git a/mem0/vector_stores/milvus.py b/mem0/vector_stores/milvus.py index ee0539821..711634688 100644 --- a/mem0/vector_stores/milvus.py +++ b/mem0/vector_stores/milvus.py @@ -11,7 +11,7 @@ try: except ImportError: raise ImportError("The 'pymilvus' library is required. Please install it using 'pip install pymilvus'.") -from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusClient +from pymilvus import CollectionSchema, DataType, FieldSchema, Function, FunctionType, MilvusClient logger = logging.getLogger(__name__) @@ -73,14 +73,34 @@ class MilvusDB(VectorStoreBase): FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=512), FieldSchema(name="vectors", dtype=DataType.FLOAT_VECTOR, dim=vector_size), FieldSchema(name="metadata", dtype=DataType.JSON), + # Text field for BM25 full-text search (auto-tokenized by Milvus analyzer) + FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535, enable_analyzer=True), + # Sparse vector field populated automatically by the BM25 function below + FieldSchema(name="sparse", dtype=DataType.SPARSE_FLOAT_VECTOR), ] schema = CollectionSchema(fields, enable_dynamic_field=True) - index = self.client.prepare_index_params( + # Add BM25 function so Milvus auto-generates sparse vectors from the text field + bm25_function = Function( + name="bm25", + input_field_names=["text"], + output_field_names=["sparse"], + function_type=FunctionType.BM25, + ) + schema.add_function(bm25_function) + + index_params = self.client.prepare_index_params() + index_params.add_index( field_name="vectors", metric_type=metric_type, index_type="AUTOINDEX", index_name="vector_index" ) - self.client.create_collection(collection_name=collection_name, schema=schema, index_params=index) + index_params.add_index( + field_name="sparse", + index_type="SPARSE_INVERTED_INDEX", + metric_type="BM25", + index_name="sparse_index", + ) + self.client.create_collection(collection_name=collection_name, schema=schema, index_params=index_params) def insert(self, ids, vectors, payloads, **kwargs: Optional[dict[str, any]]): """Insert vectors into a collection. @@ -92,7 +112,13 @@ class MilvusDB(VectorStoreBase): """ # Batch insert all records at once for better performance and consistency data = [ - {"id": idx, "vectors": embedding, "metadata": metadata} + { + "id": idx, + "vectors": embedding, + "metadata": metadata, + # Populate the text field for BM25 sparse search; prefer lemmatized text, fall back to raw data + "text": (metadata.get("text_lemmatized") or metadata.get("data", ""))[:65535] if metadata else "", + } for idx, embedding, metadata in zip(ids, vectors, payloads) ] self.client.insert(collection_name=self.collection_name, data=data, **kwargs) @@ -214,7 +240,10 @@ class MilvusDB(VectorStoreBase): vector (List[float], optional): Updated vector. payload (Dict, optional): Updated payload. """ - schema = {"id": vector_id, "vectors": vector, "metadata": payload} + text = "" + if payload: + text = (payload.get("text_lemmatized") or payload.get("data", ""))[:65535] + schema = {"id": vector_id, "vectors": vector, "metadata": payload, "text": text} self.client.upsert(collection_name=self.collection_name, data=schema) def get(self, vector_id): diff --git a/mem0/vector_stores/mongodb.py b/mem0/vector_stores/mongodb.py index e92da83c4..3ae1a1801 100644 --- a/mem0/vector_stores/mongodb.py +++ b/mem0/vector_stores/mongodb.py @@ -86,6 +86,41 @@ class MongoDB(VectorStoreBase): logger.info( f"Search index '{self.index_name}' created successfully for collection '{self.collection_name}'." ) + + # Create Atlas Search text index for keyword_search() + text_index_name = f"{self.collection_name}_text_search_index" + try: + found_text_indexes = list(collection.list_search_indexes(name=text_index_name)) + if not found_text_indexes: + text_search_index_model = SearchIndexModel( + name=text_index_name, + definition={ + "mappings": { + "dynamic": False, + "fields": { + "payload": { + "type": "document", + "fields": { + "data": {"type": "string"}, + "text_lemmatized": {"type": "string"}, + }, + } + }, + } + }, + ) + collection.create_search_index(text_search_index_model) + logger.info( + f"Text search index '{text_index_name}' created successfully for collection '{self.collection_name}'." + ) + else: + logger.info(f"Text search index '{text_index_name}' already exists in collection '{self.collection_name}'.") + except Exception as e: + logger.warning( + f"Could not create text search index '{text_index_name}': {e}. " + "Atlas Search may not be available. keyword_search() will not work." + ) + return collection except PyMongoError as e: logger.error(f"Error creating collection and search index: {e}") diff --git a/mem0/vector_stores/pgvector.py b/mem0/vector_stores/pgvector.py index 890ddc0b2..5bf3524d2 100644 --- a/mem0/vector_stores/pgvector.py +++ b/mem0/vector_stores/pgvector.py @@ -179,6 +179,13 @@ class PGVector(VectorStoreBase): USING hnsw (vector vector_cosine_ops) """ ) + cur.execute( + f""" + CREATE INDEX IF NOT EXISTS {self.collection_name}_text_lemmatized_idx + ON {self.collection_name} + USING gin(to_tsvector('simple', payload->>'text_lemmatized')); + """ + ) def insert(self, vectors: list[list[float]], payloads=None, ids=None) -> None: logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}") @@ -283,7 +290,8 @@ class PGVector(VectorStoreBase): results = cur.fetchall() return [OutputData(id=str(r[0]), score=float(r[1]), payload=r[2]) for r in results] - except Exception: + except Exception as e: + logger.debug(f"Keyword search failed: {e}") return None def delete(self, vector_id: str) -> None: