add async to memory and postprocessor base classes (#18438)

Author: logan-markewich

Makefile

@@ -11,7 +11,7 @@ lint:	## Run linters: pre-commit (black, ruff, codespell) and mypy
 	pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
 
 test:	## Run tests via pants
-	pants --level=info --no-local-cache --changed-since=origin/main --changed-dependents=transitive --no-test-use-coverage test
+	pants --level=error --no-local-cache --changed-since=origin/main --changed-dependents=transitive --no-test-use-coverage test
 
 test-core:	## Run tests via pants
 	pants --no-local-cache test llama-index-core/::

llama-index-core/llama_index/core/agent/workflow/codeact_agent.py

@@ -344,8 +344,7 @@ async def finalize(
         Adds all in-progress messages to memory.
         """
         scratchpad: List[ChatMessage] = await ctx.get(self.scratchpad_key, default=[])
-        for msg in scratchpad:
-            await memory.aput(msg)
+        await memory.aput_messages(scratchpad)
 
         # reset scratchpad
         await ctx.set(self.scratchpad_key, [])

llama-index-core/llama_index/core/agent/workflow/function_agent.py

@@ -122,8 +122,7 @@ async def finalize(
         Adds all in-progress messages to memory.
         """
         scratchpad: List[ChatMessage] = await ctx.get(self.scratchpad_key, default=[])
-        for msg in scratchpad:
-            await memory.aput(msg)
+        await memory.aput_messages(scratchpad)
 
         # reset scratchpad
         await ctx.set(self.scratchpad_key, [])

llama-index-core/llama_index/core/agent/workflow/multi_agent_workflow.py

@@ -308,7 +308,7 @@ async def init_run(self, ctx: Context, ev: StartEvent) -> AgentInput:
 
         # First set chat history if it exists
         if chat_history:
-            memory.set(chat_history)
+            await memory.aset(chat_history)
 
         # Then add user message if it exists
         if user_msg:
@@ -335,7 +335,7 @@ async def init_run(self, ctx: Context, ev: StartEvent) -> AgentInput:
             raise ValueError("Must provide either user_msg or chat_history")
 
         # Get all messages from memory
-        input_messages = memory.get()
+        input_messages = await memory.aget()
 
         # send to the current agent
         current_agent_name: str = await ctx.get("current_agent_name")
@@ -526,7 +526,7 @@ async def aggregate_tool_results(
                 return StopEvent(result=result)
 
         user_msg_str = await ctx.get("user_msg_str")
-        input_messages = memory.get(input=user_msg_str)
+        input_messages = await memory.aget(input=user_msg_str)
 
         # get this again, in case it changed
         agent_name = await ctx.get("current_agent_name")

llama-index-core/llama_index/core/memory/types.py

@@ -1,3 +1,4 @@
+import asyncio
 from abc import abstractmethod
 from typing import Any, List, Optional
 
@@ -11,10 +12,7 @@
 
 
 class BaseMemory(BaseComponent):
-    """Base class for all memory types.
-
-    NOTE: The interface for memory is not yet finalized and is subject to change.
-    """
+    """Base class for all memory types."""
 
     @classmethod
     def class_name(cls) -> str:
@@ -33,17 +31,27 @@ def from_defaults(
     def get(self, input: Optional[str] = None, **kwargs: Any) -> List[ChatMessage]:
         """Get chat history."""
 
+    async def aget(
+        self, input: Optional[str] = None, **kwargs: Any
+    ) -> List[ChatMessage]:
+        """Get chat history."""
+        return await asyncio.to_thread(self.get, input=input, **kwargs)
+
     @abstractmethod
     def get_all(self) -> List[ChatMessage]:
         """Get all chat history."""
 
+    async def aget_all(self) -> List[ChatMessage]:
+        """Get all chat history."""
+        return await asyncio.to_thread(self.get_all)
+
     @abstractmethod
     def put(self, message: ChatMessage) -> None:
         """Put chat history."""
 
     async def aput(self, message: ChatMessage) -> None:
         """Put chat history."""
-        self.put(message)
+        await asyncio.to_thread(self.put, message)
 
     def put_messages(self, messages: List[ChatMessage]) -> None:
         """Put chat history."""
@@ -52,23 +60,27 @@ def put_messages(self, messages: List[ChatMessage]) -> None:
 
     async def aput_messages(self, messages: List[ChatMessage]) -> None:
         """Put chat history."""
-        for message in messages:
-            await self.aput(message)
+        await asyncio.to_thread(self.put_messages, messages)
 
     @abstractmethod
     def set(self, messages: List[ChatMessage]) -> None:
         """Set chat history."""
 
+    async def aset(self, messages: List[ChatMessage]) -> None:
+        """Set chat history."""
+        await asyncio.to_thread(self.set, messages)
+
     @abstractmethod
     def reset(self) -> None:
         """Reset chat history."""
 
+    async def areset(self) -> None:
+        """Reset chat history."""
+        await asyncio.to_thread(self.reset)
+
 
 class BaseChatStoreMemory(BaseMemory):
-    """Base class for any .
-
-    NOTE: The interface for memory is not yet finalized and is subject to change.
-    """
+    """Base class for storing multi-tenant chat history."""
 
     chat_store: SerializeAsAny[BaseChatStore] = Field(default_factory=SimpleChatStore)
     chat_store_key: str = Field(default=DEFAULT_CHAT_STORE_KEY)
@@ -98,6 +110,20 @@ def get_all(self) -> List[ChatMessage]:
         """Get all chat history."""
         return self.chat_store.get_messages(self.chat_store_key)
 
+    async def aget_all(self) -> List[ChatMessage]:
+        """Get all chat history."""
+        return await self.chat_store.aget_messages(self.chat_store_key)
+
+    def get(self, input: Optional[str] = None, **kwargs: Any) -> List[ChatMessage]:
+        """Get chat history."""
+        return self.chat_store.get_messages(self.chat_store_key, **kwargs)
+
+    async def aget(
+        self, input: Optional[str] = None, **kwargs: Any
+    ) -> List[ChatMessage]:
+        """Get chat history."""
+        return await self.chat_store.aget_messages(self.chat_store_key, **kwargs)
+
     def put(self, message: ChatMessage) -> None:
         """Put chat history."""
         # ensure everything is serialized
@@ -112,6 +138,15 @@ def set(self, messages: List[ChatMessage]) -> None:
         """Set chat history."""
         self.chat_store.set_messages(self.chat_store_key, messages)
 
+    async def aset(self, messages: List[ChatMessage]) -> None:
+        """Set chat history."""
+        # ensure everything is serialized
+        await self.chat_store.aset_messages(self.chat_store_key, messages)
+
     def reset(self) -> None:
         """Reset chat history."""
         self.chat_store.delete_messages(self.chat_store_key)
+
+    async def areset(self) -> None:
+        """Reset chat history."""
+        await self.chat_store.adelete_messages(self.chat_store_key)

llama-index-core/llama_index/core/postprocessor/types.py

@@ -1,3 +1,4 @@
+import asyncio
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional
 
@@ -61,6 +62,29 @@ def _postprocess_nodes(
     ) -> List[NodeWithScore]:
         """Postprocess nodes."""
 
+    async def apostprocess_nodes(
+        self,
+        nodes: List[NodeWithScore],
+        query_bundle: Optional[QueryBundle] = None,
+        query_str: Optional[str] = None,
+    ) -> List[NodeWithScore]:
+        """Postprocess nodes (async)."""
+        if query_str is not None and query_bundle is not None:
+            raise ValueError("Cannot specify both query_str and query_bundle")
+        elif query_str is not None:
+            query_bundle = QueryBundle(query_str)
+        else:
+            pass
+        return await self._apostprocess_nodes(nodes, query_bundle)
+
+    async def _apostprocess_nodes(
+        self,
+        nodes: List[NodeWithScore],
+        query_bundle: Optional[QueryBundle] = None,
+    ) -> List[NodeWithScore]:
+        """Postprocess nodes (async)."""
+        return await asyncio.to_thread(self._postprocess_nodes, nodes, query_bundle)
+
     def _as_query_component(self, **kwargs: Any) -> QueryComponent:
         """As query component."""
         return PostprocessorComponent(postprocessor=self)

llama-index-core/llama_index/core/storage/chat_store/base.py

@@ -1,4 +1,5 @@
 """Base interface class for storing chat history per user."""
+import asyncio
 from abc import abstractmethod
 from typing import List, Optional
 
@@ -49,28 +50,28 @@ def get_keys(self) -> List[str]:
 
     async def aset_messages(self, key: str, messages: List[ChatMessage]) -> None:
         """Async version of Get messages for a key."""
-        self.set_messages(key, messages)
+        await asyncio.to_thread(self.set_messages, key, messages)
 
     async def aget_messages(self, key: str) -> List[ChatMessage]:
         """Async version of Get messages for a key."""
-        return self.get_messages(key)
+        return await asyncio.to_thread(self.get_messages, key)
 
     async def async_add_message(self, key: str, message: ChatMessage) -> None:
         """Async version of Add a message for a key."""
-        self.add_message(key, message)
+        await asyncio.to_thread(self.add_message, key, message)
 
     async def adelete_messages(self, key: str) -> Optional[List[ChatMessage]]:
         """Async version of Delete messages for a key."""
-        return self.delete_messages(key)
+        return await asyncio.to_thread(self.delete_messages, key)
 
     async def adelete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
         """Async version of Delete specific message for a key."""
-        return self.delete_message(key, idx)
+        return await asyncio.to_thread(self.delete_message, key, idx)
 
     async def adelete_last_message(self, key: str) -> Optional[ChatMessage]:
         """Async version of Delete last message for a key."""
-        return self.delete_last_message(key)
+        return await asyncio.to_thread(self.delete_last_message, key)
 
     async def aget_keys(self) -> List[str]:
         """Async version of Get all keys."""
-        return self.get_keys()
+        return await asyncio.to_thread(self.get_keys)

llama-index-core/tests/agent/workflow/test_code_act_agent.py

@@ -166,4 +166,4 @@ async def test_code_act_agent_tool_handling(
     # Finalize
     final_output = await agent.finalize(ctx, output, mock_memory)
     assert isinstance(final_output, AgentOutput)
-    assert mock_memory.aput.called  # Verify memory was updated
+    assert mock_memory.aput_messages.called  # Verify memory was updated