feat(agent-runner): enrich plugin runner host context

This commit is contained in:
huanghuoguoguo
2026-05-17 23:26:52 +08:00
parent 19557c3227
commit 036affe01f
9 changed files with 806 additions and 114 deletions

View File

@@ -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

View 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']

View File

@@ -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}