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:
Soumil Rathi
2026-04-02 19:32:29 -07:00
parent 41532d8850
commit cb64d978fc
3 changed files with 116 additions and 1 deletions
+82 -1
View File
@@ -679,10 +679,91 @@ class Memory(MemoryBase):
try:
all_texts = [r[1] for r in records]
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):
entities = all_entities[idx] if idx < len(all_entities) else []
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:
logger.warning(f"Batch entity linking failed: {e}")
+17
View File
@@ -73,3 +73,20 @@ class VectorStoreBase(ABC):
List of search results with id, score, payload, or None if not supported.
"""
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)]
+17
View File
@@ -229,6 +229,23 @@ class Qdrant(VectorStoreBase):
)
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):
"""
Search using BM25 sparse vectors for keyword-based retrieval.