langroid 0.23.3__py3-none-any.whl → 0.25.0__py3-none-any.whl

langroid/agent/base.py

@@ -142,12 +142,19 @@ class Agent(ABC):
         self.llm_tools_handled: Set[str] = set()
         self.llm_tools_usable: Set[str] = set()
         self.llm_tools_known: Set[str] = set()  # all known tools, handled/used or not
+        # Indicates which tool-names are allowed to be inferred when
+        # the LLM "forgets" to include the request field in its
+        # tool-call.
+        self.enabled_requests_for_inference: Optional[Set[str]] = (
+            None  # If None, we allow all
+        )
         self.interactive: bool = True  # may be modified by Task wrapper
         self.token_stats_str = ""
         self.default_human_response: Optional[str] = None
         self._indent = ""
         self.llm = LanguageModel.create(config.llm)
         self.vecdb = VectorStore.create(config.vecdb) if config.vecdb else None
+        self.tool_error = False
         if config.parsing is not None and self.config.llm is not None:
             # token_encoding_model is used to obtain the tokenizer,
             # so in case it's an OpenAI model, we ensure that the tokenizer
@@ -168,6 +175,7 @@ class Agent(ABC):
             show_llm_response=noop_fn,
             show_agent_response=noop_fn,
             get_user_response=None,
+            get_user_response_async=None,
             get_last_step=noop_fn,
             set_parent_agent=noop_fn,
             show_error_message=noop_fn,
@@ -315,6 +323,52 @@ class Agent(ABC):
                 lambda msg: message_class.handle_message_fallback(self, msg),
             )
 
+        async_tool_name = f"{tool}_async"
+        if (
+            hasattr(message_class, "handle_async")
+            and inspect.isfunction(message_class.handle_async)
+            and not hasattr(self, async_tool_name)
+        ):
+            has_chat_doc_arg = (
+                len(inspect.signature(message_class.handle_async).parameters) > 1
+            )
+
+            if has_chat_doc_arg:
+
+                @no_type_check
+                async def handler(obj, chat_doc):
+                    return await obj.handle_async(chat_doc)
+
+            else:
+
+                @no_type_check
+                async def handler(obj):
+                    return await obj.handle_async()
+
+            setattr(self, async_tool_name, handler)
+        elif (
+            hasattr(message_class, "response_async")
+            and inspect.isfunction(message_class.response_async)
+            and not hasattr(self, async_tool_name)
+        ):
+            has_chat_doc_arg = (
+                len(inspect.signature(message_class.response_async).parameters) > 2
+            )
+
+            if has_chat_doc_arg:
+
+                @no_type_check
+                async def handler(obj, chat_doc):
+                    return await obj.response_async(self, chat_doc)
+
+            else:
+
+                @no_type_check
+                async def handler(obj):
+                    return await obj.response_async(self)
+
+            setattr(self, async_tool_name, handler)
+
         return [tool]
 
     def enable_message_handling(
@@ -386,32 +440,14 @@ class Agent(ABC):
             recipient=recipient,
         )
 
