velocia 0.3.1

velocia – production-ready AI agent framework using ADK-Rust, A2A protocol, and AWS DynamoDB
//! `ToolFactory` – creates `DynTool` instances from YAML configuration.
//!
//! Mirrors Python's `ToolFactory` class.

use std::collections::HashMap;

use serde_json::Value;
use tracing::error;

use crate::config::agent::AgentConfig;
use crate::config::tool::McpToolConfig;
use crate::error::{AgentKitError, Result};
use crate::tools::loader::DynTool;
use crate::tools::mcp::McpToolLoader;

pub struct ToolFactory {
    #[allow(dead_code)]
    agent_config: AgentConfig,
    /// Registry for `function` and `class_method` tool types.
    function_registry: HashMap<String, DynTool>,
}

impl ToolFactory {
    pub fn new(agent_config: &AgentConfig, registry: HashMap<String, DynTool>) -> Self {
        Self {
            agent_config: agent_config.clone(),
            function_registry: registry,
        }
    }

    /// Parse `raw` JSON (one entry from the `tools:` YAML list) and return
    /// the corresponding `DynTool`, or `None` if the entry should be skipped.
    pub async fn create_tool(&mut self, raw: Value) -> Result<Option<DynTool>> {
        let type_str = raw["type"]
            .as_str()
            .ok_or_else(|| AgentKitError::ToolConfig("Missing 'type' field".into()))?;

        match type_str {
            "mcp" => {
                let cfg: McpToolConfig = serde_json::from_value(raw)
                    .map_err(|e| AgentKitError::ToolConfig(e.to_string()))?;
                McpToolLoader::load(&cfg).await.map(Some)
            }

            "function" => {
                let name = raw["function_name"]
                    .as_str()
                    .or_else(|| raw["name"].as_str())
                    .unwrap_or("")
                    .to_string();

                match self.function_registry.remove(&name) {
                    Some(tool) => Ok(Some(tool)),
                    None => {
                        error!("Function tool '{name}' not found in registry. Register it via AgentFactory::register_tool().");
                        Ok(None)
                    }
                }
            }

            "class_method" => {
                let class_name = raw["class_name"].as_str().unwrap_or("");
                let method_name = raw["method_name"].as_str().unwrap_or("");
                let key = format!("{class_name}::{method_name}");

                match self.function_registry.remove(&key) {
                    Some(tool) => Ok(Some(tool)),
                    None => {
                        error!("Class method tool '{key}' not found in registry.");
                        Ok(None)
                    }
                }
            }

            other => Err(AgentKitError::UnknownToolType(other.to_string())),
        }
    }
}