traiy_core 0.0.13

//! Agents module
//!
//! This module defines the `Agent` struct, which is responsible for interacting with
//! a Language Model (LLM) to generate recommendations or enhance existing content.
//!
//! The `Agent` uses a `Handbook` to manage guidelines, input files, and LLM configurations.
//! It supports actions such as recommending improvements or enhancing content
//! based on provided guidelines.

use crate::builder::build_model;
use crate::cli::Action;
use crate::handb::Handbook;
use crate::prompts::RECOMMENDATION_PROMPT_TEMPLATE;
use anyhow::{Result, anyhow};
use futures::future::join_all;
use llm::LLMProvider as ExternalLLMProvider;
use llm::chat::ChatMessage;
use log::{debug, error, info};
use std::env::var;
use std::fs;
use std::path::{Path, PathBuf};

/// Represents an agent capable of interacting with a Language Model.
///
/// The agent uses a `Handbook` to manage guidelines, input files, and LLM configurations.
pub struct Agent {
    /// The 'Action' instance representing the type of action: Recommend, Enhance.
    /// It determines what the agent is supposed to do.
    action: Action,
    /// The `Handbook` instance containing the configuration for the agent.
    /// It includes guidelines, input files, and LLM configurations.
    handbook: Handbook,
}

impl Agent {
    /// Creates a new Agent instance.
    ///
    /// # Arguments
    ///
    /// * `action` - The action to be performed by the agent (e.g., Recommend, Enhance).
    /// * `handbook` - The handbook containing guidelines, input files, and LLM configurations.
    ///
    /// # Returns
    ///
    /// A `Result` containing the new `Agent` instance, or an error if the creation fails.
    ///
    /// # Errors
    ///
    /// This function does not return a `Result` or throw errors, but the Agent
    /// initialization may fail if the provided Handbook is invalid.
    pub fn new(action: Action, handbook: Handbook) -> Agent {
        debug!("Creating Agent.");
        Self { action, handbook }
    }

    /// Builds the language model based on the configuration specified in the `Action` and `Handbook`.
    ///
    /// This function constructs the LLM (Language Model) based on the provided configurations
    /// (compatibility checked at cli.rs).
    /// It supports different LLM providers (e.g., Google) and configures parameters
    /// such as API keys, model names, max tokens, and temperature.
    ///
    /// # Returns
    ///
    /// A `Result` containing a boxed `ExternalLLMProvider` if the model is built successfully,
    /// or an error if the model cannot be built due to configuration issues or API failures.
    ///
    /// # Errors
    ///
    /// This function may return an error if:
    /// - The LLM_API_KEY environment variable is not set.
    /// - The model fails to build due to invalid configurations.
    fn build_model(&self) -> Result<Box<dyn ExternalLLMProvider>> {
        debug!("Creating Model.");
        match &self.action {
            Action::Recommend {
                num_recommendations,
                model,
                max_tokens,
                temperature,
                llm_provider,
                ..
            } => {
                let api_key = var("LLM_API_KEY").expect("LLM_API_KEY environment variable must be set for traiy to work with your llm backend.");
                let model = model.to_string();
                let max_tokens = max_tokens.unwrap_or(9999);
                let temperature = temperature.unwrap_or(0.75);
                let num_recomm = num_recommendations.unwrap_or(3);
                let system_msg = format!("{} {}", RECOMMENDATION_PROMPT_TEMPLATE, num_recomm);
                let llm = build_model(
                    &self.action,
                    api_key,
                    model,
                    max_tokens,
                    temperature,
                    system_msg,
                    llm_provider,
                )?;
                Ok(llm)
            }
            Action::Enhance { .. } => {
                unimplemented!(" ==> In progress of implementation <== ")
            }
        }
    }