-    async def agent_response_async(
+    def _agent_response_final(
         self,
-        msg: Optional[str | ChatDocument] = None,
-    ) -> Optional[ChatDocument]:
-        return self.agent_response(msg)
-
-    def agent_response(
-        self,
-        msg: Optional[str | ChatDocument] = None,
+        msg: Optional[str | ChatDocument],
+        results: Optional[str | OrderedDict[str, str] | ChatDocument],
     ) -> Optional[ChatDocument]:
         """
-        Response from the "agent itself", typically (but not only)
-        used to handle LLM's "tool message" or `function_call`
-        (e.g. OpenAI `function_call`).
-        Args:
-            msg (str|ChatDocument): the input to respond to: if msg is a string,
-                and it contains a valid JSON-structured "tool message", or
-                if msg is a ChatDocument, and it contains a `function_call`.
-        Returns:
-            Optional[ChatDocument]: the response, packaged as a ChatDocument
-
+        Convert results to final response.
         """
-        if msg is None:
-            return None
-
-        results = self.handle_message(msg)
         if results is None:
             return None
         if not settings.quiet:
@@ -431,7 +467,7 @@ class Agent(ABC):
         if isinstance(results, ChatDocument):
             # Preserve trail of tool_ids for OpenAI Assistant fn-calls
             results.metadata.tool_ids = (
-                [] if isinstance(msg, str) else msg.metadata.tool_ids
+                [] if msg is None or isinstance(msg, str) else msg.metadata.tool_ids
             )
             return results
         sender_name = self.config.name
@@ -454,10 +490,49 @@ class Agent(ABC):
                 sender_name=sender_name,
                 oai_tool_id=oai_tool_id,
                 # preserve trail of tool_ids for OpenAI Assistant fn-calls
-                tool_ids=[] if isinstance(msg, str) else msg.metadata.tool_ids,
+                tool_ids=(
+                    [] if msg is None or isinstance(msg, str) else msg.metadata.tool_ids
+                ),
             ),
         )
 
+    async def agent_response_async(
+        self,
+        msg: Optional[str | ChatDocument] = None,
+    ) -> Optional[ChatDocument]:
+        """
+        Asynch version of `agent_response`. See there for details.
+        """
+        if msg is None:
+            return None
+
+        results = await self.handle_message_async(msg)
+
+        return self._agent_response_final(msg, results)
+
+    def agent_response(
+        self,
+        msg: Optional[str | ChatDocument] = None,
+    ) -> Optional[ChatDocument]:
+        """
+        Response from the "agent itself", typically (but not only)
+        used to handle LLM's "tool message" or `function_call`
+        (e.g. OpenAI `function_call`).
+        Args:
+            msg (str|ChatDocument): the input to respond to: if msg is a string,
+                and it contains a valid JSON-structured "tool message", or
+                if msg is a ChatDocument, and it contains a `function_call`.
+        Returns:
+            Optional[ChatDocument]: the response, packaged as a ChatDocument
+
+        """
+        if msg is None:
+            return None
+
+        results = self.handle_message(msg)
+
+        return self._agent_response_final(msg, results)
+
     def process_tool_results(
         self,
         results: str,
@@ -618,60 +693,46 @@ class Agent(ABC):
             recipient=recipient,
         )
 
-    async def user_response_async(
-        self,
-        msg: Optional[str | ChatDocument] = None,
-    ) -> Optional[ChatDocument]:
-        return self.user_response(msg)
-
-    def user_response(
-        self,
-        msg: Optional[str | ChatDocument] = None,
-    ) -> Optional[ChatDocument]:
+    def user_can_respond(self, msg: Optional[str | ChatDocument] = None) -> bool:
         """
-        Get user response to current message. Could allow (human) user to intervene
-        with an actual answer, or quit using "q" or "x"
+        Whether the user can respond to a message.
 
         Args:
             msg (str|ChatDocument): the string to respond to.
 
         Returns:
-            (str) User response, packaged as a ChatDocument
 
         """
-
         # When msg explicitly addressed to user, this means an actual human response
         # is being sought.
         need_human_response = (
             isinstance(msg, ChatDocument) and msg.metadata.recipient == Entity.USER
         )
-        default_user_msg = (
-            (self.default_human_response or "null") if need_human_response else ""
-        )
 
         if not self.interactive and not need_human_response:
