max auto function calls

Author: howieleung

sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py

@@ -13,6 +13,7 @@
 import logging
 import os
 import time
+import json
 from pathlib import Path
 from typing import (
     IO,
@@ -664,6 +665,7 @@ class AgentsOperations(AgentsOperationsGenerated):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self._function_tool = _models.AsyncFunctionTool(set())
+        self._function_tool_max_retry = 10
 
     # pylint: disable=arguments-differ
     @overload
@@ -1622,6 +1624,7 @@ async def create_and_process_run(
         )
 
         # Monitor and process the run status
+        current_retry = 0
         while run.status in [
             RunStatus.QUEUED,
             RunStatus.IN_PROGRESS,
@@ -1643,6 +1646,17 @@ async def create_and_process_run(
                     toolset.add(self._function_tool)
                     tool_outputs = await toolset.execute_tool_calls(tool_calls)
 
+                    if self._has_errors_in_toolcalls_output(tool_outputs):
+                        if current_retry >= self._function_tool_max_retry:
+                            logging.warning(
+                                f"Tool outputs contain errors - reaching max retry {self._function_tool_max_retry}"
+                            )
+                            await self.cancel_run(thread_id=thread_id, run_id=run.id)
+                            break
+                        else:
+                            logging.warning(f"Tool outputs contain errors - retrying")
+                            current_retry += 1
+
                     logging.info("Tool outputs: %s", tool_outputs)
                     if tool_outputs:
                         await self.submit_tool_outputs_to_run(
@@ -1653,6 +1667,25 @@ async def create_and_process_run(
 
         return run
 
+    def _has_errors_in_toolcalls_output(self, tool_outputs: List[Dict]) -> bool:
+        """
+        Check if any tool output contains an error.
+
+        :param List[Dict] tool_outputs: A list of tool outputs to check.
+        :return: True if any output contains an error, False otherwise.
+        :rtype: bool
+        """
+        for tool_output in tool_outputs:
+            output = tool_output.get("output")
+            if isinstance(output, str):
+                try:
+                    output_json = json.loads(output)
+                    if "error" in output_json:
+                        return True
+                except json.JSONDecodeError:
+                    continue
+        return False
+
     @overload
     async def create_stream(
         self,
@@ -2090,6 +2123,8 @@ async def create_stream(  # pyright: ignore[reportInconsistentOverload]
 
         if not event_handler:
             event_handler = cast(_models.BaseAsyncAgentEventHandlerT, _models.AsyncAgentEventHandler())
+        if isinstance(event_handler, _models.AsyncAgentEventHandler):
+            event_handler.set_max_retry(self._function_tool_max_retry)
 
         return _models.AsyncAgentRunStream(response_iterator, self._handle_submit_tool_outputs, event_handler)
 
@@ -2320,13 +2355,14 @@ async def submit_tool_outputs_to_stream(  # pyright: ignore[reportInconsistentOv
         event_handler.initialize(response_iterator, self._handle_submit_tool_outputs)
 
     async def _handle_submit_tool_outputs(
-        self, run: _models.ThreadRun, event_handler: _models.BaseAsyncAgentEventHandler
-    ) -> None:
+        self, run: _models.ThreadRun, event_handler: _models.BaseAsyncAgentEventHandler, submit_with_error: bool
+    ) -> Any:
+        tool_outputs = []
         if isinstance(run.required_action, _models.SubmitToolOutputsAction):
             tool_calls = run.required_action.submit_tool_outputs.tool_calls
             if not tool_calls:
                 logger.debug("No tool calls to execute.")
-                return
+                return tool_outputs
 
             # We need tool set only if we are executing local function. In case if
             # the tool is azure_function we just need to wait when it will be finished.
@@ -2338,11 +2374,20 @@ async def _handle_submit_tool_outputs(
                 toolset.add(self._function_tool)
                 tool_outputs = await toolset.execute_tool_calls(tool_calls)
 
+                if self._has_errors_in_toolcalls_output(tool_outputs):
+                    if submit_with_error:
+                        logging.warning(f"Tool outputs contain errors - retrying")
+                    else:
+                        logging.warning(f"Tool outputs contain errors - reaching max retry limit")
+                        await self.cancel_run(thread_id=run.thread_id, run_id=run.id)
+                        return tool_outputs
+
                 logger.info("Tool outputs: %s", tool_outputs)
                 if tool_outputs:
                     await self.submit_tool_outputs_to_stream(
                         thread_id=run.thread_id, run_id=run.id, tool_outputs=tool_outputs, event_handler=event_handler
                     )
+        return tool_outputs
 
     # pylint: disable=arguments-differ
     @overload
@@ -3128,25 +3173,31 @@ async def delete_agent(self, agent_id: str, **kwargs: Any) -> _models.AgentDelet
         return await super().delete_agent(agent_id, **kwargs)
 
     @overload
-    def enable_auto_function_calls(self, *, functions: Set[Callable[..., Any]]) -> None:
+    def enable_auto_function_calls(self, *, functions: Set[Callable[..., Any]], max_retry: int = 10) -> None:
         """Enables tool calls to be executed automatically during create_and_process_run or streaming.
         If this is not set, functions must be called manually.
+        If automatic function calls fail, the agents will receive error messages allowing it to retry with another
+        function call or figure out the answer with its knowledge.
         :keyword functions: A set of callable functions to be used as tools.
         :type functions: Set[Callable[..., Any]]
         """
 
     @overload
-    def enable_auto_function_calls(self, *, function_tool: _models.AsyncFunctionTool) -> None:
+    def enable_auto_function_calls(self, *, function_tool: _models.AsyncFunctionTool, max_retry: int = 10) -> None:
         """Enables tool calls to be executed automatically during create_and_process_run or streaming.
         If this is not set, functions must be called manually.
+        If automatic function calls fail, the agents will receive error messages allowing it to retry with another
+        function call or figure out the answer with its knowledge.
         :keyword function_tool: An AsyncFunctionTool object representing the tool to be used.
         :type function_tool: Optional[_models.AsyncFunctionTool]
         """
 
     @overload
-    def enable_auto_function_calls(self, *, toolset: _models.AsyncToolSet) -> None:
+    def enable_auto_function_calls(self, *, toolset: _models.AsyncToolSet, max_retry: int = 10) -> None:
         """Enables tool calls to be executed automatically during create_and_process_run or streaming.
         If this is not set, functions must be called manually.
+        If automatic function calls fail, the agents will receive error messages allowing it to retry with another
+        function call or figure out the answer with its knowledge.
         :keyword toolset: An AsyncToolSet object representing the set of tools to be used.
         :type toolset: Optional[_models.AsyncToolSet]
         """
@@ -3157,9 +3208,12 @@ def enable_auto_function_calls(
         functions: Optional[Set[Callable[..., Any]]] = None,
         function_tool: Optional[_models.AsyncFunctionTool] = None,
         toolset: Optional[_models.AsyncToolSet] = None,
+        max_retry: int = 10,
     ) -> None:
         """Enables tool calls to be executed automatically during create_and_process_run or streaming.
         If this is not set, functions must be called manually.
+        If automatic function calls fail, the agents will receive error messages allowing it to retry with another
+        function call or figure out the answer with its knowledge.
         :keyword functions: A set of callable functions to be used as tools.
         :type functions: Set[Callable[..., Any]]
         :keyword function_tool: An AsyncFunctionTool object representing the tool to be used.
@@ -3175,6 +3229,8 @@ def enable_auto_function_calls(
             tool = toolset.get_tool(_models.AsyncFunctionTool)
             self._function_tool = tool
 
+        self._function_tool_max_retry = max_retry
+
 
 class _SyncCredentialWrapper(TokenCredential):
     """

sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py

@@ -98,6 +98,26 @@
 StreamEventData = Union["MessageDeltaChunk", "ThreadMessage", ThreadRun, RunStep, str]
 
 
+def _has_errors_in_toolcalls_output(tool_outputs: List[Dict]) -> bool:
+    """
+    Check if any tool output contains an error.
+
+    :param List[Dict] tool_outputs: A list of tool outputs to check.
+    :return: True if any output contains an error, False otherwise.
+    :rtype: bool
+    """
+    for tool_output in tool_outputs:
+        output = tool_output.get("output")
+        if isinstance(output, str):
+            try:
+                output_json = json.loads(output)
+                if "error" in output_json:
+                    return True
+            except json.JSONDecodeError:
+                continue
+    return False
+
+
 def _filter_parameters(model_class: Type, parameters: Dict[str, Any]) -> Dict[str, Any]:
     """
     Remove the parameters, non present in class public fields; return shallow copy of a dictionary.
@@ -734,7 +754,7 @@ def execute(self, tool_call: RequiredFunctionToolCall) -> Any:
         try:
             function, parsed_arguments = self._get_func_and_args(tool_call)
             return function(**parsed_arguments) if parsed_arguments else function()
-        except TypeError as e:
+        except Exception as e:  # pylint: disable=broad-exception-caught
             error_message = f"Error executing function '{tool_call.function.name}': {e}"
             logging.error(error_message)
             # Return error message as JSON string back to agent in order to make possible self
@@ -745,13 +765,12 @@ def execute(self, tool_call: RequiredFunctionToolCall) -> Any:
 class AsyncFunctionTool(BaseFunctionTool):
 
     async def execute(self, tool_call: RequiredFunctionToolCall) -> Any:  # pylint: disable=invalid-overridden-method
-        function, parsed_arguments = self._get_func_and_args(tool_call)
-
         try:
+            function, parsed_arguments = self._get_func_and_args(tool_call)
             if inspect.iscoroutinefunction(function):
                 return await function(**parsed_arguments) if parsed_arguments else await function()
             return function(**parsed_arguments) if parsed_arguments else function()
-        except TypeError as e:
+        except Exception as e:  # pylint: disable=broad-exception-caught
             error_message = f"Error executing function '{tool_call.function.name}': {e}"
             logging.error(error_message)
             # Return error message as JSON string back to agent in order to make possible self correction
@@ -1511,13 +1530,13 @@ class BaseAgentEventHandler(Iterator[T]):
 
     def __init__(self) -> None:
         self.response_iterator: Optional[Iterator[bytes]] = None
-        self.submit_tool_outputs: Optional[Callable[[ThreadRun, "BaseAgentEventHandler[T]"], None]] = None
+        self.submit_tool_outputs: Optional[Callable[[ThreadRun, "BaseAgentEventHandler[T]", bool], Any]]
         self.buffer: Optional[bytes] = None
 
     def initialize(
         self,
         response_iterator: Iterator[bytes],
-        submit_tool_outputs: Callable[[ThreadRun, "BaseAgentEventHandler[T]"], None],
+        submit_tool_outputs: Callable[[ThreadRun, "BaseAgentEventHandler[T]", bool], Any],
     ) -> None:
         self.response_iterator = (
             itertools.chain(self.response_iterator, response_iterator) if self.response_iterator else response_iterator
@@ -1569,17 +1588,33 @@ def until_done(self) -> None:
 
 
 class AsyncAgentEventHandler(BaseAsyncAgentEventHandler[Tuple[str, StreamEventData, Optional[EventFunctionReturnT]]]):
+    def __init__(self) -> None:
+        super().__init__()
+        self._max_retry = 10
+        self.current_retry = 0
+
+    def set_max_retry(self, max_retry: int) -> None:
+        """
+        Set the maximum number of retries for tool output submission.
+
+        :param int max_retry: The maximum number of retries.
+        """
+        self._max_retry = max_retry
 
     async def _process_event(self, event_data_str: str) -> Tuple[str, StreamEventData, Optional[EventFunctionReturnT]]:
+
         event_type, event_data_obj = _parse_event(event_data_str)
         if (
             isinstance(event_data_obj, ThreadRun)
             and event_data_obj.status == "requires_action"
             and isinstance(event_data_obj.required_action, SubmitToolOutputsAction)
         ):
-            await cast(Callable[[ThreadRun, "BaseAsyncAgentEventHandler"], Awaitable[None]], self.submit_tool_outputs)(
-                event_data_obj, self
-            )
+            tool_output = await cast(
+                Callable[[ThreadRun, "BaseAsyncAgentEventHandler", bool], Awaitable[Any]], self.submit_tool_outputs
+            )(event_data_obj, self, self.current_retry < self._max_retry)
+
+            if _has_errors_in_toolcalls_output(tool_output):
+                self.current_retry += 1
 
         func_rt: Optional[EventFunctionReturnT] = None
         try:
@@ -1682,6 +1717,18 @@ async def on_unhandled_event(
 
 
 class AgentEventHandler(BaseAgentEventHandler[Tuple[str, StreamEventData, Optional[EventFunctionReturnT]]]):
+    def __init__(self) -> None:
+        super().__init__()
+        self._max_retry = 10
+        self.current_retry = 0
+
+    def set_max_retry(self, max_retry: int) -> None:
+        """
+        Set the maximum number of retries for tool output submission.
+
+        :param int max_retry: The maximum number of retries.
+        """
+        self._max_retry = max_retry
 
     def _process_event(self, event_data_str: str) -> Tuple[str, StreamEventData, Optional[EventFunctionReturnT]]:
 
@@ -1691,10 +1738,13 @@ def _process_event(self, event_data_str: str) -> Tuple[str, StreamEventData, Opt
             and event_data_obj.status == "requires_action"
             and isinstance(event_data_obj.required_action, SubmitToolOutputsAction)
         ):
-            cast(Callable[[ThreadRun, "BaseAgentEventHandler"], Awaitable[None]], self.submit_tool_outputs)(
-                event_data_obj, self
+            tool_output = cast(Callable[[ThreadRun, "BaseAgentEventHandler", bool], Any], self.submit_tool_outputs)(
+                event_data_obj, self, self.current_retry < self._max_retry
             )
 
+            if _has_errors_in_toolcalls_output(tool_output):
+                self.current_retry += 1
+
         func_rt: Optional[EventFunctionReturnT] = None
         try:
             if isinstance(event_data_obj, MessageDeltaChunk):
@@ -1792,15 +1842,15 @@ class AsyncAgentRunStream(Generic[BaseAsyncAgentEventHandlerT]):
     def __init__(
         self,
         response_iterator: AsyncIterator[bytes],
-        submit_tool_outputs: Callable[[ThreadRun, BaseAsyncAgentEventHandlerT], Awaitable[None]],
+        submit_tool_outputs: Callable[[ThreadRun, BaseAsyncAgentEventHandlerT, bool], Awaitable[Any]],
         event_handler: BaseAsyncAgentEventHandlerT,
     ):
         self.response_iterator = response_iterator
         self.event_handler = event_handler
         self.submit_tool_outputs = submit_tool_outputs
         self.event_handler.initialize(
             self.response_iterator,
-            cast(Callable[[ThreadRun, BaseAsyncAgentEventHandler], Awaitable[None]], submit_tool_outputs),
+            cast(Callable[[ThreadRun, BaseAsyncAgentEventHandler], Awaitable[Any]], submit_tool_outputs),
         )
 
     async def __aenter__(self):
@@ -1818,15 +1868,15 @@ class AgentRunStream(Generic[BaseAgentEventHandlerT]):
     def __init__(
         self,
         response_iterator: Iterator[bytes],
-        submit_tool_outputs: Callable[[ThreadRun, BaseAgentEventHandlerT], None],
+        submit_tool_outputs: Callable[[ThreadRun, BaseAgentEventHandlerT, bool], Any],
         event_handler: BaseAgentEventHandlerT,
     ):
         self.response_iterator = response_iterator
         self.event_handler = event_handler
         self.submit_tool_outputs = submit_tool_outputs
         self.event_handler.initialize(
             self.response_iterator,
-            cast(Callable[[ThreadRun, BaseAgentEventHandler], None], submit_tool_outputs),
+            cast(Callable[[ThreadRun, BaseAgentEventHandler, bool], Any], submit_tool_outputs),
         )
 
     def __enter__(self):