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:
|
||||
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}")
|
||||
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user