mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-21 03:46:11 +00:00
test(provider): support fake token counting
This commit is contained in:
@@ -7,6 +7,8 @@ without calling real LLM APIs or network requests.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from unittest.mock import AsyncMock, Mock
|
from unittest.mock import AsyncMock, Mock
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
@@ -30,6 +32,25 @@ class FakeProviderAPIRequester(requester.ProviderAPIRequester):
|
|||||||
self._invoke_count = 0
|
self._invoke_count = 0
|
||||||
self._last_messages = None
|
self._last_messages = None
|
||||||
self._last_model = None
|
self._last_model = None
|
||||||
|
self._last_count_tokens_payload = None
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _content_to_text(content) -> str:
|
||||||
|
if content is None:
|
||||||
|
return ''
|
||||||
|
if isinstance(content, str):
|
||||||
|
return content
|
||||||
|
if isinstance(content, list):
|
||||||
|
parts: list[str] = []
|
||||||
|
for item in content:
|
||||||
|
if isinstance(item, dict):
|
||||||
|
text = item.get('text')
|
||||||
|
else:
|
||||||
|
text = getattr(item, 'text', None)
|
||||||
|
if text:
|
||||||
|
parts.append(str(text))
|
||||||
|
return ''.join(parts)
|
||||||
|
return str(content)
|
||||||
|
|
||||||
async def invoke_llm(
|
async def invoke_llm(
|
||||||
self,
|
self,
|
||||||
@@ -70,6 +91,37 @@ class FakeProviderAPIRequester(requester.ProviderAPIRequester):
|
|||||||
content=[provider_message.ContentElement(type='text', text='Fake stream chunk')],
|
content=[provider_message.ContentElement(type='text', text='Fake stream chunk')],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def count_tokens(
|
||||||
|
self,
|
||||||
|
model: requester.RuntimeLLMModel,
|
||||||
|
messages: list,
|
||||||
|
funcs=None,
|
||||||
|
extra_args={},
|
||||||
|
) -> int:
|
||||||
|
"""Return deterministic token estimates for token-free integration tests."""
|
||||||
|
payload: list[dict] = []
|
||||||
|
for message in messages:
|
||||||
|
payload.append(
|
||||||
|
{
|
||||||
|
'role': getattr(message, 'role', ''),
|
||||||
|
'content': self._content_to_text(getattr(message, 'content', None)),
|
||||||
|
'tool_calls': getattr(message, 'tool_calls', None),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
for func in funcs or []:
|
||||||
|
payload.append(
|
||||||
|
{
|
||||||
|
'name': getattr(func, 'name', ''),
|
||||||
|
'description': getattr(func, 'description', ''),
|
||||||
|
'parameters': getattr(func, 'parameters', {}),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
self._last_count_tokens_payload = payload
|
||||||
|
text = json.dumps(payload, ensure_ascii=False, sort_keys=True, default=str)
|
||||||
|
return max(1, (len(text) + 3) // 4)
|
||||||
|
|
||||||
async def invoke_embedding(self, model, input_text: list, extra_args={}):
|
async def invoke_embedding(self, model, input_text: list, extra_args={}):
|
||||||
"""Return fake embedding vectors."""
|
"""Return fake embedding vectors."""
|
||||||
return [[0.1, 0.2, 0.3] for _ in input_text]
|
return [[0.1, 0.2, 0.3] for _ in input_text]
|
||||||
|
|||||||
@@ -0,0 +1,57 @@
|
|||||||
|
"""Tests for provider test doubles used by integration paths."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from langbot.pkg.entity.persistence import model as persistence_model
|
||||||
|
from langbot.pkg.provider.modelmgr import requester
|
||||||
|
from langbot_plugin.api.entities.builtin.provider import message as provider_message
|
||||||
|
from langbot_plugin.api.entities.builtin.resource import tool as resource_tool
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fake_requester_counts_messages_and_tools(runtime_provider):
|
||||||
|
"""Fake requester should support token-free AgentRunner context budgeting."""
|
||||||
|
runtime_model = requester.RuntimeLLMModel(
|
||||||
|
model_entity=persistence_model.LLMModel(
|
||||||
|
uuid='fake-count-model',
|
||||||
|
name='fake-count-model',
|
||||||
|
provider_uuid=runtime_provider.provider_entity.uuid,
|
||||||
|
abilities=['func_call'],
|
||||||
|
extra_args={},
|
||||||
|
),
|
||||||
|
provider=runtime_provider,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _placeholder_func(**kwargs):
|
||||||
|
return kwargs
|
||||||
|
|
||||||
|
messages = [
|
||||||
|
provider_message.Message(role='system', content='You are a test assistant.'),
|
||||||
|
provider_message.Message(
|
||||||
|
role='user',
|
||||||
|
content=[
|
||||||
|
provider_message.ContentElement(type='text', text='hello'),
|
||||||
|
provider_message.ContentElement(type='text', text=' world'),
|
||||||
|
],
|
||||||
|
),
|
||||||
|
]
|
||||||
|
tools = [
|
||||||
|
resource_tool.LLMTool(
|
||||||
|
name='lookup',
|
||||||
|
human_desc='Lookup',
|
||||||
|
description='Lookup a value',
|
||||||
|
parameters={'type': 'object', 'properties': {'query': {'type': 'string'}}},
|
||||||
|
func=_placeholder_func,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
requester_inst = runtime_provider.requester
|
||||||
|
message_tokens = await requester_inst.count_tokens(runtime_model, messages, funcs=[])
|
||||||
|
message_and_tool_tokens = await requester_inst.count_tokens(runtime_model, messages, funcs=tools)
|
||||||
|
|
||||||
|
assert message_tokens > 0
|
||||||
|
assert message_and_tool_tokens > message_tokens
|
||||||
|
assert requester_inst._last_count_tokens_payload[-1]['name'] == 'lookup'
|
||||||
|
assert requester_inst._last_count_tokens_payload[1]['content'] == 'hello world'
|
||||||
Reference in New Issue
Block a user