mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-09 23:36:02 +00:00
feat(agent-runner): enrich plugin runner host context
This commit is contained in:
@@ -66,6 +66,21 @@ class FakeMessage:
|
||||
self.content = content
|
||||
self.role = 'user'
|
||||
|
||||
def model_dump(self, mode='json'):
|
||||
return {'role': self.role, 'content': self.content}
|
||||
|
||||
|
||||
class FakePrompt:
|
||||
"""Fake prompt container."""
|
||||
def __init__(self, messages=None):
|
||||
self.messages = messages or []
|
||||
|
||||
|
||||
class FakeAdapter:
|
||||
"""Fake adapter with streaming capability."""
|
||||
async def is_stream_output_supported(self):
|
||||
return True
|
||||
|
||||
|
||||
class TestBuildParams:
|
||||
"""Tests for _build_params filtering."""
|
||||
@@ -446,4 +461,35 @@ class TestBuildParamsInContext:
|
||||
|
||||
# state should have seeded conversation_id
|
||||
assert 'state' in context
|
||||
assert context['state']['conversation']['external.conversation_id'] == 'conv_abc'
|
||||
assert context['state']['conversation']['external.conversation_id'] == 'conv_abc'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_includes_effective_prompt_and_runtime_capabilities(self):
|
||||
"""Context should expose host-preprocessed prompt and adapter capabilities."""
|
||||
reset_state_store()
|
||||
ap = FakeApplication()
|
||||
builder = AgentRunContextBuilder(ap)
|
||||
descriptor = make_descriptor()
|
||||
resources = make_resources()
|
||||
|
||||
session = FakeSession()
|
||||
query = type('Query', (), {
|
||||
'query_id': 1,
|
||||
'bot_uuid': 'bot_001',
|
||||
'pipeline_uuid': 'pipeline_001',
|
||||
'sender_id': 'user_001',
|
||||
'session': session,
|
||||
'user_message': None,
|
||||
'message_chain': None,
|
||||
'messages': [],
|
||||
'prompt': FakePrompt([FakeMessage('Effective prompt')]),
|
||||
'adapter': FakeAdapter(),
|
||||
'pipeline_config': {'output': {'misc': {'remove-think': True}}},
|
||||
'variables': {},
|
||||
})()
|
||||
|
||||
context = await builder.build_context(query, descriptor, resources)
|
||||
|
||||
assert context['prompt'][0]['content'] == 'Effective prompt'
|
||||
assert context['runtime']['metadata']['streaming_supported'] is True
|
||||
assert context['runtime']['metadata']['remove_think'] is True
|
||||
|
||||
148
tests/unit_tests/agent/test_resource_builder.py
Normal file
148
tests/unit_tests/agent/test_resource_builder.py
Normal file
@@ -0,0 +1,148 @@
|
||||
"""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']
|
||||
@@ -7,6 +7,7 @@ from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
from langbot_plugin.api.entities.builtin.provider import message as provider_message
|
||||
from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction, RuntimeToLangBotAction
|
||||
|
||||
|
||||
@@ -27,6 +28,22 @@ def compiled_params(statement):
|
||||
return statement.compile().params
|
||||
|
||||
|
||||
def make_agent_resources(
|
||||
models: list[dict] | None = None,
|
||||
tools: list[dict] | None = None,
|
||||
knowledge_bases: list[dict] | None = None,
|
||||
):
|
||||
"""Create a minimal AgentRun resources payload for run-scoped action tests."""
|
||||
return {
|
||||
'models': models or [],
|
||||
'tools': tools or [],
|
||||
'knowledge_bases': knowledge_bases or [],
|
||||
'files': [],
|
||||
'storage': {'plugin_storage': False, 'workspace_storage': False},
|
||||
'platform_capabilities': {},
|
||||
}
|
||||
|
||||
|
||||
class TestInitializePluginSettings:
|
||||
"""Tests for initialize_plugin_settings action handler."""
|
||||
|
||||
@@ -349,3 +366,231 @@ class TestHandlerQueryLookup:
|
||||
|
||||
assert response.code == 0
|
||||
assert response.data == {'bot_uuid': 'test-bot-uuid'}
|
||||
|
||||
|
||||
class TestAgentRunProxyActions:
|
||||
"""Tests for AgentRunner proxy actions that need host Query semantics."""
|
||||
|
||||
@pytest.fixture
|
||||
def app(self):
|
||||
mock_app = Mock()
|
||||
mock_app.logger = Mock()
|
||||
mock_app.query_pool = Mock()
|
||||
mock_app.query_pool.cached_queries = {}
|
||||
mock_app.model_mgr = Mock()
|
||||
mock_app.model_mgr.get_model_by_uuid = AsyncMock()
|
||||
mock_app.model_mgr.get_rerank_model_by_uuid = AsyncMock()
|
||||
mock_app.tool_mgr = Mock()
|
||||
mock_app.tool_mgr.execute_func_call = AsyncMock(return_value={'ok': True})
|
||||
return mock_app
|
||||
|
||||
@staticmethod
|
||||
def query(remove_think=True):
|
||||
return SimpleNamespace(
|
||||
pipeline_config={'output': {'misc': {'remove-think': remove_think}}},
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoke_llm_restores_query_and_model_options(self, app):
|
||||
"""INVOKE_LLM passes Query, model extra_args and remove-think to provider."""
|
||||
from langbot.pkg.agent.runner.session_registry import get_session_registry
|
||||
|
||||
run_id = 'run_proxy_invoke_llm_options'
|
||||
query = self.query(remove_think=True)
|
||||
app.query_pool.cached_queries[901] = query
|
||||
|
||||
registry = get_session_registry()
|
||||
await registry.unregister(run_id)
|
||||
await registry.register(
|
||||
run_id=run_id,
|
||||
runner_id='plugin:test/runner/default',
|
||||
query_id=901,
|
||||
plugin_identity='test/runner',
|
||||
resources=make_agent_resources(models=[{'model_id': 'llm_001'}]),
|
||||
)
|
||||
|
||||
provider = SimpleNamespace(
|
||||
invoke_llm=AsyncMock(return_value=provider_message.Message(role='assistant', content='ok')),
|
||||
)
|
||||
model = SimpleNamespace(
|
||||
model_entity=SimpleNamespace(
|
||||
abilities=['func_call'],
|
||||
extra_args={'temperature': 0.2, 'top_p': 0.8},
|
||||
),
|
||||
provider=provider,
|
||||
)
|
||||
app.model_mgr.get_model_by_uuid.return_value = model
|
||||
runtime_handler = make_handler(app)
|
||||
|
||||
try:
|
||||
response = await runtime_handler.actions[PluginToRuntimeAction.INVOKE_LLM.value]({
|
||||
'run_id': run_id,
|
||||
'llm_model_uuid': 'llm_001',
|
||||
'messages': [{'role': 'user', 'content': 'hello'}],
|
||||
'funcs': [{
|
||||
'name': 'search',
|
||||
'human_desc': 'Search',
|
||||
'description': 'Search',
|
||||
'parameters': {'type': 'object'},
|
||||
}],
|
||||
'extra_args': {'temperature': 0.7, 'presence_penalty': 0.1},
|
||||
})
|
||||
finally:
|
||||
await registry.unregister(run_id)
|
||||
|
||||
assert response.code == 0
|
||||
provider.invoke_llm.assert_awaited_once()
|
||||
kwargs = provider.invoke_llm.await_args.kwargs
|
||||
assert kwargs['query'] is query
|
||||
assert kwargs['extra_args'] == {
|
||||
'temperature': 0.7,
|
||||
'top_p': 0.8,
|
||||
'presence_penalty': 0.1,
|
||||
}
|
||||
assert kwargs['remove_think'] is True
|
||||
assert [tool.name for tool in kwargs['funcs']] == ['search']
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoke_llm_stream_restores_query_and_options(self, app):
|
||||
"""INVOKE_LLM_STREAM applies the same host context as non-streaming calls."""
|
||||
from langbot.pkg.agent.runner.session_registry import get_session_registry
|
||||
|
||||
class StreamProvider:
|
||||
def __init__(self):
|
||||
self.kwargs = None
|
||||
|
||||
async def invoke_llm_stream(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
yield provider_message.MessageChunk(role='assistant', content='hi')
|
||||
|
||||
run_id = 'run_proxy_invoke_llm_stream_options'
|
||||
query = self.query(remove_think=False)
|
||||
app.query_pool.cached_queries[902] = query
|
||||
|
||||
registry = get_session_registry()
|
||||
await registry.unregister(run_id)
|
||||
await registry.register(
|
||||
run_id=run_id,
|
||||
runner_id='plugin:test/runner/default',
|
||||
query_id=902,
|
||||
plugin_identity='test/runner',
|
||||
resources=make_agent_resources(models=[{'model_id': 'llm_stream_001'}]),
|
||||
)
|
||||
|
||||
provider = StreamProvider()
|
||||
model = SimpleNamespace(
|
||||
model_entity=SimpleNamespace(abilities=[], extra_args={'max_tokens': 128}),
|
||||
provider=provider,
|
||||
)
|
||||
app.model_mgr.get_model_by_uuid.return_value = model
|
||||
runtime_handler = make_handler(app)
|
||||
|
||||
responses = []
|
||||
try:
|
||||
stream = runtime_handler.actions[PluginToRuntimeAction.INVOKE_LLM_STREAM.value]({
|
||||
'run_id': run_id,
|
||||
'llm_model_uuid': 'llm_stream_001',
|
||||
'messages': [{'role': 'user', 'content': 'hello'}],
|
||||
'funcs': [{
|
||||
'name': 'search',
|
||||
'human_desc': 'Search',
|
||||
'description': 'Search',
|
||||
'parameters': {'type': 'object'},
|
||||
}],
|
||||
'extra_args': {'max_tokens': 256},
|
||||
'remove_think': True,
|
||||
})
|
||||
async for response in stream:
|
||||
responses.append(response)
|
||||
finally:
|
||||
await registry.unregister(run_id)
|
||||
|
||||
assert [response.code for response in responses] == [0]
|
||||
assert provider.kwargs['query'] is query
|
||||
assert provider.kwargs['extra_args'] == {'max_tokens': 256}
|
||||
assert provider.kwargs['remove_think'] is True
|
||||
assert provider.kwargs['funcs'] == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_tool_passes_current_query(self, app):
|
||||
"""CALL_TOOL passes the current Query back into tool execution."""
|
||||
from langbot.pkg.agent.runner.session_registry import get_session_registry
|
||||
|
||||
run_id = 'run_proxy_call_tool_query'
|
||||
query = self.query()
|
||||
app.query_pool.cached_queries[903] = query
|
||||
|
||||
registry = get_session_registry()
|
||||
await registry.unregister(run_id)
|
||||
await registry.register(
|
||||
run_id=run_id,
|
||||
runner_id='plugin:test/runner/default',
|
||||
query_id=903,
|
||||
plugin_identity='test/runner',
|
||||
resources=make_agent_resources(tools=[{'tool_name': 'test/search'}]),
|
||||
)
|
||||
|
||||
runtime_handler = make_handler(app)
|
||||
|
||||
try:
|
||||
response = await runtime_handler.actions[PluginToRuntimeAction.CALL_TOOL.value]({
|
||||
'run_id': run_id,
|
||||
'tool_name': 'test/search',
|
||||
'parameters': {'q': 'langbot'},
|
||||
})
|
||||
finally:
|
||||
await registry.unregister(run_id)
|
||||
|
||||
assert response.code == 0
|
||||
app.tool_mgr.execute_func_call.assert_awaited_once_with(
|
||||
name='test/search',
|
||||
parameters={'q': 'langbot'},
|
||||
query=query,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoke_rerank_uses_authorized_model_and_extra_args(self, app):
|
||||
"""INVOKE_RERANK validates run-scoped model access and merges model extra_args."""
|
||||
from langbot.pkg.agent.runner.session_registry import get_session_registry
|
||||
|
||||
run_id = 'run_proxy_rerank_options'
|
||||
registry = get_session_registry()
|
||||
await registry.unregister(run_id)
|
||||
await registry.register(
|
||||
run_id=run_id,
|
||||
runner_id='plugin:test/runner/default',
|
||||
query_id=904,
|
||||
plugin_identity='test/runner',
|
||||
resources=make_agent_resources(models=[{'model_id': 'rerank_001'}]),
|
||||
)
|
||||
|
||||
provider = SimpleNamespace(
|
||||
invoke_rerank=AsyncMock(return_value=[
|
||||
{'index': 0, 'relevance_score': 0.2},
|
||||
{'index': 1, 'relevance_score': 0.9},
|
||||
]),
|
||||
)
|
||||
rerank_model = SimpleNamespace(
|
||||
model_entity=SimpleNamespace(extra_args={'top_n': 5, 'return_documents': False}),
|
||||
provider=provider,
|
||||
)
|
||||
app.model_mgr.get_rerank_model_by_uuid.return_value = rerank_model
|
||||
runtime_handler = make_handler(app)
|
||||
|
||||
try:
|
||||
response = await runtime_handler.actions[PluginToRuntimeAction.INVOKE_RERANK.value]({
|
||||
'run_id': run_id,
|
||||
'rerank_model_uuid': 'rerank_001',
|
||||
'query': 'hello',
|
||||
'documents': ['a', 'b'],
|
||||
'top_k': 1,
|
||||
'extra_args': {'top_n': 2},
|
||||
})
|
||||
finally:
|
||||
await registry.unregister(run_id)
|
||||
|
||||
assert response.code == 0
|
||||
assert response.data['results'] == [{'index': 1, 'relevance_score': 0.9}]
|
||||
provider.invoke_rerank.assert_awaited_once()
|
||||
kwargs = provider.invoke_rerank.await_args.kwargs
|
||||
assert kwargs['extra_args'] == {'top_n': 2, 'return_documents': False}
|
||||
|
||||
Reference in New Issue
Block a user