feat(agent-runner): support scoped token counting

This commit is contained in:
huanghuoguoguo
2026-06-27 01:31:08 +08:00
parent ae49753f74
commit d0f6fe2cec
10 changed files with 302 additions and 15 deletions
+1 -1
View File
@@ -77,7 +77,7 @@ def make_session(
}
authorized_operations: dict[str, dict[str, set[str]]] = {
'model': {
m.get('model_id'): set(m.get('operations') or ['invoke', 'stream', 'rerank'])
m.get('model_id'): set(m.get('operations') or ['invoke', 'stream', 'rerank', 'count_tokens'])
for m in res.get('models', [])
if m.get('model_id')
},
@@ -14,7 +14,7 @@ from langbot.pkg.agent.runner.resource_builder import AgentResourceBuilder
RUNNER_ID = 'plugin:test/runner/default'
FULL_PERMISSIONS = {
'models': ['invoke', 'stream', 'rerank'],
'models': ['count_tokens', 'invoke', 'stream', 'rerank'],
'tools': ['detail', 'call'],
'knowledge_bases': ['list', 'retrieve'],
'history': ['page', 'search'],
@@ -139,9 +139,24 @@ async def test_build_models_authorizes_config_declared_llm_and_rerank_models(app
resources = await build_resources(app, query, descriptor)
assert resources['models'] == [
{'model_id': 'primary', 'model_type': 'llm', 'provider': 'test-provider', 'operations': ['invoke', 'stream']},
{'model_id': 'fallback', 'model_type': 'llm', 'provider': 'test-provider', 'operations': ['invoke', 'stream']},
{'model_id': 'aux', 'model_type': 'llm', 'provider': 'aux-provider', 'operations': ['invoke', 'stream']},
{
'model_id': 'primary',
'model_type': 'llm',
'provider': 'test-provider',
'operations': ['invoke', 'stream', 'count_tokens'],
},
{
'model_id': 'fallback',
'model_type': 'llm',
'provider': 'test-provider',
'operations': ['invoke', 'stream', 'count_tokens'],
},
{
'model_id': 'aux',
'model_type': 'llm',
'provider': 'aux-provider',
'operations': ['invoke', 'stream', 'count_tokens'],
},
{'model_id': 'rerank', 'model_type': 'rerank', 'provider': 'rerank-provider', 'operations': ['rerank']},
]
@@ -189,7 +204,12 @@ async def test_build_models_authorizes_rerank_and_llm_refs_from_config(app):
resources = await build_resources(app, query, descriptor)
assert resources['models'] == [
{'model_id': 'llm', 'model_type': 'llm', 'provider': 'test-provider', 'operations': ['invoke', 'stream']},
{
'model_id': 'llm',
'model_type': 'llm',
'provider': 'test-provider',
'operations': ['invoke', 'stream', 'count_tokens'],
},
{'model_id': 'rerank', 'model_type': 'rerank', 'provider': 'rerank-provider', 'operations': ['rerank']},
]
@@ -222,7 +242,12 @@ async def test_build_resources_accepts_dynamic_form_type_aliases(app):
resources = await build_resources(app, query, descriptor)
assert resources['models'] == [
{'model_id': 'llm_alias', 'model_type': 'llm', 'provider': 'test-provider', 'operations': ['invoke', 'stream']},
{
'model_id': 'llm_alias',
'model_type': 'llm',
'provider': 'test-provider',
'operations': ['invoke', 'stream', 'count_tokens'],
},
]
assert resources['knowledge_bases'] == [
{'kb_id': 'kb_alias', 'kb_name': 'name-kb_alias', 'kb_type': 'default', 'operations': ['list', 'retrieve']},