perf: skip ThreadPoolExecutor when graph store is disabled

Avoids spawning 2 threads per add() call when graph is not configured.
Fixes thread leakage under high concurrency (was accumulating 3500+ threads
per worker). Calls _add_to_vector_store directly in the current thread.

Applied to both sync Memory and async AsyncMemory classes.
This commit is contained in:
Soumil Rathi
2026-04-03 07:42:32 -07:00
parent cb64d978fc
commit a6130c5436
+4 -2
View File
@@ -454,6 +454,7 @@ class Memory(MemoryBase):
else: else:
messages = parse_vision_messages(messages) messages = parse_vision_messages(messages)
if self.enable_graph:
with concurrent.futures.ThreadPoolExecutor() as executor: with concurrent.futures.ThreadPoolExecutor() as executor:
future1 = executor.submit(self._add_to_vector_store, messages, processed_metadata, effective_filters, infer, timestamp) future1 = executor.submit(self._add_to_vector_store, messages, processed_metadata, effective_filters, infer, timestamp)
future2 = executor.submit(self._add_to_graph, messages, effective_filters) future2 = executor.submit(self._add_to_graph, messages, effective_filters)
@@ -463,12 +464,12 @@ class Memory(MemoryBase):
vector_store_result = future1.result() vector_store_result = future1.result()
graph_result = future2.result() graph_result = future2.result()
if self.enable_graph:
return { return {
"results": vector_store_result, "results": vector_store_result,
"relations": graph_result, "relations": graph_result,
} }
vector_store_result = self._add_to_vector_store(messages, processed_metadata, effective_filters, infer, timestamp)
return {"results": vector_store_result} return {"results": vector_store_result}
def _add_to_vector_store(self, messages, metadata, filters, infer, timestamp=None): def _add_to_vector_store(self, messages, metadata, filters, infer, timestamp=None):
@@ -1737,6 +1738,7 @@ class AsyncMemory(MemoryBase):
else: else:
messages = parse_vision_messages(messages) messages = parse_vision_messages(messages)
if self.enable_graph:
vector_store_task = asyncio.create_task( vector_store_task = asyncio.create_task(
self._add_to_vector_store(messages, processed_metadata, effective_filters, infer) self._add_to_vector_store(messages, processed_metadata, effective_filters, infer)
) )
@@ -1744,12 +1746,12 @@ class AsyncMemory(MemoryBase):
vector_store_result, graph_result = await asyncio.gather(vector_store_task, graph_task) vector_store_result, graph_result = await asyncio.gather(vector_store_task, graph_task)
if self.enable_graph:
return { return {
"results": vector_store_result, "results": vector_store_result,
"relations": graph_result, "relations": graph_result,
} }
vector_store_result = await self._add_to_vector_store(messages, processed_metadata, effective_filters, infer)
return {"results": vector_store_result} return {"results": vector_store_result}
async def _add_to_vector_store( async def _add_to_vector_store(