wehub-resource-sync
279 lines
11 KiB
Python
279 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from langchain.agents import create_agent
|
|
from langchain_core.language_models import BaseChatModel
|
|
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
|
from langchain_core.outputs import ChatGeneration, ChatResult
|
|
from langchain_core.tools import tool
|
|
|
|
from deerflow.agents.middlewares.terminal_response_middleware import TerminalResponseMiddleware
|
|
from deerflow.runtime.runs.worker import _extract_llm_error_fallback_message
|
|
|
|
|
|
@tool
|
|
def lookup_status() -> str:
|
|
"""Return a deterministic tool result."""
|
|
return "tool completed"
|
|
|
|
|
|
class _PostToolResponseModel(BaseChatModel):
|
|
responses: list[str]
|
|
call_count: int = 0
|
|
observed_messages: list[list[Any]] = []
|
|
|
|
@property
|
|
def _llm_type(self) -> str:
|
|
return "post-tool-response"
|
|
|
|
def bind_tools(self, tools, **kwargs):
|
|
return self
|
|
|
|
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
|
self.observed_messages.append(list(messages))
|
|
self.call_count += 1
|
|
if self.call_count == 1:
|
|
message = AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "call-1", "name": "lookup_status", "args": {}}],
|
|
response_metadata={"finish_reason": "tool_calls"},
|
|
)
|
|
else:
|
|
message = AIMessage(
|
|
content=self.responses[self.call_count - 2],
|
|
response_metadata={"finish_reason": "stop"},
|
|
)
|
|
return ChatResult(generations=[ChatGeneration(message=message)])
|
|
|
|
async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs):
|
|
return self._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
|
|
|
|
|
|
class _PerRunRetryBudgetModel(BaseChatModel):
|
|
call_count: int = 0
|
|
observed_messages: list[list[Any]] = []
|
|
|
|
@property
|
|
def _llm_type(self) -> str:
|
|
return "per-run-retry-budget"
|
|
|
|
def bind_tools(self, tools, **kwargs):
|
|
return self
|
|
|
|
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
|
self.observed_messages.append(list(messages))
|
|
self.call_count += 1
|
|
if self.call_count == 1:
|
|
message = AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "call-budget-1", "name": "lookup_status", "args": {}}],
|
|
response_metadata={"finish_reason": "tool_calls"},
|
|
)
|
|
elif self.call_count == 2:
|
|
message = AIMessage(content="", response_metadata={"finish_reason": "stop"})
|
|
elif self.call_count == 3:
|
|
message = AIMessage(
|
|
content="I need one more status check.",
|
|
tool_calls=[{"id": "call-budget-2", "name": "lookup_status", "args": {}}],
|
|
response_metadata={"finish_reason": "tool_calls"},
|
|
)
|
|
else:
|
|
message = AIMessage(content="", response_metadata={"finish_reason": "stop"})
|
|
return ChatResult(generations=[ChatGeneration(message=message)])
|
|
|
|
async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs):
|
|
return self._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
|
|
|
|
|
|
def _agent(model: BaseChatModel):
|
|
return create_agent(
|
|
model=model,
|
|
tools=[lookup_status],
|
|
middleware=[TerminalResponseMiddleware()],
|
|
)
|
|
|
|
|
|
def _empty_terminal_messages(messages: list[Any]) -> list[AIMessage]:
|
|
return [message for message in messages if isinstance(message, AIMessage) and not message.tool_calls and not message.invalid_tool_calls and not str(message.content).strip()]
|
|
|
|
|
|
def test_retries_empty_post_tool_response_once_and_returns_model_answer():
|
|
model = _PostToolResponseModel(responses=["", "The tool completed successfully."])
|
|
|
|
result = _agent(model).invoke(
|
|
{"messages": [HumanMessage(content="Check the status")]},
|
|
context={"thread_id": "thread-1", "run_id": "run-1"},
|
|
)
|
|
|
|
assert model.call_count == 3
|
|
final = result["messages"][-1]
|
|
assert isinstance(final, AIMessage)
|
|
assert final.content == "The tool completed successfully."
|
|
assert _empty_terminal_messages(result["messages"]) == []
|
|
assert any(isinstance(message, HumanMessage) and message.name == "terminal_response_recovery" and message.additional_kwargs.get("hide_from_ui") is True for message in model.observed_messages[-1])
|
|
assert not any(isinstance(message, HumanMessage) and message.name == "terminal_response_recovery" for message in result["messages"])
|
|
|
|
|
|
def test_second_empty_post_tool_response_becomes_visible_error_fallback():
|
|
model = _PostToolResponseModel(responses=["", ""])
|
|
|
|
result = _agent(model).invoke(
|
|
{"messages": [HumanMessage(content="Check the status")]},
|
|
context={"thread_id": "thread-2", "run_id": "run-2"},
|
|
)
|
|
|
|
assert model.call_count == 3
|
|
final = result["messages"][-1]
|
|
assert isinstance(final, AIMessage)
|
|
assert "returned no final response" in str(final.content)
|
|
assert final.additional_kwargs["deerflow_error_fallback"] is True
|
|
assert _empty_terminal_messages(result["messages"]) == []
|
|
assert _extract_llm_error_fallback_message(result) == ("Model returned an empty terminal response after one retry")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_graph_retries_empty_post_tool_response_once():
|
|
model = _PostToolResponseModel(responses=["", "Recovered asynchronously."])
|
|
|
|
result = await _agent(model).ainvoke(
|
|
{"messages": [HumanMessage(content="Check the status")]},
|
|
context={"thread_id": "thread-async", "run_id": "run-async"},
|
|
)
|
|
|
|
assert model.call_count == 3
|
|
assert result["messages"][-1].content == "Recovered asynchronously."
|
|
assert _empty_terminal_messages(result["messages"]) == []
|
|
|
|
|
|
def test_graph_with_thread_id_only_keeps_recovery_state_across_model_loop():
|
|
model = _PostToolResponseModel(responses=["", "Recovered without a run id."])
|
|
|
|
result = _agent(model).invoke(
|
|
{"messages": [HumanMessage(content="Check the status")]},
|
|
context={"thread_id": "thread-only"},
|
|
)
|
|
|
|
assert model.call_count == 3
|
|
assert result["messages"][-1].content == "Recovered without a run id."
|
|
assert _empty_terminal_messages(result["messages"]) == []
|
|
|
|
|
|
def test_recovery_budget_is_once_per_run_even_when_retry_calls_another_tool():
|
|
model = _PerRunRetryBudgetModel()
|
|
|
|
result = _agent(model).invoke(
|
|
{"messages": [HumanMessage(content="Check the status twice")]},
|
|
context={"thread_id": "thread-budget", "run_id": "run-budget"},
|
|
)
|
|
|
|
assert model.call_count == 4
|
|
final = result["messages"][-1]
|
|
assert final.additional_kwargs["deerflow_error_fallback"] is True
|
|
assert _empty_terminal_messages(result["messages"]) == []
|
|
recovery_prompt_count = sum(1 for request_messages in model.observed_messages for message in request_messages if isinstance(message, HumanMessage) and message.name == "terminal_response_recovery")
|
|
assert recovery_prompt_count == 1
|
|
|
|
|
|
def test_empty_response_without_tool_result_is_not_retried():
|
|
middleware = TerminalResponseMiddleware()
|
|
message = AIMessage(content="", response_metadata={"finish_reason": "stop"})
|
|
state = {"messages": [HumanMessage(content="Hello"), message]}
|
|
runtime = type("RuntimeStub", (), {"context": {"thread_id": "thread-3", "run_id": "run-3"}})()
|
|
|
|
assert middleware.after_model(state, runtime) is None
|
|
|
|
|
|
def test_tool_call_intent_is_not_treated_as_empty_terminal_response():
|
|
middleware = TerminalResponseMiddleware()
|
|
message = AIMessage(
|
|
content="",
|
|
tool_calls=[{"id": "call-2", "name": "lookup_status", "args": {}}],
|
|
response_metadata={"finish_reason": "tool_calls"},
|
|
)
|
|
state = {"messages": [HumanMessage(content="Hello"), message]}
|
|
runtime = type("RuntimeStub", (), {"context": {"thread_id": "thread-4", "run_id": "run-4"}})()
|
|
|
|
assert middleware.after_model(state, runtime) is None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"message",
|
|
[
|
|
AIMessage(content="", invalid_tool_calls=[{"id": "bad-1", "name": "lookup_status", "args": "{"}]),
|
|
AIMessage(content="", additional_kwargs={"function_call": {"name": "lookup_status", "arguments": "{}"}}),
|
|
AIMessage(content="", response_metadata={"finish_reason": "function_call"}),
|
|
],
|
|
)
|
|
def test_invalid_or_legacy_tool_call_intent_is_not_treated_as_empty_terminal_response(message):
|
|
middleware = TerminalResponseMiddleware()
|
|
state = {"messages": [HumanMessage(content="Hello"), message]}
|
|
runtime = type("RuntimeStub", (), {"context": {"thread_id": "thread-5", "run_id": "run-5"}})()
|
|
|
|
assert middleware.after_model(state, runtime) is None
|
|
|
|
|
|
def test_after_agent_clears_retry_state_for_the_run():
|
|
middleware = TerminalResponseMiddleware()
|
|
runtime = type("RuntimeStub", (), {"context": {"thread_id": "thread-6", "run_id": "run-6"}})()
|
|
empty_after_tool = {
|
|
"messages": [
|
|
HumanMessage(content="Check the status"),
|
|
ToolMessage(content="tool completed", tool_call_id="call-6"),
|
|
AIMessage(content="", response_metadata={"finish_reason": "stop"}),
|
|
]
|
|
}
|
|
|
|
first = middleware.after_model(empty_after_tool, runtime)
|
|
assert first is not None and first["jump_to"] == "model"
|
|
middleware.after_agent(empty_after_tool, runtime)
|
|
second = middleware.after_model(empty_after_tool, runtime)
|
|
assert second is not None and second["jump_to"] == "model"
|
|
|
|
|
|
def test_before_agent_clears_same_run_state_for_resumed_invocation():
|
|
middleware = TerminalResponseMiddleware()
|
|
runtime = type("RuntimeStub", (), {"context": {"thread_id": "thread-7", "run_id": "run-7"}})()
|
|
empty_after_tool = {
|
|
"messages": [
|
|
HumanMessage(content="Check the status"),
|
|
ToolMessage(content="tool completed", tool_call_id="call-7"),
|
|
AIMessage(content="", response_metadata={"finish_reason": "stop"}),
|
|
]
|
|
}
|
|
|
|
first = middleware.after_model(empty_after_tool, runtime)
|
|
assert first is not None and first["jump_to"] == "model"
|
|
middleware.before_agent(empty_after_tool, runtime)
|
|
resumed = middleware.after_model(empty_after_tool, runtime)
|
|
assert resumed is not None and resumed["jump_to"] == "model"
|
|
|
|
|
|
def test_tool_history_without_real_user_message_does_not_trigger_recovery():
|
|
middleware = TerminalResponseMiddleware()
|
|
runtime = type("RuntimeStub", (), {"context": {"thread_id": "thread-8", "run_id": "run-8"}})()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="internal", additional_kwargs={"hide_from_ui": True}),
|
|
ToolMessage(content="tool completed", tool_call_id="call-8"),
|
|
AIMessage(content="", response_metadata={"finish_reason": "stop"}),
|
|
]
|
|
}
|
|
|
|
assert middleware.after_model(state, runtime) is None
|
|
|
|
|
|
def test_abandoned_run_state_is_bounded():
|
|
middleware = TerminalResponseMiddleware()
|
|
|
|
for index in range(1001):
|
|
key = (f"thread-{index}", f"run-{index}")
|
|
middleware._retry_counts[key] = 1
|
|
middleware._pending_prompts[key] = True
|
|
|
|
assert len(middleware._retry_counts) == 1000
|
|
assert len(middleware._pending_prompts) == 1000
|
|
assert ("thread-0", "run-0") not in middleware._retry_counts
|
|
assert ("thread-0", "run-0") not in middleware._pending_prompts
|