max retry for auto toolcalls

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}"
+                            )
+                            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,
@@ -3128,25 +3161,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 +3196,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 +3217,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/operations/_patch.py

@@ -12,6 +12,7 @@
 import os
 import sys
 import time
+import json
 from pathlib import Path
 from typing import (
     IO,
@@ -844,6 +845,7 @@ class AgentsOperations(AgentsOperationsGenerated):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self._function_tool = _models.FunctionTool(set())
+        self._function_tool_max_retry = 10
 
     # pylint: disable=arguments-differ
     @overload
@@ -1803,6 +1805,7 @@ def create_and_process_run(
         )
 
         # Monitor and process the run status
+        current_retry = 0
         while run.status in [
             RunStatus.QUEUED,
             RunStatus.IN_PROGRESS,
@@ -1826,14 +1829,47 @@ def create_and_process_run(
                     toolset.add(self._function_tool)
                     tool_outputs = 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}"
+                            )
+                            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:
-                        self.submit_tool_outputs_to_run(thread_id=thread_id, run_id=run.id, tool_outputs=tool_outputs)
+                        run2 = self.submit_tool_outputs_to_run(
+                            thread_id=thread_id, run_id=run.id, tool_outputs=tool_outputs
+                        )
+                        logging.info("Tool outputs submitted to run: %s", run2.id)
 
             logging.info("Current run status: %s", run.status)
 
         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
     def create_stream(
         self,
@@ -3309,27 +3345,39 @@ def delete_agent(self, agent_id: str, **kwargs: Any) -> _models.AgentDeletionSta
         return 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]]
+        :keyword max_retry: Maximum number of errors allowed and retry per run or stream. Default value is 10.
+        :type max_retry: int
         """
 
     @overload
-    def enable_auto_function_calls(self, *, function_tool: _models.FunctionTool) -> None:
+    def enable_auto_function_calls(self, *, function_tool: _models.FunctionTool, 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: A FunctionTool object representing the tool to be used.
         :type function_tool: Optional[_models.FunctionTool]
+        :keyword max_retry: Maximum number of errors allowed and retry per run or stream. Default value is 10.
+        :type max_retry: int
         """
 
     @overload
