own_assist/src/main.rs

mod config;
mod human_interface;
mod i18n;
mod model;

use crate::config::Config;
use crate::human_interface::HumanInterface;
use crate::i18n::translate;
use crate::model::create_model_from_config;
use base64::Engine;
use clap::Parser;
use cpal::traits::HostTrait;
use own_mcp::AgentChat;
use std::fmt::Display;

/// partial mcp client with cli and voice interaction
#[derive(Parser, Debug)]
#[command(version, long_about = None)]
struct Args {
    /// generate an encoded basic auth header
    #[clap(long, short, action)]
    basic_auth: bool,
}

fn basic_auth_tool() {
    let username: String = dialoguer::Input::new()
        .with_prompt(tlt!("username"))
        .interact_text()
        .unwrap();

    let password: String = dialoguer::Password::new()
        .with_prompt(tlt!("password"))
        .interact()
        .unwrap();

    let encoded = base64::prelude::BASE64_STANDARD.encode(format!("{}:{}", username, password));
    println!("Basic {}", encoded);
}

async fn chat_loop(mut human_interface: impl HumanInterface, mut agent_chat: AgentChat) {
    loop {
        let user_message = human_interface.expect_user_message().await.unwrap();
        let immutable_interface = &human_interface;
        let agent_message = agent_chat
            .message(user_message, async |mcp_server_name, tool_name| {
                immutable_interface
                    .ask_for_permission(mcp_server_name, tool_name)
                    .await
                    .inspect_err(exit_msg!(
                        "Human interface error while asking for permission"
                    ))
                    .unwrap()
            })
            .await
            .inspect_err(exit_msg!("Failed communicating with agent."))
            .unwrap();
        human_interface
            .agent_message(agent_message.content)
            .await
            .inspect_err(exit_msg!(
                "Human interface error while trying to display message"
            ))
            .unwrap();
    }
}

#[tokio::main]
async fn main() {
    let args = Args::parse();
    if args.basic_auth {
        basic_auth_tool();
        return;
    }

    env_logger::init();

    let config = Config::from_file()
        .inspect_err(|e| {
            log::error!("error loading assist.toml: {}", e);
            std::process::exit(1);
        })
        .unwrap();

    /*let audio_client = config.audio_client();
    log::debug!("Audio server status {:?}", audio_client.status().await.inspect_err(|e|{
        log::error!("audio server error {}", e);
        std::process::exit(1);
    }).unwrap());*/

    let ollama = config.ollama_instance();
    let model_name = &config.ollama_config().model.name;

    create_model_from_config(&ollama, &config.ollama_config().model)
        .await
        .inspect_err(exit_msg!(format!(
            "failed creating ollama model `{model_name}`"
        )))
        .unwrap();

    let mcp_clients = config
        .mcp_clients()
        .await
        .inspect_err(exit_msg!("failed creating MCP clients"))
        .unwrap();

    let mut agent_chat = own_mcp::AgentChat::new(ollama, model_name.clone(), mcp_clients)
        .await
        .inspect_err(exit_msg!("failed to create agent"))
        .unwrap();
    config.set_tool_permissions(&mut agent_chat).await.unwrap();

    log::info!(
        "all tools: {:#?}",
        agent_chat.get_all_tools().collect::<Vec<_>>()
    );

    match config.interface {
        config::HumanInterface::Cli => {
            let human_interface = human_interface::cli::CommandLine::new();
            chat_loop(human_interface, agent_chat).await;
        }
        config::HumanInterface::Audio => {
            match config.audio_client() {
                Some(audio_client) => {
                    let human_interface = human_interface::audio::Audio::new(
                        audio_client,
                        cpal::Host::default()
                            .default_input_device()
                            .expect("no default device"),
                    )
                    .await
                    .inspect_err(exit_msg!("failed creating audio client")).unwrap();

                    chat_loop(human_interface, agent_chat).await;
                }
                None => {
                    eprintln!("audio client was not configured");
                }
            };
        }
    };
}

fn exit_with_error_message(error: impl Display, message: impl Display) {
    log::error!("{message}: {error}");
    eprintln!("{}", message); // makes sure that message is printed even if logging is deactivated
    std::process::exit(1);
}

#[macro_export]
macro_rules! exit_msg {
    ($message:expr) => {
        |e| exit_with_error_message(e, $message)
    };
}