Xi Zhang
c407d2e20f
* feat(middleware): add ConfigurableModelMiddleware for dynamic model resolution - Introduced ConfigurableModelMiddleware to resolve chat models from RunnableConfig.configurable on each call. - Updated middleware initialization to include ConfigurableModelMiddleware. - Enhanced context editing middleware tests to verify presence of ConfigurableModelMiddleware. - Implemented tests for ConfigurableModelMiddleware to ensure correct model overriding and caching behavior. - Added tests for deepagents model-passthrough patch to verify configuration injection in async tasks. * feat(async-subagent): update middleware handling to prevent deadlocks in async sub-agents * style: Refactor code formatting for improved readability in patches and test files * refactor: streamline middleware construction and improve async handling in ConfigurableModelMiddleware * fix: remove unused request parameter from _read_model_override function * refactor: improve async handling in _ClientProxy and enhance logging in ConfigurableModelMiddleware test: add behavior test to ensure AskUserMiddleware is excluded in async subagent mode
163 lines
5.5 KiB
Python
163 lines
5.5 KiB
Python
"""Tests for ContextEditingMiddleware integration and compute_context_editing_trigger."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from langchain.agents.middleware import ContextEditingMiddleware
|
|
|
|
from EvoScientist.middleware.context_editing import compute_context_editing_trigger
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# compute_context_editing_trigger tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_compute_trigger_with_profile():
|
|
model = MagicMock()
|
|
model.profile = {"max_input_tokens": 200_000}
|
|
assert compute_context_editing_trigger(model) == 100_000 # 50%
|
|
|
|
|
|
def test_compute_trigger_with_1m_profile():
|
|
model = MagicMock()
|
|
model.profile = {"max_input_tokens": 1_000_000}
|
|
assert compute_context_editing_trigger(model) == 500_000 # 50%
|
|
|
|
|
|
def test_compute_trigger_with_context_length_attr():
|
|
model = MagicMock(spec=["context_length", "profile"])
|
|
model.context_length = 1_000_000
|
|
model.profile = None
|
|
assert compute_context_editing_trigger(model) == 500_000 # 50%
|
|
|
|
|
|
def test_compute_trigger_with_num_ctx():
|
|
model = MagicMock(spec=["num_ctx", "profile"])
|
|
model.num_ctx = 32_768
|
|
model.profile = None
|
|
assert compute_context_editing_trigger(model) == 16_384 # 50%
|
|
|
|
|
|
def test_compute_trigger_without_profile():
|
|
model = MagicMock()
|
|
model.profile = None
|
|
assert compute_context_editing_trigger(model) == 100_000 # fallback
|
|
|
|
|
|
def test_compute_trigger_no_profile_attr():
|
|
model = MagicMock(spec=[]) # no attributes at all
|
|
assert compute_context_editing_trigger(model) == 100_000 # fallback
|
|
|
|
|
|
def test_compute_trigger_empty_profile():
|
|
model = MagicMock()
|
|
model.profile = {}
|
|
assert compute_context_editing_trigger(model) == 100_000 # fallback
|
|
|
|
|
|
def test_compute_trigger_custom_fraction():
|
|
model = MagicMock()
|
|
model.profile = {"max_input_tokens": 200_000}
|
|
assert compute_context_editing_trigger(model, fraction=0.30) == 60_000
|
|
|
|
|
|
def test_compute_trigger_custom_fallback():
|
|
model = MagicMock()
|
|
model.profile = None
|
|
assert compute_context_editing_trigger(model, fallback=50_000) == 50_000
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# create_context_editing_middleware tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_create_middleware_configuration():
|
|
from EvoScientist.middleware.context_editing import (
|
|
create_context_editing_middleware,
|
|
)
|
|
|
|
model = MagicMock()
|
|
model.profile = {"max_input_tokens": 200_000}
|
|
mw = create_context_editing_middleware(model)
|
|
edit = mw.edits[0]
|
|
assert edit.trigger == 100_000
|
|
assert edit.keep == 5
|
|
assert "think_tool" in edit.exclude_tools
|
|
|
|
|
|
@patch("EvoScientist.EvoScientist._ensure_chat_model")
|
|
def test_create_middleware_model_none_fallback(mock_model):
|
|
from EvoScientist.middleware.context_editing import (
|
|
create_context_editing_middleware,
|
|
)
|
|
|
|
mock_model.return_value = MagicMock(profile=None)
|
|
mw = create_context_editing_middleware(None)
|
|
edit = mw.edits[0]
|
|
assert edit.trigger == 100_000 # fallback
|
|
mock_model.assert_called_once()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Middleware list integration tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@patch(
|
|
"EvoScientist.middleware.create_tool_selector_middleware",
|
|
return_value=[MagicMock(), MagicMock()],
|
|
)
|
|
@patch("EvoScientist.EvoScientist._ensure_chat_model")
|
|
@patch("EvoScientist.EvoScientist._ensure_config")
|
|
def test_default_middleware_includes_context_editing(mock_config, mock_model, mock_ts):
|
|
mock_model.return_value = MagicMock(profile={"max_input_tokens": 200_000})
|
|
cfg = MagicMock()
|
|
cfg.enable_ask_user = False
|
|
cfg.auto_approve = False
|
|
mock_config.return_value = cfg
|
|
|
|
from EvoScientist.EvoScientist import _get_default_middleware
|
|
|
|
mw = _get_default_middleware()
|
|
# ContextEditingMiddleware is present (its absolute position depends on
|
|
# other leading middlewares like ConfigurableModelMiddleware).
|
|
assert any(isinstance(m, ContextEditingMiddleware) for m in mw)
|
|
|
|
|
|
@patch("EvoScientist.EvoScientist._ensure_chat_model")
|
|
def test_inject_subagent_includes_context_editing(mock_model):
|
|
mock_model.return_value = MagicMock(profile={"max_input_tokens": 200_000})
|
|
|
|
from EvoScientist.EvoScientist import _inject_subagent_middleware
|
|
|
|
subs = [{"name": "test-agent"}]
|
|
_inject_subagent_middleware(subs)
|
|
|
|
middleware_types = [type(m) for m in subs[0]["middleware"]]
|
|
assert ContextEditingMiddleware in middleware_types
|
|
|
|
|
|
@patch(
|
|
"EvoScientist.middleware.create_tool_selector_middleware",
|
|
return_value=[MagicMock(), MagicMock()],
|
|
)
|
|
@patch("EvoScientist.EvoScientist._ensure_chat_model")
|
|
@patch("EvoScientist.EvoScientist._ensure_config")
|
|
def test_context_editing_before_overflow_mapper(mock_config, mock_model, mock_ts):
|
|
mock_model.return_value = MagicMock(profile={"max_input_tokens": 200_000})
|
|
cfg = MagicMock()
|
|
cfg.enable_ask_user = False
|
|
cfg.auto_approve = False
|
|
mock_config.return_value = cfg
|
|
|
|
from EvoScientist.EvoScientist import _get_default_middleware
|
|
|
|
mw = _get_default_middleware()
|
|
type_names = [type(m).__name__ for m in mw]
|
|
|
|
ce_idx = type_names.index("ContextEditingMiddleware")
|
|
co_idx = type_names.index("ContextOverflowMapperMiddleware")
|
|
assert ce_idx < co_idx, (
|
|
"ContextEditingMiddleware should come before ContextOverflowMapperMiddleware"
|
|
)
|