Proper async invocation for Langchain and LlamaIndex agents

motleycrew/agents/crewai/crewai.py

@@ -8,8 +8,7 @@
 
 from motleycrew.agents.crewai import CrewAIAgentWithConfig
 from motleycrew.agents.parent import MotleyAgentParent
-from motleycrew.common import MotleyAgentFactory
-from motleycrew.common import MotleySupportedTool
+from motleycrew.common import MotleyAgentFactory, MotleySupportedTool
 from motleycrew.common.utils import ensure_module_is_installed
 from motleycrew.tools import MotleyTool
 from motleycrew.tracking import add_default_callbacks_to_langchain_config
@@ -92,7 +91,7 @@ def invoke(
         config: Optional[RunnableConfig] = None,
         **kwargs: Any,
     ) -> Any:
-        prompt = self.prepare_for_invocation(input=input)
+        prompt = self._prepare_for_invocation(input=input)
 
         langchain_tools = [tool.to_langchain_tool() for tool in self.tools.values()]
         config = add_default_callbacks_to_langchain_config(config)

motleycrew/agents/langchain/langchain.py

@@ -1,18 +1,21 @@
 from __future__ import annotations
 
+import asyncio
 from typing import Any, Optional, Sequence
 
 from langchain.agents import AgentExecutor
 from langchain_core.chat_history import InMemoryChatMessageHistory
+from langchain_core.prompts.chat import ChatPromptTemplate
 from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables.config import merge_configs
-from langchain_core.runnables.history import RunnableWithMessageHistory, GetSessionHistoryCallable
-from langchain_core.prompts.chat import ChatPromptTemplate
+from langchain_core.runnables.history import (
+    GetSessionHistoryCallable,
+    RunnableWithMessageHistory,
+)
 
 from motleycrew.agents.mixins import LangchainOutputHandlingAgentMixin
 from motleycrew.agents.parent import MotleyAgentParent
-from motleycrew.common import MotleyAgentFactory
-from motleycrew.common import MotleySupportedTool, logger
+from motleycrew.common import MotleyAgentFactory, MotleySupportedTool, logger
 from motleycrew.tracking import add_default_callbacks_to_langchain_config
 
 
@@ -146,26 +149,46 @@ def materialize(self):
                 history_messages_key="chat_history",
             )
 
-    def invoke(
-        self,
-        input: dict,
-        config: Optional[RunnableConfig] = None,
-        **kwargs: Any,
-    ) -> Any:
+    def _prepare_config(self, config: RunnableConfig) -> RunnableConfig:
         config = merge_configs(self.runnable_config, config)
-        prompt = self.prepare_for_invocation(input=input, prompt_as_messages=self.input_as_messages)
-
         config = add_default_callbacks_to_langchain_config(config)
         if self.get_session_history_callable:
             config["configurable"] = config.get("configurable") or {}
             config["configurable"]["session_id"] = (
                 config["configurable"].get("session_id") or "default"
             )
+        return config
+
+    def invoke(
+        self,
+        input: dict,
+        config: Optional[RunnableConfig] = None,
+        **kwargs: Any,
+    ) -> Any:
+        config = self._prepare_config(config)
+        prompt = self._prepare_for_invocation(
+            input=input, prompt_as_messages=self.input_as_messages
+        )
 
         output = self.agent.invoke({"input": prompt}, config, **kwargs)
         output = output.get("output")
         return output
 
+    async def ainvoke(
+        self,
+        input: dict,
+        config: Optional[RunnableConfig] = None,
+        **kwargs: Any,
+    ) -> Any:
+        config = self._prepare_config(config)
+        prompt = await asyncio.to_thread(
+            self._prepare_for_invocation, input=input, prompt_as_messages=self.input_as_messages
+        )
+
+        output = await self.agent.ainvoke({"input": prompt}, config, **kwargs)
+        output = output.get("output")
+        return output
+
     @staticmethod
     def from_agent(
         agent: AgentExecutor,

motleycrew/agents/llama_index/llama_index.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
+import asyncio
 import uuid
 from typing import Any, Optional, Sequence
 
 try:
     from llama_index.core.agent import AgentRunner
-    from llama_index.core.chat_engine.types import ChatResponseMode
     from llama_index.core.agent.types import TaskStep, TaskStepOutput
-    from llama_index.core.chat_engine.types import AgentChatResponse
+    from llama_index.core.chat_engine.types import AgentChatResponse, ChatResponseMode
 except ImportError:
     AgentRunner = None
     ChatResponseMode = None
@@ -18,7 +18,7 @@
 from langchain_core.runnables import RunnableConfig
 
 from motleycrew.agents.parent import MotleyAgentParent
-from motleycrew.common import MotleySupportedTool, MotleyAgentFactory, AuxPrompts
+from motleycrew.common import AuxPrompts, MotleyAgentFactory, MotleySupportedTool
 from motleycrew.common.utils import ensure_module_is_installed
 from motleycrew.tools import DirectOutput
 
@@ -154,7 +154,7 @@ def invoke(
         config: Optional[RunnableConfig] = None,
         **kwargs: Any,
     ) -> Any:
-        prompt = self.prepare_for_invocation(input=input)
+        prompt = self._prepare_for_invocation(input=input)
 
         output = self.agent.chat(prompt)
 
@@ -163,6 +163,21 @@ def invoke(
 
         return output.response
 
+    async def ainvoke(
+        self,
+        input: dict,
+        config: Optional[RunnableConfig] = None,
+        **kwargs: Any,
+    ) -> Any:
+        prompt = await asyncio.to_thread(self._prepare_for_invocation, input=input)
+
+        output = await self.agent.achat(prompt)
+
+        if self.direct_output is not None:
+            return self.direct_output.output
+
+        return output.response
+
     @staticmethod
     def from_agent(
         agent: AgentRunner,

motleycrew/agents/parent.py

@@ -175,7 +175,7 @@ def materialize(self):
 
         self._agent = self.agent_factory(tools=self.tools)
 
-    def prepare_for_invocation(self, input: dict, prompt_as_messages: bool = False) -> str:
+    def _prepare_for_invocation(self, input: dict, prompt_as_messages: bool = False) -> str:
         """Prepare the agent for invocation by materializing it and composing the prompt.
 
         Should be called in the beginning of the agent's invoke method.

motleycrew/crew/crew.py

@@ -1,10 +1,10 @@
 import asyncio
 import threading
 import time
-from typing import Collection, Generator, Optional, Any
+from typing import Any, Collection, Generator, Optional
 
 from motleycrew.agents.parent import MotleyAgentParent
-from motleycrew.common import logger, AsyncBackend, Defaults
+from motleycrew.common import AsyncBackend, Defaults, logger
 from motleycrew.crew.crew_threads import TaskUnitThreadPool
 from motleycrew.storage import MotleyGraphStore
 from motleycrew.storage.graph_store_utils import init_graph_store
@@ -112,7 +112,7 @@ def _prepare_next_unit_for_dispatch(
             Agent, task, unit to be dispatched.
         """
         available_tasks = self.get_available_tasks()
-        logger.info("Available tasks: %s", available_tasks)
+        logger.debug("Available tasks: %s", available_tasks)
 
         for task in available_tasks:
             if not task.allow_async_units and task in running_sync_tasks: