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
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
+3
View File
@@ -1320,6 +1320,7 @@ class Memory(MemoryBase):
metadata["data"] = data metadata["data"] = data
metadata["hash"] = hashlib.md5(data.encode()).hexdigest() metadata["hash"] = hashlib.md5(data.encode()).hexdigest()
metadata["created_at"] = datetime.now(pytz.timezone("US/Pacific")).isoformat() metadata["created_at"] = datetime.now(pytz.timezone("US/Pacific")).isoformat()
metadata["text_lemmatized"] = lemmatize_for_bm25(data)
self.vector_store.insert( self.vector_store.insert(
vectors=[embeddings], vectors=[embeddings],
@@ -1391,6 +1392,7 @@ class Memory(MemoryBase):
new_metadata["data"] = data new_metadata["data"] = data
new_metadata["hash"] = hashlib.md5(data.encode()).hexdigest() 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["created_at"] = existing_memory.payload.get("created_at")
new_metadata["updated_at"] = datetime.now(pytz.timezone("US/Pacific")).isoformat() new_metadata["updated_at"] = datetime.now(pytz.timezone("US/Pacific")).isoformat()
@@ -2610,6 +2612,7 @@ class AsyncMemory(MemoryBase):
new_metadata["data"] = data new_metadata["data"] = data
new_metadata["hash"] = hashlib.md5(data.encode()).hexdigest() 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["created_at"] = existing_memory.payload.get("created_at")
new_metadata["updated_at"] = datetime.now(pytz.timezone("US/Pacific")).isoformat() new_metadata["updated_at"] = datetime.now(pytz.timezone("US/Pacific")).isoformat()
+16 -1
View File
@@ -178,17 +178,32 @@ class AzureMySQL(VectorStoreBase):
dims = vector_size or self.embedding_model_dims dims = vector_size or self.embedding_model_dims
with self._get_cursor(commit=True) as cur: 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""" cur.execute(f"""
CREATE TABLE IF NOT EXISTS `{table_name}` ( CREATE TABLE IF NOT EXISTS `{table_name}` (
id VARCHAR(255) PRIMARY KEY, id VARCHAR(255) PRIMARY KEY,
vector JSON, vector JSON,
payload 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)) INDEX idx_payload_keys ((CAST(payload AS CHAR(255)) ARRAY))
) )
""") """)
logger.info(f"Created collection '{table_name}' with vector dimension {dims}") 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): def insert(self, vectors: List[List[float]], payloads: Optional[List[Dict]] = None, ids: Optional[List[str]] = None):
""" """
Insert vectors into the collection. Insert vectors into the collection.
+34 -5
View File
@@ -11,7 +11,7 @@ try:
except ImportError: except ImportError:
raise ImportError("The 'pymilvus' library is required. Please install it using 'pip install pymilvus'.") 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__) 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="id", dtype=DataType.VARCHAR, is_primary=True, max_length=512),
FieldSchema(name="vectors", dtype=DataType.FLOAT_VECTOR, dim=vector_size), FieldSchema(name="vectors", dtype=DataType.FLOAT_VECTOR, dim=vector_size),
FieldSchema(name="metadata", dtype=DataType.JSON), 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) 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" 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]]): def insert(self, ids, vectors, payloads, **kwargs: Optional[dict[str, any]]):
"""Insert vectors into a collection. """Insert vectors into a collection.
@@ -92,7 +112,13 @@ class MilvusDB(VectorStoreBase):
""" """
# Batch insert all records at once for better performance and consistency # Batch insert all records at once for better performance and consistency
data = [ 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) for idx, embedding, metadata in zip(ids, vectors, payloads)
] ]
self.client.insert(collection_name=self.collection_name, data=data, **kwargs) self.client.insert(collection_name=self.collection_name, data=data, **kwargs)
@@ -214,7 +240,10 @@ class MilvusDB(VectorStoreBase):
vector (List[float], optional): Updated vector. vector (List[float], optional): Updated vector.
payload (Dict, optional): Updated payload. 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) self.client.upsert(collection_name=self.collection_name, data=schema)
def get(self, vector_id): def get(self, vector_id):
+35
View File
@@ -86,6 +86,41 @@ class MongoDB(VectorStoreBase):
logger.info( logger.info(
f"Search index '{self.index_name}' created successfully for collection '{self.collection_name}'." 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 return collection
except PyMongoError as e: except PyMongoError as e:
logger.error(f"Error creating collection and search index: {e}") logger.error(f"Error creating collection and search index: {e}")
+9 -1
View File
@@ -179,6 +179,13 @@ class PGVector(VectorStoreBase):
USING hnsw (vector vector_cosine_ops) 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: def insert(self, vectors: list[list[float]], payloads=None, ids=None) -> None:
logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}") logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}")
@@ -283,7 +290,8 @@ class PGVector(VectorStoreBase):
results = cur.fetchall() results = cur.fetchall()
return [OutputData(id=str(r[0]), score=float(r[1]), payload=r[2]) for r in results] 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 return None
def delete(self, vector_id: str) -> None: def delete(self, vector_id: str) -> None: