v2.11.1 (#721)
Co-authored-by: Clément Drouin <clement.drouin@mistral.ai> Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai> Co-authored-by: Guillaume LE GOFF <guillaume.lgf@gmail.com> Co-authored-by: Mert Unsal <mertunsal1905@gmail.com> Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com> Co-authored-by: Quentin <quentin.torroba@mistral.ai> Co-authored-by: Val <102326092+vdeva@users.noreply.github.com> Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com> Co-authored-by: maximevoisin-pm <maxime.voisin@mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
parent
adb1ca74ce
commit
cf3f4ca58f
143 changed files with 3457 additions and 1351 deletions
|
|
@ -14,6 +14,7 @@ from vibe.core.middleware import (
|
|||
MiddlewarePipeline,
|
||||
ReadOnlyAgentMiddleware,
|
||||
ResetReason,
|
||||
TokenLimitMiddleware,
|
||||
make_plan_agent_reminder,
|
||||
)
|
||||
from vibe.core.types import AgentStats, MessageList
|
||||
|
|
@ -634,3 +635,31 @@ class TestReadOnlyAgentMiddlewareIntegration:
|
|||
# 8. Stay in default: no injection
|
||||
r = await plan_middleware.before_turn(_ctx())
|
||||
assert r.action == MiddlewareAction.CONTINUE
|
||||
|
||||
|
||||
class TestTokenLimitMiddleware:
|
||||
@pytest.mark.asyncio
|
||||
async def test_stops_when_session_total_tokens_exceeds_limit(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
middleware = TokenLimitMiddleware(14)
|
||||
ctx.stats.session_prompt_tokens = 10
|
||||
ctx.stats.session_completion_tokens = 5
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
|
||||
assert result.action == MiddlewareAction.STOP
|
||||
assert result.reason == "Token limit exceeded: 15 > 14"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_allows_when_session_total_tokens_matches_limit(
|
||||
self, ctx: ConversationContext
|
||||
) -> None:
|
||||
middleware = TokenLimitMiddleware(15)
|
||||
ctx.stats.session_prompt_tokens = 10
|
||||
ctx.stats.session_completion_tokens = 5
|
||||
|
||||
result = await middleware.before_turn(ctx)
|
||||
|
||||
assert result.action == MiddlewareAction.CONTINUE
|
||||
assert result.reason is None
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue