Files
AgentLab/backend/tests/test_model_budget.py
T
ThearchyHelios 5afa7a9a70 核查、证据裁判退回小额度时多认两种拒绝,被嫌大过的接入记在进程里
复核第 1 轮的发现:
- 自建服务(text-generation-inference 一类)嫌额度大时回 422,原话说的是 max_new_tokens,不含 max_tokens,认不出就不退回;核查、裁判的额度从 4096 提到 16384 之后才容易撞上。新增 WIDER_BUDGET_WORDS 和 budget_refused,只给核查、裁判用;助手改图仍按 budget_rejected 认,BUDGET_WORDS 不动(data/recipe_ai.py 的测试拿它比对)
- Bearer 鉴权的 Anthropic 接入(factory 换上的 SDK 客户端没设超时)调 Opus 4、4.1 这一代的 id 时,16384 的非流式请求 anthropic SDK 在本地就抛 ValueError(Streaming is required…),请求根本没发出去。budget_refused 把它也认成额度被挡下,退回小额度;重来不会重复输出,也没花钱
- 退回的开关原来只在一组调用里共用:双盲复算的几路盲算、反方审稿的每一句、每次质询、每次点击裁判各自先按 16384 撞一次。BudgetFallback 加 key:被嫌大一次就记进进程里的 _REJECTED,之后同一个 key 新建的一开始就用小的(同 factory 的 _NO_STREAM_USAGE,重启就忘);不给 key 的(助手改图)不记也不查,行为不变
2026-10-06 14:29:48 +08:00

