mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-02 20:14:36 +00:00
264 lines
8.5 KiB
Python
264 lines
8.5 KiB
Python
"""Tests for AgentResourceBuilder."""
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
import pytest
|
|
|
|
from langbot.pkg.agent.runner.descriptor import AgentRunnerDescriptor
|
|
from langbot.pkg.agent.runner.pipeline_adapter import PipelineAdapter
|
|
from langbot.pkg.agent.runner.resource_builder import AgentResourceBuilder
|
|
|
|
|
|
RUNNER_ID = 'plugin:test/runner/default'
|
|
|
|
|
|
def make_descriptor(
|
|
*,
|
|
permissions: dict | None = None,
|
|
config_schema: list[dict] | None = None,
|
|
capabilities: dict | None = None,
|
|
) -> AgentRunnerDescriptor:
|
|
return AgentRunnerDescriptor(
|
|
id=RUNNER_ID,
|
|
source='plugin',
|
|
label={'en_US': 'Test Runner'},
|
|
plugin_author='test',
|
|
plugin_name='runner',
|
|
runner_name='default',
|
|
capabilities=capabilities or {},
|
|
permissions=permissions or {'models': ['invoke', 'stream']},
|
|
config_schema=config_schema or [],
|
|
)
|
|
|
|
|
|
def make_model(model_type='llm', provider='test-provider'):
|
|
return SimpleNamespace(
|
|
model_entity=SimpleNamespace(model_type=model_type),
|
|
provider_entity=SimpleNamespace(name=provider),
|
|
)
|
|
|
|
|
|
def make_query(
|
|
runner_config: dict,
|
|
*,
|
|
variables: dict | None = None,
|
|
use_llm_model_uuid=None,
|
|
use_funcs: list | None = None,
|
|
):
|
|
return SimpleNamespace(
|
|
pipeline_config={
|
|
'ai': {
|
|
'runner': {'id': RUNNER_ID},
|
|
'runner_config': {RUNNER_ID: runner_config},
|
|
},
|
|
},
|
|
variables=variables or {},
|
|
use_llm_model_uuid=use_llm_model_uuid,
|
|
use_funcs=use_funcs or [],
|
|
pipeline_uuid='pipeline_001',
|
|
)
|
|
|
|
|
|
async def build_resources(app, query, descriptor):
|
|
binding = PipelineAdapter.pipeline_config_to_binding(query, descriptor.id)
|
|
return await AgentResourceBuilder(app).build_resources_from_binding(
|
|
event=Mock(),
|
|
binding=binding,
|
|
descriptor=descriptor,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def app():
|
|
mock_app = Mock()
|
|
mock_app.logger = Mock()
|
|
mock_app.model_mgr = Mock()
|
|
mock_app.rag_mgr = Mock()
|
|
mock_app.rag_mgr.get_knowledge_base_by_uuid = AsyncMock(return_value=None)
|
|
return mock_app
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_build_models_authorizes_config_declared_llm_and_rerank_models(app):
|
|
"""DynamicForm model selectors should become run-scoped authorized models."""
|
|
llm_models = {
|
|
'primary': make_model(),
|
|
'fallback': make_model(),
|
|
'aux': make_model(provider='aux-provider'),
|
|
}
|
|
rerank_models = {
|
|
'rerank': make_model(model_type='rerank', provider='rerank-provider'),
|
|
}
|
|
|
|
async def get_model_by_uuid(model_uuid):
|
|
return llm_models.get(model_uuid)
|
|
|
|
async def get_rerank_model_by_uuid(model_uuid):
|
|
return rerank_models.get(model_uuid)
|
|
|
|
app.model_mgr.get_model_by_uuid = AsyncMock(side_effect=get_model_by_uuid)
|
|
app.model_mgr.get_rerank_model_by_uuid = AsyncMock(side_effect=get_rerank_model_by_uuid)
|
|
descriptor = make_descriptor(
|
|
permissions={'models': ['invoke', 'stream', 'rerank']},
|
|
config_schema=[
|
|
{'name': 'model', 'type': 'model-fallback-selector'},
|
|
{'name': 'aux-model', 'type': 'llm-model-selector'},
|
|
{'name': 'rerank-model', 'type': 'rerank-model-selector'},
|
|
],
|
|
)
|
|
query = make_query({
|
|
'model': {'primary': 'primary', 'fallbacks': ['fallback', 'primary']},
|
|
'aux-model': 'aux',
|
|
'rerank-model': 'rerank',
|
|
})
|
|
|
|
resources = await build_resources(app, query, descriptor)
|
|
|
|
assert resources['models'] == [
|
|
{'model_id': 'primary', 'model_type': 'llm', 'provider': 'test-provider'},
|
|
{'model_id': 'fallback', 'model_type': 'llm', 'provider': 'test-provider'},
|
|
{'model_id': 'aux', 'model_type': 'llm', 'provider': 'aux-provider'},
|
|
{'model_id': 'rerank', 'model_type': 'rerank', 'provider': 'rerank-provider'},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_build_models_still_honors_manifest_permissions(app):
|
|
"""Config-selected models should not bypass runner manifest permissions."""
|
|
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=make_model())
|
|
app.model_mgr.get_rerank_model_by_uuid = AsyncMock(return_value=make_model(model_type='rerank'))
|
|
descriptor = make_descriptor(
|
|
permissions={'models': []},
|
|
config_schema=[
|
|
{'name': 'model', 'type': 'model-fallback-selector'},
|
|
{'name': 'rerank-model', 'type': 'rerank-model-selector'},
|
|
],
|
|
)
|
|
query = make_query({
|
|
'model': {'primary': 'primary', 'fallbacks': ['fallback']},
|
|
'rerank-model': 'rerank',
|
|
})
|
|
|
|
resources = await build_resources(app, query, descriptor)
|
|
|
|
assert resources['models'] == []
|
|
app.model_mgr.get_model_by_uuid.assert_not_awaited()
|
|
app.model_mgr.get_rerank_model_by_uuid.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_build_models_authorizes_rerank_only_runner(app):
|
|
"""A rerank-only runner should receive config-selected rerank models."""
|
|
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=make_model())
|
|
app.model_mgr.get_rerank_model_by_uuid = AsyncMock(
|
|
return_value=make_model(model_type='rerank', provider='rerank-provider')
|
|
)
|
|
descriptor = make_descriptor(
|
|
permissions={'models': ['rerank']},
|
|
config_schema=[
|
|
{'name': 'model', 'type': 'llm-model-selector'},
|
|
{'name': 'rerank-model', 'type': 'rerank-model-selector'},
|
|
],
|
|
)
|
|
query = make_query({
|
|
'model': 'llm',
|
|
'rerank-model': 'rerank',
|
|
})
|
|
|
|
resources = await build_resources(app, query, descriptor)
|
|
|
|
assert resources['models'] == [
|
|
{'model_id': 'rerank', 'model_type': 'rerank', 'provider': 'rerank-provider'},
|
|
]
|
|
app.model_mgr.get_model_by_uuid.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_build_models_deduplicates_query_and_config_models(app):
|
|
"""A model selected by both preproc and runner config should appear once."""
|
|
app.model_mgr.get_model_by_uuid = AsyncMock(return_value=make_model())
|
|
app.model_mgr.get_rerank_model_by_uuid = AsyncMock(return_value=None)
|
|
descriptor = make_descriptor(
|
|
config_schema=[
|
|
{'name': 'model', 'type': 'model-fallback-selector'},
|
|
],
|
|
)
|
|
query = make_query(
|
|
{'model': {'primary': 'primary', 'fallbacks': ['fallback']}},
|
|
variables={'_fallback_model_uuids': ['fallback']},
|
|
use_llm_model_uuid='primary',
|
|
)
|
|
|
|
resources = await build_resources(app, query, descriptor)
|
|
|
|
assert [model['model_id'] for model in resources['models']] == ['primary', 'fallback']
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_build_tools_authorizes_query_declared_tools(app):
|
|
"""Tools discovered by Pipeline preprocessing become run-scoped authorized resources."""
|
|
descriptor = make_descriptor(
|
|
permissions={
|
|
'models': [],
|
|
'tools': ['detail', 'call'],
|
|
},
|
|
)
|
|
query = make_query(
|
|
{},
|
|
use_funcs=[
|
|
{'name': 'qa_plugin_echo', 'description': 'Echo test tool'},
|
|
SimpleNamespace(name='qa_mcp_echo'),
|
|
],
|
|
)
|
|
|
|
resources = await build_resources(app, query, descriptor)
|
|
|
|
assert resources['tools'] == [
|
|
{
|
|
'tool_name': 'qa_plugin_echo',
|
|
'tool_type': None,
|
|
'description': None,
|
|
},
|
|
{
|
|
'tool_name': 'qa_mcp_echo',
|
|
'tool_type': None,
|
|
'description': None,
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_build_knowledge_bases_unions_config_and_policy_grants(app):
|
|
descriptor = make_descriptor(
|
|
capabilities={'knowledge_retrieval': True},
|
|
permissions={
|
|
'models': [],
|
|
'knowledge_bases': ['retrieve'],
|
|
},
|
|
config_schema=[
|
|
{'name': 'knowledge-bases', 'type': 'knowledge-base-multi-selector'},
|
|
],
|
|
)
|
|
query = make_query(
|
|
{'knowledge-bases': ['kb_config']},
|
|
variables={'_knowledge_base_uuids': ['kb_policy']},
|
|
)
|
|
|
|
async def get_kb(kb_uuid):
|
|
return SimpleNamespace(
|
|
uuid=kb_uuid,
|
|
get_name=lambda: f'name-{kb_uuid}',
|
|
knowledge_base_entity=SimpleNamespace(kb_type='default'),
|
|
)
|
|
|
|
app.rag_mgr.get_knowledge_base_by_uuid = AsyncMock(side_effect=get_kb)
|
|
|
|
resources = await build_resources(app, query, descriptor)
|
|
|
|
assert resources['knowledge_bases'] == [
|
|
{'kb_id': 'kb_config', 'kb_name': 'name-kb_config', 'kb_type': 'default'},
|
|
{'kb_id': 'kb_policy', 'kb_name': 'name-kb_policy', 'kb_type': 'default'},
|
|
]
|