-            return None
-        elif self.default_human_response is not None:
-            user_msg = self.default_human_response
-        else:
-            if self.callbacks.get_user_response is not None:
-                # ask user with empty prompt: no need for prompt
-                # since user has seen the conversation so far.
-                # But non-empty prompt can be useful when Agent
-                # uses a tool that requires user input, or in other scenarios.
-                user_msg = self.callbacks.get_user_response(prompt="")
-            else:
-                user_msg = Prompt.ask(
-                    f"[blue]{self.indent}"
-                    + self.config.human_prompt
-                    + f"\n{self.indent}"
-                ).strip()
+            return False
+
+        return True
+
+    def _user_response_final(
+        self, msg: Optional[str | ChatDocument], user_msg: str
+    ) -> Optional[ChatDocument]:
+        """
+        Convert user_msg to final response.
+        """
+        if not user_msg:
+            need_human_response = (
+                isinstance(msg, ChatDocument) and msg.metadata.recipient == Entity.USER
+            )
+            user_msg = (
+                (self.default_human_response or "null") if need_human_response else ""
+            )
+        user_msg = user_msg.strip()
 
         tool_ids = []
         if msg is not None and isinstance(msg, ChatDocument):
             tool_ids = msg.metadata.tool_ids
 
-        user_msg = user_msg.strip() or default_user_msg.strip()
         # only return non-None result if user_msg not empty
         if not user_msg:
             return None
@@ -693,6 +754,72 @@ class Agent(ABC):
                 ),
             )
 
+    async def user_response_async(
+        self,
+        msg: Optional[str | ChatDocument] = None,
+    ) -> Optional[ChatDocument]:
+        """
+        Asynch version of `user_response`. See there for details.
+        """
+        if not self.user_can_respond(msg):
+            return None
+
+        if self.default_human_response is not None:
+            user_msg = self.default_human_response
+        else:
+            if (
+                self.callbacks.get_user_response_async is not None
+                and self.callbacks.get_user_response_async is not async_noop_fn
+            ):
+                user_msg = await self.callbacks.get_user_response_async(prompt="")
+            elif self.callbacks.get_user_response is not None:
+                user_msg = self.callbacks.get_user_response(prompt="")
+            else:
+                user_msg = Prompt.ask(
+                    f"[blue]{self.indent}"
+                    + self.config.human_prompt
+                    + f"\n{self.indent}"
+                )
+
+        return self._user_response_final(msg, user_msg)
+
+    def user_response(
+        self,
+        msg: Optional[str | ChatDocument] = None,
+    ) -> Optional[ChatDocument]:
+        """
+        Get user response to current message. Could allow (human) user to intervene
+        with an actual answer, or quit using "q" or "x"
+
+        Args:
+            msg (str|ChatDocument): the string to respond to.
+
+        Returns:
+            (str) User response, packaged as a ChatDocument
+
+        """
+
+        if not self.user_can_respond(msg):
+            return None
+
+        if self.default_human_response is not None:
+            user_msg = self.default_human_response
+        else:
+            if self.callbacks.get_user_response is not None:
+                # ask user with empty prompt: no need for prompt
+                # since user has seen the conversation so far.
+                # But non-empty prompt can be useful when Agent
+                # uses a tool that requires user input, or in other scenarios.
+                user_msg = self.callbacks.get_user_response(prompt="")
+            else:
+                user_msg = Prompt.ask(
+                    f"[blue]{self.indent}"
+                    + self.config.human_prompt
+                    + f"\n{self.indent}"
+                )
+
+        return self._user_response_final(msg, user_msg)
+
     @no_type_check
     def llm_can_respond(self, message: Optional[str | ChatDocument] = None) -> bool:
         """
@@ -784,7 +911,7 @@ class Agent(ABC):
                     f"""
                 Requested output length has been shortened to {output_len}
                 so that the total length of Prompt + Output is less than
-                the completion context length of the LLM. 
+                the completion context length of the LLM.
                 """
                 )
 
@@ -858,7 +985,7 @@ class Agent(ABC):
                         f"""
                     Requested output length has been shortened to {output_len}
                     so that the total length of Prompt + Output is less than
-                    the completion context length of the LLM. 
+                    the completion context length of the LLM.
                     """
                     )
             if self.llm.get_stream() and not settings.quiet:
@@ -1040,6 +1167,7 @@ class Agent(ABC):
         Returns:
             List[ToolMessage]: list of ToolMessage objects
         """
