mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-30 04:15:45 +08:00
82 lines
3.1 KiB
Python
82 lines
3.1 KiB
Python
"""Run-local token counts retain provider semantics without duplicate Host RPCs."""
|
|
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from langbot_plugin.api.entities.builtin.provider.message import Message
|
|
from langbot_plugin.api.entities.builtin.runner.resources import ModelResource
|
|
|
|
from pkg.context_pipeline import ContextTokenCounterRequiredError, HostContextTokenCounter
|
|
from pkg.run_assembly import AgentRunAssembler
|
|
from tests.test_runner import FakeRunnerAPIProxy, make_context
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_input_keeps_provider_count_and_reuse_is_run_local():
|
|
api = FakeRunnerAPIProxy()
|
|
api.count_tokens = AsyncMock(return_value=7)
|
|
counter = HostContextTokenCounter(api, "model-1")
|
|
|
|
assert await counter.count([]) == 7
|
|
assert await counter.count([]) == 7
|
|
api.count_tokens.assert_awaited_once_with(
|
|
llm_model_uuid="model-1", messages=[], funcs=[], reasoning_level="provider_default"
|
|
)
|
|
|
|
assert await HostContextTokenCounter(api, "model-1").count([]) == 7
|
|
assert api.count_tokens.await_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_count_reuse_tracks_message_content_model_and_reasoning():
|
|
api = FakeRunnerAPIProxy()
|
|
api.count_tokens = AsyncMock(side_effect=[11, 12, 13, 14, 15])
|
|
counter = HostContextTokenCounter(api, "model-1")
|
|
message = Message(role="user", content="original")
|
|
|
|
assert await counter.count([message]) == 11
|
|
assert await counter.count([message.model_copy(deep=True)]) == 11
|
|
message.content = "changed"
|
|
assert await counter.count([message]) == 12
|
|
message.provider_specific_fields = {"signature": "changed"}
|
|
assert await counter.count([message]) == 13
|
|
counter.model_id = "model-2"
|
|
assert await counter.count([message]) == 14
|
|
counter.reasoning_level = "high"
|
|
assert await counter.count([message]) == 15
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("failure", [True, -1, "7", RuntimeError("Host unavailable")])
|
|
async def test_failed_or_invalid_counts_are_not_reused(failure):
|
|
api = FakeRunnerAPIProxy()
|
|
api.count_tokens = AsyncMock(side_effect=[failure, 9])
|
|
counter = HostContextTokenCounter(api, "model-1")
|
|
|
|
with pytest.raises((ContextTokenCounterRequiredError, RuntimeError)):
|
|
await counter.count([])
|
|
assert await counter.count([]) == 9
|
|
assert api.count_tokens.await_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initial_assembly_and_model_hook_share_counts():
|
|
api = FakeRunnerAPIProxy(models=[ModelResource(model_id="model-1")])
|
|
context = make_context(
|
|
config={
|
|
"model": {"primary": "model-1"},
|
|
"prompt": [{"role": "system", "content": "Keep the sentinel"}],
|
|
"date-grounding": False,
|
|
},
|
|
input_text="request sentinel",
|
|
)
|
|
assembly = await AgentRunAssembler(api, context).assemble()
|
|
prepared = await assembly.hooks.prepare_model_turn(assembly.messages)
|
|
|
|
assert [(message.role, message.content) for message in prepared] == [
|
|
("system", "Keep the sentinel"),
|
|
("user", "request sentinel"),
|
|
]
|
|
assert api.count_tokens.await_count <= 4
|
|
assert any(call.kwargs["messages"] == [] for call in api.count_tokens.await_args_list)
|