Genarrative/rust/crates/agent-runtime-core/src/capability.rs

//! 能力目录。目录只描述能力和稳定函数名，不包含执行权限。

use std::collections::BTreeMap;

use serde_json::Value;

use crate::contract::{ContractError, validate_identifier, validate_non_empty, validate_object};

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CapabilityDefinition<D> {
    id: String,
    function_name: String,
    description: String,
    input_schema: Value,
    dispatch: D,
}

impl<D> CapabilityDefinition<D> {
    pub fn try_new(
        id: impl Into<String>,
        function_name: impl Into<String>,
        description: impl Into<String>,
        input_schema: Value,
        dispatch: D,
    ) -> Result<Self, ContractError> {
        let id = id.into();
        let function_name = function_name.into();
        let description = description.into();
        validate_identifier(&id, "capability id")?;
        validate_identifier(&function_name, "capability function name")?;
        validate_non_empty(&description, "capability description")?;
        validate_object(&input_schema, "capability inputSchema")?;
        Ok(Self {
            id,
            function_name,
            description,
            input_schema,
            dispatch,
        })
    }

    pub fn id(&self) -> &str {
        &self.id
    }
    pub fn function_name(&self) -> &str {
        &self.function_name
    }
    pub fn description(&self) -> &str {
        &self.description
    }
    pub fn input_schema(&self) -> &Value {
        &self.input_schema
    }
    pub fn dispatch(&self) -> &D {
        &self.dispatch
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CapabilityRegistryError {
    message: String,
}

impl CapabilityRegistryError {
    pub fn message(&self) -> &str {
        &self.message
    }
}
impl std::fmt::Display for CapabilityRegistryError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.message)
    }
}
impl std::error::Error for CapabilityRegistryError {}
impl From<ContractError> for CapabilityRegistryError {
    fn from(value: ContractError) -> Self {
        Self {
            message: value.message().to_owned(),
        }
    }
}

#[derive(Clone, Debug)]
pub struct CapabilityRegistry<D> {
    definitions: Vec<CapabilityDefinition<D>>,
    by_id: BTreeMap<String, usize>,
    by_function_name: BTreeMap<String, usize>,
}

impl<D> CapabilityRegistry<D> {
    pub fn try_new(
        definitions: impl IntoIterator<Item = CapabilityDefinition<D>>,
    ) -> Result<Self, CapabilityRegistryError> {
        let definitions = definitions.into_iter().collect::<Vec<_>>();
        let mut by_id = BTreeMap::new();
        let mut by_function_name = BTreeMap::new();
        for (index, definition) in definitions.iter().enumerate() {
            if by_id.insert(definition.id.clone(), index).is_some() {
                return Err(CapabilityRegistryError {
                    message: format!("capability id 重复：{}", definition.id),
                });
            }
            if by_function_name
                .insert(definition.function_name.clone(), index)
                .is_some()
            {
                return Err(CapabilityRegistryError {
                    message: format!(
                        "capability function name 重复：{}",
                        definition.function_name
                    ),
                });
            }
        }
        Ok(Self {
            definitions,
            by_id,
            by_function_name,
        })
    }

    pub fn get(&self, id: &str) -> Option<&CapabilityDefinition<D>> {
        self.by_id
            .get(id)
            .and_then(|index| self.definitions.get(*index))
    }

    pub fn get_by_function_name(&self, name: &str) -> Option<&CapabilityDefinition<D>> {
        self.by_function_name
            .get(name)
            .and_then(|index| self.definitions.get(*index))
    }

    pub fn iter(&self) -> impl ExactSizeIterator<Item = &CapabilityDefinition<D>> {
        self.definitions.iter()
    }
    pub fn len(&self) -> usize {
        self.definitions.len()
    }
    pub fn is_empty(&self) -> bool {
        self.definitions.is_empty()
    }
}