+        self.tool_error = False
         substrings = XMLToolMessage.find_candidates(input_str)
         is_json = False
         if len(substrings) == 0:
@@ -1049,7 +1177,11 @@ class Agent(ABC):
                 return []
 
         results = [self._get_one_tool_message(j, is_json) for j in substrings]
-        return [r for r in results if r is not None]
+        valid_results = [r for r in results if r is not None]
+        # If any tool is correctly formed we do not set the flag
+        if len(valid_results) > 0:
+            self.tool_error = False
+        return valid_results
 
     def get_function_call_class(self, msg: ChatDocument) -> Optional[ToolMessage]:
         """
@@ -1063,7 +1195,7 @@ class Agent(ABC):
         if tool_name not in self.llm_tools_handled:
             logger.warning(
                 f"""
-                The function_call '{tool_name}' is not handled 
+                The function_call '{tool_name}' is not handled
                 by the agent named '{self.config.name}'!
                 If you intended this agent to handle this function_call,
                 either the fn-call name is incorrectly generated by the LLM,
@@ -1071,7 +1203,10 @@ class Agent(ABC):
                 or you need to enable this agent to handle this fn-call.
                 """
             )
+            if tool_name not in self.all_llm_tools_known:
+                self.tool_error = True
             return None
+        self.tool_error = False
         tool_class = self.llm_tools_map[tool_name]
         tool_msg.update(dict(request=tool_name))
         tool = tool_class.parse_obj(tool_msg)
@@ -1086,6 +1221,7 @@ class Agent(ABC):
         if msg.oai_tool_calls is None:
             return []
         tools = []
+        all_errors = True
         for tc in msg.oai_tool_calls:
             if tc.function is None:
                 continue
@@ -1094,7 +1230,7 @@ class Agent(ABC):
             if tool_name not in self.llm_tools_handled:
                 logger.warning(
                     f"""
-                    The tool_call '{tool_name}' is not handled 
+                    The tool_call '{tool_name}' is not handled
                     by the agent named '{self.config.name}'!
                     If you intended this agent to handle this function_call,
                     either the fn-call name is incorrectly generated by the LLM,
@@ -1103,11 +1239,14 @@ class Agent(ABC):
                     """
                 )
                 continue
+            all_errors = False
             tool_class = self.llm_tools_map[tool_name]
             tool_msg.update(dict(request=tool_name))
             tool = tool_class.parse_obj(tool_msg)
             tool.id = tc.id or ""
             tools.append(tool)
+        # When no tool is valid, set the recovery flag
+        self.tool_error = all_errors
         return tools
 
     def tool_validation_error(self, ve: ValidationError) -> str:
@@ -1126,65 +1265,18 @@ class Agent(ABC):
             [f"{e['loc']}: {e['msg']}" for e in ve.errors() if "loc" in e]
         )
         return f"""
-        There were one or more errors in your attempt to use the 
-        TOOL or function_call named '{tool_name}': 
+        There were one or more errors in your attempt to use the
+        TOOL or function_call named '{tool_name}':
         {bad_field_errors}
         Please write your message again, correcting the errors.
         """
 
-    def handle_message(
-        self, msg: str | ChatDocument
-    ) -> None | str | OrderedDict[str, str] | ChatDocument:
+    def _get_multiple_orch_tool_errs(
+        self, tools: List[ToolMessage]
+    ) -> List[str | ChatDocument | None]:
         """
