Added ability to add custom cache breakpoints for anthropic models

Author: amanjaiswal73892

src/agentlab/agents/tool_use_agent/multi_tool_agent.py

@@ -65,7 +65,7 @@ class StructuredDiscussion:
 
     def __init__(self, keep_last_n_obs=None):
         self.groups: list[MsgGroup] = []
-        self.keep_last_n_obs = keep_last_n_obs
+        self.keep_last_n_obs: int| None = keep_last_n_obs
 
     def append(self, message: MessageBuilder):
         """Append a message to the last group."""
@@ -87,11 +87,14 @@ def flatten(self) -> list[MessageBuilder]:
             print(
                 f"Processing group {i} ({group.name}), is_tail={is_tail}, len(greoup)={len(group.messages)}"
             )
+            # Include only summary if group not in last n groups.
             if not is_tail and group.summary is not None:
                 messages.append(group.summary)
             else:
                 messages.extend(group.messages)
-
+            # Mark all summarized messages for caching   
+            if i == len(self.groups) - keep_last_n_obs:
+                messages[i].mark_all_previous_msg_for_caching()
         return messages
 
     def set_last_summary(self, summary: MessageBuilder):
@@ -453,7 +456,8 @@ def get_action(self, obs: Any) -> float:
             messages=messages,
             tool_choice="any",
             cache_tool_definition=True,
-            cache_complete_prompt=True,
+            cache_complete_prompt=False,
+            use_cache_breakpoints=True,
         )
 
         action = response.action

src/agentlab/llm/response_api.py

@@ -107,6 +107,11 @@ def add_image_url(self, image_url: str) -> "MessageBuilder":
         self.content.append({"image": image_to_png_base64_url(image_url)})
         return self
 
+    def mark_all_previous_msg_for_caching(self):
+        """Insert a cache breakpoint in the message content."""
+        # This is a placeholder for future implementation.
+        raise NotImplementedError
+
 
 # TODO: Support parallel tool calls.
 
@@ -216,6 +221,10 @@ def transform_content(self, content: ContentItem) -> ContentItem:
         else:
             raise ValueError(f"Unsupported content type: {content}")
 
+    def mark_all_previous_msg_for_caching(self) -> List[Message]:
+        """Insert a cache breakpoint in the message content to mark all previous messages for caching."""
+        self._cache_breakpoint = True
+
 
 class OpenAIChatCompletionAPIMessageBuilder(MessageBuilder):
 
@@ -521,7 +530,10 @@ def _call_api(
         sys_msg, other_msgs = self.filter_system_messages(messages)
         sys_msg_text = "\n".join(c["text"] for m in sys_msg for c in m.content)
         for msg in other_msgs:
-            input.extend(msg.prepare_message() if isinstance(msg, MessageBuilder) else [msg])
+            temp = msg.prepare_message() if isinstance(msg, MessageBuilder) else [msg]
+            if kwargs.pop("use_cache_breakpoints", False):
+                temp = self.apply_cache_breakpoints(msg, temp)
+            input.extend(temp)
 
         api_params: Dict[str, Any] = {
             "model": self.model_name,
@@ -581,6 +593,16 @@ def _parse_response(self, response: dict) -> dict:
                 result.think += output.text
         return result
 
+    # def ensure_cache_conditions(self, msgs: List[Message]) -> bool:
+    #     """Ensure API specific cache conditions are met."""
+    #     assert sum(getattr(msg, "_cache_breakpoint", 0) for msg in msgs) <= 4, "Too many cache breakpoints in the message."
+
+    def apply_cache_breakpoints(self, msg: Message, prepared_msg: dict) -> List[Message]:
+        """Apply cache breakpoints to the messages."""
+        if getattr(msg, "_cache_breakpoint", False):
+                prepared_msg[-1]["content"][-1]["cache_control"] = {"type": "ephemeral"}
+        return prepared_msg
+
 
 def cua_response_to_text(action):
     """