traiy_core 0.0.13

//! # Config module
//!
//! This module defines the configuration and loading mechanism for Language Model (LLM)
//! provider model compatibility.
//!
//! It includes functionalities to load compatibility
//! settings from a YAML file and provides a data structure (`CompatibilityConfig`)
//! to represent the supported models and their configurations for different providers
//! such as Google and OpenAI.

use anyhow::Result;
use serde::Deserialize;
use std::collections::HashMap;

/// Enumerates the available Language Model Providers.
#[derive(Deserialize, Debug)]
pub struct CompatibilityConfig {
    pub google: HashMap<String, String>,
    pub openai: HashMap<String, String>,
}

/// Loads the LLM provider model compatibility configuration from a YAML file.\
///
/// Returns: A `Result` containing the `CompatibilityConfig` on success, or
/// an error if loading/parsing fails.
pub fn load_model_compatibility() -> Result<CompatibilityConfig> {
    let compat_bytes = include_bytes!("../assets/model_compatibility.yaml");
    let contents = String::from_utf8_lossy(compat_bytes);
    let compat: CompatibilityConfig = serde_yaml::from_str(&contents)
        .map_err(|e| anyhow::anyhow!("Failed to parse model compatibility configuration: {}", e))?;
    Ok(compat)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_load_model_compatibility_success() -> Result<()> {
        let config = load_model_compatibility()?;
        // Add assertions here to validate the contents of the config, e.g.:
        assert!(config.google.contains_key("gemini-2.0-flash"));
        Ok(())
    }
}