-        Handle a "tool" message either a string containing one or more
-        valid "tool" JSON substrings,  or a
-        ChatDocument containing a `function_call` attribute.
-        Handle with the corresponding handler method, and return
-        the results as a combined string.
-
-        Args:
-            msg (str | ChatDocument): The string or ChatDocument to handle
-
-        Returns:
-            The result of the handler method can be:
-             - None if no tools successfully handled, or no tools present
-             - str if langroid-native JSON tools were handled, and results concatenated,
-                 OR there's a SINGLE OpenAI tool-call.
-                (We do this so the common scenario of a single tool/fn-call
-                 has a simple behavior).
-             - Dict[str, str] if multiple OpenAI tool-calls were handled
-                 (dict is an id->result map)
-             - ChatDocument if a handler returned a ChatDocument, intended to be the
-                 final response of the `agent_response` method.
+        Return error document if the message contains multiple orchestration tools
         """
-        try:
-            tools = self.get_tool_messages(msg)
-            tools = [t for t in tools if self._tool_recipient_match(t)]
-        except ValidationError as ve:
-            # correct tool name but bad fields
-            return self.tool_validation_error(ve)
-        except XMLException as xe:  # from XMLToolMessage parsing
-            return str(xe)
-        except ValueError:
-            # invalid tool name
-            # We return None since returning "invalid tool name" would
-            # be considered a valid result in task loop, and would be treated
-            # as a response to the tool message even though the tool was not intended
-            # for this agent.
-            return None
-        if len(tools) > 1 and not self.config.allow_multiple_tools:
-            return self.to_ChatDocument("ERROR: Use ONE tool at a time!")
-        if len(tools) == 0:
-            fallback_result = self.handle_message_fallback(msg)
-            if fallback_result is None:
-                return None
-            return self.to_ChatDocument(
-                fallback_result,
-                chat_doc=msg if isinstance(msg, ChatDocument) else None,
-            )
-        has_ids = all([t.id != "" for t in tools])
-        chat_doc = msg if isinstance(msg, ChatDocument) else None
-
         # check whether there are multiple orchestration-tools (e.g. DoneTool etc),
         # in which case set result to error-string since we don't yet support
         # multi-tools with one or more orch tools.
@@ -1211,20 +1303,24 @@ class Agent(ABC):
         )
 
         has_orch = any(isinstance(t, ORCHESTRATION_TOOLS) for t in tools)
-        results: List[str | ChatDocument | None]
         if has_orch and len(tools) > 1:
             err_str = "ERROR: Use ONE tool at a time!"
-            results = [err_str for _ in tools]
-        else:
-            results = [self.handle_tool_message(t, chat_doc=chat_doc) for t in tools]
-            # if there's a solitary ChatDocument|str result, return it as is
-            if len(results) == 1 and isinstance(results[0], (str, ChatDocument)):
-                return results[0]
-            # extract content from ChatDocument results so we have all str|None
-            results = [r.content if isinstance(r, ChatDocument) else r for r in results]
+            return [err_str for _ in tools]
+
+        return []
+
+    def _handle_message_final(
+        self, tools: List[ToolMessage], results: List[str | ChatDocument | None]
+    ) -> None | str | OrderedDict[str, str] | ChatDocument:
+        """
+        Convert results to final response
+        """
+        # extract content from ChatDocument results so we have all str|None
+        results = [r.content if isinstance(r, ChatDocument) else r for r in results]
 
-        # now all results are str|None
         tool_names = [t.default_value("request") for t in tools]
+
+        has_ids = all([t.id != "" for t in tools])
         if has_ids:
             id2result = OrderedDict(
                 (t.id, r)
@@ -1259,6 +1355,117 @@ class Agent(ABC):
         final = "\n\n".join(str_results)
         return final
 
+    async def handle_message_async(
+        self, msg: str | ChatDocument
+    ) -> None | str | OrderedDict[str, str] | ChatDocument:
+        """
+        Asynch version of `handle_message`. See there for details.
+        """
+        try:
+            tools = self.get_tool_messages(msg)
+            tools = [t for t in tools if self._tool_recipient_match(t)]
+        except ValidationError as ve:
+            # correct tool name but bad fields
+            return self.tool_validation_error(ve)
+        except XMLException as xe:  # from XMLToolMessage parsing
+            return str(xe)
+        except ValueError:
+            # invalid tool name
+            # We return None since returning "invalid tool name" would
+            # be considered a valid result in task loop, and would be treated
+            # as a response to the tool message even though the tool was not intended
+            # for this agent.
+            return None
+        if len(tools) > 1 and not self.config.allow_multiple_tools:
+            return self.to_ChatDocument("ERROR: Use ONE tool at a time!")
+        if len(tools) == 0:
+            fallback_result = self.handle_message_fallback(msg)
+            if fallback_result is None:
+                return None
+            return self.to_ChatDocument(
+                fallback_result,
+                chat_doc=msg if isinstance(msg, ChatDocument) else None,
+            )
+        chat_doc = msg if isinstance(msg, ChatDocument) else None
+
+        results = self._get_multiple_orch_tool_errs(tools)
+        if not results:
+            results = [
+                await self.handle_tool_message_async(t, chat_doc=chat_doc)
+                for t in tools
+            ]
+            # if there's a solitary ChatDocument|str result, return it as is
+            if len(results) == 1 and isinstance(results[0], (str, ChatDocument)):
+                return results[0]
+
+        return self._handle_message_final(tools, results)
+
+    def handle_message(
+        self, msg: str | ChatDocument
+    ) -> None | str | OrderedDict[str, str] | ChatDocument:
+        """
+        Handle a "tool" message either a string containing one or more
+        valid "tool" JSON substrings,  or a
+        ChatDocument containing a `function_call` attribute.
+        Handle with the corresponding handler method, and return
+        the results as a combined string.
+
+        Args:
+            msg (str | ChatDocument): The string or ChatDocument to handle
+
+        Returns:
+            The result of the handler method can be:
+             - None if no tools successfully handled, or no tools present
+             - str if langroid-native JSON tools were handled, and results concatenated,
+                 OR there's a SINGLE OpenAI tool-call.
+                (We do this so the common scenario of a single tool/fn-call
+                 has a simple behavior).
+             - Dict[str, str] if multiple OpenAI tool-calls were handled
+                 (dict is an id->result map)
+             - ChatDocument if a handler returned a ChatDocument, intended to be the
+                 final response of the `agent_response` method.
+        """
+        try:
+            tools = self.get_tool_messages(msg)
+            tools = [t for t in tools if self._tool_recipient_match(t)]
+        except ValidationError as ve:
+            # correct tool name but bad fields
+            return self.tool_validation_error(ve)
+        except XMLException as xe:  # from XMLToolMessage parsing
+            return str(xe)
+        except ValueError:
+            # invalid tool name
+            # We return None since returning "invalid tool name" would
+            # be considered a valid result in task loop, and would be treated
+            # as a response to the tool message even though the tool was not intended
+            # for this agent.
+            return None
+        if len(tools) > 1 and not self.config.allow_multiple_tools:
+            return self.to_ChatDocument("ERROR: Use ONE tool at a time!")
+        if len(tools) == 0:
+            fallback_result = self.handle_message_fallback(msg)
+            if fallback_result is None:
+                return None
+            return self.to_ChatDocument(
+                fallback_result,
+                chat_doc=msg if isinstance(msg, ChatDocument) else None,
+            )
+        chat_doc = msg if isinstance(msg, ChatDocument) else None
+
+        results = self._get_multiple_orch_tool_errs(tools)
+        if not results:
+            results = [self.handle_tool_message(t, chat_doc=chat_doc) for t in tools]
+            # if there's a solitary ChatDocument|str result, return it as is
+            if len(results) == 1 and isinstance(results[0], (str, ChatDocument)):
+                return results[0]
+
+        return self._handle_message_final(tools, results)
+
+    @property
+    def all_llm_tools_known(self) -> set[str]:
+        """All known tools; this may extend self.llm_tools_known."""
+        return self.llm_tools_known
+
     def handle_message_fallback(self, msg: str | ChatDocument) -> Any:
         """
         Fallback method for the "no-tools" scenario.
