Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions src/langbot/pkg/provider/runners/localagent.py
Original file line number Diff line number Diff line change
Expand Up @@ -523,7 +523,6 @@ async def run(
final_msg = stream_accumulator.final_message()

pending_tool_calls = final_msg.tool_calls
first_content = final_msg.content
if isinstance(final_msg, provider_message.MessageChunk):
first_end_sequence = final_msg.msg_sequence

Expand Down Expand Up @@ -606,9 +605,13 @@ async def run(
)

if is_stream:
# Do NOT re-seed the accumulator with first_content:
# the previous round's text was already pushed to the
# platform adapter. Re-seeding would cause every
# subsequent round to repeat the entire opening line,
# which the platform then forwards as a duplicate message.
stream_accumulator = _StreamAccumulator(
msg_sequence=first_end_sequence,
initial_content=first_content,
remove_think=remove_think,
)

Expand Down
217 changes: 217 additions & 0 deletions tests/unit_tests/provider/test_localagent_no_duplicate.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,217 @@
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


class RecordingProvider:
"""Non-streaming provider that returns a tool-call on round 1 and plain text on round 2."""

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 check that.',
tool_calls=[
provider_message.ToolCall(
id='call-1',
type='function',
function=provider_message.FunctionCall(
name='exec',
arguments=json.dumps({'command': "python -c 'print(42)'"}),
),
)
],
)

return provider_message.Message(
role='assistant',
content='The result is 42.',
)


class RecordingStreamProvider:
"""Streaming provider that returns a tool-call on round 1 and plain text on round 2."""

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',
content='Let me check that.',
tool_calls=[
provider_message.ToolCall(
id='call-1',
type='function',
function=provider_message.FunctionCall(
name='exec',
arguments=json.dumps({'command': "python -c 'print(42)'"}),
),
)
],
is_final=True,
)
return

yield provider_message.MessageChunk(
role='assistant',
content='The result is 42.',
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='no-dup-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='What is the answer?',
),
use_funcs=[SimpleNamespace(name='exec')],
use_llm_model_uuid='test-model-uuid',
variables={},
)


def _make_app(provider) -> SimpleNamespace:
model = SimpleNamespace(
provider=provider,
model_entity=SimpleNamespace(
uuid='test-model-uuid',
name='test-model',
abilities=['func_call'],
extra_args={},
),
)
return SimpleNamespace(
logger=Mock(),
model_mgr=SimpleNamespace(get_model_by_uuid=AsyncMock(return_value=model)),
tool_mgr=SimpleNamespace(
execute_func_call=AsyncMock(
return_value={
'session_id': 'no-dup-query',
'backend': 'podman',
'status': 'completed',
'ok': True,
'exit_code': 0,
'stdout': '42',
'stderr': '',
'duration_ms': 10,
}
)
),
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),
),
)


@pytest.mark.asyncio
async def test_localagent_non_streaming_no_duplicate():
"""Non-streaming: round-2 content must not contain round-1 text."""
provider = RecordingProvider()
app = _make_app(provider)

runner = LocalAgentRunner(app, pipeline_config={})
query = make_query()

results = [message async for message in runner.run(query)]

# Expect: assistant (tool call) -> tool -> assistant (final answer)
assert [message.role for message in results] == ['assistant', 'tool', 'assistant']

final_message = results[-1]
assert final_message.content == 'The result is 42.'
assert 'Let me check that.' not in final_message.content


@pytest.mark.asyncio
async def test_localagent_streaming_no_duplicate():
"""Streaming: round-2 content must not be re-seeded with round-1 text.

Regression test for the bug where _StreamAccumulator was initialized with
initial_content=first_content, causing every subsequent round to repeat
the entire opening line.
"""
provider = RecordingStreamProvider()
app = _make_app(provider)

adapter = AsyncMock()
adapter.is_stream_output_supported = AsyncMock(return_value=True)

query = make_query()
query.adapter = adapter

runner = LocalAgentRunner(app, pipeline_config={})

results = [message async for message in runner.run(query)]

# All yielded messages should be MessageChunk in streaming mode
assert all(isinstance(message, provider_message.MessageChunk) for message in results)

# The last assistant chunk is the final answer for round 2
assistant_chunks = [m for m in results if m.role == 'assistant']
assert len(assistant_chunks) >= 2

final_chunk = assistant_chunks[-1]
assert final_chunk.content == 'The result is 42.'
assert 'Let me check that.' not in final_chunk.content
Loading