perf: batch entity linking — embed, search, insert in bulk instead of per-entity
- Replace sequential _upsert_entity() loop with batched pipeline: 1. Global dedup: collect unique entities across all memories 2. Single embed_batch() call for all entity texts 3. Batch search via search_batch() (Qdrant query_batch_points) 4. Batch insert for new entities, individual updates for existing - Add search_batch() to VectorStoreBase (sequential fallback) - Add native search_batch() to Qdrant using query_batch_points - ~6-10x reduction in entity linking network round trips
This commit is contained in:
+82
-1
@@ -679,10 +679,91 @@ class Memory(MemoryBase):
|
|||||||
try:
|
try:
|
||||||
all_texts = [r[1] for r in records]
|
all_texts = [r[1] for r in records]
|
||||||
all_entities = extract_entities_batch(all_texts)
|
all_entities = extract_entities_batch(all_texts)
|
||||||
|
|
||||||
|
# 7a: Global dedup — collect unique entities across all memories
|
||||||
|
global_entities = {} # normalized_key -> (entity_type, entity_text, set of memory_ids)
|
||||||
for idx, (memory_id, text, embedding, payload) in enumerate(records):
|
for idx, (memory_id, text, embedding, payload) in enumerate(records):
|
||||||
entities = all_entities[idx] if idx < len(all_entities) else []
|
entities = all_entities[idx] if idx < len(all_entities) else []
|
||||||
for entity_type, entity_text in entities:
|
for entity_type, entity_text in entities:
|
||||||
self._upsert_entity(entity_text, entity_type, memory_id, search_filters)
|
key = entity_text.strip().lower()
|
||||||
|
if key in global_entities:
|
||||||
|
global_entities[key][2].add(memory_id)
|
||||||
|
else:
|
||||||
|
global_entities[key] = [entity_type, entity_text, {memory_id}]
|
||||||
|
|
||||||
|
if global_entities:
|
||||||
|
ordered_keys = list(global_entities.keys())
|
||||||
|
entity_texts = [global_entities[k][1] for k in ordered_keys]
|
||||||
|
|
||||||
|
# 7b: Single batch embed for all unique entities
|
||||||
|
try:
|
||||||
|
entity_embeddings = self.embedding_model.embed_batch(entity_texts, "add")
|
||||||
|
except Exception:
|
||||||
|
# Fallback: embed individually, use None for failures
|
||||||
|
entity_embeddings = []
|
||||||
|
for t in entity_texts:
|
||||||
|
try:
|
||||||
|
entity_embeddings.append(self.embedding_model.embed(t, "add"))
|
||||||
|
except Exception:
|
||||||
|
entity_embeddings.append(None)
|
||||||
|
|
||||||
|
# Filter out entities with failed embeddings
|
||||||
|
valid = [(i, k) for i, k in enumerate(ordered_keys) if entity_embeddings[i] is not None]
|
||||||
|
if valid:
|
||||||
|
valid_indices, valid_keys = zip(*valid)
|
||||||
|
valid_vectors = [entity_embeddings[i] for i in valid_indices]
|
||||||
|
|
||||||
|
# 7c: Batch search for existing entities
|
||||||
|
valid_texts = [global_entities[k][1] for k in valid_keys]
|
||||||
|
existing_matches = self.entity_store.search_batch(
|
||||||
|
queries=valid_texts,
|
||||||
|
vectors_list=valid_vectors,
|
||||||
|
limit=1,
|
||||||
|
filters=search_filters,
|
||||||
|
)
|
||||||
|
|
||||||
|
# 7d: Separate into inserts vs updates
|
||||||
|
to_insert_vectors, to_insert_ids, to_insert_payloads = [], [], []
|
||||||
|
for j, key in enumerate(valid_keys):
|
||||||
|
entity_type, entity_text, memory_ids = global_entities[key]
|
||||||
|
matches = existing_matches[j] if j < len(existing_matches) else []
|
||||||
|
|
||||||
|
if matches and matches[0].score >= 0.95:
|
||||||
|
# Update existing entity
|
||||||
|
match = matches[0]
|
||||||
|
payload = match.payload or {}
|
||||||
|
linked = set(payload.get("linked_memory_ids", []))
|
||||||
|
linked |= memory_ids
|
||||||
|
payload["linked_memory_ids"] = sorted(linked)
|
||||||
|
try:
|
||||||
|
self.entity_store.update(
|
||||||
|
vector_id=match.id,
|
||||||
|
vector=None,
|
||||||
|
payload=payload,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.debug(f"Entity update failed for '{entity_text}': {e}")
|
||||||
|
else:
|
||||||
|
# New entity — collect for batch insert
|
||||||
|
to_insert_vectors.append(valid_vectors[j])
|
||||||
|
to_insert_ids.append(str(uuid.uuid4()))
|
||||||
|
to_insert_payloads.append({
|
||||||
|
"data": entity_text,
|
||||||
|
"entity_type": entity_type,
|
||||||
|
"linked_memory_ids": sorted(memory_ids),
|
||||||
|
**search_filters,
|
||||||
|
})
|
||||||
|
|
||||||
|
# 7e: Single batch insert for all new entities
|
||||||
|
if to_insert_vectors:
|
||||||
|
try:
|
||||||
|
self.entity_store.insert(
|
||||||
|
vectors=to_insert_vectors,
|
||||||
|
ids=to_insert_ids,
|
||||||
|
payloads=to_insert_payloads,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Batch entity insert failed: {e}")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"Batch entity linking failed: {e}")
|
logger.warning(f"Batch entity linking failed: {e}")
|
||||||
|
|
||||||
|
|||||||
@@ -73,3 +73,20 @@ class VectorStoreBase(ABC):
|
|||||||
List of search results with id, score, payload, or None if not supported.
|
List of search results with id, score, payload, or None if not supported.
|
||||||
"""
|
"""
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
def search_batch(self, queries: list, vectors_list: list, limit: int = 1, filters: dict = None):
|
||||||
|
"""Batch search for multiple queries at once.
|
||||||
|
|
||||||
|
Default implementation calls search() sequentially. Override in subclasses
|
||||||
|
with native batch support (e.g., Qdrant query_batch_points).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
queries: List of query texts.
|
||||||
|
vectors_list: List of query vectors (one per query).
|
||||||
|
limit: Maximum results per query.
|
||||||
|
filters: Optional metadata filters applied to all queries.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of result lists, one per query.
|
||||||
|
"""
|
||||||
|
return [self.search(q, v, limit=limit, filters=filters) for q, v in zip(queries, vectors_list)]
|
||||||
|
|||||||
@@ -229,6 +229,23 @@ class Qdrant(VectorStoreBase):
|
|||||||
)
|
)
|
||||||
return hits.points
|
return hits.points
|
||||||
|
|
||||||
|
def search_batch(self, queries: list, vectors_list: list, limit: int = 1, filters: dict = None):
|
||||||
|
"""Batch search using Qdrant's query_batch_points for efficiency."""
|
||||||
|
query_filter = self._create_filter(filters) if filters else None
|
||||||
|
requests = [
|
||||||
|
models.QueryRequest(query=vec, filter=query_filter, limit=limit)
|
||||||
|
for vec in vectors_list
|
||||||
|
]
|
||||||
|
try:
|
||||||
|
results = self.client.query_batch_points(
|
||||||
|
collection_name=self.collection_name,
|
||||||
|
requests=requests,
|
||||||
|
)
|
||||||
|
return [r.points for r in results]
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Batch search failed, falling back to sequential: {e}")
|
||||||
|
return [self.search(q, v, limit=limit, filters=filters) for q, v in zip(queries, vectors_list)]
|
||||||
|
|
||||||
def keyword_search(self, query, limit=5, filters=None):
|
def keyword_search(self, query, limit=5, filters=None):
|
||||||
"""
|
"""
|
||||||
Search using BM25 sparse vectors for keyword-based retrieval.
|
Search using BM25 sparse vectors for keyword-based retrieval.
|
||||||
|
|||||||
Reference in New Issue
Block a user