agent/mod.rs

use std::str::FromStr;

use ollama_rs::{
    Ollama,
    coordinator::Coordinator,
    generation::{
        chat::ChatMessage,
        //tools::implementations::{Calculator, DDGSearcher, Scraper},
    },
    models::ModelOptions,
};
use reqwest::Url;

use crate::agents::config::Config;

pub mod fenrir;
pub mod freya;
pub mod loki;
pub mod odin;
pub mod thor;
pub mod tyr;

pub mod config;
pub mod scheduler;

const AGENT_DISCLAIMER: &str = "* messages from and to you are encapsulated by a program:
        DO NOT include subjects, to or from fields in yout answer!";

pub struct Agent {
    name: String,
    subject: String,
    sys_prompt: String,
    coordinator: ollama_rs::coordinator::Coordinator<Vec<ChatMessage>>,
    sent_messages: u64,
}

impl Agent {
    pub fn new(
        name: String,
        subject: String,
        config: &Config,
    ) -> Result<Self, crate::error::BalubaitecError> {
        let ollama = match config.host(&name) {
            Some(url) => Ollama::from_url(Url::from_str(url.as_str()).unwrap()),
            None => Ollama::default(),
        };

        let model = config.model(&name)?;
        let tools = config.tools(&name)?;

        let history = vec![];

        let model_options = ModelOptions::default().num_ctx(config.ctx_window(&name));
        let mut coordinator = Coordinator::new(ollama, model, history)
            .options(model_options)
            .add_tool(crate::tools::internal::request_send)
            .add_tool(crate::tools::internal::reply_check);

        for tool in tools.into_iter() {
            coordinator = match tool.trim().to_lowercase().as_str() {
                "time" => coordinator.add_tool(crate::tools::general::get_time),
                "weather" => coordinator.add_tool(crate::tools::general::get_weather),
                "bash" => coordinator.add_tool(crate::tools::development::shell_command),
                "read_file" => coordinator.add_tool(crate::tools::development::read_file),
                "dir" => coordinator.add_tool(crate::tools::development::list_directory_content),
                _ => return Err(crate::error::BalubaitecError::ToolNotAvailableError(tool)),
            };
        }

        let sent_messages = 0;

        let sys_prompt = config.system_prompt_self(&name)?;

        Ok(Self {
            name,
            subject,
            sys_prompt,
            coordinator,
            sent_messages,
        })
    }

    pub async fn chat(&mut self, from: String, message: String) -> String {
        let name = self.name.clone();
        let subject = self.subject.clone();

        let msg = format!(
            "
From: {from}\n
Subject: {subject}\n
\n
{message}
"
        );

        let mut messages = vec![ChatMessage::user(msg)];
        if self.sent_messages == 0 {
            messages.push(ChatMessage::system(self.sys_prompt.clone()));
        }

        match self.coordinator.chat(messages).await {
            Ok(resp) => {
                println!("Execution of agent {name} has ended successfully.");
                resp.message.content
            }
            Err(err) => {
                eprintln!("Ollama error on agent {name}: {err}");
                format!("There was an error contacting the agent {name}: {err}")
            }
        }
    }
}