Add a max_actions_per_run variable on AgentService as a backstop (#535)

An observed failure mode of a reasoning Agent is when it spirals into
endlessly calling tools.

In a development loop, the programmer can quickly break the loop, but in
a hosted situation this can result in surprising costs if each tool
involves LLM usage.

This PR adds a simple backstop in which the `max_actions_per_run`
variable on AgentService (default = 5) acts as a limit on how many
actions the agent is permitted to run before being stopped.

If the AgentService attempts to run an action after
`max_actions_per_run` has been reached, it is treated as an error which
will be returned to the chat invoker.

Author: eob

src/steamship/agents/examples/example_assistant_with_caching.py

@@ -4,6 +4,9 @@
 from steamship.agents.service.agent_service import AgentService
 from steamship.agents.tools.image_generation import DalleTool
 from steamship.agents.tools.search import SearchTool
+from steamship.data import TagValueKey
+from steamship.invocable import post
+from steamship.utils.kv_store import KeyValueStore
 from steamship.utils.repl import AgentREPL
 
 
@@ -15,6 +18,11 @@ class MyCachingAssistant(AgentService):
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs, use_llm_cache=True, use_action_cache=True)
+
+        # Load the max_actions_per_run from the saved store for use in testing.
+        self.kv = KeyValueStore(self.client)
+        self.max_actions_per_run = self.get_max_actions_per_run()
+
         self.set_default_agent(
             FunctionsBasedAgent(
                 tools=[
@@ -26,6 +34,19 @@ def __init__(self, **kwargs):
             )
         )
 
+    @post("set_max_actions_per_run")
+    def set_max_actions_per_run(self, value: int):
+        """Save the max_actions_per_run value so that it will be reloaded upon next request."""
+        self.max_actions_per_run = self.kv.set(
+            "max_actions_per_run", {TagValueKey.NUMBER_VALUE: value}
+        )
+        return value
+
+    @post("get_max_actions_per_run")
+    def get_max_actions_per_run(self) -> int:
+        """Save the max_actions_per_run value so that it will be reloaded upon next request."""
+        return (self.kv.get("max_actions_per_run") or {}).get(TagValueKey.NUMBER_VALUE, 5)
+
 
 if __name__ == "__main__":
     # AgentREPL provides a mechanism for local execution of an AgentService method.

src/steamship/agents/service/agent_service.py

@@ -24,15 +24,24 @@ class AgentService(PackageService):
     use_action_cache: bool
     """Whether or not to cache agent Actions (for tool runs) by default."""
 
+    max_actions_per_run: int
+    """The maximum number of actions to permit while the agent is reasoning.
+
+    This is intended primarily to act as a backstop to prevent a condition in which the Agent decides to loop endlessly
+    on tool runs that consume resources with a cost-basis (e.g. prompt completions, embedding operations, vector lookups)
+    """
+
     def __init__(
         self,
         use_llm_cache: Optional[bool] = False,
         use_action_cache: Optional[bool] = False,
+        max_actions_per_run: Optional[int] = 5,
         agent: Optional[Agent] = None,
         **kwargs,
     ):
         self.use_llm_cache = use_llm_cache
         self.use_action_cache = use_action_cache
+        self.max_actions_per_run = max_actions_per_run
         self.agent = agent
         super().__init__(**kwargs)
 
@@ -144,8 +153,22 @@ def run_agent(self, agent: Agent, context: AgentContext):
             agent=agent, input_blocks=[context.chat_history.last_user_message], context=context
         )
 
+        number_of_actions_run = 0
+
         while not isinstance(action, FinishAction):
+            # Throw an error if we've exceeded our action budget.
+            if number_of_actions_run >= self.max_actions_per_run:
+                raise SteamshipError(
+                    message=(
+                        f"Agent reached its Action budget of {self.max_actions_per_run} without arriving at a response. If you are the developer, checking the logs may reveal it was selecting unhelpful tools or receiving unhelpful responses from them."
+                    )
+                )
+
+            # Run the next action
             self.run_action(agent=agent, action=action, context=context)
+            number_of_actions_run += 1
+
+            # Select a new next_action and log it
             action = self.next_action(agent=agent, input_blocks=action.output, context=context)
 
             # TODO: Arrive at a solid design for the details of this structured log object

tests/steamship_tests/agents/test_agent_service.py

@@ -42,6 +42,17 @@ def test_example_with_caching_service(client: Steamship):
         assert "image" in blocks[0].text
         assert blocks[1].is_image()
 
+        # attempt with a max_actions_per_run budget of 0 (should fail!)
+        assert caching_agent.invoke("get_max_actions_per_run") == 5
+        caching_agent.invoke("set_max_actions_per_run", value=0)
+        assert caching_agent.invoke("get_max_actions_per_run") == 0
+
+        with pytest.raises(SteamshipError, match="budget"):
+            blocks_json = caching_agent.invoke("prompt", prompt="draw a cat", context_id=context_id)
+
+        caching_agent.invoke("set_max_actions_per_run", value=5)
+        assert caching_agent.invoke("get_max_actions_per_run") == 5
+
         agent_context = AgentContext.get_or_create(
             client=client, context_keys=context_keys, use_llm_cache=True, use_action_cache=True
         )