bep/
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 {
31                    role: "user".into(),
32                    content: input.into(),
33                });
34                chat_log.push(Message {
35                    role: "assistant".into(),
36                    content: response.clone(),
37                });
38
39                println!("========================== Response ============================");
40                println!("{response}");
41                println!("================================================================\n\n");
42
43                tracing::info!("Response:\n{}\n", response);
44            }
45            Err(error) => println!("Error reading input: {}", error),
46        }
47    }
48
49    Ok(())
50}