LangBot/tests/unit_tests/agent/test_resource_builder.py

"""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.resource_builder import AgentResourceBuilder


RUNNER_ID = 'plugin:test/runner/default'


def make_descriptor(
    *,
    permissions: dict | None = None,
    config_schema: list[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',
        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):
    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,
    )


@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(
        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 AgentResourceBuilder(app).build_resources(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 AgentResourceBuilder(app).build_resources(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_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 AgentResourceBuilder(app).build_resources(query, descriptor)

    assert [model['model_id'] for model in resources['models']] == ['primary', 'fallback']