mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/cli: assistant/reasoning on new lines; flush markdown on source change
Labels end with a newline before content. Switching reasoning↔assistant or hitting tools/errors flushes the assistant markdown buffer so streams do not mix.
This commit is contained in:
+75
-48
@@ -4,7 +4,7 @@ import contextlib
|
|||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
from contextvars import ContextVar
|
from contextvars import ContextVar
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING, Literal
|
||||||
|
|
||||||
import click
|
import click
|
||||||
|
|
||||||
@@ -32,6 +32,8 @@ _TOOL_ARGS_PREVIEW = 80
|
|||||||
_verbose_tool_results: ContextVar[bool] = ContextVar("verbose_tool_results", default=False)
|
_verbose_tool_results: ContextVar[bool] = ContextVar("verbose_tool_results", default=False)
|
||||||
_markdown_enabled: ContextVar[bool] = ContextVar("markdown_enabled", default=True)
|
_markdown_enabled: ContextVar[bool] = ContextVar("markdown_enabled", default=True)
|
||||||
|
|
||||||
|
type StreamSource = Literal["reasoning", "assistant"]
|
||||||
|
|
||||||
|
|
||||||
def set_verbose_tool_results(enabled: bool) -> None: # noqa: FBT001
|
def set_verbose_tool_results(enabled: bool) -> None: # noqa: FBT001
|
||||||
"""Set whether tool results print in full (True) or as a short preview."""
|
"""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:
|
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:
|
if not body and not label:
|
||||||
return 0
|
return 0
|
||||||
text = label + body
|
# Label is on its own line; body may contain newlines.
|
||||||
# +1 for trailing newline after the stream block in render_events.
|
text = f"{label}\n{body}" if label else body
|
||||||
return text.count("\n") + 1
|
return text.count("\n") + 1
|
||||||
|
|
||||||
|
|
||||||
def print_markdown(text: str, *, label: str = "assistant:") -> None:
|
def print_markdown(text: str, *, label: str = "assistant:") -> None:
|
||||||
"""Render *text* as markdown via Rich, with a cyan label."""
|
"""Render *text* as markdown via Rich; *label* on its own line when set."""
|
||||||
from rich.console import Console
|
from rich.console import Console
|
||||||
from rich.markdown import Markdown
|
from rich.markdown import Markdown
|
||||||
from rich.text import Text
|
from rich.text import Text
|
||||||
|
|
||||||
console = Console(file=sys.stdout, highlight=False)
|
console = Console(file=sys.stdout, highlight=False)
|
||||||
if label:
|
if label:
|
||||||
console.print(Text(label, style="cyan"), end="")
|
console.print(Text(label, style="cyan"))
|
||||||
console.print(Markdown(text))
|
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
|
async def render_events( # noqa: C901, PLR0912, PLR0915
|
||||||
events: AsyncIterator[AgentEvent],
|
events: AsyncIterator[AgentEvent],
|
||||||
*,
|
*,
|
||||||
@@ -115,38 +131,65 @@ async def render_events( # noqa: C901, PLR0912, PLR0915
|
|||||||
) -> None:
|
) -> None:
|
||||||
"""Print agent events to the terminal (text deltas stream as they arrive).
|
"""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
|
Assistant and reasoning each start on a new line after their label. When the
|
||||||
text is replaced at end-of-turn with a Rich markdown render.
|
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
|
show_full = get_verbose_tool_results() if verbose is None else verbose
|
||||||
use_markdown = get_markdown_enabled() if markdown is None else markdown
|
use_markdown = get_markdown_enabled() if markdown is None else markdown
|
||||||
pretty = bool(use_markdown and markdown_render_available())
|
pretty = bool(use_markdown and markdown_render_available())
|
||||||
|
|
||||||
printed_reasoning = False
|
source: StreamSource | None = None
|
||||||
printed_text = False
|
|
||||||
assistant_buf: list[str] = []
|
assistant_buf: list[str] = []
|
||||||
# Tools/errors after assistant text: skip replace (would erase tool lines).
|
printed_reasoning = False
|
||||||
interrupted_by_other = 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:
|
async for event in events:
|
||||||
if isinstance(event, ReasoningDeltaEvent):
|
if isinstance(event, ReasoningDeltaEvent):
|
||||||
if not printed_reasoning:
|
begin_reasoning()
|
||||||
click.echo()
|
|
||||||
click.secho("reasoning: ", fg="bright_black", nl=False)
|
|
||||||
printed_reasoning = True
|
|
||||||
_echo_stream(event.content)
|
_echo_stream(event.content)
|
||||||
elif isinstance(event, TextDeltaEvent):
|
elif isinstance(event, TextDeltaEvent):
|
||||||
if printed_reasoning and not printed_text:
|
begin_assistant()
|
||||||
click.echo()
|
|
||||||
if not printed_text:
|
|
||||||
click.echo()
|
|
||||||
click.secho("assistant: ", fg="cyan", nl=False)
|
|
||||||
printed_text = True
|
|
||||||
assistant_buf.append(event.content)
|
assistant_buf.append(event.content)
|
||||||
_echo_stream(event.content)
|
_echo_stream(event.content)
|
||||||
elif isinstance(event, ToolCallEvent):
|
elif isinstance(event, ToolCallEvent):
|
||||||
if printed_text:
|
flush_assistant()
|
||||||
interrupted_by_other = True
|
|
||||||
call = event.tool_call
|
call = event.tool_call
|
||||||
if isinstance(call, AssistantFunctionToolCall):
|
if isinstance(call, AssistantFunctionToolCall):
|
||||||
args = _preview(call.function.arguments, _TOOL_ARGS_PREVIEW)
|
args = _preview(call.function.arguments, _TOOL_ARGS_PREVIEW)
|
||||||
@@ -154,8 +197,7 @@ async def render_events( # noqa: C901, PLR0912, PLR0915
|
|||||||
else:
|
else:
|
||||||
click.secho(f"\n[tool] custom id={call.id}", fg="yellow")
|
click.secho(f"\n[tool] custom id={call.id}", fg="yellow")
|
||||||
elif isinstance(event, ToolResultEvent):
|
elif isinstance(event, ToolResultEvent):
|
||||||
if printed_text:
|
flush_assistant()
|
||||||
interrupted_by_other = True
|
|
||||||
content = event.message.content
|
content = event.message.content
|
||||||
if show_full:
|
if show_full:
|
||||||
click.secho(f"[tool ok]\n{content}", fg="magenta")
|
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", " ")
|
one_line = preview.replace("\n", " ")
|
||||||
click.secho(f"[tool ok] {one_line}", fg="magenta")
|
click.secho(f"[tool ok] {one_line}", fg="magenta")
|
||||||
elif isinstance(event, ErrorEvent):
|
elif isinstance(event, ErrorEvent):
|
||||||
if printed_text:
|
flush_assistant()
|
||||||
interrupted_by_other = True
|
|
||||||
suffix = ""
|
suffix = ""
|
||||||
if event.source:
|
if event.source:
|
||||||
suffix += f" source={event.source}"
|
suffix += f" source={event.source}"
|
||||||
@@ -173,15 +214,13 @@ async def render_events( # noqa: C901, PLR0912, PLR0915
|
|||||||
suffix += " (fatal)"
|
suffix += " (fatal)"
|
||||||
click.secho(f"\n[error]{suffix} {event.message}", fg="bright_red")
|
click.secho(f"\n[error]{suffix} {event.message}", fg="bright_red")
|
||||||
elif isinstance(event, CancelledEvent):
|
elif isinstance(event, CancelledEvent):
|
||||||
if printed_text:
|
flush_assistant()
|
||||||
interrupted_by_other = True
|
|
||||||
if event.reason:
|
if event.reason:
|
||||||
click.secho(f"\n[cancelled] {event.reason}", fg="yellow")
|
click.secho(f"\n[cancelled] {event.reason}", fg="yellow")
|
||||||
else:
|
else:
|
||||||
click.secho("\n[cancelled]", fg="yellow")
|
click.secho("\n[cancelled]", fg="yellow")
|
||||||
elif isinstance(event, MaxRoundsEvent):
|
elif isinstance(event, MaxRoundsEvent):
|
||||||
if printed_text:
|
flush_assistant()
|
||||||
interrupted_by_other = True
|
|
||||||
if event.continued:
|
if event.continued:
|
||||||
click.secho(
|
click.secho(
|
||||||
f"\n[max rounds {event.rounds} reached — continuing with a higher allowance]",
|
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.
|
# AssistantMessageEvent — text already shown via TextDeltaEvent.
|
||||||
_ = event
|
_ = event
|
||||||
|
|
||||||
full_assistant = "".join(assistant_buf)
|
# End-of-turn: flush any open assistant segment; close reasoning stream.
|
||||||
if pretty and printed_text and full_assistant.strip() and not interrupted_by_other:
|
if assistant_buf or printed_assistant:
|
||||||
# Replace plain stream with markdown (assistant block only).
|
flush_assistant()
|
||||||
lines = _line_count_for_clear("assistant: ", full_assistant)
|
elif printed_reasoning:
|
||||||
# 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:
|
|
||||||
click.echo()
|
click.echo()
|
||||||
click.echo()
|
click.echo()
|
||||||
|
|||||||
@@ -41,9 +41,43 @@ async def test_render_reasoning_and_text(capsys: pytest.CaptureFixture[str]) ->
|
|||||||
assert "think" in out
|
assert "think" in out
|
||||||
assert "assistant:" in out
|
assert "assistant:" in out
|
||||||
assert "hello" 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)
|
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:
|
async def test_tool_result_preview_vs_verbose(capsys: pytest.CaptureFixture[str]) -> None:
|
||||||
long = "x" * 200
|
long = "x" * 200
|
||||||
msg = ToolChatMessage(content=long, tool_call_id="1")
|
msg = ToolChatMessage(content=long, tool_call_id="1")
|
||||||
|
|||||||
Reference in New Issue
Block a user