Coverage for src / osiris_cli / tui_app.py: 0%

1import asyncio 

2from typing import Optional, List, Dict, Any 

3from prompt_toolkit.application import Application 

4from prompt_toolkit.buffer import Buffer 

5from prompt_toolkit.document import Document 

6from prompt_toolkit.key_binding import KeyBindings 

7from prompt_toolkit.layout.containers import HSplit, Window, WindowAlign 

8from prompt_toolkit.layout.controls import BufferControl, FormattedTextControl 

9from prompt_toolkit.layout.layout import Layout 

10from prompt_toolkit.layout.dimension import Dimension 

11from prompt_toolkit.lexers import Lexer 

12from prompt_toolkit.styles import Style 

13from prompt_toolkit.formatted_text import ANSI 

14import io 

15from rich.console import Console 

16 

17from .input_normalizer import normalize_user_input 

18from .chat_engine import ChatEngine, StreamingState 

19 

20class AnsiLexer(Lexer): 

21 def lex_document(self, document): 

22 def get_line(lineno): 

23 line_text = document.lines[lineno] 

24 # ANSI class can be used as formatted text directly 

25 return ANSI(line_text)._formatted_text 

26 return get_line 

27 

28class OsirisTUI: 

29 def __init__(self, cli_context: Any, user_header: str = "OSIRIS CLI"): 

30 self.cli = cli_context # The BeautifulCLI instance 

31 self.chat_engine = cli_context.chat_engine 

32 self.user_header = user_header 

33 

34 # State 

35 self.streaming = False 

36 self.follow_up_queue: List[str] = [] 

37 self.stream_task: Optional[asyncio.Task] = None 

38 

39 # Buffers 

40 self.transcript_buffer = Buffer(read_only=True) 

41 self.input_buffer = Buffer(multiline=True) 

42 

43 # UI Components 

44 self.header_control = FormattedTextControl(self._get_header_text) 

45 self.status_control = FormattedTextControl(self._get_status_text) 

46 

47 # Capture console for Rich -> ANSI conversion 

48 self._capture_io = io.StringIO() 

49 self._rich_capture_console = Console( 

50 file=self._capture_io, 

51 force_terminal=True, 

52 width=100, 

53 color_system="truecolor" 

54 ) 

55 

56 self.layout = Layout( 

57 HSplit([ 

58 Window(content=self.header_control, height=1, align=WindowAlign.CENTER, style="class:header"), 

59 Window(content=BufferControl( 

60 buffer=self.transcript_buffer, 

61 lexer=AnsiLexer(), 

62 ), wrap_lines=True, style="class:transcript"), 

63 Window(height=1, char='─', style="class:separator"), 

64 # Status line for working indicators 

65 Window(content=self.status_control, height=1, style="class:status"), 

66 # Input box 

67 Window(content=BufferControl(buffer=self.input_buffer), height=Dimension(min=2, max=5), style="class:input"), 

68 ]) 

69 ) 

70 

71 # Key bindings 

72 self.kb = KeyBindings() 

73 

74 @self.kb.add('c-c') 

75 def _(event): 

76 self._handle_ctrl_c(event) 

77 

78 @self.kb.add('enter') 

79 def _(event): 

80 self._handle_enter(event) 

81 

82 # Style 

83 self.style = Style.from_dict({ 

84 'header': 'bg:#00D9FF #000000 bold', 

85 'transcript': '#ffffff', 

86 'input': '#ffffff', 

87 'status': 'bg:#1a1a1a #00D9FF italic', 

88 'separator': '#444444', 

89 }) 

90 

91 self.app = Application( 

92 layout=self.layout, 

93 key_bindings=self.kb, 

94 style=self.style, 

95 full_screen=True, 

96 mouse_support=True, 

97 ) 

98 

99 def _get_header_text(self): 

100 return f" {self.user_header} " 

101 

102 def _get_status_text(self): 

103 if self.streaming: 

