rust-sdk/examples/rig-integration/src/chat.rs

use futures::StreamExt;
use rig::{
    agent::{Agent, MultiTurnStreamItem},
    completion::CompletionModel,
    message::{Message, Text},
    streaming::{StreamedAssistantContent, StreamingChat},
};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};

pub async fn cli_chatbot<M>(chatbot: Agent<M>) -> anyhow::Result<()>
where
    M: CompletionModel + 'static,
    M::StreamingResponse: Send,
{
    let mut chat_log = vec![];

    let mut output = BufWriter::new(tokio::io::stdout());
    let mut input = BufReader::new(tokio::io::stdin());
    output.write_all(b"Enter :q to quit\n").await?;
    loop {
        output.write_all(b"\x1b[32muser>\x1b[0m ").await?;
        // Flush stdout to ensure the prompt appears before input
        output.flush().await?;
        let mut input_buf = String::new();
        input.read_line(&mut input_buf).await?;
        // Remove the newline character from the input
        let input = input_buf.trim();
        // Check for a command to exit
        if input == ":q" {
            break;
        }

        tracing::info!(%input);
        chat_log.push(Message::user(input));

        let mut response = chatbot.stream_chat(input, chat_log.clone()).await;
        stream_output_agent_start(&mut output).await?;
        let mut message_buf = String::new();

        while let Some(message) = response.next().await {
            match message {
                Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(
                    Text { text },
                ))) => {
                    message_buf.push_str(&text);
                    output_agent(&text, &mut output).await?;
                }
                Ok(MultiTurnStreamItem::StreamAssistantItem(
                    StreamedAssistantContent::ToolCall(tool_call),
                )) => {
                    let name = &tool_call.function.name;
                    let arguments = &tool_call.function.arguments;
                    stream_output_toolcall(
                        format!("Calling tool: {name} with args: {arguments}"),
                        &mut output,
                    )
                    .await?;
                }
                Ok(MultiTurnStreamItem::StreamUserItem(user_content)) => {
                    // Tool results are streamed back as user items
                    stream_output_toolcall(format!("Tool result: {:?}", user_content), &mut output)
                        .await?;
                }
                Ok(MultiTurnStreamItem::FinalResponse(final_response)) => {
                    tracing::info!("Final response received: {:?}", final_response);
                }
                Ok(_) => {
                    // Handle other stream items (reasoning, deltas, etc.)
                }
                Err(error) => {
                    output_error(error, &mut output).await?;
                }
            }
        }

        chat_log.push(Message::assistant(message_buf));
        stream_output_agent_finished(&mut output).await?;
    }

    Ok(())
}

pub async fn output_error(
    e: impl std::fmt::Display,
    output: &mut BufWriter<tokio::io::Stdout>,
) -> std::io::Result<()> {
    output
        .write_all(b"\x1b[1;31m\xE2\x9D\x8C ERROR: \x1b[0m")
        .await?;
    output.write_all(e.to_string().as_bytes()).await?;
    output.write_all(b"\n").await?;
    output.flush().await?;
    Ok(())
}

pub async fn output_agent(
    content: impl std::fmt::Display,
    output: &mut BufWriter<tokio::io::Stdout>,
) -> std::io::Result<()> {
    output.write_all(content.to_string().as_bytes()).await?;
    output.flush().await?;
    Ok(())
}

pub async fn stream_output_toolcall(
    content: impl std::fmt::Display,
    output: &mut BufWriter<tokio::io::Stdout>,
) -> std::io::Result<()> {
    output
        .write_all(b"\x1b[1;33m\xF0\x9F\x9B\xA0 Tool Call: \x1b[0m")
        .await?;
    output.write_all(content.to_string().as_bytes()).await?;
    output.write_all(b"\n").await?;
    output.flush().await?;
    Ok(())
}

pub async fn stream_output_agent_start(
    output: &mut BufWriter<tokio::io::Stdout>,
) -> std::io::Result<()> {
    output
        .write_all(b"\x1b[1;34m\xF0\x9F\xA4\x96 Agent: \x1b[0m")
        .await?;
    output.flush().await?;
    Ok(())
}

pub async fn stream_output_agent_finished(
    output: &mut BufWriter<tokio::io::Stdout>,
) -> std::io::Result<()> {
    output.write_all(b"\n").await?;
    output.flush().await?;
    Ok(())
}