diff --git a/mem0/memory/main.py b/mem0/memory/main.py index bef176699..035c75a22 100644 --- a/mem0/memory/main.py +++ b/mem0/memory/main.py @@ -1384,21 +1384,33 @@ class AsyncMemory(MemoryBase): "mem0.get_all", self, {"limit": limit, "keys": keys, "encoded_ids": encoded_ids, "sync_type": "async"} ) - with concurrent.futures.ThreadPoolExecutor() as executor: - future_memories = executor.submit(self._get_all_from_vector_store, effective_filters, limit) - future_graph_entities = ( - executor.submit(self.graph.get_all, effective_filters, limit) if self.enable_graph else None - ) - - concurrent.futures.wait( - [future_memories, future_graph_entities] if future_graph_entities else [future_memories] - ) - - all_memories_result = future_memories.result() - graph_entities_result = future_graph_entities.result() if future_graph_entities else None + vector_store_task = asyncio.create_task( + self._get_all_from_vector_store(effective_filters, limit) + ) + graph_task = None if self.enable_graph: - return {"results": all_memories_result, "relations": graph_entities_result} + graph_get_all = getattr(self.graph, "get_all", None) + if callable(graph_get_all): + if asyncio.iscoroutinefunction(graph_get_all): + graph_task = asyncio.create_task( + graph_get_all(effective_filters, limit) + ) + else: + graph_task = asyncio.create_task( + asyncio.to_thread(graph_get_all, effective_filters, limit) + ) + + results_dict = {} + if graph_task: + vector_store_result, graph_entities_result = await asyncio.gather( + vector_store_task, graph_task + ) + results_dict.update( + {"results": vector_store_result, "relations": graph_entities_result} + ) + else: + results_dict.update({"results": await vector_store_task}) if self.api_version == "v1.0": warnings.warn( @@ -1408,9 +1420,9 @@ class AsyncMemory(MemoryBase): category=DeprecationWarning, stacklevel=2, ) - return all_memories_result - else: - return {"results": all_memories_result} + return results_dict["results"] + + return results_dict async def _get_all_from_vector_store(self, filters, limit): memories_result = await asyncio.to_thread(self.vector_store.list, filters=filters, limit=limit)