wanwu/deer-flow
1
0
Fork 0
mirror of https://gitee.com/wanwujie/deer-flow synced 2026-04-25 23:14:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
Xun
2026-01-23 08:31:48 +08:00
committed by GitHub
parent 65cdc182d3
commit c0849af37e
5 changed files with 120 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
tests/unit/tools/test_search.py
View File

@@ -17,7 +17,7 @@ class TestGetWebSearchTool:
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 @@ class TestGetWebSearchTool:
"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 @@ class TestGetWebSearchTool:
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 @@ class TestGetWebSearchTool:
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 @@ class TestGetWebSearchTool:
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 @@ class TestGetWebSearchTool:
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 @@ class TestGetWebSearchTool:
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 @@ class TestGetWebSearchTool:
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 @@ class TestGetWebSearchTool:
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 @@ class TestGetWebSearchTool:
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 @@ class TestGetWebSearchTool:
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

42
tests/unit/utils/test_context_manager.py
View File

@@ -85,8 +85,8 @@ class TestContextManager:
# 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)
@@ -94,15 +94,26 @@ class TestContextManager:
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 @@ class TestContextManager:
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