mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 08:13:11 +03:00
fix(agent): complete native reasoning lifecycle
This commit is contained in:
+14
-1
@@ -949,6 +949,7 @@ class AgentRunner:
|
|||||||
wants_streaming = hook.wants_streaming()
|
wants_streaming = hook.wants_streaming()
|
||||||
|
|
||||||
active_hosted_tools: dict[str, dict[str, Any]] = {}
|
active_hosted_tools: dict[str, dict[str, Any]] = {}
|
||||||
|
native_reasoning_open = False
|
||||||
request_started_at = 0.0
|
request_started_at = 0.0
|
||||||
first_output_at: float | None = None
|
first_output_at: float | None = None
|
||||||
generation_started_at: float | None = None
|
generation_started_at: float | None = None
|
||||||
@@ -971,9 +972,17 @@ class AgentRunner:
|
|||||||
generation_elapsed_s += max(0.0, time.perf_counter() - generation_started_at)
|
generation_elapsed_s += max(0.0, time.perf_counter() - generation_started_at)
|
||||||
generation_started_at = None
|
generation_started_at = None
|
||||||
|
|
||||||
|
async def _close_native_reasoning() -> None:
|
||||||
|
nonlocal native_reasoning_open
|
||||||
|
if not native_reasoning_open:
|
||||||
|
return
|
||||||
|
native_reasoning_open = False
|
||||||
|
await hook.emit_reasoning_end()
|
||||||
|
|
||||||
async def _provider_tool_event(event: dict[str, Any]) -> None:
|
async def _provider_tool_event(event: dict[str, Any]) -> None:
|
||||||
if event.get("kind") != "hosted_tool":
|
if event.get("kind") != "hosted_tool":
|
||||||
return
|
return
|
||||||
|
await _close_native_reasoning()
|
||||||
await hook.on_provider_tool_event(context, event)
|
await hook.on_provider_tool_event(context, event)
|
||||||
call_id = event.get("call_id")
|
call_id = event.get("call_id")
|
||||||
if not call_id:
|
if not call_id:
|
||||||
@@ -991,10 +1000,11 @@ class AgentRunner:
|
|||||||
_generation_delta(delta)
|
_generation_delta(delta)
|
||||||
if delta:
|
if delta:
|
||||||
context.streamed_content = True
|
context.streamed_content = True
|
||||||
|
await _close_native_reasoning()
|
||||||
await hook.on_stream(context, delta)
|
await hook.on_stream(context, delta)
|
||||||
|
|
||||||
async def _thinking(delta: str) -> None:
|
async def _thinking(delta: str) -> None:
|
||||||
nonlocal thinking_buf
|
nonlocal native_reasoning_open, thinking_buf
|
||||||
if not delta:
|
if not delta:
|
||||||
return
|
return
|
||||||
_generation_delta(delta)
|
_generation_delta(delta)
|
||||||
@@ -1004,10 +1014,12 @@ class AgentRunner:
|
|||||||
incremental = new_clean[len(prev_clean):]
|
incremental = new_clean[len(prev_clean):]
|
||||||
if incremental:
|
if incremental:
|
||||||
context.streamed_reasoning = True
|
context.streamed_reasoning = True
|
||||||
|
native_reasoning_open = True
|
||||||
await hook.emit_reasoning(incremental)
|
await hook.emit_reasoning(incremental)
|
||||||
|
|
||||||
async def _stream_recover() -> None:
|
async def _stream_recover() -> None:
|
||||||
_pause_generation()
|
_pause_generation()
|
||||||
|
await _close_native_reasoning()
|
||||||
await hook.on_stream_end(context, resuming=True)
|
await hook.on_stream_end(context, resuming=True)
|
||||||
|
|
||||||
coro = spec.runtime.provider.chat_stream_with_retry(
|
coro = spec.runtime.provider.chat_stream_with_retry(
|
||||||
@@ -1054,6 +1066,7 @@ class AgentRunner:
|
|||||||
error_kind="timeout",
|
error_kind="timeout",
|
||||||
)
|
)
|
||||||
_pause_generation()
|
_pause_generation()
|
||||||
|
await _close_native_reasoning()
|
||||||
if first_output_at is not None:
|
if first_output_at is not None:
|
||||||
response.ttft_ms = max(0, round((first_output_at - request_started_at) * 1000))
|
response.ttft_ms = max(0, round((first_output_at - request_started_at) * 1000))
|
||||||
if generation_elapsed_s > 0:
|
if generation_elapsed_s > 0:
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ channels, gated by ``context.streamed_reasoning`` rather than
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Any
|
||||||
from unittest.mock import AsyncMock, MagicMock
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
@@ -48,6 +49,39 @@ class _StreamRecordingHook(_RecordingHook):
|
|||||||
self.streamed.append(delta)
|
self.streamed.append(delta)
|
||||||
|
|
||||||
|
|
||||||
|
class _LifecycleRecordingHook(AgentHook):
|
||||||
|
def __init__(self) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.events: list[str] = []
|
||||||
|
|
||||||
|
def wants_streaming(self) -> bool:
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def emit_reasoning(self, reasoning_content: str | None) -> None:
|
||||||
|
if reasoning_content:
|
||||||
|
self.events.append(f"reasoning:{reasoning_content}")
|
||||||
|
|
||||||
|
async def emit_reasoning_end(self) -> None:
|
||||||
|
self.events.append("reasoning_end")
|
||||||
|
|
||||||
|
async def on_stream(self, _ctx: AgentHookContext, delta: str) -> None:
|
||||||
|
self.events.append(f"content:{delta}")
|
||||||
|
|
||||||
|
async def on_stream_end(self, _ctx: AgentHookContext, *, resuming: bool) -> None:
|
||||||
|
self.events.append(f"stream_end:{resuming}")
|
||||||
|
|
||||||
|
async def before_execute_tools(self, context: AgentHookContext) -> None:
|
||||||
|
names = ",".join(call.name for call in context.tool_calls)
|
||||||
|
self.events.append(f"local_tools:{names}")
|
||||||
|
|
||||||
|
async def on_provider_tool_event(
|
||||||
|
self,
|
||||||
|
_context: AgentHookContext,
|
||||||
|
event: dict[str, Any],
|
||||||
|
) -> None:
|
||||||
|
self.events.append(f"hosted_tool:{event.get('phase')}")
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_runner_preserves_reasoning_fields_in_assistant_history():
|
async def test_runner_preserves_reasoning_fields_in_assistant_history():
|
||||||
"""Reasoning fields ride along on the persisted assistant message so
|
"""Reasoning fields ride along on the persisted assistant message so
|
||||||
@@ -371,6 +405,155 @@ async def test_runner_streams_native_thinking_deltas_without_post_hoc_dup():
|
|||||||
assert hook.emitted == ["part1", "part2"]
|
assert hook.emitted == ["part1", "part2"]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_runner_closes_native_reasoning_before_streaming_answer():
|
||||||
|
from nanobot.agent.runner import AgentRunner
|
||||||
|
|
||||||
|
provider = MagicMock()
|
||||||
|
|
||||||
|
async def chat_stream_with_retry(
|
||||||
|
*, on_content_delta=None, on_thinking_delta=None, **kwargs
|
||||||
|
):
|
||||||
|
if on_thinking_delta:
|
||||||
|
await on_thinking_delta("inspect")
|
||||||
|
if on_content_delta:
|
||||||
|
await on_content_delta("done")
|
||||||
|
return LLMResponse(content="done", tool_calls=[], usage=None)
|
||||||
|
|
||||||
|
provider.chat_stream_with_retry = chat_stream_with_retry
|
||||||
|
tools = MagicMock()
|
||||||
|
tools.get_definitions.return_value = []
|
||||||
|
hook = _LifecycleRecordingHook()
|
||||||
|
|
||||||
|
result = await AgentRunner().run(make_run_spec(
|
||||||
|
provider,
|
||||||
|
initial_messages=[{"role": "user", "content": "q"}],
|
||||||
|
tools=tools,
|
||||||
|
model="test-model",
|
||||||
|
max_iterations=1,
|
||||||
|
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
|
||||||
|
hook=hook,
|
||||||
|
))
|
||||||
|
|
||||||
|
assert result.final_content == "done"
|
||||||
|
assert hook.events == [
|
||||||
|
"reasoning:inspect",
|
||||||
|
"reasoning_end",
|
||||||
|
"content:done",
|
||||||
|
"stream_end:False",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_runner_closes_native_reasoning_before_local_tool_execution():
|
||||||
|
from nanobot.agent.runner import AgentRunner
|
||||||
|
|
||||||
|
provider = MagicMock()
|
||||||
|
responses = iter([
|
||||||
|
LLMResponse(
|
||||||
|
content="",
|
||||||
|
finish_reason="tool_calls",
|
||||||
|
tool_calls=[ToolCallRequest(id="call-1", name="list_dir", arguments={"path": "."})],
|
||||||
|
usage=None,
|
||||||
|
),
|
||||||
|
LLMResponse(content="done", tool_calls=[], usage=None),
|
||||||
|
])
|
||||||
|
|
||||||
|
async def chat_stream_with_retry(
|
||||||
|
*, on_content_delta=None, on_thinking_delta=None, **kwargs
|
||||||
|
):
|
||||||
|
response = next(responses)
|
||||||
|
if response.tool_calls:
|
||||||
|
if on_thinking_delta:
|
||||||
|
await on_thinking_delta("inspect")
|
||||||
|
elif on_content_delta:
|
||||||
|
await on_content_delta("done")
|
||||||
|
return response
|
||||||
|
|
||||||
|
provider.chat_stream_with_retry = chat_stream_with_retry
|
||||||
|
tools = MagicMock()
|
||||||
|
tools.get_definitions.return_value = []
|
||||||
|
tools.execute = AsyncMock(return_value="tool result")
|
||||||
|
hook = _LifecycleRecordingHook()
|
||||||
|
|
||||||
|
result = await AgentRunner().run(make_run_spec(
|
||||||
|
provider,
|
||||||
|
initial_messages=[{"role": "user", "content": "inspect"}],
|
||||||
|
tools=tools,
|
||||||
|
model="test-model",
|
||||||
|
max_iterations=2,
|
||||||
|
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
|
||||||
|
hook=hook,
|
||||||
|
))
|
||||||
|
|
||||||
|
assert result.final_content == "done"
|
||||||
|
assert hook.events == [
|
||||||
|
"reasoning:inspect",
|
||||||
|
"reasoning_end",
|
||||||
|
"stream_end:True",
|
||||||
|
"local_tools:list_dir",
|
||||||
|
"content:done",
|
||||||
|
"stream_end:False",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_runner_closes_native_reasoning_before_hosted_tool_event():
|
||||||
|
from nanobot.agent.runner import AgentRunner
|
||||||
|
|
||||||
|
provider = MagicMock()
|
||||||
|
|
||||||
|
async def chat_stream_with_retry(
|
||||||
|
*, on_content_delta=None, on_thinking_delta=None, on_tool_call_delta=None, **kwargs
|
||||||
|
):
|
||||||
|
if on_thinking_delta:
|
||||||
|
await on_thinking_delta("search")
|
||||||
|
if on_tool_call_delta:
|
||||||
|
await on_tool_call_delta({
|
||||||
|
"kind": "hosted_tool",
|
||||||
|
"phase": "start",
|
||||||
|
"call_id": "search-1",
|
||||||
|
"name": "web_search",
|
||||||
|
"arguments": {"query": "nanobot"},
|
||||||
|
})
|
||||||
|
await on_tool_call_delta({
|
||||||
|
"kind": "hosted_tool",
|
||||||
|
"phase": "end",
|
||||||
|
"call_id": "search-1",
|
||||||
|
"name": "web_search",
|
||||||
|
"arguments": {"query": "nanobot"},
|
||||||
|
"result": {"count": 1},
|
||||||
|
})
|
||||||
|
if on_content_delta:
|
||||||
|
await on_content_delta("done")
|
||||||
|
return LLMResponse(content="done", tool_calls=[], usage=None)
|
||||||
|
|
||||||
|
provider.chat_stream_with_retry = chat_stream_with_retry
|
||||||
|
tools = MagicMock()
|
||||||
|
tools.get_definitions.return_value = []
|
||||||
|
hook = _LifecycleRecordingHook()
|
||||||
|
|
||||||
|
result = await AgentRunner().run(make_run_spec(
|
||||||
|
provider,
|
||||||
|
initial_messages=[{"role": "user", "content": "search"}],
|
||||||
|
tools=tools,
|
||||||
|
model="test-model",
|
||||||
|
max_iterations=1,
|
||||||
|
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
|
||||||
|
hook=hook,
|
||||||
|
))
|
||||||
|
|
||||||
|
assert result.final_content == "done"
|
||||||
|
assert hook.events == [
|
||||||
|
"reasoning:search",
|
||||||
|
"reasoning_end",
|
||||||
|
"hosted_tool:start",
|
||||||
|
"hosted_tool:end",
|
||||||
|
"content:done",
|
||||||
|
"stream_end:False",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_runner_strips_thinking_tags_from_native_thinking_deltas():
|
async def test_runner_strips_thinking_tags_from_native_thinking_deltas():
|
||||||
from nanobot.agent.runner import AgentRunner
|
from nanobot.agent.runner import AgentRunner
|
||||||
|
|||||||
Reference in New Issue
Block a user