    /// Retrieves a response from the language model.
    ///
    /// # Arguments
    ///
    /// * `task` - The task or query to be sent to the language model.
    ///
    /// # Returns
    ///
    /// A `Result` containing an `Option<String>` with the LLM's response.
    /// Returns `Ok(Some(String))` with the LLM's response if the request is successful.
    /// Returns `Ok(None)` if the LLM does not provide a response.
    /// Returns an `Err` if the request fails.
    ///
    /// # Errors
    ///
    /// This function may return an error if:
    /// - The language model fails to build.
    /// - The chat request to the language model fails.
    async fn get_llm_response(&self, task: &str) -> Result<Option<String>> {
        info!("Task: {:?}", &task);
        // debug!("Asking model.");
        let llm = self.build_model()?;

        ///////////////////
        // TODO
        // Diverge by action
        ///////////////////

        // Construct the input messages for the LLM.
        let mut input_messages = Vec::with_capacity(6);
        input_messages.extend([
            ChatMessage::user().content("GUIDELINE:").build(),
            ChatMessage::user().content(task).build(),
            ChatMessage::user().content("DOCUMENT:").build(),
            ChatMessage::user()
                .content(&self.handbook.input_csv)
                .build(),
        ]);

        if let Some(context) = &self.handbook.context_csv {
            if !context.trim().is_empty() {
                input_messages.insert(0, ChatMessage::user().content("CONTEXT:").build());
                input_messages.insert(1, ChatMessage::user().content(context).build());
            }
        }

        let response = llm.chat(&input_messages).await?;
        let response = response.text();

        Ok(response)
    }

    /// Writes the LLM's response to a file.
    ///
    /// # Arguments
    ///
    /// * `file_path` - The path to the file where the response will be written.
    /// * `content` - The content to write to the file.
    ///
    /// # Returns
    ///
    /// A `Result` indicating success or an error if writing fails.
    ///
    /// # Errors
    ///
    /// This function may return an error if the file cannot be written to,
    /// for example, due to insufficient permissions or the file path
    /// not being found.
    fn write_response_to_file(&self, file_path: &Path, content: &str) -> Result<()> {
        debug!("Writing to output path.");
        let file_str = file_path.to_string_lossy();
        fs::write(file_path, content)
            .map_err(|e| anyhow!("Error writing to file '{}': {}", file_str, e))
    }

    /// Executes the language model for each guideline provided.
    ///
    /// This function iterates over each guideline, invokes the language model,
    /// and collects the responses. It handles potential errors during the
    /// execution of each task.
    ///
    /// # Returns
    ///
    /// A `Result` containing a `Vec<Option<String>>` with the LLM's responses for each guideline.
    ///
    /// # Errors
    ///
    /// This function may return an error if:
    /// - An error occurs while retrieving the LLM's response for a task.
    ///   The error message will provide details about the specific task that failed.
    pub async fn loop_guidelines(&self) -> Result<Option<String>> {
        let guidelines = &self.handbook.guidelines_csv;
        let future = self.get_llm_response(&guidelines);
        let result = future.await?;

        // let result = join_all(futures)
        //     .await
        //     .into_iter()
        //     .map(|x| {
        //         x.map_err(|e| {
        //             error!("An error happened while looping this task: {}", e);
        //             anyhow!("An error happened while looping this task: {}", e)
        //         })
        //         .ok()
        //         .flatten()
        //     })
        //     .collect();
        // info!("Guidelines looped.");
        // Ok(result)

        Ok(result)
    }

    /// Extracts the file stem and original path from the input CSV file specified
    /// in the action.
    ///
    /// This function retrieves the path to the input CSV file from the `Action` enum
    /// (either `Recommend` or `Enhance`) and extracts the file's stem (the filename
    /// without the extension).
    ///
    /// # Returns
    ///
    /// A `Result` containing a tuple:
    /// - The `String` of the file stem.
    /// - The `PathBuf` of the original input file path.
    ///
    /// # Errors
    ///
    /// This function may return an error if:
    /// - The file stem cannot be extracted from the path. This might occur if the
    /// path is malformed or does not point to a valid file.
    ///
    /// # Remarks
    ///
    /// The function currently uses `unwrap()` to handle the case where the file stem cannot be extracted.
    /// Consider adding more robust error handling to manage cases where the file stem is not present.
    fn get_file_stem_and_original_path(&self) -> Result<(String, PathBuf)> {
        let input_csv = match &self.action {
            Action::Recommend { input_csv, .. } => input_csv,
            Action::Enhance { input_csv, .. } => input_csv,
        };
        let original_path = PathBuf::from(input_csv);
        let file_stem = original_path.file_stem().and_then(|s| s.to_str()).unwrap(); // Handle this error ? Maybe not as the file has already been checked
        Ok((file_stem.to_string(), original_path))
    }

