Skip to content

Commit

Permalink
fix: accept text through pipes in chat (#259)
Browse files Browse the repository at this point in the history 
* 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]>
  • Loading branch information
@Rahel-A
Rahel-A and Rahel Ahmed authored Dec 19, 2024
1 parent 496c04a commit a7b55ae
Show file tree
Hide file tree
Showing 2 changed files with 169 additions and 92 deletions.
226 changes: 134 additions & 92 deletions crates/q_cli/src/cli/chat/mod.rs
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
mod api;
mod parse;
mod prompt;
mod terminal;

use std::io::{
Stderr,
IsTerminal,
Read,
Write,
stderr,
stdin,
};
use std::process::ExitCode;
use std::time::Duration;
Expand All @@ -17,11 +19,10 @@ use crossterm::style::{
Print,
};
use crossterm::{
ExecutableCommand,
QueueableCommand,
cursor,
execute,
queue,
style,
terminal,
};
use eyre::{
Result,
Expand Down Expand Up @@ -58,7 +59,7 @@ enum ApiResponse {
End,
}

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

region_check("chat")?;

let mut stderr = stderr();
let result = try_chat(&mut stderr, input).await;
let stdin = stdin();
let is_interactive = stdin.is_terminal();

stderr
.queue(style::SetAttribute(Attribute::Reset))?
.queue(style::ResetColor)?
.flush()
.ok();
if !is_interactive {
// append to input string any extra info that was provided.
stdin.lock().read_to_string(&mut input).unwrap();
}
let mut output = terminal::new(is_interactive);
let result = try_chat(&mut output, input, is_interactive).await;

if is_interactive {
queue!(output, style::SetAttribute(Attribute::Reset), style::ResetColor).ok();
}
output.flush().ok();

result.map(|_| ExitCode::SUCCESS)
}

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

