Skip to content

Commit a7b55ae

Browse files
Rahel-ARahel Ahmed
Rahel-A
and
Rahel Ahmed
authored
fix: accept text through pipes in chat (#259)
* fix: accept text through pipes in chat For the `q chat` command, also allow piping text into the q chat command. * Update mod.rs * Print to stdout if input is piped * Combine the 2 try_chat functions * code format --------- Co-authored-by: Rahel Ahmed <[email protected]>
1 parent 496c04a commit a7b55ae

File tree

2 files changed

+169
-92
lines changed

2 files changed

+169
-92
lines changed

‎crates/q_cli/src/cli/chat/mod.rs‎

Lines changed: 134 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
mod api;
22
mod parse;
33
mod prompt;
4+
mod terminal;
45

56
use std::io::{
6-
Stderr,
7+
IsTerminal,
8+
Read,
79
Write,
8-
stderr,
10+
stdin,
911
};
1012
use std::process::ExitCode;
1113
use std::time::Duration;
@@ -17,11 +19,10 @@ use crossterm::style::{
1719
Print,
1820
};
1921
use crossterm::{
20-
ExecutableCommand,
21-
QueueableCommand,
2222
cursor,
23+
execute,
24+
queue,
2325
style,
24-
terminal,
2526
};
2627
use eyre::{
2728
Result,
@@ -58,7 +59,7 @@ enum ApiResponse {
5859
End,
5960
}
6061

61-
pub async fn chat(input: String) -> Result<ExitCode> {
62+
pub async fn chat(mut input: String) -> Result<ExitCode> {
6263
if !fig_util::system_info::in_cloudshell() && !fig_auth::is_logged_in().await {
6364
bail!(
6465
"You are not logged in, please log in with {}",
@@ -68,76 +69,92 @@ pub async fn chat(input: String) -> Result<ExitCode> {
6869

6970
region_check("chat")?;
7071

71-
let mut stderr = stderr();
72-
let result = try_chat(&mut stderr, input).await;
72+
let stdin = stdin();
73+
let is_interactive = stdin.is_terminal();
7374

74-
stderr
75-
.queue(style::SetAttribute(Attribute::Reset))?
76-
.queue(style::ResetColor)?
77-
.flush()
78-
.ok();
75+
if !is_interactive {
76+
// append to input string any extra info that was provided.
77+
stdin.lock().read_to_string(&mut input).unwrap();
78+
}
79+
let mut output = terminal::new(is_interactive);
80+
let result = try_chat(&mut output, input, is_interactive).await;
81+
82+
if is_interactive {
83+
queue!(output, style::SetAttribute(Attribute::Reset), style::ResetColor).ok();
84+
}
85+
output.flush().ok();
7986

8087
result.map(|_| ExitCode::SUCCESS)
8188
}
8289

83-
async fn try_chat(stderr: &mut Stderr, mut input: String) -> Result<()> {
84-
let mut rl = rl()?;
90+
async fn try_chat<W: Write>(output: &mut W, mut input: String, interactive: bool) -> Result<()> {
91+
let mut rl = if interactive { Some(rl()?) } else { None };
8592
let client = StreamingClient::new().await?;
8693
let mut rx = None;
8794
let mut conversation_id: Option<String> = None;
8895
let mut message_id = None;
8996

9097
loop {
91-
// Make request with input
92-
match input.trim() {
93-
"exit" | "quit" => {
94-
if let Some(conversation_id) = conversation_id {
95-
fig_telemetry::send_end_chat(conversation_id.clone()).await;
96-
}
97-
return Ok(());
98-
},
99-
_ => (),
100-
}
101-
102-
if !input.is_empty() {
103-
stderr.queue(style::SetForegroundColor(Color::Magenta))?;
104-
if input.contains("@history") {
105-
stderr.queue(style::Print("Using shell history\n"))?;
98+
// Make request with input, otherwise used already provided buffer for input
99+
if !interactive {
100+
rx = Some(send_message(client.clone(), input.clone(), conversation_id.clone()).await?);
101+
} else {
102+
match input.trim() {
103+
"exit" | "quit" => {
104+
if let Some(conversation_id) = conversation_id {
105+
fig_telemetry::send_end_chat(conversation_id.clone()).await;
106+
}
107+
return Ok(());
108+
},
109+
_ => (),
106110
}
107111

108-
if input.contains("@git") {
109-
stderr.queue(style::Print("Using git context\n"))?;
110-
}
112+
if !input.is_empty() {
113+
queue!(output, style::SetForegroundColor(Color::Magenta))?;
114+
if input.contains("@history") {
115+
queue!(output, style::Print("Using shell history\n"))?;
116+
}
111117

112-
if input.contains("@env") {
113-
stderr.queue(style::Print("Using environment\n"))?;
114-
}
118+
if input.contains("@git") {
119+
queue!(output, style::Print("Using git context\n"))?;
120+
}
121+
122+
if input.contains("@env") {
123+
queue!(output, style::Print("Using environment\n"))?;
124+
}
115125

116-
rx = Some(send_message(client.clone(), input, conversation_id.clone()).await?);
117-
stderr
118-
.queue(style::SetForegroundColor(Color::Reset))?
119-
.execute(style::Print("\n"))?;
120-
} else if fig_settings::settings::get_bool_or("chat.greeting.enabled", true) {
121-
stderr.execute(style::Print(format!(
122-
"
126+
rx = Some(send_message(client.clone(), input.clone(), conversation_id.clone()).await?);
127+
queue!(output, style::SetForegroundColor(Color::Reset))?;
128+
execute!(output, style::Print("\n"))?;
129+
} else if fig_settings::settings::get_bool_or("chat.greeting.enabled", true) {
130+
execute!(
131+
output,
132+
style::Print(format!(
133+
"
123134
Hi, I'm Amazon Q. I can answer questions about your shell and CLI tools!
124135
You can include additional context by adding the following to your prompt:
125136
126137
{} to pass your shell history
127138
{} to pass information about your current git repository
128139
{} to pass your shell environment
129-
130140
",
131-
"@history".bold(),
132-
"@git".bold(),
133-
"@env".bold()
134-
)))?;
141+
"@history".bold(),
142+
"@git".bold(),
143+
"@env".bold()
144+
))
145+
)?;
146+
}
135147
}
136148

137149
// Print response as we receive it
138150
if let Some(rx) = &mut rx {
139-
stderr.queue(cursor::Hide)?;
140-
let mut spinner = Some(Spinner::new(Spinners::Dots, "Generating your answer...".to_owned()));
151+
// compiler complains about unused variable for spinner (bad global state usage)
152+
let mut _spinner = if interactive {
153+
queue!(output, cursor::Hide)?;
154+
Some(Spinner::new(Spinners::Dots, "Generating your answer...".to_owned()))
155+
} else {
156+
None
157+
};
141158

142159
let mut buf = String::new();
143160
let mut offset = 0;
@@ -175,23 +192,34 @@ You can include additional context by adding the following to your prompt:
175192
ended = true;
176193
},
177194
ApiResponse::Error(error) => {
178-
drop(spinner.take());
179-
stderr.queue(cursor::MoveToColumn(0))?;
180-
181-
match error {
182-
Some(error) => stderr
183-
.queue(style::SetForegroundColor(Color::Red))?
184-
.queue(style::SetAttribute(Attribute::Bold))?
185-
.queue(style::Print("error"))?
186-
.queue(style::SetForegroundColor(Color::Reset))?
187-
.queue(style::SetAttribute(Attribute::Reset))?
188-
.queue(style::Print(format!(": {error}\n")))?,
189-
None => stderr.queue(style::Print(
190-
"Amazon Q is having trouble responding right now. Try again later.",
191-
))?,
192-
};
193-
194-
stderr.flush()?;
195+
if interactive {
196+
_spinner = None;
197+
queue!(output, cursor::MoveToColumn(0))?;
198+
199+
match error {
200+
Some(error) => {
201+
queue!(
202+
output,
203+
style::SetForegroundColor(Color::Red),
204+
style::SetAttribute(Attribute::Bold),
205+
style::Print("error"),
206+
style::SetForegroundColor(Color::Reset),
207+
style::SetAttribute(Attribute::Reset),
208+
style::Print(format!(": {error}\n"))
209+
)?;
210+
},
211+
None => {
212+
queue!(
213+
output,
214+
style::Print(
215+
"Amazon Q is having trouble responding right now. Try again later.",
216+
)
217+
)?;
218+
},
219+
};
220+
}
221+
222+
output.flush()?;
195223
ended = true;
196224
},
197225
}
@@ -203,19 +231,24 @@ You can include additional context by adding the following to your prompt:
203231
buf.push('\n');
204232
}
205233

206-
if !buf.is_empty() && spinner.take().is_some() {
207-
stderr
208-
.queue(terminal::Clear(terminal::ClearType::CurrentLine))?
209-
.queue(cursor::MoveToColumn(0))?
210-
.queue(cursor::Show)?;
234+
if !buf.is_empty() && interactive {
235+
_spinner = None;
236+
queue!(
237+
output,
238+
crossterm::terminal::Clear(crossterm::terminal::ClearType::CurrentLine),
239+
cursor::MoveToColumn(0),
240+
cursor::Show
241+
)?;
211242
}
212243

213244
loop {
214245
let input = Partial::new(&buf[offset..]);
215-
match interpret_markdown(input, stderr as &mut Stderr, &mut state) {
246+
// fresh reborrow required on output
247+
match interpret_markdown(input, &mut *output, &mut state) {
216248
Ok(parsed) => {
217249
offset += parsed.offset_from(&input);
218-
stderr.lock().flush()?;
250+
output.flush()?;
251+
// output.lock().flush()?;
219252
state.newline = state.set_newline;
220253
state.set_newline = false;
221254
},
@@ -229,22 +262,28 @@ You can include additional context by adding the following to your prompt:
229262
}
230263

231264
if ended {
232-
stderr
233-
.queue(style::ResetColor)?
234-
.queue(style::SetAttribute(Attribute::Reset))?
235-
.queue(Print("\n"))?;
236-
237-
for (i, citation) in &state.citations {
238-
stderr
239-
.queue(style::SetForegroundColor(Color::Blue))?
240-
.queue(style::Print(format!("{i} ")))?
241-
.queue(style::SetForegroundColor(Color::DarkGrey))?
242-
.queue(style::Print(format!("{citation}\n")))?
243-
.queue(style::SetForegroundColor(Color::Reset))?;
244-
}
245-
246-
if !state.citations.is_empty() {
247-
stderr.execute(Print("\n"))?;
265+
if interactive {
266+
queue!(
267+
output,
268+
style::ResetColor,
269+
style::SetAttribute(Attribute::Reset),
270+
Print("\n")
271+
)?;
272+
273+
for (i, citation) in &state.citations {
274+
queue!(
275+
output,
276+
style::SetForegroundColor(Color::Blue),
277+
style::Print(format!("{i} ")),
278+
style::SetForegroundColor(Color::DarkGrey),
279+
style::Print(format!("{citation}\n")),
280+
style::SetForegroundColor(Color::Reset)
281+
)?;
282+
}
283+
284+
if !state.citations.is_empty() {
285+
execute!(output, Print("\n"))?;
286+
}
248287
}
249288

250289
if let (Some(conversation_id), Some(message_id)) = (&conversation_id, &message_id) {
@@ -256,14 +295,17 @@ You can include additional context by adding the following to your prompt:
256295
}
257296
}
258297

298+
if !interactive {
299+
break Ok(());
300+
}
259301
loop {
260-
let readline = rl.readline(PROMPT);
302+
let readline = rl.as_mut().unwrap().readline(PROMPT);
261303
match readline {
262304
Ok(line) => {
263305
if line.trim().is_empty() {
264306
continue;
265307
}
266-
let _ = rl.add_history_entry(line.as_str());
308+
let _ = rl.as_mut().unwrap().add_history_entry(line.as_str());
267309
input = line;
268310
break;
269311
},

‎crates/q_cli/src/cli/chat/terminal.rs‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
use std::io::{
2+
Stderr,
3+
Stdout,
4+
Write,
5+
};
6+
7+
pub enum WriteOutput {
8+
Interactive(Stderr),
9+
NonInteractive(Stdout),
10+
}
11+
12+
// wraps the stderr/stdout handle with a enum type.
13+
pub fn new(is_interactive: bool) -> WriteOutput {
14+
if is_interactive {
15+
WriteOutput::Interactive(std::io::stderr())
16+
} else {
17+
WriteOutput::NonInteractive(std::io::stdout())
18+
}
19+
}
20+
21+
impl Write for WriteOutput {
22+
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
23+
match self {
24+
WriteOutput::Interactive(stderr) => stderr.write(buf),
25+
WriteOutput::NonInteractive(stdout) => stdout.write(buf),
26+
}
27+
}
28+
29+
fn flush(&mut self) -> std::io::Result<()> {
30+
match self {
31+
WriteOutput::Interactive(stderr) => stderr.flush(),
32+
WriteOutput::NonInteractive(stdout) => stdout.flush(),
33+
}
34+
}
35+
}

0 commit comments

Comments
 (0)