mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-08 23:06:03 +00:00
feat(agent-runner): enrich plugin runner host context
This commit is contained in:
huanghuoguoguo
@@ -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