From b8ed9e0ee9b6dd0cba430a31e7302f5fe9018888 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Fri, 17 Jul 2026 15:36:29 +0800 Subject: [PATCH] core/cli: assistant/reasoning on new lines; flush markdown on source change MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Labels end with a newline before content. Switching reasoning↔assistant or hitting tools/errors flushes the assistant markdown buffer so streams do not mix. --- src/plyngent/cli/display.py | 125 ++++++++++++++++++++------------- tests/test_cli/test_display.py | 38 +++++++++- 2 files changed, 112 insertions(+), 51 deletions(-) diff --git a/src/plyngent/cli/display.py b/src/plyngent/cli/display.py index 2a19051..348c6f2 100644 --- a/src/plyngent/cli/display.py +++ b/src/plyngent/cli/display.py @@ -4,7 +4,7 @@ import contextlib import os import sys from contextvars import ContextVar -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import click @@ -32,6 +32,8 @@ _TOOL_ARGS_PREVIEW = 80 _verbose_tool_results: ContextVar[bool] = ContextVar("verbose_tool_results", default=False) _markdown_enabled: ContextVar[bool] = ContextVar("markdown_enabled", default=True) +type StreamSource = Literal["reasoning", "assistant"] + def set_verbose_tool_results(enabled: bool) -> None: # noqa: FBT001 """Set whether tool results print in full (True) or as a short preview.""" @@ -87,26 +89,40 @@ def _clear_streamed_lines(line_count: int) -> None: def _line_count_for_clear(label: str, body: str) -> int: - """Approximate terminal lines used by ``label + body`` for cursor erase.""" + """Approximate terminal lines used by ``label\\n + body`` for cursor erase.""" if not body and not label: return 0 - text = label + body - # +1 for trailing newline after the stream block in render_events. + # Label is on its own line; body may contain newlines. + text = f"{label}\n{body}" if label else body return text.count("\n") + 1 -def print_markdown(text: str, *, label: str = "assistant: ") -> None: - """Render *text* as markdown via Rich, with a cyan label.""" +def print_markdown(text: str, *, label: str = "assistant:") -> None: + """Render *text* as markdown via Rich; *label* on its own line when set.""" from rich.console import Console from rich.markdown import Markdown from rich.text import Text console = Console(file=sys.stdout, highlight=False) if label: - console.print(Text(label, style="cyan"), end="") + console.print(Text(label, style="cyan")) console.print(Markdown(text)) +def _flush_assistant_markdown(body: str, *, pretty: bool) -> None: + """Replace the plain assistant stream with markdown when enabled.""" + if not body.strip(): + click.echo() + return + if pretty: + lines = _line_count_for_clear("assistant:", body) + _clear_streamed_lines(lines) + print_markdown(body, label="assistant:") + click.echo() + else: + click.echo() + + async def render_events( # noqa: C901, PLR0912, PLR0915 events: AsyncIterator[AgentEvent], *, @@ -115,38 +131,65 @@ async def render_events( # noqa: C901, PLR0912, PLR0915 ) -> None: """Print agent events to the terminal (text deltas stream as they arrive). - When markdown is enabled and stdout is a TTY, the plain streamed assistant - text is replaced at end-of-turn with a Rich markdown render. + Assistant and reasoning each start on a new line after their label. When the + content source changes (reasoning ↔ assistant, or tools/errors), the + assistant markdown buffer is flushed so streams do not mix and Rich can + re-render completed assistant segments. """ show_full = get_verbose_tool_results() if verbose is None else verbose use_markdown = get_markdown_enabled() if markdown is None else markdown pretty = bool(use_markdown and markdown_render_available()) - printed_reasoning = False - printed_text = False + source: StreamSource | None = None assistant_buf: list[str] = [] - # Tools/errors after assistant text: skip replace (would erase tool lines). - interrupted_by_other = False + printed_reasoning = False + printed_assistant = False + + def flush_assistant() -> None: + nonlocal source, assistant_buf, printed_assistant + if source != "assistant" and not assistant_buf: + return + body = "".join(assistant_buf) + assistant_buf = [] + if printed_assistant: + _flush_assistant_markdown(body, pretty=pretty) + printed_assistant = False + if source == "assistant": + source = None + + def begin_reasoning() -> None: + nonlocal source, printed_reasoning + if source == "reasoning": + return + if source == "assistant": + flush_assistant() + click.echo() + click.secho("reasoning:", fg="bright_black") + source = "reasoning" + printed_reasoning = True + + def begin_assistant() -> None: + nonlocal source, printed_assistant + if source == "assistant": + return + if source == "reasoning": + click.echo() # end reasoning stream line + source = None + click.echo() + click.secho("assistant:", fg="cyan") + source = "assistant" + printed_assistant = True async for event in events: if isinstance(event, ReasoningDeltaEvent): - if not printed_reasoning: - click.echo() - click.secho("reasoning: ", fg="bright_black", nl=False) - printed_reasoning = True + begin_reasoning() _echo_stream(event.content) elif isinstance(event, TextDeltaEvent): - if printed_reasoning and not printed_text: - click.echo() - if not printed_text: - click.echo() - click.secho("assistant: ", fg="cyan", nl=False) - printed_text = True + begin_assistant() assistant_buf.append(event.content) _echo_stream(event.content) elif isinstance(event, ToolCallEvent): - if printed_text: - interrupted_by_other = True + flush_assistant() call = event.tool_call if isinstance(call, AssistantFunctionToolCall): args = _preview(call.function.arguments, _TOOL_ARGS_PREVIEW) @@ -154,8 +197,7 @@ async def render_events( # noqa: C901, PLR0912, PLR0915 else: click.secho(f"\n[tool] custom id={call.id}", fg="yellow") elif isinstance(event, ToolResultEvent): - if printed_text: - interrupted_by_other = True + flush_assistant() content = event.message.content if show_full: click.secho(f"[tool ok]\n{content}", fg="magenta") @@ -164,8 +206,7 @@ async def render_events( # noqa: C901, PLR0912, PLR0915 one_line = preview.replace("\n", " ") click.secho(f"[tool ok] {one_line}", fg="magenta") elif isinstance(event, ErrorEvent): - if printed_text: - interrupted_by_other = True + flush_assistant() suffix = "" if event.source: suffix += f" source={event.source}" @@ -173,15 +214,13 @@ async def render_events( # noqa: C901, PLR0912, PLR0915 suffix += " (fatal)" click.secho(f"\n[error]{suffix} {event.message}", fg="bright_red") elif isinstance(event, CancelledEvent): - if printed_text: - interrupted_by_other = True + flush_assistant() if event.reason: click.secho(f"\n[cancelled] {event.reason}", fg="yellow") else: click.secho("\n[cancelled]", fg="yellow") elif isinstance(event, MaxRoundsEvent): - if printed_text: - interrupted_by_other = True + flush_assistant() if event.continued: click.secho( f"\n[max rounds {event.rounds} reached — continuing with a higher allowance]", @@ -195,21 +234,9 @@ async def render_events( # noqa: C901, PLR0912, PLR0915 # AssistantMessageEvent — text already shown via TextDeltaEvent. _ = event - full_assistant = "".join(assistant_buf) - if pretty and printed_text and full_assistant.strip() and not interrupted_by_other: - # Replace plain stream with markdown (assistant block only). - lines = _line_count_for_clear("assistant: ", full_assistant) - # Also account for blank line before label when reasoning was shown. - if printed_reasoning: - lines += 1 - _clear_streamed_lines(lines) - if printed_reasoning: - # Reasoning already printed above; leave it and only re-print assistant. - pass - print_markdown(full_assistant, label="assistant: ") - click.echo() - return - - if printed_text or printed_reasoning: + # End-of-turn: flush any open assistant segment; close reasoning stream. + if assistant_buf or printed_assistant: + flush_assistant() + elif printed_reasoning: click.echo() click.echo() diff --git a/tests/test_cli/test_display.py b/tests/test_cli/test_display.py index e3c3faa..140f725 100644 --- a/tests/test_cli/test_display.py +++ b/tests/test_cli/test_display.py @@ -37,13 +37,47 @@ async def test_render_reasoning_and_text(capsys: pytest.CaptureFixture[str]) -> ) ) out = capsys.readouterr().out - assert "reasoning: " in out + assert "reasoning:" in out assert "think" in out - assert "assistant: " in out + assert "assistant:" in out assert "hello" in out + # Labels on their own lines (content begins after newline). + assert "assistant:\nhello" in out or "assistant:\r\nhello" in out set_markdown_enabled(True) +async def test_flush_markdown_on_source_change( + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Assistant segment before tools is flushed so later text is a new segment.""" + monkeypatch.setattr("plyngent.cli.display.markdown_render_available", lambda: True) + set_markdown_enabled(True) + from plyngent.agent import ToolCallEvent + from plyngent.lmproto.openai_compatible.model import ( + AssistantFunctionTool, + AssistantFunctionToolCall, + ) + + call = AssistantFunctionToolCall( + id="1", + function=AssistantFunctionTool(name="read_file", arguments="{}"), + ) + await render_events( + _aiter( + [ + TextDeltaEvent(content="before **tool**"), + ToolCallEvent(tool_call=call), + TextDeltaEvent(content="after"), + ] + ), + markdown=True, + ) + out = capsys.readouterr().out + assert "[tool]" in out + assert "after" in out + + async def test_tool_result_preview_vs_verbose(capsys: pytest.CaptureFixture[str]) -> None: long = "x" * 200 msg = ToolChatMessage(content=long, tool_call_id="1")