Skip to main content

graphrag_cli/
config.rs

1//! Configuration loading and management
2
3use crate::handlers::FileOperations;
4use color_eyre::eyre::{eyre, Result};
5use graphrag_core::config::json5_loader::{detect_config_format, ConfigFormat};
6use graphrag_core::config::setconfig::SetConfig;
7use graphrag_core::Config as GraphRAGConfig;
8use std::path::Path;
9
10/// Load GraphRAG configuration from file (supports JSON5, JSON, TOML)
11pub async fn load_config(path: &Path) -> Result<GraphRAGConfig> {
12    // Validate file
13    FileOperations::validate_file(path).await?;
14
15    // Read file
16    let content = FileOperations::read_to_string(path).await?;
17
18    // Detect format from file extension
19    let format = detect_config_format(path)
20        .ok_or_else(|| eyre!("Unsupported config file format: {:?}", path.extension()))?;
21
22    // Parse based on detected format - always parse as SetConfig first, then convert
23    let set_config: SetConfig = match format {
24        ConfigFormat::Json5 => {
25            #[cfg(feature = "json5-support")]
26            {
27                json5::from_str(&content)
28                    .map_err(|e| eyre!("Failed to parse JSON5 config: {}", e))?
29            }
30            #[cfg(not(feature = "json5-support"))]
31            {
32                return Err(eyre!(
33                    "JSON5 support not enabled. Recompile with json5-support feature."
34                ));
35            }
36        },
37        ConfigFormat::Json => serde_json::from_str(&content)
38            .map_err(|e| eyre!("Failed to parse JSON config: {}", e))?,
39        ConfigFormat::Toml => {
40            toml::from_str(&content).map_err(|e| eyre!("Failed to parse TOML config: {}", e))?
41        },
42        ConfigFormat::Yaml => {
43            #[cfg(feature = "yaml-support")]
44            {
45                serde_yaml::from_str(&content)
46                    .map_err(|e| eyre!("Failed to parse YAML config: {}", e))?
47            }
48            #[cfg(not(feature = "yaml-support"))]
49            {
50                return Err(eyre!(
51                    "YAML support not enabled. Recompile with yaml-support feature."
52                ));
53            }
54        },
55    };
56
57    // Convert SetConfig to Config
58    let config = set_config.to_graphrag_config();
59
60    // Log configuration details for debugging
61    tracing::info!("Loaded {:?} configuration from: {}", format, path.display());
62    tracing::info!("Mode approach: {}", set_config.mode.approach);
63    tracing::info!(
64        "Entity extraction use_gleaning: {}",
65        config.entities.use_gleaning
66    );
67    tracing::info!(
68        "Entity extraction max_gleaning_rounds: {}",
69        config.entities.max_gleaning_rounds
70    );
71    tracing::info!("Ollama enabled: {}", config.ollama.enabled);
72    tracing::info!("Ollama chat_model: {}", config.ollama.chat_model);
73
74    Ok(config)
75}
76
77/// Get default configuration
78#[allow(dead_code)]
79pub fn default_config() -> GraphRAGConfig {
80    GraphRAGConfig::default()
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use std::io::Write;
87    use tempfile::NamedTempFile;
88    use tokio;
89
90    #[tokio::test]
91    async fn test_load_valid_toml_config() {
92        let mut temp_file = NamedTempFile::with_suffix(".toml").unwrap();
93        writeln!(temp_file, "[general]").unwrap();
94        writeln!(temp_file, "log_level = \"info\"").unwrap();
95
96        let result = load_config(temp_file.path()).await;
97        assert!(result.is_ok());
98    }
99
100    #[tokio::test]
101    #[cfg(feature = "json5-support")]
102    async fn test_load_valid_json5_config() {
103        let mut temp_file = NamedTempFile::with_suffix(".json5").unwrap();
104        writeln!(temp_file, "{{").unwrap();
105        writeln!(temp_file, "  // Comment in JSON5").unwrap();
106        writeln!(temp_file, "  general: {{").unwrap();
107        writeln!(temp_file, "    log_level: \"info\",").unwrap();
108        writeln!(temp_file, "  }},").unwrap();
109        writeln!(temp_file, "}}").unwrap();
110
111        let result = load_config(temp_file.path()).await;
112        assert!(
113            result.is_ok(),
114            "Failed to load JSON5 config: {:?}",
115            result.err()
116        );
117    }
118
119    #[tokio::test]
120    async fn test_load_valid_json_config() {
121        let mut temp_file = NamedTempFile::with_suffix(".json").unwrap();
122        writeln!(temp_file, "{{").unwrap();
123        writeln!(temp_file, "  \"general\": {{").unwrap();
124        writeln!(temp_file, "    \"log_level\": \"info\"").unwrap();
125        writeln!(temp_file, "  }}").unwrap();
126        writeln!(temp_file, "}}").unwrap();
127
128        let result = load_config(temp_file.path()).await;
129        assert!(result.is_ok());
130    }
131
132    #[tokio::test]
133    async fn test_load_invalid_toml() {
134        let mut temp_file = NamedTempFile::with_suffix(".toml").unwrap();
135        writeln!(temp_file, "invalid toml content {{{{").unwrap();
136
137        let result = load_config(temp_file.path()).await;
138        assert!(result.is_err());
139    }
140
141    #[tokio::test]
142    async fn test_load_unsupported_format() {
143        let result = load_config(Path::new("/nonexistent/file.txt")).await;
144        assert!(result.is_err());
145    }
146
147    #[tokio::test]
148    async fn test_load_nonexistent_file() {
149        let result = load_config(Path::new("/nonexistent/file.toml")).await;
150        assert!(result.is_err());
151    }
152}