loop {
// Make request with input
match input.trim() {
"exit" | "quit" => {
if let Some(conversation_id) = conversation_id {
fig_telemetry::send_end_chat(conversation_id.clone()).await;
}
return Ok(());
},
_ => (),
}

if !input.is_empty() {
stderr.queue(style::SetForegroundColor(Color::Magenta))?;
if input.contains("@history") {
stderr.queue(style::Print("Using shell history\n"))?;
// Make request with input, otherwise used already provided buffer for input
if !interactive {
rx = Some(send_message(client.clone(), input.clone(), conversation_id.clone()).await?);
} else {
match input.trim() {
"exit" | "quit" => {
if let Some(conversation_id) = conversation_id {
fig_telemetry::send_end_chat(conversation_id.clone()).await;
}
return Ok(());
},
_ => (),
}

if input.contains("@git") {
stderr.queue(style::Print("Using git context\n"))?;
}
if !input.is_empty() {
queue!(output, style::SetForegroundColor(Color::Magenta))?;
if input.contains("@history") {
queue!(output, style::Print("Using shell history\n"))?;
}

if input.contains("@env") {
stderr.queue(style::Print("Using environment\n"))?;
}
if input.contains("@git") {
queue!(output, style::Print("Using git context\n"))?;
}

if input.contains("@env") {
queue!(output, style::Print("Using environment\n"))?;
}

rx = Some(send_message(client.clone(), input, conversation_id.clone()).await?);
stderr
.queue(style::SetForegroundColor(Color::Reset))?
.execute(style::Print("\n"))?;
} else if fig_settings::settings::get_bool_or("chat.greeting.enabled", true) {
stderr.execute(style::Print(format!(
"
rx = Some(send_message(client.clone(), input.clone(), conversation_id.clone()).await?);
queue!(output, style::SetForegroundColor(Color::Reset))?;
execute!(output, style::Print("\n"))?;
} else if fig_settings::settings::get_bool_or("chat.greeting.enabled", true) {
execute!(
output,
style::Print(format!(
"
Hi, I'm Amazon Q. I can answer questions about your shell and CLI tools!
You can include additional context by adding the following to your prompt:
{} to pass your shell history
{} to pass information about your current git repository
{} to pass your shell environment
",
"@history".bold(),
"@git".bold(),
"@env".bold()
)))?;
"@history".bold(),
"@git".bold(),
"@env".bold()
))
)?;
}
}

// Print response as we receive it
if let Some(rx) = &mut rx {
stderr.queue(cursor::Hide)?;
let mut spinner = Some(Spinner::new(Spinners::Dots, "Generating your answer...".to_owned()));
// compiler complains about unused variable for spinner (bad global state usage)
let mut _spinner = if interactive {
queue!(output, cursor::Hide)?;
Some(Spinner::new(Spinners::Dots, "Generating your answer...".to_owned()))
} else {
None
};

let mut buf = String::new();
let mut offset = 0;
Expand Down Expand Up @@ -175,23 +192,34 @@ You can include additional context by adding the following to your prompt:
ended = true;
},
ApiResponse::Error(error) => {
drop(spinner.take());
stderr.queue(cursor::MoveToColumn(0))?;

match error {
Some(error) => stderr
.queue(style::SetForegroundColor(Color::Red))?
.queue(style::SetAttribute(Attribute::Bold))?
.queue(style::Print("error"))?
.queue(style::SetForegroundColor(Color::Reset))?
.queue(style::SetAttribute(Attribute::Reset))?
.queue(style::Print(format!(": {error}\n")))?,
None => stderr.queue(style::Print(
"Amazon Q is having trouble responding right now. Try again later.",
))?,
};

stderr.flush()?;
if interactive {
_spinner = None;
queue!(output, cursor::MoveToColumn(0))?;

match error {
Some(error) => {
queue!(
output,
style::SetForegroundColor(Color::Red),
style::SetAttribute(Attribute::Bold),
style::Print("error"),
style::SetForegroundColor(Color::Reset),
style::SetAttribute(Attribute::Reset),
style::Print(format!(": {error}\n"))
)?;
},
None => {
queue!(
output,
style::Print(
"Amazon Q is having trouble responding right now. Try again later.",
)
)?;
},
};
}

output.flush()?;
ended = true;
},
}
Expand All @@ -203,19 +231,24 @@ You can include additional context by adding the following to your prompt:
buf.push('\n');
}

if !buf.is_empty() && spinner.take().is_some() {
stderr
.queue(terminal::Clear(terminal::ClearType::CurrentLine))?
.queue(cursor::MoveToColumn(0))?
.queue(cursor::Show)?;
if !buf.is_empty() && interactive {
_spinner = None;
queue!(
output,
crossterm::terminal::Clear(crossterm::terminal::ClearType::CurrentLine),
cursor::MoveToColumn(0),
cursor::Show
)?;
}

loop {
let input = Partial::new(&buf[offset..]);
match interpret_markdown(input, stderr as &mut Stderr, &mut state) {
// fresh reborrow required on output
match interpret_markdown(input, &mut *output, &mut state) {
Ok(parsed) => {
offset += parsed.offset_from(&input);
stderr.lock().flush()?;
output.flush()?;
// output.lock().flush()?;
state.newline = state.set_newline;
state.set_newline = false;
},
Expand All @@ -229,22 +262,28 @@ You can include additional context by adding the following to your prompt:
}

if ended {
stderr
.queue(style::ResetColor)?
.queue(style::SetAttribute(Attribute::Reset))?
.queue(Print("\n"))?;

for (i, citation) in &state.citations {
stderr
.queue(style::SetForegroundColor(Color::Blue))?
.queue(style::Print(format!("{i} ")))?
.queue(style::SetForegroundColor(Color::DarkGrey))?
.queue(style::Print(format!("{citation}\n")))?
.queue(style::SetForegroundColor(Color::Reset))?;
}

if !state.citations.is_empty() {
stderr.execute(Print("\n"))?;
if interactive {
queue!(
output,
style::ResetColor,
style::SetAttribute(Attribute::Reset),
Print("\n")
)?;

for (i, citation) in &state.citations {
queue!(
output,
style::SetForegroundColor(Color::Blue),
style::Print(format!("{i} ")),
style::SetForegroundColor(Color::DarkGrey),
style::Print(format!("{citation}\n")),
style::SetForegroundColor(Color::Reset)
)?;
}

if !state.citations.is_empty() {
execute!(output, Print("\n"))?;
}
}

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

if !interactive {
break Ok(());
}
loop {
let readline = rl.readline(PROMPT);
let readline = rl.as_mut().unwrap().readline(PROMPT);
match readline {
Ok(line) => {
if line.trim().is_empty() {
continue;
}
let _ = rl.add_history_entry(line.as_str());
let _ = rl.as_mut().unwrap().add_history_entry(line.as_str());
input = line;
break;
},
Expand Down
35 changes: 35 additions & 0 deletions crates/q_cli/src/cli/chat/terminal.rs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
use std::io::{
Stderr,
Stdout,
Write,
};

pub enum WriteOutput {
Interactive(Stderr),
NonInteractive(Stdout),
}

// wraps the stderr/stdout handle with a enum type.
pub fn new(is_interactive: bool) -> WriteOutput {
if is_interactive {
WriteOutput::Interactive(std::io::stderr())
} else {
WriteOutput::NonInteractive(std::io::stdout())
}
}

impl Write for WriteOutput {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
match self {
WriteOutput::Interactive(stderr) => stderr.write(buf),
WriteOutput::NonInteractive(stdout) => stdout.write(buf),
}
}

fn flush(&mut self) -> std::io::Result<()> {
match self {
WriteOutput::Interactive(stderr) => stderr.flush(),
WriteOutput::NonInteractive(stdout) => stdout.flush(),
}
}
}

0 comments on commit a7b55ae

Please sign in to comment.