mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:41:59 +00:00
输出额度的通用逻辑挪到 providers/budget.py：认截断（finish_reason=length、stop_reason=max_tokens、openai 的 LengthFinishReasonError）、认服务方嫌额度太大（400/422 且报错提到 max_tokens 一类）、截断那次的用量、给人看的「某某模型的输出达到上限（N token）被截断」，以及退回小额度的 BudgetFallback
BudgetFallback 在助手原来的做法上补齐核查和裁判要用的：非流式 ainvoke 也认；bind_tools、with_structured_output 包出来的和原来的共用同一个开关，另一份配置建的模型（核查收尾时关掉思考的那个）经 shared_with 也共用，工具循环里退回过，收尾那次直接用小的；小额度的模型可以要退回时才建，平常一个都不多建。只在还没吐出任何东西时重来，和原来一样。 助手改为引用它，_cut_off、_BudgetFallback、_BUDGET_WORDS 留着老名字，行为不变，助手和配方助手的测试照样通过。
This commit is contained in:
ThearchyHelios
@@ -27,6 +27,10 @@ from app.db.models import Workflow
|
||||
from app.engine.layout import CORRIDOR_MIN, MIN_ROW_GAP, NODE_W, _height, auto_layout
|
||||
from app.engine.schema import GraphSpec, NodeType, ValidationIssue, text_parse_issues, validate_graph
|
||||
from app.engine.state import message_text
|
||||
# 额度用完的认法、退回小额度的做法挪到了 providers/budget.py（核查、证据裁判也用），这里留着老名字；
|
||||
# _BUDGET_WORDS 另有 data/recipe_ai.py 的测试拿来比对
|
||||
from app.providers.budget import BUDGET_WORDS as _BUDGET_WORDS # noqa: F401
|
||||
from app.providers.budget import BudgetFallback as _BudgetFallback, cut_off as _cut_off
|
||||
from app.providers.factory import ModelSpec, ProviderNotConfigured, get_chat_model
|
||||
from app.tools.registry import all_specs
|
||||
|
||||
@@ -1053,54 +1057,12 @@ class CopilotIncomplete(Exception):
|
||||
self.detail = detail
|
||||
|
||||
|
||||
def _cut_off(chunk: Any) -> str:
|
||||
"""这一块是不是在说「额度用完了」：OpenAI 系是 finish_reason=length，Anthropic 是 stop_reason=max_tokens。"""
|
||||
meta = getattr(chunk, "response_metadata", None) or {}
|
||||
reason = meta.get("finish_reason") or meta.get("stop_reason")
|
||||
return str(reason) if reason in ("length", "max_tokens") else ""
|
||||
|
||||
|
||||
#: 流式改图的输出额度。会思考的模型（DeepSeek、Claude 的 adaptive thinking）思考也算在里面：
|
||||
#: 8192 时复杂的改图常常想完就没额度了，正文一个字没写（2026-09 部署实例上同类请求 5 轮坏了 4 轮）
|
||||
COPILOT_MAX_TOKENS = 32768
|
||||
#: 服务方不收大额度时退回这个（原来的默认值）
|
||||
#: 服务方不收大额度时退回这个（原来的默认值）。退回的做法（_BudgetFallback：只在还没吐出任何东西时重来，
|
||||
#: 之后的调用即自查修正轮都用小的）和截断的认法（_cut_off）在 providers/budget.py，核查、证据裁判也用它
|
||||
_FALLBACK_MAX_TOKENS = 8192
|
||||
#: 服务方嫌额度太大时报错里会出现的字样：超过模型的输出上限（DeepSeek chat 系只到 8K），
|
||||
#: 或者输入加额度超过上下文长度（本地 vLLM）
|
||||
_BUDGET_WORDS = ("max_tokens", "max_completion_tokens", "max_output_tokens", "maximum context length")
|
||||
|
||||
|
||||
def _budget_rejected(e: BaseException) -> bool:
|
||||
status = getattr(e, "status_code", None) or getattr(getattr(e, "response", None), "status_code", None)
|
||||
return status in (400, 422) and any(w in str(e).lower() for w in _BUDGET_WORDS)
|
||||
|
||||
|
||||
class _BudgetFallback:
|
||||
"""先按大额度调；服务方一开口就嫌额度太大时按原来的额度重来，之后的调用（自查修正轮）都用小的。
|
||||
|
||||
只在还没吐出任何东西时重来：这类拒绝都发生在请求刚发出去的时候，重来不会重复输出。
|
||||
"""
|
||||
|
||||
def __init__(self, big: Any, small: Any) -> None:
|
||||
self.model, self._small = big, small
|
||||
|
||||
@property
|
||||
def max_tokens(self) -> Any:
|
||||
return getattr(self.model, "max_tokens", None)
|
||||
|
||||
async def astream(self, messages: list[Any]):
|
||||
started = False
|
||||
try:
|
||||
async for chunk in self.model.astream(messages):
|
||||
started = True
|
||||
yield chunk
|
||||
return
|
||||
except Exception as e:
|
||||
if started or self.model is self._small or not _budget_rejected(e):
|
||||
raise
|
||||
self.model = self._small
|
||||
async for chunk in self.model.astream(messages):
|
||||
yield chunk
|
||||
|
||||
|
||||
async def _op_model(session: AsyncSession, payload: GenerateIn) -> tuple[Any, str]:
|
||||
|
||||
@@ -0,0 +1,168 @@
|
||||
"""模型调用的输出额度：认出「输出到了上限被截断」，服务方嫌额度太大时退回小额度重来。
|
||||
|
||||
助手改图（api/copilot.py）、核查（engine/verify_agent.py：双盲复算、反方审稿、审批质询共用）和证据裁判（engine/judge.py）
|
||||
共用这一份。会思考的模型思考也算在输出额度里：经 OpenAI 兼容接入的推理模型（DeepSeek、GLM、Kimi 这一类）收不到
|
||||
thinking=off，Claude 开了 adaptive thinking 也一样。额度小了，想完就没额度写正文——回复是空的，或者 JSON 只写了一半。
|
||||
这时要照实说是输出额度不够，不能当成「模型没按格式交」「没查到」。
|
||||
|
||||
- cut_off：这一块（流式）或这条回复是不是因为额度用完停下的
|
||||
- truncated_error：这个异常是不是在说同一件事（openai 结构化输出解析时抛的 LengthFinishReasonError）
|
||||
- budget_rejected：服务方一开口就嫌额度太大
|
||||
- error_usage：截断的那次调用花掉的用量
|
||||
- cut_words：给人看的「某某模型的输出达到上限（N token）被截断」
|
||||
- BudgetFallback：先按大额度调，被嫌大时按小额度重来，之后都用小的
|
||||
|
||||
engine/nodes/llm.py（cite_fields 的抽取）、data/recipe_ai.py 各有一份同样的判断，是在这里之前写的，行为相同。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
#: 服务方嫌额度太大时报错里会出现的字样：超过模型的输出上限（DeepSeek chat 系只到 8K），
|
||||
#: 或者输入加额度超过上下文长度（本地 vLLM）
|
||||
BUDGET_WORDS = ("max_tokens", "max_completion_tokens", "max_output_tokens", "maximum context length")
|
||||
#: 截断类报错的原话：openai 的 LengthFinishReasonError（结构化输出解析时 finish_reason=length）
|
||||
LENGTH_WORDS = ("length limit was reached", "finish_reason=length", "finish_reason='length'")
|
||||
|
||||
|
||||
def cut_off(chunk: Any) -> str:
|
||||
"""这一块是不是在说「额度用完了」：OpenAI 系是 finish_reason=length，Anthropic 是 stop_reason=max_tokens。"""
|
||||
meta = getattr(chunk, "response_metadata", None) or {}
|
||||
reason = meta.get("finish_reason") or meta.get("stop_reason")
|
||||
return str(reason) if reason in ("length", "max_tokens") else ""
|
||||
|
||||
|
||||
def truncated_error(error: BaseException | None) -> bool:
|
||||
"""这个异常是不是在说「输出被截断」：openai 的 LengthFinishReasonError，或者 langchain、网关转述的同一件事。
|
||||
|
||||
服务方嫌额度太大的 400（报错里也有 max_tokens）不算：那是请求被拒，模型一个字都没写。
|
||||
"""
|
||||
if error is None:
|
||||
return False
|
||||
if type(error).__name__ == "LengthFinishReasonError":
|
||||
return True
|
||||
text = str(error).lower()
|
||||
return any(word in text for word in LENGTH_WORDS)
|
||||
|
||||
|
||||
def budget_rejected(e: BaseException) -> bool:
|
||||
status = getattr(e, "status_code", None) or getattr(getattr(e, "response", None), "status_code", None)
|
||||
return status in (400, 422) and any(w in str(e).lower() for w in BUDGET_WORDS)
|
||||
|
||||
|
||||
def error_usage(error: BaseException | None) -> tuple[int, int]:
|
||||
"""失败的那次调用花掉的 (输入, 输出) token。截断时 openai 的异常带着 completion.usage：额度是真花出去了，要记账，
|
||||
不然每日上限会漏掉最贵的那几次。别的异常没有用量，是 (0, 0)。"""
|
||||
usage = getattr(getattr(error, "completion", None), "usage", None)
|
||||
return int(getattr(usage, "prompt_tokens", 0) or 0), int(getattr(usage, "completion_tokens", 0) or 0)
|
||||
|
||||
|
||||
def cut_words(who: str, limit: int | None) -> str:
|
||||
"""「核查模型的输出达到上限（16384 token）被截断」。额度不知道（没给 max_tokens、按接入的默认值）就不写数。"""
|
||||
return f"{who}的输出达到上限{f'（{limit} token）' if limit else ''}被截断"
|
||||
|
||||
|
||||
class _Switch:
|
||||
"""几次调用共用的「退回小额度了没有」：核查的工具循环和收尾、裁判的各批。"""
|
||||
|
||||
__slots__ = ("small",)
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.small = False
|
||||
|
||||
|
||||
class BudgetFallback:
|
||||
"""先按大额度调；服务方一开口就嫌额度太大时按小额度重来，之后的调用都用小的。
|
||||
|
||||
只在还没吐出任何东西时重来：这类拒绝都发生在请求刚发出去的时候，重来不会重复输出。流式（astream）、非流式
|
||||
（ainvoke）都认。bind_tools、with_structured_output 照样包一层，包出来的和原来的共用同一个开关；shared_with 给了，
|
||||
另一份配置建的模型（核查收尾时关掉思考的那个）也和它共用：工具循环里退回过，收尾那次直接用小的，不再先撞一次。
|
||||
|
||||
small 是现成的小额度模型；也可以不给，给 build（无参的协程函数，返回小额度的模型），要退回时才建，平常一个都不多建。
|
||||
limits 是 (大额度, 小额度)：给了，max_tokens 就按它报（包过一层的模型身上不一定读得到额度）；不给就读眼下那个模型的
|
||||
max_tokens。
|
||||
"""
|
||||
|
||||
def __init__(self, big: Any, small: Any = None, *, build: Callable[[], Awaitable[Any]] | None = None,
|
||||
limits: tuple[int | None, int | None] | None = None,
|
||||
shared_with: BudgetFallback | None = None) -> None:
|
||||
self._big = big
|
||||
self._small = small
|
||||
self._build = build
|
||||
self._limits = limits
|
||||
self._switch = shared_with._switch if shared_with is not None else _Switch()
|
||||
|
||||
@property
|
||||
def model(self) -> Any:
|
||||
"""眼下用的那个：退回过、小额度的也建好了就是小的，否则是大的。"""
|
||||
return self._small if self._switch.small and self._small is not None else self._big
|
||||
|
||||
@property
|
||||
def max_tokens(self) -> Any:
|
||||
if self._limits is not None:
|
||||
return self._limits[1] if self._switch.small else self._limits[0]
|
||||
return getattr(self.model, "max_tokens", None)
|
||||
|
||||
def _can_fall_back(self, used: Any) -> bool:
|
||||
"""刚才那次用的是大额度、又有小额度可退。大小是同一个模型时重来也是同样的额度，不算可退。"""
|
||||
if used is not self._big:
|
||||
return False
|
||||
return self._build is not None if self._small is None else self._small is not self._big
|
||||
|
||||
async def _small_model(self) -> Any:
|
||||
if self._small is None and self._build is not None:
|
||||
built = await self._build()
|
||||
if self._small is None: # 并发的几次调用同时要建时，留先建好的那个
|
||||
self._small = built
|
||||
return self._small
|
||||
|
||||
async def _current(self) -> Any:
|
||||
if self._switch.small and (small := await self._small_model()) is not None:
|
||||
return small
|
||||
return self._big
|
||||
|
||||
async def ainvoke(self, messages: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
used = await self._current()
|
||||
try:
|
||||
return await used.ainvoke(messages, *args, **kwargs)
|
||||
except Exception as e:
|
||||
if not self._can_fall_back(used) or not budget_rejected(e):
|
||||
raise
|
||||
self._switch.small = True
|
||||
return await (await self._current()).ainvoke(messages, *args, **kwargs)
|
||||
|
||||
async def astream(self, messages: Any, *args: Any, **kwargs: Any):
|
||||
used = await self._current()
|
||||
started = False
|
||||
try:
|
||||
async for chunk in used.astream(messages, *args, **kwargs):
|
||||
started = True
|
||||
yield chunk
|
||||
return
|
||||
except Exception as e:
|
||||
if started or not self._can_fall_back(used) or not budget_rejected(e):
|
||||
raise
|
||||
self._switch.small = True
|
||||
async for chunk in (await self._current()).astream(messages, *args, **kwargs):
|
||||
yield chunk
|
||||
|
||||
def _derived(self, make: Callable[[Any], Any]) -> BudgetFallback:
|
||||
"""对大小两个模型做同一件事（绑工具、要结构化输出），包成共用开关的一层。小的还没建时，建好再做。"""
|
||||
small, build = None, None
|
||||
if self._small is not None:
|
||||
small = make(self._small)
|
||||
elif self._build is not None:
|
||||
async def build() -> Any:
|
||||
return make(await self._small_model())
|
||||
return BudgetFallback(make(self._big), small, build=build, limits=self._limits, shared_with=self)
|
||||
|
||||
def bind_tools(self, tools: Any, **kwargs: Any) -> BudgetFallback:
|
||||
return self._derived(lambda m: m.bind_tools(tools, **kwargs))
|
||||
|
||||
def with_structured_output(self, schema: Any, **kwargs: Any) -> BudgetFallback:
|
||||
return self._derived(lambda m: m.with_structured_output(schema, **kwargs))
|
||||
|
||||
|
||||
__all__ = ["BUDGET_WORDS", "LENGTH_WORDS", "BudgetFallback", "budget_rejected", "cut_off", "cut_words", "error_usage",
|
||||
"truncated_error"]
|
||||
@@ -0,0 +1,244 @@
|
||||
"""模型调用的输出额度（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 包出来的和原来的共用同一个开关；只在还没吐出任何东西时重来；小额度的模型要退回时才建
|
||||
|
||||
模型都是剧本，不碰真接口。
|
||||
"""
|
||||
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.budget import (
|
||||
BUDGET_WORDS,
|
||||
BudgetFallback,
|
||||
budget_rejected,
|
||||
cut_off,
|
||||
cut_words,
|
||||
error_usage,
|
||||
truncated_error,
|
||||
)
|
||||
|
||||
|
||||
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"
|
||||
|
||||
|
||||
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_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
|
||||
Reference in New Issue
Block a user