add parameter to limit the size of icl context

Author: BartekCupial

balrog/agents/__init__.py

@@ -49,7 +49,7 @@ def create_agent(self):
         elif self.config.agent.type == "custom":
             return CustomAgent(client_factory, prompt_builder)
         elif self.config.agent.type == "few_shot":
-            return FewShotAgent(client_factory, prompt_builder)
+            return FewShotAgent(client_factory, prompt_builder, self.config.agent.max_icl_history)
 
         else:
             raise ValueError(f"Unknown agent type: {self.config.agent}")

balrog/agents/few_shot.py

@@ -16,12 +16,13 @@ def __repr__(self):
 
 
 class FewShotAgent(BaseAgent):
-    def __init__(self, client_factory, prompt_builder):
+    def __init__(self, client_factory, prompt_builder, max_icl_history):
         """Initialize the FewShotAgent with a client and prompt builder."""
         super().__init__(client_factory, prompt_builder)
         self.client = client_factory()
         self.icl_episodes = []
         self.icl_events = []
+        self.max_icl_history = max_icl_history
         self.cached_icl = False
 
     def update_icl_observation(self, obs: dict):
@@ -76,8 +77,19 @@ def get_icl_prompt(self) -> List[Message]:
 
         # unroll the wrapped icl episodes messages
         icl_messages = [icl_instruction]
+        i = 0
         for icl_episode in self.icl_episodes:
-            icl_messages.extend(icl_episode)
+            episode_steps = len(icl_episode) - 2  # not count start and end messages
+            if i + episode_steps <= self.max_icl_history:
+                icl_messages.extend(icl_episode)
+                i += episode_steps
+            else:
+                icl_episode = icl_episode[: self.max_icl_history - i + 1] + [
+                    icl_episode[-1]
+                ]  # +1 for start message -1 for end message
+                icl_messages.extend(icl_episode)
+                i += len(icl_episode) - 2  # not count start and end messages
+                break
 
         end_demo_message = Message(
             role="user",

balrog/config/config.yaml

@@ -4,6 +4,7 @@ agent:
   max_history: 16        # Maximum number of previous turns to keep in the dialogue history
   max_image_history: 0   # Maximum number of images to keep in the history
   max_cot_history: 1     # Maximum number of chain-of-thought steps to keep in history (if using 'cot' type of agent)
+  max_icl_history: 200   # Maximum number of ICL steps to keep in history (if using 'few_shot' type of agent)
   cache_icl: False
 
 eval: