From a8bbfb2da3f8c28869b12c8a9bb21209b0d03089 Mon Sep 17 00:00:00 2001
From: Nir Gazit <nirga@users.noreply.github.com>
Date: Thu, 13 Jul 2023 09:09:25 +0300
Subject: [PATCH] feat(agents): allow trimming of intermediate steps to last N
 (#6476)

Added an option to trim intermediate steps to last N steps. This is
especially useful for long-running agents. Users can explicitly specify
N or provide a function that does custom trimming/manipulation on
intermediate steps. I've mimicked the API of the `handle_parsing_errors`
parameter.
---
 langchain/agents/agent.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/langchain/agents/agent.py b/langchain/agents/agent.py
index 807101e074..7c3a354b5c 100644
--- a/langchain/agents/agent.py
+++ b/langchain/agents/agent.py
@@ -659,6 +659,9 @@ s
      as an argument, and the result of that function will be passed to the agent
       as an observation.
     """
+    trim_intermediate_steps: Union[
+        int, Callable[[List[Tuple[AgentAction, str]]], List[Tuple[AgentAction, str]]]
+    ] = -1
 
     @classmethod
     def from_agent_and_tools(
@@ -788,6 +791,8 @@ s
         Override this to take control of how the agent makes and acts on choices.
         """
         try:
+            intermediate_steps = self._prepare_intermediate_steps(intermediate_steps)
+
             # Call the LLM to see what to do.
             output = self.agent.plan(
                 intermediate_steps,
@@ -879,6 +884,8 @@ s
         Override this to take control of how the agent makes and acts on choices.
         """
         try:
+            intermediate_steps = self._prepare_intermediate_steps(intermediate_steps)
+
             # Call the LLM to see what to do.
             output = await self.agent.aplan(
                 intermediate_steps,
@@ -1088,3 +1095,16 @@ s
                     "",
                 )
         return None
+
+    def _prepare_intermediate_steps(
+        self, intermediate_steps: List[Tuple[AgentAction, str]]
+    ) -> List[Tuple[AgentAction, str]]:
+        if (
+            isinstance(self.trim_intermediate_steps, int)
+            and self.trim_intermediate_steps > 0
+        ):
+            return intermediate_steps[-self.trim_intermediate_steps :]
+        elif callable(self.trim_intermediate_steps):
+            return self.trim_intermediate_steps(intermediate_steps)
+        else:
+            return intermediate_steps