refactor: Improve async handling in AsyncMemory class for better performance (#3250)

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:
+        vector_store_task = asyncio.create_task(
-            future_memories = executor.submit(self._get_all_from_vector_store, effective_filters, limit)
+            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
 
+        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
+            return results_dict["results"]
-        else:
+
-            return {"results": all_memories_result}
+        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)