feat(context): decrease token in web_search AIMessage (bytedance#827)

This PR addresses token limit issues when web_search is enabled with include_raw_content by implementing a two-pronged approach: changing the default behavior to exclude raw content and adding compression logic for when raw content is included.

Author: xunliu

src/tools/crawl.py

@@ -37,13 +37,13 @@ def crawl_tool(
             "error": "PDF files cannot be crawled directly. Please download and view the PDF manually.",
             "crawled_content": None,
             "is_pdf": True
-        })
+        }, ensure_ascii=False)
         return pdf_message
 
     try:
         crawler = Crawler()
         article = crawler.crawl(url)
-        return json.dumps({"url": url, "crawled_content": article.to_markdown()[:1000]})
+        return json.dumps({"url": url, "crawled_content": article.to_markdown()[:1000]}, ensure_ascii=False)
     except BaseException as e:
         error_msg = f"Failed to crawl. Error: {repr(e)}"
         logger.error(error_msg)

src/tools/search.py

@@ -57,7 +57,7 @@ def get_web_search_tool(max_search_results: int):
         exclude_domains: Optional[List[str]] = search_config.get("exclude_domains", [])
         include_answer: bool = search_config.get("include_answer", False)
         search_depth: str = search_config.get("search_depth", "advanced")
-        include_raw_content: bool = search_config.get("include_raw_content", True)
+        include_raw_content: bool = search_config.get("include_raw_content", False)
         include_images: bool = search_config.get("include_images", True)
         include_image_descriptions: bool = include_images and search_config.get(
             "include_image_descriptions", True

src/utils/context_manager.py

@@ -188,77 +188,86 @@ def compress_messages(self, state: dict, runtime: Runtime | None = None) -> dict
 
     def _compress_messages(self, messages: List[BaseMessage]) -> List[BaseMessage]:
         """
-        Compress compressible messages
-
+        Compress messages to fit within token limit through two strategies:
+        1. First, compress web_search ToolMessage raw_content by truncating to 1024 chars
+        2. If still over limit, drop oldest messages while preserving prefix messages and system messages
+        
         Args:
             messages: List of messages to compress
-
-        Returns:
-            Compressed message list
-        """
-
-        available_token = self.token_limit
-        prefix_messages = []
-
-        # 1. Preserve head messages of specified length to retain system prompts and user input
-        for i in range(min(self.preserve_prefix_message_count, len(messages))):
-            cur_token_cnt = self._count_message_tokens(messages[i])
-            if available_token > 0 and available_token >= cur_token_cnt:
-                prefix_messages.append(messages[i])
-                available_token -= cur_token_cnt
-            elif available_token > 0:
-                # Truncate content to fit available tokens
-                truncated_message = self._truncate_message_content(
-                    messages[i], available_token
-                )
-                prefix_messages.append(truncated_message)
-                return prefix_messages
-            else:
-                break
-
-        # 2. Compress subsequent messages from the tail, some messages may be discarded
-        messages = messages[len(prefix_messages) :]
-        suffix_messages = []
-        for i in range(len(messages) - 1, -1, -1):
-            cur_token_cnt = self._count_message_tokens(messages[i])
-
-            if cur_token_cnt > 0 and available_token >= cur_token_cnt:
-                suffix_messages = [messages[i]] + suffix_messages
-                available_token -= cur_token_cnt
-            elif available_token > 0:
-                # Truncate content to fit available tokens
-                truncated_message = self._truncate_message_content(
-                    messages[i], available_token
-                )
-                suffix_messages = [truncated_message] + suffix_messages
-                return prefix_messages + suffix_messages
-            else:
-                break
-
-        return prefix_messages + suffix_messages
-
-    def _truncate_message_content(
-        self, message: BaseMessage, max_tokens: int
-    ) -> BaseMessage:
-        """
-        Truncate message content while preserving all other attributes by copying the original message
-        and only modifying its content attribute.
-
-        Args:
-            message: The message to truncate
-            max_tokens: Maximum number of tokens to keep
-
         Returns:
-            New message instance with truncated content
+            List of messages with compressed content and/or dropped messages
         """
-
-        # Create a deep copy of the original message to preserve all attributes
-        truncated_message = copy.deepcopy(message)
-
-        # Truncate only the content attribute
-        truncated_message.content = message.content[:max_tokens]
-
-        return truncated_message
+        # Create a deep copy to avoid mutating original messages
+        compressed = copy.deepcopy(messages)
+        
+        # Step 1: Compress raw_content in web_search ToolMessages
+        for msg in compressed:
+            # Only compress ToolMessage with name 'web_search'
+            if isinstance(msg, ToolMessage) and getattr(msg, "name", None) == "web_search":
+                try:
+                    # Determine content type and check if compression is needed
+                    if isinstance(msg.content, str):
+                        # Early exit if content is small enough (avoid JSON parsing overhead)
+                        # A heuristic: if string is less than 2KB, raw_content likely doesn't need truncation
+                        if len(msg.content) < 2048:
+                            continue
+                        
+                        try:
+                            content_data = json.loads(msg.content)
+                        except json.JSONDecodeError as e:
+                            logger.error(f"Failed to parse JSON content in web_search ToolMessage: {e}. Content: {msg.content[:200]}")
+                            continue
+                    elif isinstance(msg.content, list):
+                        content_data = copy.deepcopy(msg.content)
+                    else:
+                        continue
+
+                    # Compress raw_content in the content (item by item processing)
+                    # Track if any modifications were made
+                    modified = False
+                    if isinstance(content_data, list):
+                        for item in content_data:
+                            if isinstance(item, dict) and "raw_content" in item:
+                                raw_content = item.get("raw_content")
+                                if raw_content and isinstance(raw_content, str) and len(raw_content) > 1024:
+                                    item["raw_content"] = raw_content[:1024]
+                                    modified = True
+                        
+                        # Update message content with modified data only if changes were made
+                        if modified:
+                            msg.content = json.dumps(content_data, ensure_ascii=False)
+                except Exception as e:
+                    logger.error(f"Unexpected error during message compression: {e}")
+                    continue
+
+        # Step 2: If still over limit after raw_content compression, drop oldest messages
+        # while preserving prefix messages (e.g., system message) and recent messages
+        if self.is_over_limit(compressed):
+            # Identify messages to preserve at the beginning
+            preserved_count = self.preserve_prefix_message_count
+            preserved_messages = compressed[:preserved_count]
+            remaining_messages = compressed[preserved_count:]
+            
+            # Drop messages from the middle, keeping the most recent ones
+            result_messages = preserved_messages
+            for msg in reversed(remaining_messages):
+                result_messages.insert(len(preserved_messages), msg)
+                if not self.is_over_limit(result_messages):
+                    break
+            
+            compressed = result_messages
+
+        # Step 3: Verify that compression was successful and log warning if needed
+        if self.is_over_limit(compressed):
+            current_tokens = self.count_tokens(compressed)
+            logger.warning(
+                f"Message compression failed to bring tokens below limit: "
+                f"{current_tokens} > {self.token_limit} tokens. "
+                f"Total messages: {len(compressed)}. "
+                f"Consider increasing token_limit or preserve_prefix_message_count."
+            )
+
+        return compressed
 
     def _create_summary_message(self, messages: List[BaseMessage]) -> BaseMessage:
         """

tests/unit/tools/test_search.py

@@ -17,7 +17,7 @@ def test_get_web_search_tool_tavily(self):
         tool = get_web_search_tool(max_search_results=5)
         assert tool.name == "web_search"
         assert tool.max_results == 5
-        assert tool.include_raw_content is True
+        assert tool.include_raw_content is False
         assert tool.include_images is True
         assert tool.include_image_descriptions is True
         assert tool.include_answer is False
@@ -79,7 +79,7 @@ def test_get_web_search_tool_tavily_with_custom_config(self, mock_config):
             "SEARCH_ENGINE": {
                 "include_answer": True,
                 "search_depth": "basic",
-                "include_raw_content": False,
+                "include_raw_content": True,
                 "include_images": False,
                 "include_image_descriptions": True,
                 "include_domains": ["example.com"],
@@ -91,7 +91,7 @@ def test_get_web_search_tool_tavily_with_custom_config(self, mock_config):
         assert tool.max_results == 5
         assert tool.include_answer is True
         assert tool.search_depth == "basic"
-        assert tool.include_raw_content is False
+        assert tool.include_raw_content is True
         assert tool.include_images is False
         # include_image_descriptions should be False because include_images is False
         assert tool.include_image_descriptions is False
@@ -108,7 +108,7 @@ def test_get_web_search_tool_tavily_with_empty_config(self, mock_config):
         assert tool.max_results == 10
         assert tool.include_answer is False
         assert tool.search_depth == "advanced"
-        assert tool.include_raw_content is True
+        assert tool.include_raw_content is False
         assert tool.include_images is True
         assert tool.include_image_descriptions is True
         assert tool.include_domains == []
@@ -143,7 +143,7 @@ def test_get_web_search_tool_tavily_partial_config(self, mock_config):
         tool = get_web_search_tool(max_search_results=3)
         assert tool.include_answer is True
         assert tool.search_depth == "advanced"  # default
-        assert tool.include_raw_content is True  # default
+        assert tool.include_raw_content is False  # default
         assert tool.include_domains == ["trusted.com"]
         assert tool.exclude_domains == []  # default
 
@@ -157,7 +157,7 @@ def test_get_web_search_tool_tavily_with_no_config_file(self, mock_config):
         assert tool.max_results == 5
         assert tool.include_answer is False
         assert tool.search_depth == "advanced"
-        assert tool.include_raw_content is True
+        assert tool.include_raw_content is False
         assert tool.include_images is True
 
     @patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
@@ -184,7 +184,7 @@ def test_tavily_with_no_search_engine_section(self, mock_config):
         assert tool.max_results == 5
         assert tool.include_answer is False
         assert tool.search_depth == "advanced"
-        assert tool.include_raw_content is True
+        assert tool.include_raw_content is False
         assert tool.include_images is True
         assert tool.include_domains == []
         assert tool.exclude_domains == []
@@ -199,7 +199,7 @@ def test_tavily_with_completely_empty_config(self, mock_config):
         assert tool.max_results == 5
         assert tool.include_answer is False
         assert tool.search_depth == "advanced"
-        assert tool.include_raw_content is True
+        assert tool.include_raw_content is False
         assert tool.include_images is True
 
     @patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
@@ -210,7 +210,7 @@ def test_tavily_with_only_include_answer_param(self, mock_config):
         tool = get_web_search_tool(max_search_results=5)
         assert tool.include_answer is True
         assert tool.search_depth == "advanced"
-        assert tool.include_raw_content is True
+        assert tool.include_raw_content is False
         assert tool.include_images is True
 
     @patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
@@ -221,7 +221,7 @@ def test_tavily_with_only_search_depth_param(self, mock_config):
         tool = get_web_search_tool(max_search_results=5)
         assert tool.search_depth == "basic"
         assert tool.include_answer is False
-        assert tool.include_raw_content is True
+        assert tool.include_raw_content is False
         assert tool.include_images is True
 
     @patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
@@ -286,6 +286,6 @@ def test_tavily_all_parameters_optional_mix(self, mock_config):
             tool.include_image_descriptions is False
         )  # should be False since include_images is False
         assert tool.search_depth == "advanced"  # default
-        assert tool.include_raw_content is True  # default
+        assert tool.include_raw_content is False  # default
         assert tool.include_domains == []  # default
         assert tool.exclude_domains == []  # default

tests/unit/utils/test_context_manager.py

@@ -85,24 +85,35 @@ def test_compress_messages_when_not_over_limit(self):
         # Should return the same messages when not over limit
         assert len(compressed["messages"]) == len(messages)
 
-    def test_compress_messages_with_system_message(self):
-        """Test compress_messages preserves system message"""
+    def test_compress_messages_with_tool_message(self):
+        """Test compress_messages preserves system message and compresses raw_content"""
         # Create a context manager with limited token capacity
         limited_cm = ContextManager(token_limit=200)
 
         messages = [
             SystemMessage(content="You are a helpful assistant."),
             HumanMessage(content="Hello"),
             AIMessage(content="Hi there!"),
-            HumanMessage(
-                content="Can you tell me a very long story that would exceed token limits? "
-                * 100
-            ),
+            ToolMessage(
+                name="web_search",
+                content='[{"title": "Test Result", "url": "https://example.com", "raw_content": "' + ("This is a test content that should be compressed if it exceeds 1024 characters. " * 2000) + '"}]',
+                tool_call_id="test_search",
+            )
         ]
 
         compressed = limited_cm.compress_messages({"messages": messages})
         # Should preserve system message and some recent messages
-        assert len(compressed["messages"]) == 1
+        assert len(compressed["messages"]) == 4
+        
+        # Verify raw_content was compressed to 1024 characters
+        import json
+        for msg in compressed["messages"]:
+            if isinstance(msg, ToolMessage) and getattr(msg, "name", None) == "web_search":
+                content_data = json.loads(msg.content)
+                if isinstance(content_data, list):
+                    for item in content_data:
+                        if isinstance(item, dict) and "raw_content" in item:
+                            assert len(item["raw_content"]) == 1024
 
     def test_compress_messages_with_preserve_prefix_message(self):
         """Test compress_messages when no system message is present"""
@@ -201,9 +212,24 @@ def test_compress_messages_with_runtime_when_over_limit(self):
             HumanMessage(
                 content="Can you tell me a very long story that would exceed token limits? " * 100
             ),
+            ToolMessage(
+                name="web_search",
+                content='[{"title": "Test Result", "url": "https://example.com", "raw_content": "' + ("This is a test content that should be compressed if it exceeds 1024 characters. " * 2000) + '"}]',
+                tool_call_id="test_search",
+            )
         ]
         compressed = limited_cm.compress_messages({"messages": messages}, runtime=object())
         assert isinstance(compressed, dict)
         assert "messages" in compressed
         # Should preserve only what fits; with this setup we expect heavy compression
-        assert len(compressed["messages"]) == 1
+        assert len(compressed["messages"]) == 5
+        
+        # Verify raw_content was compressed to 1024 characters
+        import json
+        for msg in compressed["messages"]:
+            if isinstance(msg, ToolMessage) and getattr(msg, "name", None) == "web_search":
+                content_data = json.loads(msg.content)
+                if isinstance(content_data, list):
+                    for item in content_data:
+                        if isinstance(item, dict) and "raw_content" in item:
+                            assert len(item["raw_content"]) == 1024