@@ -1278,7 +1485,7 @@ class Agent(ABC):
     ) -> Optional[ToolMessage]:
         """
         Parse the tool_candidate_str into ANY ToolMessage KNOWN to agent --
-        This includes non-used/handled tools, i.e. any tool in self.llm_tools_known.
+        This includes non-used/handled tools, i.e. any tool in self.all_llm_tools_known.
         The exception to this is below where we try our best to infer the tool
         when the LLM has "forgotten" to include the "request" field in the tool str ---
         in this case we ONLY look at the possible set of HANDLED tools, i.e.
@@ -1311,6 +1518,7 @@ class Agent(ABC):
         # }
 
         if not isinstance(maybe_tool_dict, dict):
+            self.tool_error = True
             return None
 
         properties = maybe_tool_dict.get("properties")
@@ -1318,7 +1526,14 @@ class Agent(ABC):
             maybe_tool_dict = properties
         request = maybe_tool_dict.get("request")
         if request is None:
-            possible = [self.llm_tools_map[r] for r in self.llm_tools_handled]
+            if self.enabled_requests_for_inference is None:
+                possible = [self.llm_tools_map[r] for r in self.llm_tools_handled]
+            else:
+                allowable = self.enabled_requests_for_inference.intersection(
+                    self.llm_tools_handled
+                )
+                possible = [self.llm_tools_map[r] for r in allowable]
+
             default_keys = set(ToolMessage.__fields__.keys())
             request_keys = set(maybe_tool_dict.keys())
 
@@ -1351,19 +1566,23 @@ class Agent(ABC):
             if len(candidate_tools) == 1:
                 return candidate_tools[0]
             else:
+                self.tool_error = True
                 return None
 
-        if not isinstance(request, str) or request not in self.llm_tools_known:
+        if not isinstance(request, str) or request not in self.all_llm_tools_known:
+            self.tool_error = True
             return None
 
         message_class = self.llm_tools_map.get(request)
         if message_class is None:
             logger.warning(f"No message class found for request '{request}'")
+            self.tool_error = True
             return None
 
         try:
             message = message_class.parse_obj(maybe_tool_dict)
         except ValidationError as ve:
+            self.tool_error = True
             raise ve
         return message
 
@@ -1474,7 +1693,6 @@ class Agent(ABC):
             value = tool.get_value_of_type(output_type)
             if value is not None:
                 return cast(T, value)
-
         return None
 
     def _maybe_truncate_result(
@@ -1511,6 +1729,37 @@ class Agent(ABC):
             ) + truncate_warning
             return result
 
+    async def handle_tool_message_async(
+        self,
+        tool: ToolMessage,
+        chat_doc: Optional[ChatDocument] = None,
+    ) -> None | str | ChatDocument:
+        """
+        Asynch version of `handle_tool_message`. See there for details.
+        """
+        tool_name = tool.default_value("request")
+        handler_method = getattr(self, tool_name + "_async", None)
+        if handler_method is None:
+            return self.handle_tool_message(tool, chat_doc=chat_doc)
+        has_chat_doc_arg = (
+            chat_doc is not None
+            and "chat_doc" in inspect.signature(handler_method).parameters
+        )
+        try:
+            if has_chat_doc_arg:
+                maybe_result = await handler_method(tool, chat_doc=chat_doc)
+            else:
+                maybe_result = await handler_method(tool)
+            result = self.to_ChatDocument(maybe_result, tool_name, chat_doc)
+        except Exception as e:
+            # raise the error here since we are sure it's
+            # not a pydantic validation error,
+            # which we check in `handle_message`
+            raise e
+        return self._maybe_truncate_result(
+            result, tool._max_result_tokens
+        )  # type: ignore
+
     def handle_tool_message(
         self,
         tool: ToolMessage,