返回 JoyAI-Echo
stream.py
1 """Streaming renderer for CLI output.
2
3 Uses Rich Live with auto_refresh=False for stable, flicker-free
4 markdown rendering during streaming. Ellipsis mode handles overflow.
5 """
6
7 from __future__ import annotations
8
9 import sys
10 import time
11
12 from rich.console import Console
13 from rich.live import Live
14 from rich.markdown import Markdown
15 from rich.text import Text
16
17 from nanobot import __logo__
18
19
20 def _make_console() -> Console:
21 """Create a Console that emits plain text when stdout is not a TTY.
22
23 Rich's spinner, Live render, and cursor-visibility escape codes all
24 key off ``Console.is_terminal``. Forcing ``force_terminal=True`` overrode
25 the ``isatty()`` check and caused control sequences (``\\x1b[?25l``,
26 braille spinner frames) to pollute programmatic consumers such as
27 ``docker exec -i`` or pipes, even with ``NO_COLOR`` or ``TERM=dumb``.
28 Deferring to ``isatty()`` keeps Rich output in interactive terminals
29 and plain text everywhere else (#3265).
30 """
31 return Console(file=sys.stdout, force_terminal=sys.stdout.isatty())
32
33
34 class ThinkingSpinner:
35 """Spinner that shows 'nanobot is thinking...' with pause support."""
36
37 def __init__(self, console: Console | None = None):
38 c = console or _make_console()
39 self._spinner = c.status("[dim]nanobot is thinking...[/dim]", spinner="dots")
40 self._active = False
41
42 def __enter__(self):
43 self._spinner.start()
44 self._active = True
45 return self
46
47 def __exit__(self, *exc):
48 self._active = False
49 self._spinner.stop()
50 return False
51
52 def pause(self):
53 """Context manager: temporarily stop spinner for clean output."""
54 from contextlib import contextmanager
55
56 @contextmanager
57 def _ctx():
58 if self._spinner and self._active:
59 self._spinner.stop()
60 try:
61 yield
62 finally:
63 if self._spinner and self._active:
64 self._spinner.start()
65
66 return _ctx()
67
68
69 class StreamRenderer:
70 """Rich Live streaming with markdown. auto_refresh=False avoids render races.
71
72 Deltas arrive pre-filtered (no <think> tags) from the agent loop.
73
74 Flow per round:
75 spinner -> first visible delta -> header + Live renders ->
76 on_end -> Live stops (content stays on screen)
77 """
78
79 def __init__(self, render_markdown: bool = True, show_spinner: bool = True):
80 self._md = render_markdown
81 self._show_spinner = show_spinner
82 self._buf = ""
83 self._live: Live | None = None
84 self._t = 0.0
85 self.streamed = False
86 self._spinner: ThinkingSpinner | None = None
87 self._start_spinner()
88
89 def _render(self):
90 return Markdown(self._buf) if self._md and self._buf else Text(self._buf or "")
91
92 def _start_spinner(self) -> None:
93 if self._show_spinner:
94 self._spinner = ThinkingSpinner()
95 self._spinner.__enter__()
96
97 def _stop_spinner(self) -> None:
98 if self._spinner:
99 self._spinner.__exit__(None, None, None)
100 self._spinner = None
101
102 async def on_delta(self, delta: str) -> None:
103 self.streamed = True
104 self._buf += delta
105 if self._live is None:
106 if not self._buf.strip():
107 return
108 self._stop_spinner()
109 c = _make_console()
110 c.print()
111 c.print(f"[cyan]{__logo__} nanobot[/cyan]")
112 self._live = Live(self._render(), console=c, auto_refresh=False)
113 self._live.start()
114 now = time.monotonic()
115 if (now - self._t) > 0.15:
116 self._live.update(self._render())
117 self._live.refresh()
118 self._t = now
119
120 async def on_end(self, *, resuming: bool = False) -> None:
121 if self._live:
122 self._live.update(self._render())
123 self._live.refresh()
124 self._live.stop()
125 self._live = None
126 self._stop_spinner()
127 if resuming:
128 self._buf = ""
129 self._start_spinner()
130 else:
131 _make_console().print()
132
133 def stop_for_input(self) -> None:
134 """Stop spinner before user input to avoid prompt_toolkit conflicts."""
135 self._stop_spinner()
136
137 async def close(self) -> None:
138 """Stop spinner/live without rendering a final streamed round."""
139 if self._live:
140 self._live.stop()
141 self._live = None
142 self._stop_spinner()
143
143 lines PYTHON