104 queue_count = f" ({len(self.follow_up_queue)} queued)" if self.follow_up_queue else "" 

105 return f" [Streaming...] Ctrl+C to cancel | Enter to queue follow-up{queue_count} " 

106 return " [Idle] Enter to send | Ctrl+C to exit " 

107 

108 def _append_rich(self, rich_obj: Any): 

109 """Convert Rich object to ANSI and append to transcript""" 

110 self._capture_io.seek(0) 

111 self._capture_io.truncate(0) 

112 self._rich_capture_console.print(rich_obj) 

113 ansi_text = self._capture_io.getvalue() 

114 self._append_text(ansi_text) 

115 

116 def _append_text(self, text: str): 

117 new_text = self.transcript_buffer.text + text 

118 self.transcript_buffer.set_document( 

119 Document(text=new_text, cursor_position=len(new_text)), 

120 bypass_readonly=True 

121 ) 

122 self.app.invalidate() 

123 

124 def _handle_ctrl_c(self, event): 

125 if self.streaming: 

126 self._cancel_stream() 

127 else: 

128 event.app.exit() 

129 

130 def _handle_enter(self, event): 

131 text = self.input_buffer.text 

132 if not text.strip(): 

133 return 

134 

135 normalized = normalize_user_input(text) 

136 self.input_buffer.text = "" # Clear 

137 

138 if self.streaming: 

139 self.follow_up_queue.append(normalized) 

140 self._append_text(f"\n\x1b[38;5;244m[Queued]: {normalized.strip()[:40]}...\x1b[0m\n") 

141 else: 

142 self._start_chat(normalized) 

143 

144 def _cancel_stream(self): 

145 if self.stream_task and not self.stream_task.done(): 

146 self.stream_task.cancel() 

147 self.streaming = False 

148 self._append_text("\n\x1b[31m[Stream Cancelled]\x1b[0m\n") 

149 

150 def _start_chat(self, user_input: str): 

151 self.streaming = True 

152 self.stream_task = asyncio.create_task(self._process_chat(user_input)) 

153 

154 async def _process_chat(self, user_input: str): 

155 try: 

156 # 1. Render User Message (once) 

157 self._append_rich(self.cli._build_user_panel(user_input)) 

158 

159 content_buffer = "" 

160 response_start = asyncio.get_event_loop().time() 

161 tool_active = False 

162 

163 # 2. Run Engine 

164 async for event in self.chat_engine.send_message(user_input, system_context=self.cli.system_context): 

165 event_type = event.get("type") 

166 

167 if event_type == "content": 

168 text = event.get("text", "") 

169 content_buffer += text 

170 # We don't stream tokens to transcript line by line to keep it clean, 

171 # but we could update a "Live" area. For now, we append the final panel. 

172 self.app.invalidate() 

173 

174 elif event_type == "tool_call": 

175 tool = event["tool"] 

176 tool_active = True 

177 self._append_text(f"\n\x1b[38;5;214m⚡ Tool: {tool.name}\x1b[0m") 

178 

179 elif event_type == "tool_result": 

180 tool_active = False 

181 

182 elif event_type == "state_change": 

183 if event["state"] == StreamingState.IDLE: 

184 # Stream finished 

185 elapsed = asyncio.get_event_loop().time() - response_start 

186 self._append_text("\n") 

187 self._append_rich(self.cli._render_response_panel(content_buffer, elapsed=elapsed)) 

188 content_buffer = "" 

189 

190 except asyncio.CancelledError: 

191 pass 

192 except Exception as e: 

193 self._append_text(f"\n\x1b[31m[Error]: {str(e)}\x1b[0m\n") 

194 finally: 

195 self.streaming = False 

196 # Check queue 

197 if self.follow_up_queue: 

198 next_input = self.follow_up_queue.pop(0) 

199 self._start_chat(next_input) 

200 

201 async def run(self): 

202 await self.app.run_async()