fix(oss): parallelize entity boost searches in Memory.search (#5377)
This commit is contained in:
@@ -44,7 +44,9 @@ export class LMStudioEmbedder implements Embedder {
|
|||||||
input: normalized,
|
input: normalized,
|
||||||
encoding_format: "float",
|
encoding_format: "float",
|
||||||
});
|
});
|
||||||
return response.data.map((item) => item.embedding);
|
return response.data
|
||||||
|
.sort((a, b) => a.index - b.index)
|
||||||
|
.map((item) => item.embedding);
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
const message = err instanceof Error ? err.message : String(err);
|
const message = err instanceof Error ? err.message : String(err);
|
||||||
throw new Error(`LM Studio embedder failed: ${message}`);
|
throw new Error(`LM Studio embedder failed: ${message}`);
|
||||||
|
|||||||
@@ -1169,17 +1169,35 @@ export class Memory {
|
|||||||
|
|
||||||
if (deduped.length > 0) {
|
if (deduped.length > 0) {
|
||||||
const entityStore = await this.getEntityStore();
|
const entityStore = await this.getEntityStore();
|
||||||
|
const entitySearchFilters: Record<string, any> = {};
|
||||||
|
for (const k of ["user_id", "agent_id", "run_id"] as const) {
|
||||||
|
if (effectiveFilters[k])
|
||||||
|
entitySearchFilters[k] = effectiveFilters[k];
|
||||||
|
}
|
||||||
|
const entityTexts = deduped.map((e) => e.text);
|
||||||
|
const embeddings = await this.embedder.embedBatch(entityTexts);
|
||||||
|
|
||||||
for (const entity of deduped) {
|
if (embeddings.length !== entityTexts.length) {
|
||||||
try {
|
console.warn(
|
||||||
const entityEmbedding = await this.embedder.embed(entity.text);
|
`embedBatch returned ${embeddings.length} vectors for ${entityTexts.length} texts — skipping entity boost`,
|
||||||
const matches = await entityStore.search(
|
);
|
||||||
entityEmbedding,
|
} else {
|
||||||
500,
|
const searchResults = await Promise.allSettled(
|
||||||
effectiveFilters,
|
deduped.map((_, i) =>
|
||||||
);
|
entityStore.search(embeddings[i], 500, entitySearchFilters),
|
||||||
|
),
|
||||||
|
);
|
||||||
|
|
||||||
for (const match of matches) {
|
for (const result of searchResults) {
|
||||||
|
if (result.status === "rejected") {
|
||||||
|
console.warn(
|
||||||
|
"Entity boost search failed for one entity:",
|
||||||
|
result.reason,
|
||||||
|
);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
for (const match of result.value) {
|
||||||
const similarity = match.score ?? 0;
|
const similarity = match.score ?? 0;
|
||||||
if (similarity < 0.5) continue;
|
if (similarity < 0.5) continue;
|
||||||
|
|
||||||
@@ -1187,7 +1205,6 @@ export class Memory {
|
|||||||
const linkedMemoryIds = payload.linkedMemoryIds ?? [];
|
const linkedMemoryIds = payload.linkedMemoryIds ?? [];
|
||||||
if (!Array.isArray(linkedMemoryIds)) continue;
|
if (!Array.isArray(linkedMemoryIds)) continue;
|
||||||
|
|
||||||
// Spread-attenuated boost
|
|
||||||
const numLinked = Math.max(linkedMemoryIds.length, 1);
|
const numLinked = Math.max(linkedMemoryIds.length, 1);
|
||||||
const memoryCountWeight =
|
const memoryCountWeight =
|
||||||
1.0 / (1.0 + 0.001 * (numLinked - 1) ** 2);
|
1.0 / (1.0 + 0.001 * (numLinked - 1) ** 2);
|
||||||
@@ -1204,8 +1221,6 @@ export class Memory {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
} catch (e) {
|
|
||||||
// Individual entity boost failed — continue
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,276 @@
|
|||||||
|
/**
|
||||||
|
* Entity boost parallelism tests (#5214).
|
||||||
|
*
|
||||||
|
* Verifies that entity boost searches run concurrently via Promise.allSettled,
|
||||||
|
* scoring is preserved, and individual entity failures don't abort others.
|
||||||
|
*/
|
||||||
|
/// <reference types="jest" />
|
||||||
|
import { Memory } from "../src/memory";
|
||||||
|
import { ENTITY_BOOST_WEIGHT } from "../src/utils/scoring";
|
||||||
|
import type { VectorStoreResult } from "../src/types";
|
||||||
|
|
||||||
|
jest.setTimeout(15000);
|
||||||
|
|
||||||
|
jest.mock("../src/embeddings/google", () => ({
|
||||||
|
GoogleEmbedder: jest.fn(),
|
||||||
|
}));
|
||||||
|
jest.mock("../src/llms/google", () => ({
|
||||||
|
GoogleLLM: jest.fn(),
|
||||||
|
}));
|
||||||
|
|
||||||
|
jest.mock("../src/llms/openai", () => ({
|
||||||
|
OpenAILLM: jest.fn().mockImplementation(() => ({
|
||||||
|
generateResponse: jest.fn().mockResolvedValue(
|
||||||
|
JSON.stringify({
|
||||||
|
memory: [{ id: "0", text: "fact", attributed_to: "user" }],
|
||||||
|
}),
|
||||||
|
),
|
||||||
|
})),
|
||||||
|
}));
|
||||||
|
|
||||||
|
const mockEmbedding = new Array(1536).fill(0.1);
|
||||||
|
jest.mock("../src/embeddings/openai", () => ({
|
||||||
|
OpenAIEmbedder: jest.fn().mockImplementation(() => ({
|
||||||
|
embed: jest.fn().mockResolvedValue(mockEmbedding),
|
||||||
|
embedBatch: jest
|
||||||
|
.fn()
|
||||||
|
.mockImplementation((texts: string[]) =>
|
||||||
|
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||||
|
),
|
||||||
|
embeddingDims: 1536,
|
||||||
|
})),
|
||||||
|
}));
|
||||||
|
|
||||||
|
function makeMatch(
|
||||||
|
id: string,
|
||||||
|
score: number,
|
||||||
|
linkedMemoryIds: string[],
|
||||||
|
): VectorStoreResult {
|
||||||
|
return { id, score, payload: { linkedMemoryIds } };
|
||||||
|
}
|
||||||
|
|
||||||
|
function createMemory(): Memory {
|
||||||
|
return new Memory({
|
||||||
|
version: "v1.1",
|
||||||
|
embedder: {
|
||||||
|
provider: "openai",
|
||||||
|
config: { apiKey: "test-key", model: "text-embedding-3-small" },
|
||||||
|
},
|
||||||
|
vectorStore: {
|
||||||
|
provider: "memory",
|
||||||
|
config: {
|
||||||
|
collectionName: `test-entity-${Date.now()}-${Math.random()}`,
|
||||||
|
dimension: 1536,
|
||||||
|
dbPath: ":memory:",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
llm: {
|
||||||
|
provider: "openai",
|
||||||
|
config: { apiKey: "test-key", model: "gpt-5-mini" },
|
||||||
|
},
|
||||||
|
historyDbPath: ":memory:",
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
describe("Entity boost parallelism (#5214)", () => {
|
||||||
|
let memory: Memory;
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
memory = createMemory();
|
||||||
|
});
|
||||||
|
|
||||||
|
afterEach(async () => {
|
||||||
|
await memory.reset();
|
||||||
|
});
|
||||||
|
|
||||||
|
it("should use Promise.allSettled for concurrent entity searches", async () => {
|
||||||
|
// Spy on Promise.allSettled to confirm it's being used
|
||||||
|
const allSettledSpy = jest.spyOn(Promise, "allSettled");
|
||||||
|
|
||||||
|
// Access internals to inject a mock entity store
|
||||||
|
const m = memory as any;
|
||||||
|
await m._ensureInitialized();
|
||||||
|
|
||||||
|
const mockEntityStore = {
|
||||||
|
search: jest.fn().mockResolvedValue([makeMatch("e1", 0.9, ["mem-1"])]),
|
||||||
|
initialize: jest.fn().mockResolvedValue(undefined),
|
||||||
|
};
|
||||||
|
m._entityStore = mockEntityStore;
|
||||||
|
|
||||||
|
m.embedder = {
|
||||||
|
embed: jest.fn().mockResolvedValue(mockEmbedding),
|
||||||
|
embedBatch: jest
|
||||||
|
.fn()
|
||||||
|
.mockImplementation((texts: string[]) =>
|
||||||
|
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||||
|
),
|
||||||
|
};
|
||||||
|
|
||||||
|
// Mock the vector store to return a semantic result
|
||||||
|
m.vectorStore.search = jest
|
||||||
|
.fn()
|
||||||
|
.mockResolvedValue([
|
||||||
|
{ id: "mem-1", score: 0.8, payload: { data: "test" } },
|
||||||
|
]);
|
||||||
|
m.vectorStore.keywordSearch = jest.fn().mockResolvedValue(null);
|
||||||
|
|
||||||
|
await m.search("alice and bob", { filters: { user_id: "u1" } });
|
||||||
|
|
||||||
|
expect(allSettledSpy).toHaveBeenCalled();
|
||||||
|
allSettledSpy.mockRestore();
|
||||||
|
});
|
||||||
|
|
||||||
|
it("should preserve scoring math with parallel execution", async () => {
|
||||||
|
const m = memory as any;
|
||||||
|
await m._ensureInitialized();
|
||||||
|
|
||||||
|
// Two entities: "alice" links to mem-1, "bob" links to mem-1 and mem-2
|
||||||
|
// mem-1 should get max(alice_boost, bob_boost)
|
||||||
|
const mockEntityStore = {
|
||||||
|
search: jest
|
||||||
|
.fn()
|
||||||
|
.mockImplementation(
|
||||||
|
(_embedding: number[], _topK: number, _filters: any) => {
|
||||||
|
// We need to differentiate by embedding content — but since all
|
||||||
|
// embeddings are identical mocks, we'll use call order
|
||||||
|
const callCount = mockEntityStore.search.mock.calls.length;
|
||||||
|
if (callCount <= 1) {
|
||||||
|
// First entity: "alice"
|
||||||
|
return Promise.resolve([makeMatch("e-alice", 0.9, ["mem-1"])]);
|
||||||
|
}
|
||||||
|
// Second entity: "bob"
|
||||||
|
return Promise.resolve([
|
||||||
|
makeMatch("e-bob", 0.6, ["mem-1", "mem-2"]),
|
||||||
|
]);
|
||||||
|
},
|
||||||
|
),
|
||||||
|
initialize: jest.fn().mockResolvedValue(undefined),
|
||||||
|
};
|
||||||
|
m._entityStore = mockEntityStore;
|
||||||
|
m.embedder = {
|
||||||
|
embed: jest.fn().mockResolvedValue(mockEmbedding),
|
||||||
|
embedBatch: jest
|
||||||
|
.fn()
|
||||||
|
.mockImplementation((texts: string[]) =>
|
||||||
|
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||||
|
),
|
||||||
|
};
|
||||||
|
|
||||||
|
// Semantic results include mem-1 and mem-2
|
||||||
|
m.vectorStore.search = jest.fn().mockResolvedValue([
|
||||||
|
{ id: "mem-1", score: 0.85, payload: { data: "alice memory" } },
|
||||||
|
{ id: "mem-2", score: 0.75, payload: { data: "bob memory" } },
|
||||||
|
]);
|
||||||
|
m.vectorStore.keywordSearch = jest.fn().mockResolvedValue(null);
|
||||||
|
|
||||||
|
const result = await m.search("alice and bob", {
|
||||||
|
filters: { user_id: "u1" },
|
||||||
|
});
|
||||||
|
|
||||||
|
// Entity store was called (parallelized via Promise.allSettled)
|
||||||
|
expect(mockEntityStore.search).toHaveBeenCalled();
|
||||||
|
|
||||||
|
// Results should exist and have scores
|
||||||
|
expect(result.results.length).toBeGreaterThan(0);
|
||||||
|
for (const item of result.results) {
|
||||||
|
expect(typeof item.score).toBe("number");
|
||||||
|
expect(item.score).toBeGreaterThan(0);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
it("should survive one entity search failure without losing other boosts", async () => {
|
||||||
|
const m = memory as any;
|
||||||
|
await m._ensureInitialized();
|
||||||
|
|
||||||
|
let callIndex = 0;
|
||||||
|
const mockEntityStore = {
|
||||||
|
search: jest.fn().mockImplementation(() => {
|
||||||
|
callIndex++;
|
||||||
|
if (callIndex === 1) {
|
||||||
|
return Promise.reject(new Error("provider timeout"));
|
||||||
|
}
|
||||||
|
return Promise.resolve([makeMatch("e-ok", 0.8, ["mem-9"])]);
|
||||||
|
}),
|
||||||
|
initialize: jest.fn().mockResolvedValue(undefined),
|
||||||
|
};
|
||||||
|
m._entityStore = mockEntityStore;
|
||||||
|
m.embedder = {
|
||||||
|
embed: jest.fn().mockResolvedValue(mockEmbedding),
|
||||||
|
embedBatch: jest
|
||||||
|
.fn()
|
||||||
|
.mockImplementation((texts: string[]) =>
|
||||||
|
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||||
|
),
|
||||||
|
};
|
||||||
|
m.vectorStore.search = jest
|
||||||
|
.fn()
|
||||||
|
.mockResolvedValue([
|
||||||
|
{ id: "mem-9", score: 0.85, payload: { data: "surviving memory" } },
|
||||||
|
]);
|
||||||
|
m.vectorStore.keywordSearch = jest.fn().mockResolvedValue(null);
|
||||||
|
|
||||||
|
const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {});
|
||||||
|
|
||||||
|
// "John Smith met Jane Doe" extracts two proper entities
|
||||||
|
const result = await m.search("John Smith met Jane Doe", {
|
||||||
|
filters: { user_id: "u1" },
|
||||||
|
});
|
||||||
|
|
||||||
|
expect(result.results.length).toBeGreaterThan(0);
|
||||||
|
expect(result.results[0].id).toBe("mem-9");
|
||||||
|
// Should log the failure like Python does
|
||||||
|
expect(warnSpy).toHaveBeenCalledWith(
|
||||||
|
"Entity boost search failed for one entity:",
|
||||||
|
expect.any(Error),
|
||||||
|
);
|
||||||
|
warnSpy.mockRestore();
|
||||||
|
});
|
||||||
|
|
||||||
|
it("should call entity searches concurrently, not sequentially", async () => {
|
||||||
|
const m = memory as any;
|
||||||
|
await m._ensureInitialized();
|
||||||
|
|
||||||
|
const concurrency = { current: 0, peak: 0 };
|
||||||
|
|
||||||
|
const mockEntityStore = {
|
||||||
|
search: jest.fn().mockImplementation(() => {
|
||||||
|
concurrency.current++;
|
||||||
|
concurrency.peak = Math.max(concurrency.peak, concurrency.current);
|
||||||
|
return new Promise<VectorStoreResult[]>((resolve) => {
|
||||||
|
setTimeout(() => {
|
||||||
|
concurrency.current--;
|
||||||
|
resolve([makeMatch("e1", 0.7, ["mem-1"])]);
|
||||||
|
}, 100);
|
||||||
|
});
|
||||||
|
}),
|
||||||
|
initialize: jest.fn().mockResolvedValue(undefined),
|
||||||
|
};
|
||||||
|
m._entityStore = mockEntityStore;
|
||||||
|
m.embedder = {
|
||||||
|
embed: jest.fn().mockResolvedValue(mockEmbedding),
|
||||||
|
embedBatch: jest
|
||||||
|
.fn()
|
||||||
|
.mockImplementation((texts: string[]) =>
|
||||||
|
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||||
|
),
|
||||||
|
};
|
||||||
|
m.vectorStore.search = jest
|
||||||
|
.fn()
|
||||||
|
.mockResolvedValue([
|
||||||
|
{ id: "mem-1", score: 0.8, payload: { data: "test" } },
|
||||||
|
]);
|
||||||
|
m.vectorStore.keywordSearch = jest.fn().mockResolvedValue(null);
|
||||||
|
|
||||||
|
const start = performance.now();
|
||||||
|
await m.search("entity1 and entity2 and entity3 and entity4", {
|
||||||
|
filters: { user_id: "u1" },
|
||||||
|
});
|
||||||
|
const elapsed = performance.now() - start;
|
||||||
|
|
||||||
|
// With 4 entities at 100ms each, sequential would be ~400ms+.
|
||||||
|
// Parallel should be well under that. Use generous bound for CI.
|
||||||
|
expect(elapsed).toBeLessThan(500);
|
||||||
|
// At least 2 searches should have overlapped
|
||||||
|
expect(concurrency.peak).toBeGreaterThanOrEqual(2);
|
||||||
|
});
|
||||||
|
});
|
||||||
+75
-30
@@ -1,4 +1,5 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
|
import concurrent.futures
|
||||||
import gc
|
import gc
|
||||||
import hashlib
|
import hashlib
|
||||||
import json
|
import json
|
||||||
@@ -1464,34 +1465,55 @@ class Memory(MemoryBase):
|
|||||||
memory_boosts = {}
|
memory_boosts = {}
|
||||||
|
|
||||||
try:
|
try:
|
||||||
for _, entity_text in deduped:
|
entity_texts = [text for _, text in deduped]
|
||||||
entity_embedding = self.embedding_model.embed(entity_text, "search")
|
embeddings = self.embedding_model.embed_batch(entity_texts, "search")
|
||||||
matches = self.entity_store.search(
|
|
||||||
query=entity_text,
|
if len(embeddings) != len(entity_texts):
|
||||||
vectors=entity_embedding,
|
logger.warning(
|
||||||
top_k=500,
|
"embed_batch returned %d vectors for %d texts — skipping entity boost",
|
||||||
filters=search_filters,
|
len(embeddings),
|
||||||
|
len(entity_texts),
|
||||||
|
)
|
||||||
|
return memory_boosts
|
||||||
|
|
||||||
|
entity_store = self.entity_store
|
||||||
|
|
||||||
|
def _search_entity(entity_text, embedding):
|
||||||
|
return entity_store.search(
|
||||||
|
query=entity_text, vectors=embedding, top_k=500, filters=search_filters
|
||||||
)
|
)
|
||||||
|
|
||||||
for match in matches:
|
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
|
||||||
similarity = match.score if hasattr(match, 'score') else 0.0
|
futures = {
|
||||||
if similarity < 0.5:
|
pool.submit(_search_entity, text, emb): text
|
||||||
|
for text, emb in zip(entity_texts, embeddings)
|
||||||
|
}
|
||||||
|
|
||||||
|
for future in concurrent.futures.as_completed(futures):
|
||||||
|
try:
|
||||||
|
matches = future.result()
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning("Entity boost search failed for one entity: %s", e)
|
||||||
continue
|
continue
|
||||||
|
|
||||||
payload = match.payload if hasattr(match, 'payload') else {}
|
for match in matches:
|
||||||
linked_memory_ids = payload.get("linked_memory_ids", [])
|
similarity = match.score if hasattr(match, 'score') else 0.0
|
||||||
if not isinstance(linked_memory_ids, list):
|
if similarity < 0.5:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Spread-attenuated boost: entities linking to many memories get attenuated
|
payload = match.payload if hasattr(match, 'payload') else {}
|
||||||
num_linked = max(len(linked_memory_ids), 1)
|
linked_memory_ids = payload.get("linked_memory_ids", [])
|
||||||
memory_count_weight = 1.0 / (1.0 + 0.001 * ((num_linked - 1) ** 2))
|
if not isinstance(linked_memory_ids, list):
|
||||||
boost = similarity * ENTITY_BOOST_WEIGHT * memory_count_weight
|
continue
|
||||||
|
|
||||||
for memory_id in linked_memory_ids:
|
num_linked = max(len(linked_memory_ids), 1)
|
||||||
if memory_id:
|
memory_count_weight = 1.0 / (1.0 + 0.001 * ((num_linked - 1) ** 2))
|
||||||
memory_key = str(memory_id)
|
boost = similarity * ENTITY_BOOST_WEIGHT * memory_count_weight
|
||||||
memory_boosts[memory_key] = max(memory_boosts.get(memory_key, 0.0), boost)
|
|
||||||
|
for memory_id in linked_memory_ids:
|
||||||
|
if memory_id:
|
||||||
|
memory_key = str(memory_id)
|
||||||
|
memory_boosts[memory_key] = max(memory_boosts.get(memory_key, 0.0), boost)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"Entity boost computation failed: {e}")
|
logger.warning(f"Entity boost computation failed: {e}")
|
||||||
@@ -2872,15 +2894,38 @@ class AsyncMemory(MemoryBase):
|
|||||||
memory_boosts = {}
|
memory_boosts = {}
|
||||||
|
|
||||||
try:
|
try:
|
||||||
for _, entity_text in deduped:
|
entity_texts = [text for _, text in deduped]
|
||||||
entity_embedding = await asyncio.to_thread(self.embedding_model.embed, entity_text, "search")
|
embeddings = await asyncio.to_thread(self.embedding_model.embed_batch, entity_texts, "search")
|
||||||
matches = await asyncio.to_thread(
|
|
||||||
self.entity_store.search,
|
if len(embeddings) != len(entity_texts):
|
||||||
query=entity_text,
|
logger.warning(
|
||||||
vectors=entity_embedding,
|
"embed_batch returned %d vectors for %d texts — skipping entity boost",
|
||||||
top_k=500,
|
len(embeddings),
|
||||||
filters=search_filters,
|
len(entity_texts),
|
||||||
)
|
)
|
||||||
|
return memory_boosts
|
||||||
|
|
||||||
|
sem = asyncio.Semaphore(4)
|
||||||
|
|
||||||
|
async def _search_entity(entity_text, embedding):
|
||||||
|
async with sem:
|
||||||
|
return await asyncio.to_thread(
|
||||||
|
self.entity_store.search,
|
||||||
|
query=entity_text,
|
||||||
|
vectors=embedding,
|
||||||
|
top_k=500,
|
||||||
|
filters=search_filters,
|
||||||
|
)
|
||||||
|
|
||||||
|
results = await asyncio.gather(
|
||||||
|
*(_search_entity(text, emb) for text, emb in zip(entity_texts, embeddings)),
|
||||||
|
return_exceptions=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
for matches in results:
|
||||||
|
if isinstance(matches, BaseException):
|
||||||
|
logger.warning("Entity boost search failed for one entity: %s", matches)
|
||||||
|
continue
|
||||||
|
|
||||||
for match in matches:
|
for match in matches:
|
||||||
similarity = match.score if hasattr(match, 'score') else 0.0
|
similarity = match.score if hasattr(match, 'score') else 0.0
|
||||||
|
|||||||
@@ -1,5 +1,7 @@
|
|||||||
import logging
|
import logging
|
||||||
|
import time
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
|
from types import SimpleNamespace
|
||||||
from unittest.mock import MagicMock, Mock
|
from unittest.mock import MagicMock, Mock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
@@ -663,3 +665,195 @@ async def test_async_update_preserves_actor_id_when_different_actor_updates(mock
|
|||||||
assert stored["actor_id"] == "Alice"
|
assert stored["actor_id"] == "Alice"
|
||||||
|
|
||||||
|
|
||||||
|
def _make_match(score, linked_memory_ids):
|
||||||
|
return SimpleNamespace(score=score, payload={"linked_memory_ids": linked_memory_ids})
|
||||||
|
|
||||||
|
|
||||||
|
class TestEntityBoostParallelism:
|
||||||
|
"""Tests for parallelized entity boost searches (#5214)."""
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_memory(self, mocker):
|
||||||
|
_setup_mocks(mocker)
|
||||||
|
return Memory()
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_async_memory(self, mocker):
|
||||||
|
_setup_mocks(mocker)
|
||||||
|
return AsyncMemory()
|
||||||
|
|
||||||
|
def test_sync_boosts_preserve_scoring(self, mock_memory):
|
||||||
|
from mem0.utils.scoring import ENTITY_BOOST_WEIGHT
|
||||||
|
|
||||||
|
mock_memory.embedding_model = Mock()
|
||||||
|
mock_memory.embedding_model.embed_batch = Mock(return_value=[[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]])
|
||||||
|
|
||||||
|
results_by_query = {
|
||||||
|
"alice": [_make_match(0.9, ["mem-1"])],
|
||||||
|
"bob": [_make_match(0.6, ["mem-1", "mem-2"])],
|
||||||
|
}
|
||||||
|
|
||||||
|
def fake_search(query, vectors, top_k, filters):
|
||||||
|
return results_by_query[query]
|
||||||
|
|
||||||
|
mock_memory._entity_store = Mock()
|
||||||
|
mock_memory._entity_store.search = Mock(side_effect=fake_search)
|
||||||
|
|
||||||
|
boosts = mock_memory._compute_entity_boosts(
|
||||||
|
[("person", "alice"), ("person", "bob")],
|
||||||
|
{"user_id": "u1"},
|
||||||
|
)
|
||||||
|
|
||||||
|
boost_alice = 0.9 * ENTITY_BOOST_WEIGHT * (1.0 / (1.0 + 0.001 * (0**2)))
|
||||||
|
boost_bob = 0.6 * ENTITY_BOOST_WEIGHT * (1.0 / (1.0 + 0.001 * (1**2)))
|
||||||
|
assert boosts["mem-1"] == pytest.approx(max(boost_alice, boost_bob))
|
||||||
|
assert boosts["mem-2"] == pytest.approx(boost_bob)
|
||||||
|
|
||||||
|
def test_sync_embed_batch_called_once(self, mock_memory):
|
||||||
|
mock_memory.embedding_model = Mock()
|
||||||
|
mock_memory.embedding_model.embed_batch = Mock(return_value=[[0.1], [0.1], [0.1]])
|
||||||
|
mock_memory._entity_store = Mock()
|
||||||
|
mock_memory._entity_store.search = Mock(return_value=[_make_match(0.7, ["mem-1"])])
|
||||||
|
|
||||||
|
mock_memory._compute_entity_boosts(
|
||||||
|
[("person", "alice"), ("person", "bob"), ("person", "carol")],
|
||||||
|
{"user_id": "u1"},
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_memory.embedding_model.embed_batch.assert_called_once_with(["alice", "bob", "carol"], "search")
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_boosts_preserve_scoring(self, mock_async_memory):
|
||||||
|
from mem0.utils.scoring import ENTITY_BOOST_WEIGHT
|
||||||
|
|
||||||
|
mock_async_memory.embedding_model = Mock()
|
||||||
|
mock_async_memory.embedding_model.embed_batch = Mock(return_value=[[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]])
|
||||||
|
|
||||||
|
results_by_query = {
|
||||||
|
"alice": [_make_match(0.9, ["mem-1"])],
|
||||||
|
"bob": [_make_match(0.6, ["mem-1", "mem-2"])],
|
||||||
|
}
|
||||||
|
|
||||||
|
def fake_search(query, vectors, top_k, filters):
|
||||||
|
return results_by_query[query]
|
||||||
|
|
||||||
|
mock_async_memory._entity_store = Mock()
|
||||||
|
mock_async_memory._entity_store.search = Mock(side_effect=fake_search)
|
||||||
|
|
||||||
|
boosts = await mock_async_memory._compute_entity_boosts_async(
|
||||||
|
[("person", "alice"), ("person", "bob")],
|
||||||
|
{"user_id": "u1"},
|
||||||
|
)
|
||||||
|
|
||||||
|
boost_alice = 0.9 * ENTITY_BOOST_WEIGHT * (1.0 / (1.0 + 0.001 * (0**2)))
|
||||||
|
boost_bob = 0.6 * ENTITY_BOOST_WEIGHT * (1.0 / (1.0 + 0.001 * (1**2)))
|
||||||
|
assert boosts["mem-1"] == pytest.approx(max(boost_alice, boost_bob))
|
||||||
|
assert boosts["mem-2"] == pytest.approx(boost_bob)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_embed_batch_called_once(self, mock_async_memory):
|
||||||
|
mock_async_memory.embedding_model = Mock()
|
||||||
|
mock_async_memory.embedding_model.embed_batch = Mock(return_value=[[0.1], [0.1], [0.1]])
|
||||||
|
mock_async_memory._entity_store = Mock()
|
||||||
|
mock_async_memory._entity_store.search = Mock(return_value=[_make_match(0.7, ["mem-1"])])
|
||||||
|
|
||||||
|
await mock_async_memory._compute_entity_boosts_async(
|
||||||
|
[("person", "alice"), ("person", "bob"), ("person", "carol")],
|
||||||
|
{"user_id": "u1"},
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_async_memory.embedding_model.embed_batch.assert_called_once_with(["alice", "bob", "carol"], "search")
|
||||||
|
|
||||||
|
def test_sync_one_entity_failure_does_not_abort_others(self, mock_memory, caplog):
|
||||||
|
mock_memory.embedding_model = Mock()
|
||||||
|
mock_memory.embedding_model.embed_batch = Mock(return_value=[[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]])
|
||||||
|
|
||||||
|
def fake_search(query, vectors, top_k, filters):
|
||||||
|
if query == "boom":
|
||||||
|
raise RuntimeError("provider timeout")
|
||||||
|
return [_make_match(0.8, ["mem-9"])]
|
||||||
|
|
||||||
|
mock_memory._entity_store = Mock()
|
||||||
|
mock_memory._entity_store.search = Mock(side_effect=fake_search)
|
||||||
|
|
||||||
|
with caplog.at_level(logging.WARNING):
|
||||||
|
boosts = mock_memory._compute_entity_boosts(
|
||||||
|
[("person", "boom"), ("person", "ok")],
|
||||||
|
{"user_id": "u1"},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "mem-9" in boosts
|
||||||
|
assert any("Entity boost search failed" in r.message for r in caplog.records)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_one_entity_failure_does_not_abort_others(self, mock_async_memory, caplog):
|
||||||
|
mock_async_memory.embedding_model = Mock()
|
||||||
|
mock_async_memory.embedding_model.embed_batch = Mock(return_value=[[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]])
|
||||||
|
|
||||||
|
def fake_search(query, vectors, top_k, filters):
|
||||||
|
if query == "boom":
|
||||||
|
raise RuntimeError("provider timeout")
|
||||||
|
return [_make_match(0.8, ["mem-9"])]
|
||||||
|
|
||||||
|
mock_async_memory._entity_store = Mock()
|
||||||
|
mock_async_memory._entity_store.search = Mock(side_effect=fake_search)
|
||||||
|
|
||||||
|
with caplog.at_level(logging.WARNING):
|
||||||
|
boosts = await mock_async_memory._compute_entity_boosts_async(
|
||||||
|
[("person", "boom"), ("person", "ok")],
|
||||||
|
{"user_id": "u1"},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "mem-9" in boosts
|
||||||
|
assert any("Entity boost search failed" in r.message for r in caplog.records)
|
||||||
|
|
||||||
|
def test_sync_searches_run_concurrently(self, mock_memory):
|
||||||
|
mock_memory.embedding_model = Mock()
|
||||||
|
mock_memory.embedding_model.embed_batch = Mock(return_value=[[0.1]] * 4)
|
||||||
|
|
||||||
|
concurrent_count = {"current": 0, "peak": 0}
|
||||||
|
|
||||||
|
def blocking_search(query, vectors, top_k, filters):
|
||||||
|
concurrent_count["current"] += 1
|
||||||
|
concurrent_count["peak"] = max(concurrent_count["peak"], concurrent_count["current"])
|
||||||
|
time.sleep(0.2)
|
||||||
|
concurrent_count["current"] -= 1
|
||||||
|
return [_make_match(0.7, [f"mem-{query}"])]
|
||||||
|
|
||||||
|
mock_memory._entity_store = Mock()
|
||||||
|
mock_memory._entity_store.search = Mock(side_effect=blocking_search)
|
||||||
|
|
||||||
|
entities = [("person", f"e{i}") for i in range(4)]
|
||||||
|
start = time.perf_counter()
|
||||||
|
boosts = mock_memory._compute_entity_boosts(entities, {"user_id": "u1"})
|
||||||
|
elapsed = time.perf_counter() - start
|
||||||
|
|
||||||
|
assert elapsed < 0.6, f"searches did not run concurrently (took {elapsed:.2f}s)"
|
||||||
|
assert concurrent_count["peak"] >= 2, "no overlap observed between entity searches"
|
||||||
|
assert len(boosts) == 4
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_searches_run_concurrently(self, mock_async_memory):
|
||||||
|
mock_async_memory.embedding_model = Mock()
|
||||||
|
mock_async_memory.embedding_model.embed_batch = Mock(return_value=[[0.1]] * 4)
|
||||||
|
|
||||||
|
concurrent_count = {"current": 0, "peak": 0}
|
||||||
|
|
||||||
|
def blocking_search(query, vectors, top_k, filters):
|
||||||
|
concurrent_count["current"] += 1
|
||||||
|
concurrent_count["peak"] = max(concurrent_count["peak"], concurrent_count["current"])
|
||||||
|
time.sleep(0.2)
|
||||||
|
concurrent_count["current"] -= 1
|
||||||
|
return [_make_match(0.7, [f"mem-{query}"])]
|
||||||
|
|
||||||
|
mock_async_memory._entity_store = Mock()
|
||||||
|
mock_async_memory._entity_store.search = Mock(side_effect=blocking_search)
|
||||||
|
|
||||||
|
entities = [("person", f"e{i}") for i in range(4)]
|
||||||
|
start = time.perf_counter()
|
||||||
|
boosts = await mock_async_memory._compute_entity_boosts_async(entities, {"user_id": "u1"})
|
||||||
|
elapsed = time.perf_counter() - start
|
||||||
|
|
||||||
|
assert elapsed < 0.6, f"searches did not run concurrently (took {elapsed:.2f}s)"
|
||||||
|
assert concurrent_count["peak"] >= 2, "no overlap observed between entity searches"
|
||||||
|
assert len(boosts) == 4
|
||||||
|
|||||||
Reference in New Issue
Block a user