345 lines
18 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""模型调用的输出额度(providers/budget.py):助手改图、核查、证据裁判共用的那几样。
推理模型的思考也算在输出额度里,额度小了,想完就没额度写正文:回复是空的,或者 JSON 只写了一半。这里守几件事:
- 认得出「输出到了上限被截断」:OpenAI 系 finish_reason=length、Anthropic stop_reason=max_tokens,以及 openai 结构化
输出解析时抛的 LengthFinishReasonError;服务方嫌额度太大的 400 不算截断(那是请求被拒,一个字都没写)
- 认得出「服务方嫌额度太大」:400 / 422,报错里有 max_tokens 一类字样;别的 400 不算
- 截断的那次调用花掉的用量从异常里取得到(额度是真花出去了)
- BudgetFallback:先按大额度调,被嫌大时按小额度重来一次,之后的调用都用小的;流式、非流式都认;bind_tools、
with_structured_output 包出来的和原来的共用同一个开关;只在还没吐出任何东西时重来;小额度的模型要退回时才建
- 核查、证据裁判另认两种「额度被挡下」(budget_refused):自建服务(text-generation-inference 一类)说的
max_new_tokens,和 anthropic SDK 在本地就拒的额度(请求没发出去);助手改图的认法(budget_rejected)不变
- 给了 key(接入、模型、大额度)的,被嫌大一次就记在进程里:之后新建的调用直接用小的,不再每次先撞一次;没给 key 的
(助手改图)不记也不查
模型都是剧本,不碰真接口。
"""
from __future__ import annotations
import pytest
from langchain_core.messages import AIMessage, AIMessageChunk
from openai import LengthFinishReasonError
from openai.types import CompletionUsage
from openai.types.chat import ChatCompletion
from app.providers import budget
from app.providers.budget import (
BUDGET_WORDS,
BudgetFallback,
budget_refused,
budget_rejected,
cut_off,
cut_words,
error_usage,
truncated_error,
)
@pytest.fixture(autouse=True)
def forget_rejections():
"""被嫌大过的接入记在进程里:每个测试从没记过开始,也不留给后面的测试。"""
budget._REJECTED.clear()
yield
budget._REJECTED.clear()
def length_error(prompt: int = 9000, completion: int = 16384) -> LengthFinishReasonError:
"""openai 的结构化输出在 finish_reason=length 时抛的那个错(原话和真实运行里一样)。"""
usage = CompletionUsage(completion_tokens=completion, prompt_tokens=prompt, total_tokens=prompt + completion)
return LengthFinishReasonError(completion=ChatCompletion.model_construct(
id="x", choices=[], created=0, model="m", object="chat.completion", usage=usage))
class Rejected(Exception):
"""仿 openai / anthropic SDK 的 400:一发请求就被拒,还没吐出任何东西。"""
def __init__(self, message: str, status_code: int = 400) -> None:
super().__init__(message)
self.status_code = status_code
TOO_LARGE = "max_tokens: 16384 > 8192, which is the maximum allowed number of output tokens"
#: text-generation-inference 的 OpenAI 兼容接口:输入加额度超过 max_total_tokens 时的 422(原话的样子)
TGI_TOO_LARGE = ("Input validation error: `inputs` tokens + `max_new_tokens` must be <= 8192. "
"Given: 1200 `inputs` tokens and 16384 `max_new_tokens`")
#: anthropic SDK 在本地拒绝时的原话:没设超时的客户端非流式调用、额度超过这个模型不流式能给的上限
SDK_REFUSED = ("Streaming is required for operations that may take longer than 10 minutes. See "
"https://platform.claude.com/docs/en/cli-sdks-libraries/sdks/python#long-requests for more details")
class Model:
"""记下每次调用的剧本模型。fail:每次调用都抛它;stream:流式时吐的块,吐到 boom 处抛 fail。"""
def __init__(self, name: str, *, fail: Exception | None = None, stream: list | None = None,
boom: int | None = None, max_tokens: int | None = None) -> None:
self.name = name
self.fail = fail
self.stream = stream if stream is not None else [AIMessageChunk(content=name)]
self.boom = boom
self.max_tokens = max_tokens
self.calls: list[tuple[str, object]] = []
self.bound: list[tuple[str, object, dict]] = []
async def ainvoke(self, messages, *args, **kwargs):
self.calls.append(("ainvoke", messages))
if self.fail is not None:
raise self.fail
return AIMessage(content=self.name)
async def astream(self, messages, *args, **kwargs):
self.calls.append(("astream", messages))
for i, chunk in enumerate(self.stream):
if self.boom is not None and i == self.boom:
raise self.fail
yield chunk
if self.fail is not None and self.boom is None:
raise self.fail
def bind_tools(self, tools, **kwargs):
self.bound.append(("tools", tools, kwargs))
return self
def with_structured_output(self, schema, **kwargs):
self.bound.append(("structured", schema, kwargs))
return self
# --------------------------------------------------------------------------
# 认出截断、认出嫌额度太大
# --------------------------------------------------------------------------
def test_cut_off_replies_are_recognised_for_both_families():
assert cut_off(AIMessage(content="", response_metadata={"finish_reason": "length"})) == "length"
assert cut_off(AIMessageChunk(content=[], response_metadata={"stop_reason": "max_tokens"})) == "max_tokens"
assert cut_off(AIMessage(content="好了", response_metadata={"finish_reason": "stop"})) == ""
assert cut_off(AIMessage(content="", response_metadata={"stop_reason": "end_turn"})) == ""
assert cut_off(None) == "" and cut_off({"finish_reason": "length"}) == ""
def test_length_errors_are_recognised_but_a_rejected_budget_is_not():
assert truncated_error(length_error())
assert truncated_error(RuntimeError("Could not parse response content as the length limit was reached"))
assert not truncated_error(None)
assert not truncated_error(RuntimeError("网关超时"))
assert not truncated_error(Rejected(TOO_LARGE)), "嫌额度太大是请求被拒,模型一个字都没写,不是截断"
@pytest.mark.parametrize("message", [
"Invalid max_tokens value, the valid range of max_tokens is [1, 8192]",
TOO_LARGE,
"max_completion_tokens is too large: 16384. This model supports at most 8192 completion tokens",
"This model's maximum context length is 32768 tokens. However, you requested 40369 tokens "
"(7601 in the messages, 32768 in the completion).",
])
def test_a_rejected_budget_is_recognised(message):
assert budget_rejected(Rejected(message))
assert budget_rejected(Rejected(message, status_code=422))
def test_other_rejections_are_not_a_budget_problem():
assert not budget_rejected(Rejected("model 'no-such-model' does not exist"))
assert not budget_rejected(Rejected(TOO_LARGE, status_code=500)), "服务方自己出错不是嫌额度大"
assert not budget_rejected(RuntimeError(TOO_LARGE)), "没有状态码的不认"
assert "max_tokens" in BUDGET_WORDS
def test_verify_and_judge_also_recognise_tgi_and_a_local_sdk_refusal():
"""核查、证据裁判的额度从 4096 提到了 16384,才撞得上这两种:自建服务说的是 max_new_tokens;Bearer 鉴权的 Anthropic
接入(客户端没设超时)调 Opus 4、4.1 这一代时,SDK 嫌 16384 太大、请求没发出去就抛 ValueError。"""
assert budget_refused(Rejected(TGI_TOO_LARGE, status_code=422))
assert budget_refused(ValueError(SDK_REFUSED))
assert budget_refused(Rejected(TOO_LARGE)), "原来认的照样认"
assert not budget_refused(ValueError("别的本地错误"))
assert not budget_refused(RuntimeError(SDK_REFUSED)), "只认 SDK 抛的 ValueError"
assert not budget_refused(Rejected("model 'no-such-model' does not exist"))
assert not budget_refused(Rejected(TGI_TOO_LARGE, status_code=500))
def test_the_copilot_wording_list_stays_as_it_was():
"""助手改图按 budget_rejected 认(流式调用,SDK 不在本地拦):认法不变,data/recipe_ai.py 的测试拿 BUDGET_WORDS 比对。"""
assert BUDGET_WORDS == ("max_tokens", "max_completion_tokens", "max_output_tokens", "maximum context length")
assert not budget_rejected(Rejected(TGI_TOO_LARGE, status_code=422))
assert not budget_rejected(ValueError(SDK_REFUSED))
def test_the_tokens_spent_by_a_cut_off_call_come_from_the_error():
assert error_usage(length_error(prompt=9000, completion=16384)) == (9000, 16384)
assert error_usage(RuntimeError("网关超时")) == (0, 0) and error_usage(None) == (0, 0)
def test_the_words_name_who_was_cut_off_and_the_limit():
assert cut_words("核查模型", 16384) == "核查模型的输出达到上限(16384 token)被截断"
assert cut_words("裁判模型", None) == "裁判模型的输出达到上限被截断", "额度不知道就不写数"
# --------------------------------------------------------------------------
# 退回小额度
# --------------------------------------------------------------------------
async def test_a_rejected_budget_is_retried_once_with_the_small_one_and_stays_small():
big, small = Model("big", fail=Rejected(TOO_LARGE)), Model("small")
model = BudgetFallback(big, small)
assert (await model.ainvoke(["问"])).content == "small"
assert (await model.ainvoke(["再问"])).content == "small", "之后的调用都用小的,不再先撞一次"
assert len(big.calls) == 1 and len(small.calls) == 2
assert model.model is small
async def test_other_errors_are_not_retried():
big, small = Model("big", fail=Rejected("model 'no-such-model' does not exist")), Model("small")
with pytest.raises(Rejected):
await BudgetFallback(big, small).ainvoke(["问"])
assert small.calls == []
async def test_the_small_model_is_only_built_when_needed():
built: list[str] = []
small = Model("small")
async def build():
built.append("small")
return small
fine = BudgetFallback(Model("big"), build=build, limits=(16384, 4096))
assert (await fine.ainvoke(["问"])).content == "big"
assert built == [] and fine.max_tokens == 16384, "平常一个小额度的模型都不多建"
rejecting = BudgetFallback(Model("big", fail=Rejected(TOO_LARGE)), build=build, limits=(16384, 4096))
assert (await rejecting.ainvoke(["问"])).content == "small"
assert (await rejecting.ainvoke(["再问"])).content == "small"
assert built == ["small"], "退回之后建一次,之后沿用"
assert rejecting.max_tokens == 4096
async def test_without_a_small_model_the_rejection_is_raised():
with pytest.raises(Rejected):
await BudgetFallback(Model("big", fail=Rejected(TOO_LARGE))).ainvoke(["问"])
same = Model("same", fail=Rejected(TOO_LARGE))
with pytest.raises(Rejected):
await BudgetFallback(same, same).ainvoke(["问"])
assert len(same.calls) == 1, "大小是同一个模型:重来也是同样的额度,不白撞第二次"
async def test_bound_and_structured_views_share_the_switch():
"""核查的工具循环里退回过小额度,收尾那次直接用小的;包出来的视图各自转给对应的小模型。"""
big, small = Model("big", fail=Rejected(TOO_LARGE)), Model("small")
model = BudgetFallback(big, small, limits=(16384, 4096))
tools = model.bind_tools(["db_query"], parallel_tool_calls=False)
assert big.bound == [("tools", ["db_query"], {"parallel_tool_calls": False})]
assert small.bound == [("tools", ["db_query"], {"parallel_tool_calls": False})]
assert (await tools.ainvoke(["查"])).content == "small"
structured = model.with_structured_output({"type": "object"}, include_raw=True)
assert (await structured.ainvoke(["交结果"])).content == "small"
assert len(big.calls) == 1, "收尾那次不再先按大额度撞一次"
assert structured.max_tokens == 4096 and model.max_tokens == 4096
async def test_a_sibling_shares_the_switch_and_builds_its_own_small_model():
"""收尾用的是另一份配置(关掉思考)建的模型:和工具循环的模型共用开关,要退回时按自己的配置建小的。"""
built: list[str] = []
loop_small, settle_small = Model("loop-small"), Model("settle-small")
async def build_loop():
built.append("loop")
return loop_small
async def build_settle():
built.append("settle")
return settle_small
loop = BudgetFallback(Model("loop-big", fail=Rejected(TOO_LARGE)), build=build_loop, limits=(16384, 4096))
assert (await loop.ainvoke(["查"])).content == "loop-small"
settle_big = Model("settle-big")
settle = BudgetFallback(settle_big, build=build_settle, limits=(16384, 4096), shared_with=loop)
assert (await settle.with_structured_output({"type": "object"}).ainvoke(["交结果"])).content == "settle-small"
assert settle_big.calls == [] and built == ["loop", "settle"]
async def test_streaming_falls_back_only_before_anything_came_out():
big = Model("big", fail=Rejected(TOO_LARGE), stream=[])
small = Model("small", stream=[AIMessageChunk(content="a"), AIMessageChunk(content="b")])
model = BudgetFallback(big, small)
assert [c.content async for c in model.astream(["问"])] == ["a", "b"]
assert len(big.calls) == 1 and len(small.calls) == 1
started = Model("big", fail=Rejected(TOO_LARGE), stream=[AIMessageChunk(content="x"), AIMessageChunk(content="y")],
boom=1)
unused = Model("small")
got: list[str] = []
with pytest.raises(Rejected):
async for chunk in BudgetFallback(started, unused).astream(["问"]):
got.append(chunk.content)
assert got == ["x"] and unused.calls == [], "已经吐出东西了:重来会重复输出,原样抛出"
async def test_without_limits_max_tokens_reads_the_current_model():
big = Model("big", fail=Rejected(TOO_LARGE), max_tokens=32768)
small = Model("small", max_tokens=8192)
model = BudgetFallback(big, small)
assert model.max_tokens == 32768
await model.ainvoke(["问"])
assert model.max_tokens == 8192
async def test_a_wider_check_can_be_given():
"""核查、证据裁判按 budget_refused 认:服务方说的是 max_new_tokens 也退回小额度。缺省(助手改图)不认。"""
tgi = Rejected(TGI_TOO_LARGE, status_code=422)
model = BudgetFallback(Model("big", fail=tgi), Model("small"), rejected=budget_refused)
assert (await model.ainvoke(["问"])).content == "small"
with pytest.raises(Rejected):
await BudgetFallback(Model("big", fail=tgi), Model("small")).ainvoke(["问"])
local = BudgetFallback(Model("big", fail=ValueError(SDK_REFUSED)), Model("small"), rejected=budget_refused)
structured = local.with_structured_output({"type": "object"})
assert (await structured.ainvoke(["交结果"])).content == "small", "包出来的视图照样按这个认"
async def test_a_rejection_is_remembered_for_later_calls_to_the_same_model():
"""双盲复算的几路盲算、反方审稿的每一句、每次质询、每次点击裁判,各是一组新建的调用:同一接入、模型、额度被嫌大
过一次,之后新建的直接用小的,不再先撞一次。别的模型、别的额度照旧先按大的试。"""
built: list[str] = []
def builder(name: str):
async def build():
built.append(name)
return Model(name)
return build
key = ("gw", "reasoner-x", 16384)
first = BudgetFallback(Model("big", fail=Rejected(TOO_LARGE)), build=builder("small-1"), limits=(16384, 4096),
key=key)
assert first.max_tokens == 16384
assert (await first.ainvoke(["问"])).content == "small-1"
later_big = Model("big", fail=Rejected(TOO_LARGE))
later = BudgetFallback(later_big, build=builder("small-2"), limits=(16384, 4096), key=key)
assert later.max_tokens == 4096, "还没调用就知道要用小的:截断时报的额度对得上"
assert (await later.bind_tools(["db_query"]).ainvoke(["再问"])).content == "small-2"
assert later_big.calls == [], "不再先按大额度撞一次"
other = Model("other-big")
assert (await BudgetFallback(other, build=builder("x"), key=("gw", "other", 16384)).ainvoke(["问"])).content \
== "other-big"
bigger = Model("bigger")
assert (await BudgetFallback(bigger, build=builder("y"), key=("gw", "reasoner-x", 32768)).ainvoke(["问"])).content \
== "bigger"
assert built == ["small-1", "small-2"]
async def test_without_a_key_nothing_is_remembered():
"""助手改图不给 key:每次改图照旧先按大额度试(行为和挪过来之前一样)。"""
await BudgetFallback(Model("big", fail=Rejected(TOO_LARGE)), Model("small")).ainvoke(["问"])
assert budget._REJECTED == set()
fresh = Model("big")
assert (await BudgetFallback(fresh, Model("small")).ainvoke(["问"])).content == "big"
async def test_a_streaming_rejection_is_remembered_too():
key = ("gw", "streamer", 16384)
big = Model("big", fail=Rejected(TOO_LARGE), stream=[])
model = BudgetFallback(big, Model("small", stream=[AIMessageChunk(content="a")]), key=key)
assert [c.content async for c in model.astream(["问"])] == ["a"]
again = Model("big", stream=[AIMessageChunk(content="b")])
assert [c.content async for c in BudgetFallback(again, Model("small", stream=[AIMessageChunk(content="c")]),
key=key).astream(["问"])] == ["c"]
assert again.calls == []