mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-08-29 03:56:43 +08:00
992 lines
33 KiB
Python
992 lines
33 KiB
Python
"""Agent 最终模型请求预算测试。"""
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from langchain.agents import create_agent
|
|
from langchain.agents.middleware import AgentMiddleware, SummarizationMiddleware
|
|
from langchain.agents.middleware.types import ModelRequest, ModelResponse
|
|
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
|
|
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
|
|
from langchain_core.tools import tool
|
|
|
|
from app.agent import MoviePilotAgent
|
|
from app.agent.middleware.tool_selection import ToolSelectorMiddleware
|
|
from app.agent.middleware.usage import UsageMiddleware
|
|
from app.chain.message import MessageChain
|
|
|
|
|
|
class _ToolBindingFakeModel(FakeMessagesListChatModel):
|
|
"""记录最终绑定工具，同时保留固定响应行为。"""
|
|
|
|
bound_tool_names: list[str] = []
|
|
|
|
def bind_tools(self, tools, **kwargs):
|
|
"""记录 LangChain 最终交给模型的工具集合。"""
|
|
self.bound_tool_names = [
|
|
item.get("function", {}).get("name") or item.get("name")
|
|
if isinstance(item, dict)
|
|
else item.name
|
|
for item in tools
|
|
]
|
|
return self
|
|
|
|
|
|
class _DynamicSystemMiddleware(AgentMiddleware):
|
|
"""模拟 MoviePilot 运行时动态追加系统上下文。"""
|
|
|
|
async def awrap_model_call(self, request, handler):
|
|
current = request.system_message.content if request.system_message else ""
|
|
return await handler(
|
|
request.override(
|
|
system_message=SystemMessage(
|
|
content=f"{current}\n{'动态系统上下文 ' * 200}"
|
|
)
|
|
)
|
|
)
|
|
|
|
|
|
def _request(
|
|
*,
|
|
messages=None,
|
|
system_message=None,
|
|
tools=None,
|
|
max_input_tokens=4096,
|
|
max_output_tokens=512,
|
|
model_settings=None,
|
|
) -> ModelRequest:
|
|
"""构造带模型窗口和最终请求组成的测试请求。"""
|
|
return ModelRequest(
|
|
model=SimpleNamespace(
|
|
model="small-model",
|
|
profile={
|
|
"max_input_tokens": max_input_tokens,
|
|
"max_output_tokens": max_output_tokens,
|
|
},
|
|
),
|
|
messages=list(messages or []),
|
|
system_message=system_message,
|
|
tools=list(tools or []),
|
|
state={},
|
|
runtime=None,
|
|
model_settings=model_settings,
|
|
)
|
|
|
|
|
|
def test_final_request_can_exceed_window_before_message_fraction_triggers():
|
|
"""动态系统提示词和工具定义可能在消息摘要阈值前耗尽输入窗口。"""
|
|
messages = [HumanMessage(content="用户上下文 " * 1000)]
|
|
model = SimpleNamespace(
|
|
_llm_type="test-chat",
|
|
profile={"max_input_tokens": 4096},
|
|
)
|
|
model.with_retry = lambda: model
|
|
summarizer = SummarizationMiddleware(
|
|
model=model,
|
|
trigger=("fraction", 0.85),
|
|
)
|
|
message_tokens = summarizer.token_counter(messages)
|
|
assert message_tokens < 4096 * 0.85
|
|
assert not summarizer._should_summarize(messages, message_tokens)
|
|
|
|
request = _request(
|
|
messages=messages,
|
|
system_message=SystemMessage(content="动态系统上下文 " * 1000),
|
|
tools=[
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "large_schema_tool",
|
|
"description": "工具业务说明 " * 800,
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"value": {
|
|
"type": "string",
|
|
"description": "参数约束 " * 400,
|
|
}
|
|
},
|
|
},
|
|
},
|
|
}
|
|
],
|
|
)
|
|
|
|
snapshot = UsageMiddleware.estimate_request(request)
|
|
|
|
assert snapshot["has_estimate"]
|
|
assert snapshot["message_tokens"] < 4096 * 0.85
|
|
assert snapshot["system_tokens"] > 0
|
|
assert snapshot["tool_tokens"] > 0
|
|
assert snapshot["estimated_input_tokens"] > 4096
|
|
assert snapshot["estimated_over_input_limit"] is True
|
|
assert snapshot["model_max_output_tokens"] == 512
|
|
assert snapshot["configured_output_limit_tokens"] is None
|
|
assert "output_headroom_tokens" not in snapshot
|
|
|
|
|
|
def test_request_budget_marks_same_request_over_limit_after_switch_to_small_model():
|
|
"""相同请求切换到小窗口模型后，应仅改变预算判断而不改变估算输入。"""
|
|
messages = [HumanMessage(content="x" * 12000)]
|
|
|
|
large_snapshot = UsageMiddleware.estimate_request(
|
|
_request(messages=messages, max_input_tokens=128000)
|
|
)
|
|
small_snapshot = UsageMiddleware.estimate_request(
|
|
_request(messages=messages, max_input_tokens=2048)
|
|
)
|
|
|
|
assert large_snapshot["estimated_input_tokens"] == small_snapshot["estimated_input_tokens"]
|
|
assert large_snapshot["estimated_over_input_limit"] is False
|
|
assert small_snapshot["estimated_over_input_limit"] is True
|
|
|
|
|
|
def test_request_budget_reads_only_explicit_per_call_output_limit():
|
|
"""只有最终请求显式配置的输出上限才可视为单次调用限制。"""
|
|
request = _request(
|
|
messages=[HumanMessage(content="hello")],
|
|
max_output_tokens=8192,
|
|
model_settings={"max_completion_tokens": 1024},
|
|
)
|
|
|
|
snapshot = UsageMiddleware.estimate_request(request)
|
|
|
|
assert snapshot["model_max_output_tokens"] == 8192
|
|
assert snapshot["configured_output_limit_tokens"] == 1024
|
|
assert snapshot["estimated_input_tokens"] < snapshot["context_window_tokens"]
|
|
|
|
|
|
def test_request_budget_rejects_non_integer_token_limits():
|
|
"""近似观察不得把 bool、浮点或字符串误报为有效 token 上限。"""
|
|
for invalid_value in (True, False, 1.5, 0, -1, "1024"):
|
|
request = _request(
|
|
messages=[HumanMessage(content="hello")],
|
|
max_input_tokens=invalid_value,
|
|
max_output_tokens=invalid_value,
|
|
model_settings={"max_completion_tokens": invalid_value},
|
|
)
|
|
|
|
snapshot = UsageMiddleware.estimate_request(request)
|
|
|
|
assert snapshot["context_window_tokens"] is None
|
|
assert snapshot["model_max_output_tokens"] is None
|
|
assert snapshot["configured_output_limit_tokens"] is None
|
|
assert snapshot["estimated_input_ratio"] is None
|
|
assert snapshot["estimated_over_input_limit"] is None
|
|
|
|
|
|
def test_request_budget_uses_next_valid_output_limit_alias():
|
|
"""高优先字段无效时，应继续读取同一请求中的有效兼容字段。"""
|
|
request = _request(
|
|
messages=[HumanMessage(content="hello")],
|
|
model_settings={"max_completion_tokens": True, "max_tokens": 1024},
|
|
)
|
|
|
|
snapshot = UsageMiddleware.estimate_request(request)
|
|
|
|
assert snapshot["configured_output_limit_tokens"] == 1024
|
|
|
|
|
|
def test_request_budget_counts_multimodal_input_without_storing_content():
|
|
"""图片按固定成本计入估算，快照不得保留请求正文或工具定义。"""
|
|
secret_marker = "REQUEST_BUDGET_SECRET_MARKER"
|
|
request = _request(
|
|
messages=[
|
|
HumanMessage(
|
|
content=[
|
|
{"type": "text", "text": secret_marker},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": "data:image/png;base64,hidden"},
|
|
},
|
|
{"type": "file", "file_id": secret_marker},
|
|
]
|
|
)
|
|
],
|
|
system_message=SystemMessage(content=secret_marker),
|
|
tools=[
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "secret_tool",
|
|
"description": secret_marker,
|
|
"parameters": {"type": "object", "properties": {}},
|
|
},
|
|
}
|
|
],
|
|
)
|
|
|
|
snapshot = UsageMiddleware.estimate_request(request)
|
|
|
|
assert snapshot["image_count"] == 1
|
|
assert snapshot["unknown_multimodal_count"] == 1
|
|
assert snapshot["multimodal_tokens"] == 85
|
|
assert snapshot["estimated_input_tokens"] == (
|
|
snapshot["message_tokens"]
|
|
+ snapshot["system_tokens"]
|
|
+ snapshot["tool_tokens"]
|
|
)
|
|
assert snapshot["model"] == "small-model"
|
|
assert secret_marker not in repr(snapshot)
|
|
assert all(
|
|
value is None or isinstance(value, (bool, int, float))
|
|
for key, value in snapshot.items()
|
|
if key != "model"
|
|
)
|
|
|
|
|
|
def test_request_budget_counts_each_image_cost_exactly_once():
|
|
"""LangChain 的消息估算已包含图片固定成本，汇总预算不得重复相加。"""
|
|
text_only = UsageMiddleware.estimate_request(
|
|
_request(
|
|
messages=[HumanMessage(content=[{"type": "text", "text": "hello"}])]
|
|
)
|
|
)
|
|
with_image = UsageMiddleware.estimate_request(
|
|
_request(
|
|
messages=[
|
|
HumanMessage(
|
|
content=[
|
|
{"type": "text", "text": "hello"},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": "data:image/png;base64,hidden"},
|
|
},
|
|
]
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
assert with_image["multimodal_tokens"] == 85
|
|
assert (
|
|
with_image["estimated_input_tokens"] - text_only["estimated_input_tokens"]
|
|
== 85
|
|
)
|
|
|
|
|
|
def test_request_budget_callback_failure_does_not_block_model_call():
|
|
"""预算观察失败不得改变模型请求和响应。"""
|
|
request = _request(messages=[HumanMessage(content="hello")])
|
|
response = ModelResponse(result=[AIMessage(content="ok")])
|
|
|
|
def _raise_callback(_snapshot):
|
|
raise RuntimeError("observer unavailable")
|
|
|
|
middleware = UsageMiddleware(on_request_budget=_raise_callback)
|
|
handled = []
|
|
|
|
async def _handler(received: ModelRequest):
|
|
handled.append(received)
|
|
return response
|
|
|
|
result = asyncio.run(middleware.awrap_model_call(request, _handler))
|
|
|
|
assert result is response
|
|
assert handled == [request]
|
|
|
|
|
|
def test_request_budget_callback_failure_clears_previous_estimate_state():
|
|
"""预算回调失败时，本轮 usage 不得与上一轮估算拼接。"""
|
|
agent = MoviePilotAgent(session_id="request-budget", user_id="user-1")
|
|
callback_count = 0
|
|
|
|
def _record_then_fail(snapshot):
|
|
nonlocal callback_count
|
|
callback_count += 1
|
|
if callback_count == 2:
|
|
raise RuntimeError("observer unavailable")
|
|
agent._record_request_budget(snapshot)
|
|
|
|
middleware = UsageMiddleware(
|
|
on_request_budget=_record_then_fail,
|
|
on_usage=agent._record_usage,
|
|
)
|
|
responses = iter(
|
|
[
|
|
ModelResponse(
|
|
result=[
|
|
AIMessage(
|
|
content="first",
|
|
usage_metadata={
|
|
"input_tokens": 10,
|
|
"output_tokens": 1,
|
|
"total_tokens": 11,
|
|
},
|
|
)
|
|
]
|
|
),
|
|
ModelResponse(
|
|
result=[
|
|
AIMessage(
|
|
content="second",
|
|
usage_metadata={
|
|
"input_tokens": 50,
|
|
"output_tokens": 2,
|
|
"total_tokens": 52,
|
|
},
|
|
)
|
|
]
|
|
),
|
|
]
|
|
)
|
|
|
|
async def _handler(_request):
|
|
return next(responses)
|
|
|
|
request = _request(messages=[HumanMessage(content="hello")])
|
|
asyncio.run(middleware.awrap_model_call(request, _handler))
|
|
assert agent.get_session_status()["last_request_estimate_available"] is True
|
|
|
|
asyncio.run(middleware.awrap_model_call(request, _handler))
|
|
status = agent.get_session_status()
|
|
|
|
assert status["last_request_sequence"] == 2
|
|
assert status["last_request_estimate_available"] is False
|
|
assert status["last_estimated_input_tokens"] is None
|
|
assert status["last_actual_input_tokens"] is None
|
|
assert status["last_estimate_error_tokens"] is None
|
|
assert status["model"] == "small-model"
|
|
assert status["context_window_tokens"] == 4096
|
|
assert status["last_input_tokens"] == 50
|
|
|
|
|
|
def test_request_budget_estimator_failure_reports_empty_snapshot_and_calls_model():
|
|
"""估算器异常应清除旧观测状态，并继续原模型调用。"""
|
|
budgets = []
|
|
middleware = UsageMiddleware(on_request_budget=budgets.append)
|
|
request = _request(messages=[HumanMessage(content="hello")])
|
|
response = ModelResponse(result=[AIMessage(content="ok")])
|
|
|
|
async def _handler(_request):
|
|
return response
|
|
|
|
with patch.object(
|
|
UsageMiddleware,
|
|
"estimate_request",
|
|
side_effect=RuntimeError("request content must not reach the snapshot"),
|
|
):
|
|
result = asyncio.run(middleware.awrap_model_call(request, _handler))
|
|
|
|
assert result is response
|
|
assert budgets == [
|
|
{
|
|
"request_sequence": 1,
|
|
"has_estimate": False,
|
|
"model": "small-model",
|
|
"context_window_tokens": 4096,
|
|
}
|
|
]
|
|
|
|
|
|
def test_request_budget_metadata_failure_still_calls_model():
|
|
"""模型元数据属性异常不得让预算观察器阻断真实模型调用。"""
|
|
|
|
class _BrokenMetadataModel:
|
|
@property
|
|
def model(self):
|
|
raise RuntimeError("model metadata unavailable")
|
|
|
|
@property
|
|
def model_name(self):
|
|
raise RuntimeError("model metadata unavailable")
|
|
|
|
@property
|
|
def model_id(self):
|
|
raise RuntimeError("model metadata unavailable")
|
|
|
|
@property
|
|
def profile(self):
|
|
raise RuntimeError("profile metadata unavailable")
|
|
|
|
budgets = []
|
|
middleware = UsageMiddleware(on_request_budget=budgets.append)
|
|
request = ModelRequest(
|
|
model=_BrokenMetadataModel(),
|
|
messages=[HumanMessage(content="hello")],
|
|
tools=[],
|
|
state={},
|
|
runtime=None,
|
|
)
|
|
response = ModelResponse(result=[AIMessage(content="ok")])
|
|
handled = []
|
|
|
|
async def _handler(received):
|
|
handled.append(received)
|
|
return response
|
|
|
|
with patch.object(
|
|
UsageMiddleware,
|
|
"estimate_request",
|
|
side_effect=RuntimeError("estimate unavailable"),
|
|
):
|
|
result = asyncio.run(middleware.awrap_model_call(request, _handler))
|
|
|
|
assert result is response
|
|
assert handled == [request]
|
|
assert budgets == [
|
|
{
|
|
"request_sequence": 1,
|
|
"has_estimate": False,
|
|
"model": None,
|
|
"context_window_tokens": None,
|
|
}
|
|
]
|
|
|
|
|
|
def test_request_sequence_callback_failure_still_calls_model():
|
|
"""会话序号分配异常时应放弃最近快照竞争，并继续真实模型调用。"""
|
|
budgets = []
|
|
|
|
def _raise_sequence():
|
|
raise RuntimeError("sequence unavailable")
|
|
|
|
middleware = UsageMiddleware(
|
|
on_request_budget=budgets.append,
|
|
next_request_sequence=_raise_sequence,
|
|
)
|
|
request = _request(messages=[HumanMessage(content="hello")])
|
|
response = ModelResponse(result=[AIMessage(content="ok")])
|
|
handled = []
|
|
|
|
async def _handler(received):
|
|
handled.append(received)
|
|
return response
|
|
|
|
result = asyncio.run(middleware.awrap_model_call(request, _handler))
|
|
|
|
assert result is response
|
|
assert handled == [request]
|
|
assert budgets[0]["request_sequence"] is None
|
|
|
|
|
|
def test_request_budget_and_actual_usage_share_one_request_sequence():
|
|
"""真实 usage 只能校准同一次成功模型调用产生的估算。"""
|
|
budgets = []
|
|
usages = []
|
|
middleware = UsageMiddleware(
|
|
on_request_budget=budgets.append,
|
|
on_usage=usages.append,
|
|
)
|
|
request = _request(messages=[HumanMessage(content="hello")])
|
|
response = ModelResponse(
|
|
result=[
|
|
AIMessage(
|
|
content="ok",
|
|
usage_metadata={
|
|
"input_tokens": 23,
|
|
"output_tokens": 4,
|
|
"total_tokens": 27,
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
async def _handler(_request):
|
|
return response
|
|
|
|
asyncio.run(middleware.awrap_model_call(request, _handler))
|
|
|
|
assert budgets[0]["request_sequence"] == 1
|
|
assert usages[0]["request_sequence"] == 1
|
|
assert usages[0]["request_budget_recorded"] is True
|
|
assert usages[0]["input_usage_available"] is True
|
|
assert usages[0]["estimated_input_tokens"] == budgets[0]["estimated_input_tokens"]
|
|
|
|
|
|
def test_partial_usage_without_input_does_not_calibrate_request_estimate():
|
|
"""仅有输出 usage 时，不得把缺失的真实输入误报为零。"""
|
|
agent = MoviePilotAgent(session_id="request-budget", user_id="user-1")
|
|
middleware = UsageMiddleware(
|
|
on_request_budget=agent._record_request_budget,
|
|
on_usage=agent._record_usage,
|
|
)
|
|
request = _request(messages=[HumanMessage(content="hello")])
|
|
response = ModelResponse(
|
|
result=[
|
|
AIMessage(
|
|
content="ok",
|
|
response_metadata={
|
|
"token_usage": {
|
|
"completion_tokens": 7,
|
|
"total_tokens": 7,
|
|
}
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
async def _handler(_request):
|
|
return response
|
|
|
|
asyncio.run(middleware.awrap_model_call(request, _handler))
|
|
status = agent.get_session_status()
|
|
|
|
assert status["last_request_estimate_available"] is True
|
|
assert status["last_input_usage_available"] is False
|
|
assert status["last_input_tokens"] is None
|
|
assert status["last_output_tokens"] == 7
|
|
assert status["last_actual_input_tokens"] is None
|
|
assert status["last_estimate_error_tokens"] is None
|
|
assert status["last_estimate_error_ratio"] is None
|
|
assert status["last_context_usage_ratio"] is None
|
|
status_text = MessageChain._format_session_status_text(status)
|
|
assert "最近一次上下文占用: 未知 / 4,096" in status_text
|
|
assert "最近一次 tokens: 输入 未知 / 输出 7 / 总计 7" in status_text
|
|
|
|
|
|
def test_missing_usage_clears_previous_last_call_values():
|
|
"""本轮没有 usage 时，最近一次状态不得保留上一轮实际值。"""
|
|
agent = MoviePilotAgent(session_id="request-budget", user_id="user-1")
|
|
agent._record_request_budget(
|
|
{
|
|
"request_sequence": 1,
|
|
"has_estimate": True,
|
|
"estimated_input_tokens": 10,
|
|
"context_window_tokens": 1000,
|
|
}
|
|
)
|
|
agent._record_usage(
|
|
{
|
|
"request_sequence": 1,
|
|
"request_budget_recorded": True,
|
|
"has_usage": True,
|
|
"input_usage_available": True,
|
|
"input_tokens": 12,
|
|
"output_tokens": 3,
|
|
"total_tokens": 15,
|
|
"cache_usage_available": True,
|
|
"cache_read_input_tokens": 4,
|
|
"cache_write_input_tokens": 0,
|
|
"uncached_input_tokens": 8,
|
|
}
|
|
)
|
|
agent._record_request_budget(
|
|
{
|
|
"request_sequence": 2,
|
|
"has_estimate": True,
|
|
"estimated_input_tokens": 20,
|
|
"context_window_tokens": 1000,
|
|
}
|
|
)
|
|
agent._record_usage(
|
|
{
|
|
"request_sequence": 2,
|
|
"request_budget_recorded": True,
|
|
"has_usage": False,
|
|
"input_usage_available": False,
|
|
}
|
|
)
|
|
|
|
status = agent.get_session_status()
|
|
|
|
assert status["last_request_sequence"] == 2
|
|
assert status["last_request_estimate_available"] is True
|
|
assert status["last_input_usage_available"] is False
|
|
assert status["last_input_tokens"] is None
|
|
assert status["last_output_tokens"] is None
|
|
assert status["last_total_tokens"] is None
|
|
assert status["last_context_usage_ratio"] is None
|
|
assert status["last_cache_usage_available"] is False
|
|
assert status["total_input_tokens"] == 12
|
|
assert status["total_output_tokens"] == 3
|
|
assert status["total_tokens"] == 15
|
|
status_text = MessageChain._format_session_status_text(status)
|
|
assert "最近一次 tokens: 输入 未知 / 输出 未知 / 总计 未知" in status_text
|
|
assert "最近一次缓存:" not in status_text
|
|
assert "当前会话累计缓存: 命中 4 / 写入 0 / 未命中 8" in status_text
|
|
|
|
|
|
def test_out_of_order_responses_keep_latest_request_snapshot():
|
|
"""较早请求晚返回时，只累计 usage，不得覆盖较新请求的最近状态。"""
|
|
agent = MoviePilotAgent(session_id="request-budget", user_id="user-1")
|
|
middleware = UsageMiddleware(
|
|
on_request_budget=agent._record_request_budget,
|
|
on_usage=agent._record_usage,
|
|
)
|
|
first_request = _request(
|
|
messages=[HumanMessage(content="first")],
|
|
max_input_tokens=1000,
|
|
)
|
|
first_request.model.model = "first-model"
|
|
second_request = _request(
|
|
messages=[HumanMessage(content="second")],
|
|
max_input_tokens=2000,
|
|
)
|
|
second_request.model.model = "second-model"
|
|
first_started = asyncio.Event()
|
|
release_first = asyncio.Event()
|
|
|
|
async def _handler(request):
|
|
if request.model.model == "first-model":
|
|
first_started.set()
|
|
await release_first.wait()
|
|
return ModelResponse(
|
|
result=[
|
|
AIMessage(
|
|
content="first",
|
|
usage_metadata={
|
|
"input_tokens": 10,
|
|
"output_tokens": 1,
|
|
"total_tokens": 11,
|
|
},
|
|
)
|
|
]
|
|
)
|
|
return ModelResponse(
|
|
result=[
|
|
AIMessage(
|
|
content="second",
|
|
usage_metadata={
|
|
"input_tokens": 50,
|
|
"output_tokens": 2,
|
|
"total_tokens": 52,
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
async def _run_out_of_order():
|
|
first_task = asyncio.create_task(
|
|
middleware.awrap_model_call(first_request, _handler)
|
|
)
|
|
await first_started.wait()
|
|
await middleware.awrap_model_call(second_request, _handler)
|
|
release_first.set()
|
|
await first_task
|
|
|
|
asyncio.run(_run_out_of_order())
|
|
status = agent.get_session_status()
|
|
|
|
assert status["last_request_sequence"] == 2
|
|
assert status["model"] == "second-model"
|
|
assert status["context_window_tokens"] == 2000
|
|
assert status["last_input_tokens"] == 50
|
|
assert status["last_output_tokens"] == 2
|
|
assert status["last_total_tokens"] == 52
|
|
assert status["last_actual_input_tokens"] == 50
|
|
assert status["total_input_tokens"] == 60
|
|
assert status["total_output_tokens"] == 3
|
|
assert status["total_tokens"] == 63
|
|
assert status["model_call_count"] == 2
|
|
|
|
|
|
def test_sequence_failure_cannot_overwrite_next_request_snapshot():
|
|
"""无序号请求晚返回时只能累计 usage，不能覆盖后续有序请求。"""
|
|
agent = MoviePilotAgent(session_id="request-budget", user_id="user-1")
|
|
allocation_count = 0
|
|
|
|
def _next_sequence():
|
|
nonlocal allocation_count
|
|
allocation_count += 1
|
|
if allocation_count == 1:
|
|
raise RuntimeError("sequence unavailable")
|
|
return agent._next_request_sequence()
|
|
|
|
middleware = UsageMiddleware(
|
|
on_request_budget=agent._record_request_budget,
|
|
on_usage=agent._record_usage,
|
|
next_request_sequence=_next_sequence,
|
|
)
|
|
first_request = _request(
|
|
messages=[HumanMessage(content="first")],
|
|
max_input_tokens=1000,
|
|
)
|
|
first_request.model.model = "first-model"
|
|
second_request = _request(
|
|
messages=[HumanMessage(content="second")],
|
|
max_input_tokens=2000,
|
|
)
|
|
second_request.model.model = "second-model"
|
|
first_started = asyncio.Event()
|
|
release_first = asyncio.Event()
|
|
|
|
async def _handler(request):
|
|
if request.model.model == "first-model":
|
|
first_started.set()
|
|
await release_first.wait()
|
|
input_tokens = 10
|
|
else:
|
|
input_tokens = 50
|
|
return ModelResponse(
|
|
result=[
|
|
AIMessage(
|
|
content=request.model.model,
|
|
usage_metadata={
|
|
"input_tokens": input_tokens,
|
|
"output_tokens": 1,
|
|
"total_tokens": input_tokens + 1,
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
async def _run_out_of_order():
|
|
first_task = asyncio.create_task(
|
|
middleware.awrap_model_call(first_request, _handler)
|
|
)
|
|
await first_started.wait()
|
|
await middleware.awrap_model_call(second_request, _handler)
|
|
release_first.set()
|
|
await first_task
|
|
|
|
asyncio.run(_run_out_of_order())
|
|
status = agent.get_session_status()
|
|
|
|
assert status["last_request_sequence"] == 1
|
|
assert status["model"] == "second-model"
|
|
assert status["context_window_tokens"] == 2000
|
|
assert status["last_input_tokens"] == 50
|
|
assert status["last_actual_input_tokens"] == 50
|
|
assert status["total_input_tokens"] == 60
|
|
assert status["model_call_count"] == 2
|
|
|
|
|
|
def test_failed_model_switch_keeps_model_and_window_from_same_request():
|
|
"""新模型请求即使失败，状态也不得混合旧模型名称与新窗口。"""
|
|
agent = MoviePilotAgent(session_id="request-budget", user_id="user-1")
|
|
agent._sync_model_profile(
|
|
SimpleNamespace(model="large-model", profile={"max_input_tokens": 128000})
|
|
)
|
|
middleware = UsageMiddleware(
|
|
on_request_budget=agent._record_request_budget,
|
|
on_usage=agent._record_usage,
|
|
)
|
|
request = _request(
|
|
messages=[HumanMessage(content="small")],
|
|
max_input_tokens=2048,
|
|
)
|
|
request.model.model = "small-model"
|
|
|
|
async def _failing_handler(_request):
|
|
raise RuntimeError("model unavailable")
|
|
|
|
async def _run_failure():
|
|
try:
|
|
await middleware.awrap_model_call(request, _failing_handler)
|
|
except RuntimeError:
|
|
pass
|
|
|
|
asyncio.run(_run_failure())
|
|
status = agent.get_session_status()
|
|
|
|
assert status["last_request_sequence"] == 1
|
|
assert status["model"] == "small-model"
|
|
assert status["context_window_tokens"] == 2048
|
|
|
|
|
|
def test_unknown_model_name_is_not_replaced_after_request_snapshot():
|
|
"""请求已开始后，未知模型名称不得与配置默认值拼成虚假快照。"""
|
|
agent = MoviePilotAgent(session_id="request-budget", user_id="user-1")
|
|
agent._record_request_budget(
|
|
{
|
|
"has_estimate": True,
|
|
"request_sequence": 1,
|
|
"model": None,
|
|
"estimated_input_tokens": 100,
|
|
"context_window_tokens": 2048,
|
|
}
|
|
)
|
|
|
|
status = agent.get_session_status()
|
|
|
|
assert status["model"] is None
|
|
assert status["context_window_tokens"] == 2048
|
|
|
|
|
|
def test_request_sequence_remains_monotonic_after_agent_graph_rebuild():
|
|
"""重建 Agent 图后，新观察器也必须延续当前会话的请求顺序。"""
|
|
agent = MoviePilotAgent(session_id="request-budget", user_id="user-1")
|
|
first_graph = UsageMiddleware(
|
|
on_request_budget=agent._record_request_budget,
|
|
on_usage=agent._record_usage,
|
|
next_request_sequence=agent._next_request_sequence,
|
|
)
|
|
rebuilt_graph = UsageMiddleware(
|
|
on_request_budget=agent._record_request_budget,
|
|
on_usage=agent._record_usage,
|
|
next_request_sequence=agent._next_request_sequence,
|
|
)
|
|
|
|
async def _handler(request):
|
|
return ModelResponse(
|
|
result=[
|
|
AIMessage(
|
|
content=request.model.model,
|
|
usage_metadata={
|
|
"input_tokens": request.model.profile["max_input_tokens"] // 10,
|
|
"output_tokens": 1,
|
|
"total_tokens": request.model.profile["max_input_tokens"] // 10 + 1,
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
first = _request(
|
|
messages=[HumanMessage(content="first")],
|
|
max_input_tokens=1000,
|
|
)
|
|
first.model.model = "first-model"
|
|
second = _request(
|
|
messages=[HumanMessage(content="second")],
|
|
max_input_tokens=2000,
|
|
)
|
|
second.model.model = "second-model"
|
|
|
|
asyncio.run(first_graph.awrap_model_call(first, _handler))
|
|
asyncio.run(rebuilt_graph.awrap_model_call(second, _handler))
|
|
status = agent.get_session_status()
|
|
|
|
assert status["last_request_sequence"] == 2
|
|
assert status["model"] == "second-model"
|
|
assert status["context_window_tokens"] == 2000
|
|
assert status["last_input_tokens"] == 200
|
|
|
|
|
|
def test_new_request_estimate_clears_stale_calibration_until_success():
|
|
"""新请求失败时不得沿用上一轮估算误差。"""
|
|
agent = MoviePilotAgent(session_id="request-budget", user_id="user-1")
|
|
first = {
|
|
"has_estimate": True,
|
|
"request_sequence": 1,
|
|
"estimated_input_tokens": 100,
|
|
"context_window_tokens": 1000,
|
|
}
|
|
second = {
|
|
"has_estimate": True,
|
|
"request_sequence": 2,
|
|
"estimated_input_tokens": 200,
|
|
"context_window_tokens": 1000,
|
|
}
|
|
|
|
agent._record_request_budget(first)
|
|
agent._record_usage(
|
|
{
|
|
"request_sequence": 1,
|
|
"request_budget_recorded": True,
|
|
"input_usage_available": True,
|
|
"estimated_input_tokens": 100,
|
|
"has_usage": True,
|
|
"input_tokens": 120,
|
|
"output_tokens": 1,
|
|
"total_tokens": 121,
|
|
}
|
|
)
|
|
assert agent.get_session_status()["last_estimate_error_tokens"] == 20
|
|
|
|
agent._record_request_budget(second)
|
|
status = agent.get_session_status()
|
|
|
|
assert status["last_request_sequence"] == 2
|
|
assert status["last_actual_input_tokens"] is None
|
|
assert status["last_estimate_error_tokens"] is None
|
|
assert status["last_estimate_error_ratio"] is None
|
|
|
|
|
|
def test_failed_estimate_clears_previous_request_components():
|
|
"""估算失败后状态不得继续展示上一轮的请求组成。"""
|
|
agent = MoviePilotAgent(session_id="request-budget", user_id="user-1")
|
|
agent._record_request_budget(
|
|
{
|
|
"has_estimate": True,
|
|
"request_sequence": 1,
|
|
"estimated_input_tokens": 100,
|
|
"message_tokens": 70,
|
|
"system_tokens": 20,
|
|
"tool_tokens": 10,
|
|
"message_count": 3,
|
|
"tool_count": 2,
|
|
}
|
|
)
|
|
|
|
agent._record_request_budget(
|
|
{"has_estimate": False, "request_sequence": 2}
|
|
)
|
|
status = agent.get_session_status()
|
|
|
|
assert status["last_request_sequence"] == 2
|
|
assert status["last_request_estimate_available"] is False
|
|
assert status["last_estimated_input_tokens"] is None
|
|
assert status["last_estimated_message_tokens"] is None
|
|
assert status["last_estimated_system_tokens"] is None
|
|
assert status["last_estimated_tool_tokens"] is None
|
|
assert status["last_message_count"] == 0
|
|
assert status["last_tool_count"] == 0
|
|
|
|
|
|
def test_new_request_with_unknown_window_clears_previous_model_window():
|
|
"""切换到窗口未知的模型后，不得继续展示上一模型的窗口。"""
|
|
agent = MoviePilotAgent(session_id="request-budget", user_id="user-1")
|
|
agent._sync_model_profile(
|
|
SimpleNamespace(model="large-model", profile={"max_input_tokens": 128000})
|
|
)
|
|
agent._record_request_budget(
|
|
{
|
|
"has_estimate": True,
|
|
"request_sequence": 1,
|
|
"estimated_input_tokens": 100,
|
|
"context_window_tokens": 128000,
|
|
}
|
|
)
|
|
|
|
agent._sync_model_profile(
|
|
SimpleNamespace(model="unknown-window-model", profile={"max_input_tokens": 0})
|
|
)
|
|
agent._record_request_budget(
|
|
{
|
|
"has_estimate": True,
|
|
"request_sequence": 2,
|
|
"estimated_input_tokens": 100,
|
|
"context_window_tokens": None,
|
|
}
|
|
)
|
|
status = agent.get_session_status()
|
|
|
|
assert status["model"] == "unknown-window-model"
|
|
assert status["context_window_tokens"] is None
|
|
assert status["last_estimated_input_ratio"] is None
|
|
assert status["last_estimated_over_input_limit"] is None
|
|
|
|
|
|
def test_real_agent_observer_sees_dynamic_system_and_selected_tools():
|
|
"""末尾观察器必须看到动态 system 和 ToolSelector 最终保留的工具。"""
|
|
|
|
@tool
|
|
def kept_tool(value: str) -> str:
|
|
"""应保留的测试工具。"""
|
|
return value
|
|
|
|
@tool
|
|
def removed_tool(value: str) -> str:
|
|
"""应被筛除的测试工具。"""
|
|
return value
|
|
|
|
selection_model = _ToolBindingFakeModel(
|
|
responses=[AIMessage(content='{"tools": []}')]
|
|
)
|
|
main_model = _ToolBindingFakeModel(responses=[AIMessage(content="done")])
|
|
budgets = []
|
|
graph = create_agent(
|
|
model=main_model,
|
|
tools=[kept_tool, removed_tool],
|
|
system_prompt="base",
|
|
middleware=[
|
|
_DynamicSystemMiddleware(),
|
|
ToolSelectorMiddleware(
|
|
model=selection_model,
|
|
selection_tools=[kept_tool, removed_tool],
|
|
max_tools=1,
|
|
always_include=["kept_tool"],
|
|
),
|
|
UsageMiddleware(on_request_budget=budgets.append),
|
|
],
|
|
)
|
|
|
|
asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="hello")]}))
|
|
|
|
assert len(budgets) == 1
|
|
assert budgets[0]["tool_count"] == 1
|
|
assert budgets[0]["system_tokens"] > 100
|
|
assert main_model.bound_tool_names == ["kept_tool"]
|