    /// Creates the output directory and determines the file stem for output files.
    ///
    /// This function constructs the output directory path based on the input file's name
    /// and appends "_traiy" to it, creating a unique directory for the output.
    /// It extracts the file stem (name without extension) from the input file path,
    /// which is used for naming output files within the newly created directory.
    ///
    /// # Arguments
    ///
    /// * `file_stem` - The stem of the input file name.
    /// * `original_path` - The path to the original input file.
    ///
    /// # Returns
    ///
    /// A `Result` containing a tuple:
    /// - The `String` of the file stem (name without extension).
    /// - The `PathBuf` of the output directory where the results will be stored.
    ///
    /// # Errors
    ///
    /// This function may return an error if:
    /// - The output directory cannot be created due to permission issues or other IO errors.
    ///
    /// # Remarks
    ///
    /// The output directory is created as a subdirectory of the original file's parent directory.
    /// If the directory already exists, this function does nothing.
    fn create_output_dir_and_file_stem(
        &self,
        file_stem: String,
        original_path: PathBuf,
    ) -> Result<(String, PathBuf)> {
        let output_dir_name = format!("{}_traiy", file_stem);
        let parent_dir = original_path.parent().unwrap();
        let output_dir_path = parent_dir.join(output_dir_name);

        fs::create_dir_all(&output_dir_path).map_err(|e| {
            anyhow!(
                "Failed to create output directory '{}': {}",
                output_dir_path.display(),
                e
            )
        })?;
        Ok((file_stem, output_dir_path))
    }

    /// Saves the results to individual files in the output directory.
    ///
    /// This function iterates through the results, constructs a file name for each,
    /// and writes the result content to the file.
    ///
    /// # Arguments
    ///
    /// * `results` - A vector of optional strings, where each string is the result of a task.
    /// * `file_stem` - The base name for the output files.
    /// * `output_dir_path` - The path to the directory where the files will be saved.
    ///
    /// # Errors
    ///
    /// This function may return an error if:
    /// - Writing to a file fails.
    fn save_results(
        &self,
        result: Option<String>,
        file_stem: String,
        output_dir_path: PathBuf,
    ) -> Result<()> {
        // debug
        if let Some(response_text) = result {
            let new_file_name = format!("{}_recommendations.{}", file_stem, "md");
            let output_path = output_dir_path.join(new_file_name);
            self.write_response_to_file(&output_path, &response_text)?;
        } else {
            error!("LLM response was empty for current task");
        }
        Ok(())
    }

    /// Executes the recommendation process.
    ///
    /// This function orchestrates the recommendation process by first looping through
    /// the guidelines, then creating the output directory and file stem, and finally
    /// saving the results to the output directory.
    ///
    /// # Errors
    ///
    /// This function may return an error if:
    /// - Looping through the guidelines fails.
    /// - Creating the output directory and file stem fails.
    /// - Saving the results fails.
    pub async fn recommend(&self) -> Result<()> {
        debug!("Starting recommendations...");
        let results = self.loop_guidelines().await?;
        let (file_stem, original_path) = self.get_file_stem_and_original_path()?;
        let (file_stem, output_dir_path) =
            self.create_output_dir_and_file_stem(file_stem, original_path)?;
        let _ = self.save_results(results, file_stem, output_dir_path);
        Ok(())
    }

    // Enhance
    // pub async fn enhance(&self) -> Result<()> {
    //     debug!("Consuming recommendations...");
    //     get input_file
    //     get (edited) recommendations
    //     call llm with both, exclusive prompt
    //     overwrite input file...?
    //     which wil serve as for the next recommendations
    //     and thus a loop
    //     Ok(())
    // }

    // pub fn enhance()
    // verify folder of recommendations exist
    // then prompt it to take those recommendations
    // and implement them on a new final file
}