430 lines
19 KiB
Python
430 lines
19 KiB
Python
"""Generated large-payload regressions; no credentials or external model needed."""
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from langchain.agents.middleware.types import ModelRequest, ModelResponse
|
|
from langchain_core.language_models.fake_chat_models import FakeListChatModel
|
|
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
|
from langgraph.errors import GraphBubbleUp
|
|
|
|
from deerflow.agents.middlewares.context_window_guard_middleware import ContextWindowGuardMiddleware, estimate_model_request_tokens
|
|
from deerflow.agents.middlewares.llm_error_handling_middleware import LLMErrorHandlingMiddleware, is_llm_error_message
|
|
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware
|
|
from deerflow.models.context_recovery import CONTEXT_RECOVERY_MESSAGE, ContextRecoveryExhausted, is_context_length_error
|
|
|
|
|
|
def request_for(text="问题", **kwargs):
|
|
return ModelRequest(model=FakeListChatModel(responses=["ok"]), messages=[HumanMessage(content=text)], **kwargs)
|
|
|
|
|
|
def run_guard(guard, request, handler, asynchronous):
|
|
if not asynchronous:
|
|
return guard.wrap_model_call(request, handler)
|
|
|
|
async def call(req):
|
|
return handler(req)
|
|
|
|
return asyncio.run(guard.awrap_model_call(request, call))
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
def test_unconfigured_window_does_not_trim_even_large_input(asynchronous):
|
|
request = request_for("中文材料" * 100_000)
|
|
seen = []
|
|
run_guard(ContextWindowGuardMiddleware(), request, lambda r: seen.append(r) or AIMessage(content="正常"), asynchronous)
|
|
assert seen == [request]
|
|
assert seen[0] is request
|
|
|
|
|
|
def test_unconfigured_upload_turn_keeps_manifest_and_file_tools_before_first_call():
|
|
upload_manifest = (
|
|
'<uploaded_files>{"format":"deerflow-upload-context-v2",'
|
|
'"current":[{"filename":"report.pdf","path":"/mnt/user-data/uploads/report.pdf.md"}]}'
|
|
"</uploaded_files>"
|
|
)
|
|
tool_names = ["read_file", "grep", "view_image", *[f"other_tool_{index}" for index in range(17)]]
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": name,
|
|
"description": f"Tool {name}",
|
|
"parameters": {"type": "object", "properties": {}},
|
|
},
|
|
}
|
|
for name in tool_names
|
|
]
|
|
request = request_for(f"{upload_manifest}\n请读取并总结文件。", tools=tools).override(
|
|
system_message=SystemMessage(content="S" * 66_000)
|
|
)
|
|
|
|
guarded = ContextWindowGuardMiddleware().guard_request(request)
|
|
|
|
assert guarded is request
|
|
assert guarded.messages[-1].content == request.messages[-1].content
|
|
assert upload_manifest in guarded.messages[-1].content
|
|
assert [tool["function"]["name"] for tool in guarded.tools] == tool_names
|
|
|
|
|
|
def test_bounded_upload_turn_trims_system_before_manifest_or_file_tools():
|
|
limit = 32_768
|
|
upload_manifest = (
|
|
'<uploaded_files>{"format":"deerflow-upload-context-v2",'
|
|
'"current":[{"filename":"report.pdf","path":"/mnt/user-data/uploads/report.pdf.md"}]}'
|
|
"</uploaded_files>"
|
|
)
|
|
file_tool_names = {"read_file", "grep", "view_image"}
|
|
tool_names = [*file_tool_names, *[f"other_tool_{index}" for index in range(17)]]
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": name,
|
|
"description": (f"Tool {name}. " * 10),
|
|
"parameters": {"type": "object", "properties": {}},
|
|
},
|
|
}
|
|
for name in tool_names
|
|
]
|
|
request = request_for(f"{upload_manifest}\n请读取并总结文件。", tools=tools).override(
|
|
system_message=SystemMessage(content="S" * 66_000)
|
|
)
|
|
|
|
guarded = ContextWindowGuardMiddleware().guard_request(request, limit)
|
|
|
|
assert estimate_model_request_tokens(guarded) <= limit
|
|
assert guarded.messages[-1].content == request.messages[-1].content
|
|
assert upload_manifest in guarded.messages[-1].content
|
|
kept_tool_names = {tool["function"]["name"] for tool in guarded.tools}
|
|
assert file_tool_names <= kept_tool_names
|
|
assert len(str(guarded.system_message.content)) < len(str(request.system_message.content))
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
@pytest.mark.parametrize("unit", ["English research. ", "中文长材料：", "😀<>&\n"])
|
|
@pytest.mark.parametrize("size", [100, 10_000, 100_000])
|
|
def test_generated_context_pressure_recovers(asynchronous, unit, size):
|
|
original = request_for("问题开始：" + unit * size + "问题结束：请概括。")
|
|
snapshot = original.messages[0].model_dump()
|
|
sizes = []
|
|
|
|
def provider(request):
|
|
estimate = estimate_model_request_tokens(request)
|
|
sizes.append(estimate)
|
|
assert request.messages and request.messages[-1].content
|
|
if len(sizes) == 1 or estimate > 1000:
|
|
raise RuntimeError("context_length_exceeded")
|
|
return ModelResponse(result=[AIMessage(content="这是正常回答。")])
|
|
|
|
response = run_guard(ContextWindowGuardMiddleware(), original, provider, asynchronous)
|
|
assert response.result[0].content == "这是正常回答。"
|
|
assert is_llm_error_message(response.result[0]) is None
|
|
assert 2 <= len(sizes) <= 5
|
|
assert all(b < a for a, b in zip(sizes, sizes[1:]))
|
|
assert original.messages[0].model_dump() == snapshot
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
def test_explicit_limit_applies_before_first_call(asynchronous):
|
|
sizes = []
|
|
run_guard(ContextWindowGuardMiddleware(max_input_tokens=2000), request_for("汉" * 50_000), lambda r: sizes.append(estimate_model_request_tokens(r)) or AIMessage(content="ok"), asynchronous)
|
|
assert len(sizes) == 1 and sizes[0] <= 2000
|
|
|
|
|
|
def test_profile_limit_is_used_without_inventing_default():
|
|
request = request_for("字" * 10000)
|
|
request.model.profile = {"max_input_tokens": 2000}
|
|
assert estimate_model_request_tokens(ContextWindowGuardMiddleware().guard_request(request)) <= 2000
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
def test_persistent_context_rejection_returns_clarification_not_error(asynchronous):
|
|
seen = []
|
|
|
|
def fail(request):
|
|
seen.append(request)
|
|
raise RuntimeError("maximum context length is 1")
|
|
|
|
response = run_guard(ContextWindowGuardMiddleware(), request_for("长" * 100000), fail, asynchronous)
|
|
assert response.content == CONTEXT_RECOVERY_MESSAGE
|
|
assert is_llm_error_message(response) is None
|
|
assert len(seen) <= 5
|
|
assert all(r.messages and r.messages[0].content for r in seen)
|
|
|
|
|
|
@pytest.mark.parametrize("limit", [1, 16, 32])
|
|
def test_impossibly_small_config_never_sends_empty_request(limit):
|
|
seen = []
|
|
response = ContextWindowGuardMiddleware(max_input_tokens=limit).wrap_model_call(request_for("问题" * 1000), lambda r: seen.append(r))
|
|
assert not seen
|
|
assert response.content == CONTEXT_RECOVERY_MESSAGE
|
|
|
|
|
|
def test_required_tools_and_response_schema_are_not_silently_removed():
|
|
tool = {"type": "function", "function": {"name": "save", "description": "保存", "parameters": {"type": "object"}}}
|
|
schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
|
|
request = request_for("长" * 10000, tools=[tool], tool_choice="required", response_format=schema)
|
|
guarded = ContextWindowGuardMiddleware().guard_request(request, 2000)
|
|
assert guarded.tools == [tool]
|
|
assert guarded.tool_choice == "required"
|
|
assert guarded.response_format == schema
|
|
with pytest.raises(ContextRecoveryExhausted):
|
|
ContextWindowGuardMiddleware(max_input_tokens=1).wrap_model_call(request, lambda r: None)
|
|
|
|
|
|
def test_tool_pair_and_hidden_latest_human_are_preserved_correctly():
|
|
request = request_for("分析检索结果")
|
|
request = request.override(
|
|
messages=[
|
|
HumanMessage(content="旧问题" * 1000),
|
|
AIMessage(content="旧回答" * 1000),
|
|
*request.messages,
|
|
AIMessage(content="", tool_calls=[{"id": "c1", "name": "search", "args": {}}]),
|
|
ToolMessage(content="检索结果" * 10000, tool_call_id="c1"),
|
|
HumanMessage(content="隐藏材料" * 10000, additional_kwargs={"hide_from_ui": True}),
|
|
]
|
|
)
|
|
before = [m.model_dump() for m in request.messages]
|
|
result = ContextWindowGuardMiddleware().guard_request(request, 1800)
|
|
assert result.messages[0].content == "分析检索结果"
|
|
assert [m.type for m in result.messages] == ["human", "ai", "tool"]
|
|
assert result.messages[1].tool_calls[0]["id"] == result.messages[2].tool_call_id
|
|
assert [m.model_dump() for m in request.messages] == before
|
|
|
|
|
|
def test_output_reservation_is_lowered_only_after_context_error():
|
|
request = request_for("问题" * 1000, model_settings={"max_tokens": 100000})
|
|
seen = []
|
|
|
|
def provider(req):
|
|
seen.append(req.model_settings["max_tokens"])
|
|
if seen[-1] > 3000:
|
|
raise RuntimeError("maximum context length exceeded")
|
|
return AIMessage(content="ok")
|
|
|
|
result = ContextWindowGuardMiddleware().wrap_model_call(request, provider)
|
|
assert result.content == "ok"
|
|
assert seen[0] == 100000 and seen[-1] <= 3000
|
|
assert request.model_settings["max_tokens"] == 100000
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
@pytest.mark.parametrize("exception", [ValueError("bad parameter"), GraphBubbleUp(), asyncio.CancelledError()])
|
|
def test_other_errors_and_control_flow_propagate(asynchronous, exception):
|
|
calls = []
|
|
|
|
def provider(req):
|
|
calls.append(req)
|
|
raise exception
|
|
|
|
with pytest.raises(type(exception)):
|
|
run_guard(ContextWindowGuardMiddleware(), request_for(), provider, asynchronous)
|
|
assert len(calls) == 1
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
def test_outer_error_handler_sees_success_after_inner_recovery(asynchronous):
|
|
errors = LLMErrorHandlingMiddleware(app_config=SimpleNamespace(circuit_breaker=SimpleNamespace(failure_threshold=3, recovery_timeout_sec=30)))
|
|
guard = ContextWindowGuardMiddleware()
|
|
calls = []
|
|
|
|
def provider(request):
|
|
calls.append(request)
|
|
if len(calls) == 1:
|
|
raise RuntimeError("maximum context length exceeded")
|
|
return AIMessage(content="正常回答")
|
|
|
|
request = request_for("汉字" * 10000)
|
|
if asynchronous:
|
|
|
|
async def invoke(req):
|
|
return provider(req)
|
|
|
|
async def inner(req):
|
|
return await guard.awrap_model_call(req, invoke)
|
|
|
|
response = asyncio.run(errors.awrap_model_call(request, inner))
|
|
else:
|
|
response = errors.wrap_model_call(request, lambda req: guard.wrap_model_call(req, provider))
|
|
assert response.content == "正常回答"
|
|
assert is_llm_error_message(response) is None
|
|
assert len(calls) == 2 and errors._circuit_failure_count == 0
|
|
|
|
|
|
class RejectingSummaryModel(FakeListChatModel):
|
|
sizes: list[int] = []
|
|
fail_always: bool = False
|
|
generic_error: bool = False
|
|
|
|
def _call(self, messages, stop=None, run_manager=None, **kwargs):
|
|
size = sum(len(str(m.content).encode("utf-8")) for m in messages)
|
|
self.sizes.append(size)
|
|
if self.generic_error:
|
|
raise RuntimeError("summary unavailable")
|
|
if self.fail_always or len(self.sizes) == 1 or size > 1200:
|
|
raise RuntimeError("context_length_exceeded")
|
|
return "摘要成功"
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
@pytest.mark.parametrize("trim", [None, 4000])
|
|
def test_summary_rendered_prompt_shrinks_until_success(asynchronous, trim):
|
|
model = RejectingSummaryModel(responses=["unused"])
|
|
middleware = DeerFlowSummarizationMiddleware(model, trim_tokens_to_summarize=trim)
|
|
messages = [HumanMessage(content="<&中文材料>" * 30000)]
|
|
result = asyncio.run(middleware._acreate_summary(messages)) if asynchronous else middleware._create_summary(messages)
|
|
assert result == "摘要成功"
|
|
assert 2 <= len(model.sizes) <= 5
|
|
assert all(b < a for a, b in zip(model.sizes, model.sizes[1:]))
|
|
if trim:
|
|
assert model.sizes[0] <= trim
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
@pytest.mark.parametrize("generic", [False, True])
|
|
def test_summary_failure_does_not_abort_main_answer_or_cache_error(asynchronous, generic):
|
|
model = RejectingSummaryModel(responses=["unused"], fail_always=True, generic_error=generic)
|
|
middleware = DeerFlowSummarizationMiddleware(model, trigger=("messages", 3), keep=("messages", 1))
|
|
request = request_for().override(
|
|
messages=[
|
|
HumanMessage(content="旧材料" * 10000, id="u1"),
|
|
AIMessage(content="旧回答", id="a1"),
|
|
HumanMessage(content="请继续分析", id="u2"),
|
|
]
|
|
)
|
|
seen = []
|
|
|
|
def main(req):
|
|
seen.append(req)
|
|
return ModelResponse(result=[AIMessage(content="正常主回答")])
|
|
|
|
response = run_guard(middleware, request, main, asynchronous)
|
|
assert response.result[0].content == "正常主回答"
|
|
assert seen[0].messages[-1].content == "请继续分析"
|
|
assert middleware._pending_state_cache_update is None
|
|
assert all("summary unavailable" not in str(m.content) for m in seen[0].messages)
|
|
|
|
|
|
def test_structured_context_code_detection():
|
|
error = RuntimeError("bad request")
|
|
error.body = {"code": "context_length_exceeded"}
|
|
assert is_context_length_error(error)
|
|
assert not is_context_length_error(RuntimeError("invalid request schema"))
|
|
|
|
|
|
def test_minimax_live_context_error_2013_is_recovered():
|
|
# Captured from the real configured MiniMax endpoint with 40,000 synthetic rows.
|
|
exc = RuntimeError("invalid params, context window exceeds limit (2013)")
|
|
assert is_context_length_error(exc)
|
|
assert not is_context_length_error(RuntimeError("invalid params (2013)"))
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
@pytest.mark.parametrize("content", ["", " ", "<think>仅有思考内容</think>", "<think>未结束的思考"])
|
|
def test_recovery_does_not_return_empty_or_thinking_only_answer(asynchronous, content):
|
|
calls = []
|
|
|
|
def provider(request):
|
|
calls.append(request)
|
|
if len(calls) == 1:
|
|
raise RuntimeError("context_length_exceeded")
|
|
return ModelResponse(result=[AIMessage(content=content)])
|
|
|
|
response = run_guard(ContextWindowGuardMiddleware(), request_for("长内容" * 10000), provider, asynchronous)
|
|
assert response.content == CONTEXT_RECOVERY_MESSAGE
|
|
assert is_llm_error_message(response) is None
|
|
assert len(calls) == 2
|
|
|
|
|
|
def test_small_output_budget_is_not_halved_on_context_rejection():
|
|
request = request_for("长内容" * 10000, model_settings={"max_tokens": 1024})
|
|
smaller = ContextWindowGuardMiddleware()._retry_request(request, 1)
|
|
assert smaller.model_settings["max_tokens"] == 1024
|
|
|
|
|
|
def test_emergency_does_not_forget_completed_action_while_forcing_tools():
|
|
request = request_for("保存结果", tools=[{"type": "function", "function": {"name": "save", "parameters": {"type": "object"}}}], tool_choice="required")
|
|
request = request.override(
|
|
messages=[
|
|
*request.messages,
|
|
AIMessage(content="", tool_calls=[{"id": "save1", "name": "save", "args": {}}]),
|
|
ToolMessage(content="已保存成功", tool_call_id="save1"),
|
|
]
|
|
)
|
|
with pytest.raises(ContextRecoveryExhausted):
|
|
ContextWindowGuardMiddleware()._minimal_request(request, 512)
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
def test_retry_deadline_stops_additional_attempts(asynchronous):
|
|
guard = ContextWindowGuardMiddleware()
|
|
guard.context_retry_timeout_sec = 0
|
|
seen = []
|
|
|
|
def provider(request):
|
|
seen.append(request)
|
|
raise RuntimeError("context_length_exceeded")
|
|
|
|
response = run_guard(guard, request_for("长内容" * 1000), provider, asynchronous)
|
|
assert response.content == CONTEXT_RECOVERY_MESSAGE
|
|
assert len(seen) == 1
|
|
|
|
|
|
def test_summary_bound_model_uses_explicit_smaller_window():
|
|
model = RejectingSummaryModel(responses=["unused"])
|
|
model.profile = {"max_input_tokens": 700}
|
|
middleware = DeerFlowSummarizationMiddleware(model.with_config(tags=["middleware:summarize"]), trim_tokens_to_summarize=4000)
|
|
result = middleware._create_summary([HumanMessage(content="长材料" * 10000)])
|
|
assert result == "摘要成功"
|
|
assert model.sizes[0] <= 700
|
|
|
|
|
|
def test_generated_mixed_payload_stress_preserves_input_and_protocol():
|
|
import random
|
|
|
|
rng = random.Random(20260911)
|
|
guard = ContextWindowGuardMiddleware()
|
|
for case in range(50):
|
|
request = request_for("请总结这段内容：" + "材料" * rng.randint(1, 20000))
|
|
request = request.override(
|
|
system_message=SystemMessage(content="按用户要求分析。" * rng.randint(1, 1000)),
|
|
messages=[HumanMessage(content="历史" * rng.randint(1, 3000)), AIMessage(content="答案" * rng.randint(1, 3000)), *request.messages],
|
|
)
|
|
if case % 2:
|
|
request = request.override(
|
|
messages=[
|
|
*request.messages,
|
|
AIMessage(content="", additional_kwargs={"reasoning_content": "推理" * rng.randint(1, 5000)}, tool_calls=[{"id": "t", "name": "search", "args": {"query": "q"}}]),
|
|
ToolMessage(content="检索" * rng.randint(1, 20000), tool_call_id="t"),
|
|
]
|
|
)
|
|
if case % 3 == 0:
|
|
request = request.override(messages=[*request.messages, HumanMessage(content="隐藏材料" * 1000, additional_kwargs={"hide_from_ui": True})])
|
|
before = [m.model_dump() for m in request.messages]
|
|
limit = rng.choice([512, 1024, 2048, 4096, 8192])
|
|
result = guard.guard_request(request, limit)
|
|
assert estimate_model_request_tokens(result) <= limit
|
|
assert any(isinstance(m, HumanMessage) and m.content for m in result.messages)
|
|
for index, message in enumerate(result.messages):
|
|
if isinstance(message, ToolMessage):
|
|
assert index and isinstance(result.messages[index - 1], AIMessage)
|
|
assert message.tool_call_id in {call["id"] for call in result.messages[index - 1].tool_calls}
|
|
assert [m.model_dump() for m in request.messages] == before
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
def test_context_failure_releases_circuit_probe_without_counting_outage(asynchronous):
|
|
errors = LLMErrorHandlingMiddleware(app_config=SimpleNamespace(circuit_breaker=SimpleNamespace(failure_threshold=3, recovery_timeout_sec=30)))
|
|
errors._circuit_state = "half_open"
|
|
|
|
def fail(request):
|
|
raise ContextRecoveryExhausted()
|
|
|
|
result = run_guard(errors, request_for(), fail, asynchronous)
|
|
assert is_llm_error_message(result) == "context_length"
|
|
assert not errors._circuit_probe_in_flight
|
|
assert errors._circuit_failure_count == 0
|
|
assert "上下文" not in result.content
|