Fixes for asynchronous crew execution (#81)

* Create new Lunary queue if running in a separate thread

* Agent call_as_tool explicit input schema

* 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

@@ -5,8 +5,9 @@
 
 from langchain_core.messages import BaseMessage
 from langchain_core.prompts.chat import ChatPromptTemplate, HumanMessage, SystemMessage
+from langchain_core.pydantic_v1 import BaseModel, Field
 from langchain_core.runnables import RunnableConfig
-from langchain_core.tools import Tool
+from langchain_core.tools import StructuredTool
 
 from motleycrew.agents.abstract_parent import MotleyAgentAbstractParent
 from motleycrew.common import MotleyAgentFactory, MotleySupportedTool, logger
@@ -174,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.
@@ -225,8 +226,10 @@ def as_tool(self, **kwargs) -> MotleyTool:
         if not getattr(self, "name", None) or not getattr(self, "description", None):
             raise ValueError("Agent must have a name and description to be called as a tool")
 
-        def call_as_tool(self, *args, **kwargs):
-            # TODO: this thing is hacky, we should have a better way to pass structured input
+        class CallAsToolInput(BaseModel):
+            input: str = Field(..., description="Input to the tool")
+
+        def call_as_tool(*args, **kwargs):
             if args:
                 return self.invoke({"prompt": args[0]})
             if len(kwargs) == 1:
@@ -235,12 +238,13 @@ def call_as_tool(self, *args, **kwargs):
 
         # To be specialized if we expect structured input
         return MotleyTool.from_langchain_tool(
-            Tool(
+            StructuredTool(
                 name=self.name.replace(
                     " ", "_"
                 ).lower(),  # OpenAI doesn't accept spaces in function names
                 description=self.description,
                 func=call_as_tool,
+                args_schema=CallAsToolInput,
             ),
             **kwargs,
         )

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:

motleycrew/tracking/callbacks.py

@@ -2,27 +2,28 @@
 The module contains callback handlers for sending data to the Lunary service
 """
 
-from typing import List, Dict, Optional, Any, Union
 import traceback
+from typing import Any, Dict, List, Optional, Union
 
 try:
+    from llama_index.core.base.llms.types import ChatMessage
     from llama_index.core.callbacks.base_handler import BaseCallbackHandler
     from llama_index.core.callbacks.schema import CBEventType, EventPayload
-    from llama_index.core.base.llms.types import ChatMessage
 except ImportError:
     BaseCallbackHandler = object
     CBEventType = None
     ChatMessage = None
 
 try:
-    from lunary import track_event, event_queue_ctx
+    from lunary import EventQueue, event_queue_ctx, track_event
 except ImportError:
     track_event = None
     event_queue_ctx = None
+    EventQueue = None
 
-from motleycrew.common.enums import LunaryRunType, LunaryEventName
-from motleycrew.common.utils import ensure_module_is_installed
 from motleycrew.common import logger
+from motleycrew.common.enums import LunaryEventName, LunaryRunType
+from motleycrew.common.utils import ensure_module_is_installed
 
 
 def event_delegate_decorator(f):
@@ -33,6 +34,7 @@ def event_delegate_decorator(f):
     Args:
         f (callable):
     """
+
     def wrapper(self, *args, **kwargs):
         ensure_module_is_installed("llama_index")
         run_type = "start" if "start" in f.__name__ else "end"
@@ -122,7 +124,12 @@ def __init__(
         self._track_event = track_event
         self._event_run_type_ids = []
 
-        self.queue = event_queue_ctx.get()
+        try:
+            self.queue = event_queue_ctx.get()
+        except (
+            LookupError
+        ):  # happens when running in a separate thread, so we can't use the global queue
+            self.queue = EventQueue()
 
     def _get_initial_track_event_params(
         self, run_type: LunaryRunType, event_name: LunaryEventName, run_id: str = None
@@ -312,7 +319,7 @@ def _on_agent_step_end(
 
         if payload:
             response = payload.get(EventPayload.RESPONSE)
-            output  = response.response
+            output = response.response
         else:
             output = ""
         params["output"] = output
@@ -400,7 +407,7 @@ def on_event_end(
             payload (dict): related event data
             event_id (str): event id (uuid)
             **kwargs:
-       """
+        """
         return
 
     @staticmethod