rig/
cli_chatbot.rs

1use std::io::{self, Write};
2
3use crate::completion::{Chat, Message, PromptError};
4
5/// Utility function to create a simple REPL CLI chatbot from a type that implements the
6/// `Chat` trait.
7pub async fn cli_chatbot(chatbot: impl Chat) -> Result<(), PromptError> {
8    let stdin = io::stdin();
9    let mut stdout = io::stdout();
10    let mut chat_log = vec![];
11
12    println!("Welcome to the chatbot! Type 'exit' to quit.");
13    loop {
14        print!("> ");
15        // Flush stdout to ensure the prompt appears before input
16        stdout.flush().unwrap();
17
18        let mut input = String::new();
19        match stdin.read_line(&mut input) {
20            Ok(_) => {
21                // Remove the newline character from the input
22                let input = input.trim();
23                // Check for a command to exit
24                if input == "exit" {
25                    break;
26                }
27                tracing::info!("Prompt:\n{}\n", input);
28
29                let response = chatbot.chat(input, chat_log.clone()).await?;
30                chat_log.push(Message::user(input));
31                chat_log.push(Message::assistant(response.clone()));
32
33                println!("========================== Response ============================");
34                println!("{response}");
35                println!("================================================================\n\n");
36
37                tracing::info!("Response:\n{}\n", response);
38            }
39            Err(error) => println!("Error reading input: {error}"),
40        }
41    }
42
43    Ok(())
44}
rig/cli_chatbot.rs

rig/
cli_chatbot.rs