-    def enable_auto_function_calls(self, *, toolset: _models.ToolSet) -> None:
+    def enable_auto_function_calls(self, *, toolset: _models.ToolSet, 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: A ToolSet object representing the set of tools to be used.
         :type toolset: Optional[_models.ToolSet]
+        :keyword max_retry: Maximum number of errors allowed and retry per run or stream. Default value is 10.
+        :type max_retry: int
         """
 
     @distributed_trace
@@ -3339,15 +3387,20 @@ def enable_auto_function_calls(
         functions: Optional[Set[Callable[..., Any]]] = None,
         function_tool: Optional[_models.FunctionTool] = None,
         toolset: Optional[_models.ToolSet] = 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: A FunctionTool object representing the tool to be used.
         :type function_tool: Optional[_models.FunctionTool]
         :keyword toolset: A ToolSet object representing the set of tools to be used.
         :type toolset: Optional[_models.ToolSet]
+        :keyword max_retry: Maximum number of errors allowed and retry per run or stream. Default value is 10.
+        :type max_retry: int
         """
         if functions:
             self._function_tool = _models.FunctionTool(functions)
@@ -3357,6 +3410,8 @@ def enable_auto_function_calls(
             tool = toolset.get_tool(_models.FunctionTool)
             self._function_tool = tool
 
+        self._function_tool_max_retry = max_retry
+
 
 __all__: List[str] = [
     "AgentsOperations",

sdk/ai/azure-ai-projects/samples/agents/sample_agents_run_with_toolset.py

@@ -45,7 +45,7 @@
     toolset.add(code_interpreter)
 
     # To enable tool calls executed automatically
-    project_client.agents.enable_auto_function_calls(toolset=toolset)
+    project_client.agents.enable_auto_function_calls(toolset=toolset, max_retry=3)
 
     agent = project_client.agents.create_agent(
         model=os.environ["MODEL_DEPLOYMENT_NAME"],

sdk/ai/azure-ai-projects/samples/agents/user_functions.py

@@ -28,6 +28,7 @@ def fetch_current_datetime(format: Optional[str] = None) -> str:
         time_format = "%Y-%m-%d %H:%M:%S"
 
     time_json = json.dumps({"current_time": current_time.strftime(time_format)})
+    raise ValueError("Just a minute")
     return time_json
 
 

sdk/ai/azure-ai-projects/tests/agents/test_agent_operations.py

@@ -52,6 +52,10 @@ def function2():
     return "output from the second agent"
 
 
+def function_throw_exception():
+    raise ValueError("Just a minute")
+
+
 class TestAgentsOperations:
     """Tests for agent operations"""
 
@@ -217,16 +221,17 @@ def _set_toolcalls(
         self, project_client: AgentsOperations, toolset1: Optional[ToolSet], toolset2: Optional[ToolSet]
     ) -> None:
         """Set the tool calls for the agent."""
+        max_retry = 3
         if toolset1 and toolset2:
             function_in_toolset1 = set(toolset1.get_tool(tool_type=FunctionTool)._functions.values())
             function_in_toolset2 = set(toolset2.get_tool(tool_type=FunctionTool)._functions.values())
             function_tool = FunctionTool(function_in_toolset1)
             function_tool.add_functions(function_in_toolset2)
-            project_client.enable_auto_function_calls(function_tool=function_tool)
+            project_client.enable_auto_function_calls(function_tool=function_tool, max_retry=max_retry)
         elif toolset1:
-            project_client.enable_auto_function_calls(toolset=toolset1)
+            project_client.enable_auto_function_calls(toolset=toolset1, max_retry=max_retry)
         elif toolset2:
-            project_client.enable_auto_function_calls(toolset=toolset2)
+            project_client.enable_auto_function_calls(toolset=toolset2, max_retry=max_retry)
 
     @patch("azure.ai.projects._patch.PipelineClient")
     @pytest.mark.parametrize(
@@ -296,6 +301,45 @@ def test_multiple_agents_create(
             project_client.agents.delete_agent(agent1.id)
             project_client.agents.delete_agent(agent2.id)
 
+    @patch("azure.ai.projects.operations._operations.AgentsOperations.cancel_run")
+    @patch("azure.ai.projects._patch.PipelineClient")
+    def test_auto_function_calls_retry(
+        self,
+        mock_pipeline_client_gen: MagicMock,
+        mock_cancel_run: MagicMock,
+    ) -> None:
+        """Test azure function with toolset."""
+        toolset = self.get_toolset("file_for_agent1", function_throw_exception)
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.side_effect = [
+            self._get_agent_json("first", "123", toolset),
+            self._get_run("run1", toolset),  # create_run
+            self._get_run("run2", toolset),  # get_run
+            self._get_run("run3", toolset),  # get_run
+            self._get_run("run4", toolset),  # get_run
+            self._get_run("run5", toolset),  # get_run
+        ]
+        mock_pipeline_response = MagicMock()
+        mock_pipeline_response.http_response = mock_response
+        mock_pipeline = MagicMock()
+        mock_pipeline._pipeline.run.return_value = mock_pipeline_response
+        mock_pipeline_client_gen.return_value = mock_pipeline
+        project_client = self.get_mock_client()
+        with project_client:
+            # Check that pipelines are created as expected.
+            self._set_toolcalls(project_client.agents, toolset, None)
+            agent1 = project_client.agents.create_agent(
+                model="gpt-4-1106-preview",
+                name="first",
+                instructions="You are a helpful assistant",
+                toolset=toolset,
+            )
+            # Create run with new tool set, which also can be none.
+            project_client.agents.create_and_process_run(thread_id="some_thread_id", agent_id=agent1.id)
+            assert mock_cancel_run.call_count == 1
+            assert project_client.agents.submit_tool_outputs_to_run.call_count == 3
+
     @patch("azure.ai.projects._patch.PipelineClient")
     @pytest.mark.parametrize(
         "file_agent_1,file_agent_2",