DevFlow/crates/df-ai/src/ai_tools.rs

//! AI 工具注册 — 将现有 CRUD 操作注册为 LLM 可调用的 Tool
//!
//! 风险分级:
//! - Low:  只读操作（list/get），AI 自动执行
//! - Medium: 创建操作（create），显示意图，可配置自动批准
//! - High:  破坏性操作（delete / run_workflow），必须人工批准

use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;

use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::provider::ToolDefinition;

// ============================================================
// 风险级别
// ============================================================

/// 工具调用风险级别
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum RiskLevel {
    /// 只读操作 — AI 自动执行
    Low,
    /// 创建操作 — 可配置自动批准
    Medium,
    /// 破坏性操作 — 必须人工批准
    High,
}

// ============================================================
// 工具定义
// ============================================================

/// 已注册的 AI 工具
pub struct AiTool {
    /// 工具定义（发送给 LLM 的 JSON Schema）
    pub definition: ToolDefinition,
    /// 风险级别
    pub risk_level: RiskLevel,
    /// 执行处理器
    pub handler: AiToolHandler,
}

/// 工具处理器类型 — 异步函数，接收 JSON 参数，返回 JSON 结果
pub type AiToolHandler =
    Box<dyn Fn(Value) -> Pin<Box<dyn Future<Output = anyhow::Result<Value>> + Send>> + Send + Sync>;

// ============================================================
// 工具注册表
// ============================================================

/// AI 工具注册表
pub struct AiToolRegistry {
    tools: HashMap<String, AiTool>,
}

impl AiToolRegistry {
    /// 创建空注册表
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
        }
    }

    /// 注册一个工具
    pub fn register(
        &mut self,
        name: impl Into<String>,
        description: impl Into<String>,
        parameters: Value,
        risk_level: RiskLevel,
        handler: AiToolHandler,
    ) {
        let name_str = name.into();
        let definition = ToolDefinition::function(&name_str, description, parameters);
        self.tools.insert(
            name_str,
            AiTool {
                definition,
                risk_level,
                handler,
            },
        );
    }

    /// 获取所有工具定义（发送给 LLM 的 tools 参数）
    pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
        self.tools.values().map(|t| t.definition.clone()).collect()
    }

    /// 根据名称获取工具
    pub fn get(&self, name: &str) -> Option<&AiTool> {
        self.tools.get(name)
    }

    /// 执行指定工具 — handler 是唯一执行路径（schema+risk+实现同源，消除双轨）
    pub async fn execute(
        &self,
        name: &str,
        args: serde_json::Value,
    ) -> anyhow::Result<serde_json::Value> {
        let tool = self
            .tools
            .get(name)
            .ok_or_else(|| anyhow::anyhow!("未知工具: {}", name))?;
        (tool.handler)(args).await
    }

    /// 获取所有已注册工具名称
    pub fn tool_names(&self) -> Vec<String> {
        self.tools.keys().cloned().collect()
    }

    /// 已注册工具数量
    pub fn len(&self) -> usize {
        self.tools.len()
    }

    /// 是否为空
    pub fn is_empty(&self) -> bool {
        self.tools.is_empty()
    }
}

impl Default for AiToolRegistry {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================
// 工具执行结果
// ============================================================

/// 工具执行结果
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolExecutionResult {
    /// 工具调用 ID
    pub tool_call_id: String,
    /// 工具名称
    pub tool_name: String,
    /// 执行参数
    pub arguments: Value,
    /// 执行结果
    pub result: Value,
    /// 是否需要人工批准
    pub approval_required: bool,
    /// 风险级别
    pub risk_level: RiskLevel,
}

// ============================================================
// 辅助: 构建 JSON Schema 参数
// ============================================================

/// 构建一个简单的 object JSON Schema
pub fn object_schema(properties: Vec<(&str, &str, bool)>) -> Value {
    let mut props = serde_json::Map::new();
    let mut required = Vec::new();

    for (name, type_str, is_required) in properties {
        props.insert(
            name.to_string(),
            serde_json::json!({ "type": type_str, "description": "" }),
        );
        if is_required {
            required.push(name.to_string());
        }
    }

    serde_json::json!({
        "type": "object",
        "properties": props,
        "required": required,
    })
}