mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-27 07:54:19 +00:00
feat(agent-runner): add plugin runner host integration
This commit is contained in:
huanghuoguoguo
@@ -1,169 +0,0 @@
|
||||
"""Tests for DifyServiceAPIRunner pure utility methods.
|
||||
|
||||
Tests the helper methods that don't require real Dify API calls.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class TestDifyExtractTextOutput:
|
||||
"""Tests for _extract_dify_text_output method."""
|
||||
|
||||
def _create_runner(self):
|
||||
"""Create runner instance."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langbot.pkg.provider.runners.difysvapi import DifyServiceAPIRunner
|
||||
|
||||
mock_app = MagicMock()
|
||||
pipeline_config = {
|
||||
'ai': {
|
||||
'dify-service-api': {
|
||||
'app-type': 'chat',
|
||||
'api-key': 'test-key',
|
||||
'base-url': 'https://api.dify.ai',
|
||||
}
|
||||
},
|
||||
'output': {'misc': {}},
|
||||
}
|
||||
|
||||
runner = DifyServiceAPIRunner(mock_app, pipeline_config)
|
||||
runner.dify_client = MagicMock()
|
||||
|
||||
return runner
|
||||
|
||||
def test_extract_none_value(self):
|
||||
"""None returns empty string."""
|
||||
runner = self._create_runner()
|
||||
|
||||
result = runner._extract_dify_text_output(None)
|
||||
|
||||
assert result == ''
|
||||
|
||||
def test_extract_string_value(self):
|
||||
"""Plain string is returned."""
|
||||
runner = self._create_runner()
|
||||
|
||||
result = runner._extract_dify_text_output('plain text')
|
||||
|
||||
assert result == 'plain text'
|
||||
|
||||
def test_extract_dict_with_content(self):
|
||||
"""Dict with 'content' key extracts content."""
|
||||
runner = self._create_runner()
|
||||
|
||||
result = runner._extract_dify_text_output({'content': 'extracted content'})
|
||||
|
||||
assert result == 'extracted content'
|
||||
|
||||
def test_extract_dict_without_content(self):
|
||||
"""Dict without 'content' key is JSON dumped."""
|
||||
runner = self._create_runner()
|
||||
|
||||
result = runner._extract_dify_text_output({'key': 'value'})
|
||||
|
||||
assert 'key' in result
|
||||
assert 'value' in result
|
||||
|
||||
def test_extract_json_string_with_content(self):
|
||||
"""JSON string with 'content' key extracts content."""
|
||||
runner = self._create_runner()
|
||||
|
||||
result = runner._extract_dify_text_output('{"content": "json content"}')
|
||||
|
||||
assert result == 'json content'
|
||||
|
||||
def test_extract_json_string_without_content(self):
|
||||
"""JSON string without 'content' key returns original."""
|
||||
runner = self._create_runner()
|
||||
|
||||
result = runner._extract_dify_text_output('{"other": "value"}')
|
||||
|
||||
assert '{"other": "value"}' in result
|
||||
|
||||
def test_extract_whitespace_string(self):
|
||||
"""Whitespace string returns empty."""
|
||||
runner = self._create_runner()
|
||||
|
||||
result = runner._extract_dify_text_output(' ')
|
||||
|
||||
assert result == ''
|
||||
|
||||
|
||||
class TestDifyRunnerConfigValidation:
|
||||
"""Tests for runner config validation."""
|
||||
|
||||
def test_invalid_app_type_raises(self):
|
||||
"""Invalid app-type raises DifyAPIError."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langbot.pkg.provider.runners.difysvapi import DifyServiceAPIRunner
|
||||
from langbot.libs.dify_service_api.v1.errors import DifyAPIError
|
||||
|
||||
mock_app = MagicMock()
|
||||
pipeline_config = {
|
||||
'ai': {
|
||||
'dify-service-api': {
|
||||
'app-type': 'invalid-type',
|
||||
'api-key': 'test',
|
||||
'base-url': 'https://api.dify.ai',
|
||||
}
|
||||
},
|
||||
'output': {'misc': {}},
|
||||
}
|
||||
|
||||
with pytest.raises(DifyAPIError, match='不支持'):
|
||||
DifyServiceAPIRunner(mock_app, pipeline_config)
|
||||
|
||||
def test_valid_app_types(self):
|
||||
"""Valid app-types don't raise."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langbot.pkg.provider.runners.difysvapi import DifyServiceAPIRunner
|
||||
|
||||
mock_app = MagicMock()
|
||||
|
||||
for app_type in ['chat', 'agent', 'workflow']:
|
||||
pipeline_config = {
|
||||
'ai': {
|
||||
'dify-service-api': {
|
||||
'app-type': app_type,
|
||||
'api-key': 'test',
|
||||
'base-url': 'https://api.dify.ai',
|
||||
}
|
||||
},
|
||||
'output': {'misc': {}},
|
||||
}
|
||||
|
||||
runner = DifyServiceAPIRunner(mock_app, pipeline_config)
|
||||
# Should not raise
|
||||
assert runner is not None
|
||||
|
||||
|
||||
class TestDifyRunnerInit:
|
||||
"""Tests for runner initialization."""
|
||||
|
||||
def test_runner_stores_config(self):
|
||||
"""Runner stores pipeline_config."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langbot.pkg.provider.runners.difysvapi import DifyServiceAPIRunner
|
||||
|
||||
mock_app = MagicMock()
|
||||
pipeline_config = {
|
||||
'ai': {
|
||||
'dify-service-api': {
|
||||
'app-type': 'chat',
|
||||
'api-key': 'test-key',
|
||||
'base-url': 'https://api.dify.ai',
|
||||
}
|
||||
},
|
||||
'output': {'misc': {}},
|
||||
}
|
||||
|
||||
runner = DifyServiceAPIRunner(mock_app, pipeline_config)
|
||||
|
||||
assert runner.pipeline_config == pipeline_config
|
||||
assert runner.ap == mock_app
|
||||
@@ -1,281 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
import langbot_plugin.api.entities.builtin.provider.session as provider_session
|
||||
|
||||
from langbot.pkg.provider.runners.localagent import LocalAgentRunner, _StreamAccumulator
|
||||
|
||||
|
||||
class RecordingProvider:
|
||||
def __init__(self):
|
||||
self.requests: list[dict] = []
|
||||
|
||||
async def invoke_llm(self, query, model, messages, funcs, extra_args=None, remove_think=None):
|
||||
self.requests.append(
|
||||
{
|
||||
'messages': list(messages),
|
||||
'funcs': list(funcs),
|
||||
'remove_think': remove_think,
|
||||
}
|
||||
)
|
||||
|
||||
if len(self.requests) == 1:
|
||||
return provider_message.Message(
|
||||
role='assistant',
|
||||
content='Let me calculate that exactly.',
|
||||
tool_calls=[
|
||||
provider_message.ToolCall(
|
||||
id='call-1',
|
||||
type='function',
|
||||
function=provider_message.FunctionCall(
|
||||
name='exec',
|
||||
arguments=json.dumps(
|
||||
{'command': ("python - <<'PY'\nnums = [1, 2, 3, 4]\nprint(sum(nums) / len(nums))\nPY")}
|
||||
),
|
||||
),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
tool_result = json.loads(messages[-1].content)
|
||||
return provider_message.Message(
|
||||
role='assistant',
|
||||
content=f'The average is {tool_result["stdout"]}.',
|
||||
)
|
||||
|
||||
|
||||
class RecordingStreamProvider:
|
||||
def __init__(self):
|
||||
self.stream_requests: list[dict] = []
|
||||
|
||||
def invoke_llm_stream(self, query, model, messages, funcs, extra_args=None, remove_think=None):
|
||||
self.stream_requests.append(
|
||||
{
|
||||
'messages': list(messages),
|
||||
'funcs': list(funcs),
|
||||
'remove_think': remove_think,
|
||||
}
|
||||
)
|
||||
|
||||
async def _stream():
|
||||
if len(self.stream_requests) == 1:
|
||||
yield provider_message.MessageChunk(
|
||||
role='assistant',
|
||||
tool_calls=[
|
||||
provider_message.ToolCall(
|
||||
id='call-1',
|
||||
type='function',
|
||||
function=provider_message.FunctionCall(
|
||||
name='exec',
|
||||
arguments=json.dumps({'command': "python -c 'print(1)'"}),
|
||||
),
|
||||
)
|
||||
],
|
||||
is_final=True,
|
||||
)
|
||||
return
|
||||
|
||||
yield provider_message.MessageChunk(
|
||||
role='assistant',
|
||||
content='Tool execution failed.',
|
||||
is_final=True,
|
||||
)
|
||||
|
||||
return _stream()
|
||||
|
||||
|
||||
def make_query() -> pipeline_query.Query:
|
||||
adapter = AsyncMock()
|
||||
adapter.is_stream_output_supported = AsyncMock(return_value=False)
|
||||
|
||||
return pipeline_query.Query.model_construct(
|
||||
query_id='avg-query',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_chain=[],
|
||||
message_event=None,
|
||||
adapter=adapter,
|
||||
pipeline_uuid='pipeline-uuid',
|
||||
bot_uuid='bot-uuid',
|
||||
pipeline_config={
|
||||
'ai': {
|
||||
'runner': {'runner': 'local-agent'},
|
||||
'local-agent': {'model': {'primary': 'test-model-uuid', 'fallbacks': []}, 'prompt': 'test-prompt'},
|
||||
},
|
||||
'output': {'misc': {'remove-think': False}},
|
||||
},
|
||||
prompt=SimpleNamespace(messages=[]),
|
||||
messages=[],
|
||||
user_message=provider_message.Message(
|
||||
role='user',
|
||||
content='Please calculate the average of 1, 2, 3, and 4.',
|
||||
),
|
||||
use_funcs=[SimpleNamespace(name='exec')],
|
||||
use_llm_model_uuid='test-model-uuid',
|
||||
variables={},
|
||||
)
|
||||
|
||||
|
||||
def test_stream_accumulator_merges_fragmented_tool_call_arguments():
|
||||
accumulator = _StreamAccumulator(msg_sequence=1)
|
||||
|
||||
assert (
|
||||
accumulator.add(
|
||||
provider_message.MessageChunk(
|
||||
role='assistant',
|
||||
tool_calls=[
|
||||
provider_message.ToolCall(
|
||||
id='call-1',
|
||||
type='function',
|
||||
function=provider_message.FunctionCall(name='exec', arguments='{"command":'),
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
emitted = accumulator.add(
|
||||
provider_message.MessageChunk(
|
||||
role='assistant',
|
||||
tool_calls=[
|
||||
provider_message.ToolCall(
|
||||
id='call-1',
|
||||
type='function',
|
||||
function=provider_message.FunctionCall(name='exec', arguments='"pwd"}'),
|
||||
)
|
||||
],
|
||||
is_final=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert emitted is not None
|
||||
final_msg = accumulator.final_message()
|
||||
assert final_msg.tool_calls[0].function.name == 'exec'
|
||||
assert final_msg.tool_calls[0].function.arguments == '{"command":"pwd"}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_localagent_uses_exec_for_exact_calculation():
|
||||
provider = RecordingProvider()
|
||||
model = SimpleNamespace(
|
||||
provider=provider,
|
||||
model_entity=SimpleNamespace(
|
||||
uuid='test-model-uuid',
|
||||
name='test-model',
|
||||
abilities=['func_call'],
|
||||
extra_args={},
|
||||
),
|
||||
)
|
||||
|
||||
tool_manager = SimpleNamespace(
|
||||
execute_func_call=AsyncMock(
|
||||
return_value={
|
||||
'session_id': 'avg-query',
|
||||
'backend': 'podman',
|
||||
'status': 'completed',
|
||||
'ok': True,
|
||||
'exit_code': 0,
|
||||
'stdout': '2.5',
|
||||
'stderr': '',
|
||||
'duration_ms': 18,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
app = SimpleNamespace(
|
||||
logger=Mock(),
|
||||
model_mgr=SimpleNamespace(get_model_by_uuid=AsyncMock(return_value=model)),
|
||||
tool_mgr=tool_manager,
|
||||
rag_mgr=SimpleNamespace(),
|
||||
box_service=SimpleNamespace(
|
||||
get_system_guidance=Mock(
|
||||
return_value=(
|
||||
'When the exec tool is available, use it for exact calculations, statistics, '
|
||||
'structured data parsing, and code execution instead of estimating mentally. '
|
||||
'Unless the user explicitly asks for the script, code, or implementation details, '
|
||||
'do not include the generated script in the final answer. '
|
||||
'A default workspace is mounted at /workspace for file tasks.'
|
||||
)
|
||||
),
|
||||
),
|
||||
skill_mgr=SimpleNamespace(
|
||||
get_skills_for_pipeline=AsyncMock(return_value=[]),
|
||||
detect_skill_activation=AsyncMock(return_value=None),
|
||||
build_activation_prompt=Mock(return_value=None),
|
||||
),
|
||||
)
|
||||
|
||||
runner = LocalAgentRunner(app, pipeline_config={})
|
||||
query = make_query()
|
||||
|
||||
results = [message async for message in runner.run(query)]
|
||||
|
||||
assert [message.role for message in results] == ['assistant', 'tool', 'assistant']
|
||||
assert results[-1].content == 'The average is 2.5.'
|
||||
|
||||
tool_manager.execute_func_call.assert_awaited_once()
|
||||
tool_name, tool_parameters = tool_manager.execute_func_call.await_args.args[:2]
|
||||
assert tool_name == 'exec'
|
||||
assert 'print(sum(nums) / len(nums))' in tool_parameters['command']
|
||||
|
||||
first_request = provider.requests[0]
|
||||
assert any(
|
||||
message.role == 'system'
|
||||
and 'exec' in str(message.content)
|
||||
and 'exact calculations' in str(message.content)
|
||||
and 'Unless the user explicitly asks for the script' in str(message.content)
|
||||
and '/workspace' in str(message.content)
|
||||
for message in first_request['messages']
|
||||
)
|
||||
assert [tool.name for tool in first_request['funcs']] == ['exec']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_localagent_streaming_tool_error_yields_message_chunks():
|
||||
provider = RecordingStreamProvider()
|
||||
model = SimpleNamespace(
|
||||
provider=provider,
|
||||
model_entity=SimpleNamespace(
|
||||
uuid='test-model-uuid',
|
||||
name='test-model',
|
||||
abilities=['func_call'],
|
||||
extra_args={},
|
||||
),
|
||||
)
|
||||
|
||||
adapter = AsyncMock()
|
||||
adapter.is_stream_output_supported = AsyncMock(return_value=True)
|
||||
|
||||
query = make_query()
|
||||
query.adapter = adapter
|
||||
|
||||
app = SimpleNamespace(
|
||||
logger=Mock(),
|
||||
model_mgr=SimpleNamespace(get_model_by_uuid=AsyncMock(return_value=model)),
|
||||
tool_mgr=SimpleNamespace(execute_func_call=AsyncMock(side_effect=RuntimeError('boom'))),
|
||||
rag_mgr=SimpleNamespace(),
|
||||
box_service=SimpleNamespace(
|
||||
get_system_guidance=Mock(return_value='sandbox guidance'),
|
||||
),
|
||||
skill_mgr=SimpleNamespace(
|
||||
get_skills_for_pipeline=AsyncMock(return_value=[]),
|
||||
detect_skill_activation=AsyncMock(return_value=None),
|
||||
build_activation_prompt=Mock(return_value=None),
|
||||
),
|
||||
)
|
||||
|
||||
runner = LocalAgentRunner(app, pipeline_config={})
|
||||
|
||||
results = [message async for message in runner.run(query)]
|
||||
|
||||
assert all(isinstance(message, provider_message.MessageChunk) for message in results)
|
||||
assert any(message.role == 'tool' and message.content == 'err: boom' for message in results)
|
||||
@@ -12,13 +12,39 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
import langbot_plugin.api.entities.builtin.provider.session as provider_session
|
||||
|
||||
from langbot.pkg.api.http.service.model import _runtime_model_data
|
||||
from langbot.pkg.agent.runner.descriptor import AgentRunnerDescriptor
|
||||
from langbot.pkg.api.http.service.provider import ModelProviderService
|
||||
from langbot.pkg.entity.persistence import model as persistence_model
|
||||
from langbot.pkg.pipeline.preproc.preproc import PreProcessor
|
||||
from langbot.pkg.provider.modelmgr import requester
|
||||
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
|
||||
from langbot.pkg.provider.modelmgr.token import TokenManager
|
||||
from langbot.pkg.provider.runners.localagent import LocalAgentRunner
|
||||
|
||||
|
||||
DEFAULT_RUNNER_ID = 'plugin:langbot/local-agent/default'
|
||||
|
||||
|
||||
class FakeAgentRunnerRegistry:
|
||||
async def get(self, runner_id, bound_plugins=None):
|
||||
return AgentRunnerDescriptor(
|
||||
id=runner_id,
|
||||
source='plugin',
|
||||
label={'en_US': 'Local Agent'},
|
||||
plugin_author='langbot',
|
||||
plugin_name='local-agent',
|
||||
runner_name='default',
|
||||
config_schema=[
|
||||
{'name': 'model', 'type': 'model-fallback-selector'},
|
||||
{'name': 'prompt', 'type': 'prompt-editor', 'default': []},
|
||||
{'name': 'knowledge-bases', 'type': 'knowledge-base-multi-selector', 'default': []},
|
||||
],
|
||||
capabilities={'tool_calling': True, 'knowledge_retrieval': True, 'multimodal_input': True},
|
||||
permissions={
|
||||
'models': ['invoke', 'stream'],
|
||||
'tools': ['detail', 'call'],
|
||||
'knowledge_bases': ['list', 'retrieve'],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_runtime_llm_model_data_preserves_uuid_after_update_payload_uuid_removed():
|
||||
@@ -120,6 +146,7 @@ async def test_updated_llm_model_is_immediately_usable_by_local_agent_pipeline()
|
||||
|
||||
ap = SimpleNamespace()
|
||||
ap.logger = Mock()
|
||||
ap.agent_runner_registry = FakeAgentRunnerRegistry()
|
||||
ap.persistence_mgr = SimpleNamespace(execute_async=AsyncMock())
|
||||
ap.tool_mgr = SimpleNamespace(get_all_tools=AsyncMock(return_value=[]))
|
||||
ap.skill_mgr = None # PreProcessor only uses skill_mgr for the local-agent skill-binding branch
|
||||
@@ -183,11 +210,13 @@ async def test_updated_llm_model_is_immediately_usable_by_local_agent_pipeline()
|
||||
)
|
||||
pipeline_config = {
|
||||
'ai': {
|
||||
'runner': {'runner': 'local-agent'},
|
||||
'local-agent': {
|
||||
'model': {'primary': model_uuid, 'fallbacks': []},
|
||||
'prompt': [],
|
||||
'knowledge-bases': [],
|
||||
'runner': {'id': DEFAULT_RUNNER_ID},
|
||||
'runner_config': {
|
||||
DEFAULT_RUNNER_ID: {
|
||||
'model': {'primary': model_uuid, 'fallbacks': []},
|
||||
'prompt': [],
|
||||
'knowledge-bases': [],
|
||||
},
|
||||
},
|
||||
},
|
||||
'trigger': {'misc': {'combine-quote-message': False}},
|
||||
@@ -220,8 +249,3 @@ async def test_updated_llm_model_is_immediately_usable_by_local_agent_pipeline()
|
||||
processed_query = result.new_query
|
||||
|
||||
assert processed_query.use_llm_model_uuid == model_uuid
|
||||
|
||||
runner = SimpleNamespace(ap=ap, pipeline_config=pipeline_config)
|
||||
candidates = await LocalAgentRunner._get_model_candidates(runner, processed_query)
|
||||
|
||||
assert [model.model_entity.uuid for model in candidates] == [model_uuid]
|
||||
|
||||
@@ -302,6 +302,59 @@ async def test_runtime_provider_invoke_llm_delegates(runtime_provider, runtime_l
|
||||
assert result.role == 'assistant'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_provider_invoke_llm_stashes_usage(runtime_provider, runtime_llm_model):
|
||||
"""RuntimeProvider preserves requester usage for upstream action handlers."""
|
||||
provider = runtime_provider
|
||||
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
|
||||
query = pipeline_query.Query.model_construct(
|
||||
query_id='test-query-usage',
|
||||
launcher_type='person',
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_chain=None,
|
||||
message_event=None,
|
||||
adapter=None,
|
||||
pipeline_uuid='pipeline-uuid',
|
||||
bot_uuid='bot-uuid',
|
||||
pipeline_config={'ai': {}, 'output': {}, 'trigger': {}},
|
||||
session=None,
|
||||
prompt=None,
|
||||
messages=[],
|
||||
user_message=None,
|
||||
use_funcs=[],
|
||||
use_llm_model_uuid=None,
|
||||
variables={},
|
||||
resp_messages=[],
|
||||
resp_message_chain=None,
|
||||
current_stage_name=None,
|
||||
)
|
||||
usage = {
|
||||
'prompt_tokens': 11,
|
||||
'completion_tokens': 7,
|
||||
'total_tokens': 18,
|
||||
'prompt_tokens_details': {'cached_tokens': 3},
|
||||
}
|
||||
provider.requester.invoke_llm = AsyncMock(
|
||||
return_value=(
|
||||
provider_message.Message(role='assistant', content='ok'),
|
||||
usage,
|
||||
)
|
||||
)
|
||||
|
||||
result = await provider.invoke_llm(
|
||||
query,
|
||||
runtime_llm_model,
|
||||
[provider_message.Message(role='user', content='Hello')],
|
||||
)
|
||||
|
||||
assert result.content == 'ok'
|
||||
assert query.variables[requester.LLM_USAGE_QUERY_VARIABLE] == usage
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_provider_invoke_llm_stream_yields_chunks(runtime_provider, runtime_llm_model):
|
||||
"""Test RuntimeProvider.invoke_llm_stream yields chunks from requester."""
|
||||
@@ -345,6 +398,61 @@ async def test_runtime_provider_invoke_llm_stream_yields_chunks(runtime_provider
|
||||
assert chunks[0].role == 'assistant'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_provider_invoke_llm_stream_stashes_usage(runtime_provider, runtime_llm_model):
|
||||
"""RuntimeProvider transfers captured stream usage to the public query usage key."""
|
||||
provider = runtime_provider
|
||||
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
|
||||
query = pipeline_query.Query.model_construct(
|
||||
query_id='test-stream-usage',
|
||||
launcher_type='person',
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_chain=None,
|
||||
message_event=None,
|
||||
adapter=None,
|
||||
pipeline_uuid='pipeline-uuid',
|
||||
bot_uuid='bot-uuid',
|
||||
pipeline_config={'ai': {}, 'output': {}, 'trigger': {}},
|
||||
session=None,
|
||||
prompt=None,
|
||||
messages=[],
|
||||
user_message=None,
|
||||
use_funcs=[],
|
||||
use_llm_model_uuid=None,
|
||||
variables={},
|
||||
resp_messages=[],
|
||||
resp_message_chain=None,
|
||||
current_stage_name=None,
|
||||
)
|
||||
usage = {
|
||||
'prompt_tokens': 13,
|
||||
'completion_tokens': 2,
|
||||
'total_tokens': 15,
|
||||
}
|
||||
|
||||
async def fake_stream(**kwargs):
|
||||
kwargs['query'].variables[requester.LLM_USAGE_QUERY_VARIABLE] = usage
|
||||
yield provider_message.MessageChunk(role='assistant', content='ok')
|
||||
|
||||
provider.requester.invoke_llm_stream = fake_stream
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in provider.invoke_llm_stream(
|
||||
query,
|
||||
runtime_llm_model,
|
||||
[provider_message.Message(role='user', content='Hello')],
|
||||
)
|
||||
]
|
||||
|
||||
assert len(chunks) == 1
|
||||
assert query.variables[requester.LLM_USAGE_QUERY_VARIABLE] == usage
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_provider_invoke_embedding_returns_vectors(runtime_provider, runtime_embedding_model):
|
||||
"""Test RuntimeProvider.invoke_embedding returns embedding vectors."""
|
||||
|
||||
Reference in New Issue
Block a user