iAs/ias-ai/src/workflow/core.rs

use serde_json::Value;
use tokio::fs;

use crate::workflow::flow::start_flow;
use crate::workflow::types::FlowConfig;

#[derive(Debug, Clone)]
pub struct LoadedFlowConfigs {
    pub configs: Vec<FlowConfig>,
    pub warnings: Vec<String>,
}

pub async fn call_flow(path: String, args: String) -> String {
    let all_flow = match load_configs(path).await {
        Ok(flows) => flows,
        Err(err) => return format!("工具调用失败:{:?}", err),
    };
    let args_json: Value = match serde_json::from_str(&args) {
        Ok(json) => json,
        Err(err) => return format!("工具调用失败:{:?}", err),
    };
    let flow_name = args_json["name"].as_str().unwrap_or_default();
    if let Some(config) = all_flow.iter().find(|e| flow_name == e.name) {
        let flow_result = start_flow(config, args_json["prompt"].clone()).await;
        match flow_result {
            Ok(value) => value,
            Err(err) => format!("工具调用失败:{:?}", err),
        }
    } else {
        "工具不存在，请检查工具名称".to_string()
    }
}

pub async fn load_tools_define(path: String) -> anyhow::Result<String> {
    let loaded = load_configs_with_warnings(path).await?;
    Ok(serde_json::to_string(&loaded.configs).unwrap_or_default())
}

pub async fn load_configs(path: String) -> anyhow::Result<Vec<FlowConfig>> {
    Ok(load_configs_with_warnings(path).await?.configs)
}

pub async fn load_configs_with_warnings(path: String) -> anyhow::Result<LoadedFlowConfigs> {
    let mut configs: Vec<FlowConfig> = vec![];
    let mut warnings: Vec<String> = vec![];
    let config_files = scan_config_dir(path).await?;
    if !config_files.is_empty() {
        for file in config_files {
            match load_yaml(file.clone()).await {
                Ok(config) => configs.push(config),
                Err(error) => warnings.push(format!("跳过工具配置 {file}: {error}")),
            }
        }
    }
    Ok(LoadedFlowConfigs { configs, warnings })
}

pub async fn load_yaml(path: String) -> anyhow::Result<FlowConfig> {
    let content: String = fs::read_to_string(path).await?;
    let config = serde_yaml::from_str(&content)?;
    Ok(config)
}

pub async fn scan_config_dir(path: String) -> anyhow::Result<Vec<String>> {
    let mut dir = fs::read_dir(path).await?;
    let mut files: Vec<String> = vec![];

    while let Some(entry) = dir.next_entry().await? {
        let path = entry.path();
        if !is_yaml_file(&path) {
            continue;
        }
        let abs = fs::canonicalize(&path).await?;
        files.push(abs.display().to_string());
    }

    files.sort();
    Ok(files)
}

fn is_yaml_file(path: &std::path::Path) -> bool {
    path.extension()
        .and_then(|value| value.to_str())
        .map(|value| matches!(value.to_ascii_lowercase().as_str(), "yaml" | "yml"))
        .unwrap_or(false)
}

#[tokio::test]
async fn list_file() {
    println!();

    let tools = load_tools_define("../tools".to_string()).await.unwrap();

    println!("{}", tools);

    println!();
}

#[tokio::test]
async fn load_configs_skips_invalid_yaml_files() {
    let dir = unique_test_dir("load_configs_skips_invalid_yaml_files");
    tokio::fs::create_dir_all(&dir).await.unwrap();

    tokio::fs::write(
        dir.join("good.yaml"),
        r#"
name: good_tool
desc: good config
params: []
steps: []
"#,
    )
    .await
    .unwrap();
    tokio::fs::write(
        dir.join("bad.yaml"),
        r#"
name: bad_tool
steps:
  - query:
      - key: value
"#,
    )
    .await
    .unwrap();

    let loaded = load_configs_with_warnings(dir.display().to_string())
        .await
        .unwrap();
    assert_eq!(loaded.configs.len(), 1);
    assert_eq!(loaded.configs[0].name, "good_tool");
    assert_eq!(loaded.warnings.len(), 1);
    assert!(loaded.warnings[0].contains("bad.yaml"));

    let configs = load_configs(dir.display().to_string()).await.unwrap();
    assert_eq!(configs.len(), 1);
    assert_eq!(configs[0].name, "good_tool");

    let _ = tokio::fs::remove_dir_all(&dir).await;
}

#[cfg(test)]
fn unique_test_dir(name: &str) -> std::path::PathBuf {
    let nanos = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_nanos();
    std::env::temp_dir().join(format!("ias-ai-{name}-{}-{nanos}", std::process::id()))
}