JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: add text_lemmatized to create/update memory + fix keyword search across vector stores

Browse Source 
- 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
This commit is contained in:
Soumil Rathi
2026-04-02 10:34:03 -07:00
parent 4534acdfbe
commit 70c62ea54f
5 changed files with 97 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0/memory/main.py
+3
View File
@@ -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()
mem0/vector_stores/azure_mysql.py
+16 -1
View File
@@ -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.
mem0/vector_stores/milvus.py
+34 -5
View File
@@ -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):
mem0/vector_stores/mongodb.py
+35
View File
@@ -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}")
mem0/vector_stores/pgvector.py
+9 -1
View File
@@ -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: