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

fix(memory): guard entity embed_batch count mismatch in v3 add pipeline (#5604)

Browse Source 
Co-authored-by: Kartik <kartik.labhshetwar@mem0.ai>
This commit is contained in:
Bartok
2026-06-18 17:22:44 +05:30
committed by GitHub
co-authored by Kartik
parent 1dcee153b9
commit 7ae6a8c36a
2 changed files with 152 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0/memory/main.py
+21
View File
@@ -993,6 +993,17 @@ class Memory(MemoryBase):
except Exception:
entity_embeddings.append(None)
if len(entity_embeddings) != len(ordered_keys):
logger.warning(
"embed_batch returned %d vectors for %d entity texts — "
"padding/truncating to avoid dropping entity links",
len(entity_embeddings),
len(ordered_keys),
)
entity_embeddings = list(entity_embeddings[: len(ordered_keys)])
entity_embeddings += [None] * (len(ordered_keys) - len(entity_embeddings))
# 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:
@@ -2520,6 +2531,16 @@ class AsyncMemory(MemoryBase):
except Exception:
entity_embeddings.append(None)
if len(entity_embeddings) != len(ordered_keys):
logger.warning(
"embed_batch returned %d vectors for %d entity texts — "
"padding/truncating to avoid dropping entity links",
len(entity_embeddings),
len(ordered_keys),
)
entity_embeddings = list(entity_embeddings[: len(ordered_keys)])
entity_embeddings += [None] * (len(ordered_keys) - len(entity_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)