pe_tools/

selector.rs

//! Tool selection — mechanism for filtering available tools before an LLM call.
//!
//! The library provides `ToolSelector` as the extension point and
//! `AllToolsSelector` as the default (pass-through) implementation.
//! Users implement their own selectors for context-aware tool filtering.

use async_trait::async_trait;
use pe_core::{Message, ToolSchema};

/// Async trait for selecting which tools to present to the LLM.
///
/// Implementations receive the full set of available tools and the
/// conversation history, then return the names of tools to include.
///
/// # Examples
///
/// ```
/// use pe_tools::selector::{ToolSelector, AllToolsSelector};
/// use pe_core::{ToolSchema, Message};
///
/// # tokio::runtime::Runtime::new().unwrap().block_on(async {
/// let selector = AllToolsSelector;
/// let tools = vec![
///     ToolSchema {
///         name: "search".into(),
///         description: "Search the web".into(),
///         parameters: serde_json::json!({}),
///         strict: false,
///     },
/// ];
/// let selected = selector.select(&tools, &[]).await;
/// assert_eq!(selected, vec!["search".to_string()]);
/// # });
/// ```
#[async_trait]
pub trait ToolSelector: Send + Sync + 'static {
    /// Select tool names from the available set.
    ///
    /// Returns a `Vec<String>` of tool names that should be included
    /// in the next LLM call. Names must match `ToolSchema::name` values.
    async fn select(&self, available: &[ToolSchema], messages: &[Message]) -> Vec<String>;
}

/// Default selector that returns all available tools.
///
/// This is the pass-through implementation: every registered tool
/// is presented to the LLM. Users who need filtering implement
/// their own `ToolSelector`.
///
/// # Examples
///
/// ```
/// use pe_tools::selector::{ToolSelector, AllToolsSelector};
/// use pe_core::{ToolSchema, Message};
///
/// # tokio::runtime::Runtime::new().unwrap().block_on(async {
/// let selector = AllToolsSelector;
/// let names = selector.select(&[], &[]).await;
/// assert!(names.is_empty());
/// # });
/// ```
#[derive(Debug, Clone, Default)]
pub struct AllToolsSelector;

#[async_trait]
impl ToolSelector for AllToolsSelector {
    async fn select(&self, available: &[ToolSchema], _messages: &[Message]) -> Vec<String> {
        available.iter().map(|t| t.name.clone()).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use pe_core::Message;
    use serde_json::json;

    fn make_tool(name: &str) -> ToolSchema {
        ToolSchema {
            name: name.into(),
            description: format!("{name} tool"),
            parameters: json!({}),
            strict: false,
        }
    }

    #[tokio::test]
    async fn all_tools_selector_returns_every_tool_name() {
        let selector = AllToolsSelector;
        let tools = vec![
            make_tool("search"),
            make_tool("calculator"),
            make_tool("email"),
        ];

        let selected = selector.select(&tools, &[]).await;

        assert_eq!(selected, vec!["search", "calculator", "email"]);
    }

    #[tokio::test]
    async fn all_tools_selector_empty_input_returns_empty() {
        let selector = AllToolsSelector;
        let selected = selector.select(&[], &[]).await;
        assert!(selected.is_empty());
    }

    #[tokio::test]
    async fn all_tools_selector_ignores_messages() {
        let selector = AllToolsSelector;
        let tools = vec![make_tool("read_file")];
        let messages = vec![
            Message::human("Please read my file"),
            Message::ai("I'll use the read_file tool"),
        ];

        let selected = selector.select(&tools, &messages).await;

        assert_eq!(selected, vec!["read_file"]);
    }

    #[tokio::test]
    async fn custom_selector_filters_by_conversation() {
        /// A test selector that only includes tools mentioned in messages.
        struct MentionedToolsSelector;

        #[async_trait]
        impl ToolSelector for MentionedToolsSelector {
            async fn select(&self, available: &[ToolSchema], messages: &[Message]) -> Vec<String> {
                let text: String = messages
                    .iter()
                    .filter_map(|m| match m {
                        Message::Human(h) => h.content.as_text().map(|s| s.to_owned()),
                        Message::Ai(a) => a.content.as_text().map(|s| s.to_owned()),
                        Message::System(s) => Some(s.content.clone()),
                        Message::Tool(t) => Some(t.content.clone()),
                        _ => None,
                    })
                    .collect::<Vec<_>>()
                    .join(" ");

                available
                    .iter()
                    .filter(|t| text.contains(&t.name))
                    .map(|t| t.name.clone())
                    .collect()
            }
        }

        let selector = MentionedToolsSelector;
        let tools = vec![
            make_tool("search"),
            make_tool("calculator"),
            make_tool("email"),
        ];
        let messages = vec![Message::human("I need to search for something")];

        let selected = selector.select(&tools, &messages).await;

        assert_eq!(selected, vec!["search"]);
    }

    #[tokio::test]
    async fn custom_selector_can_return_empty() {
        struct NoToolsSelector;

        #[async_trait]
        impl ToolSelector for NoToolsSelector {
            async fn select(
                &self,
                _available: &[ToolSchema],
                _messages: &[Message],
            ) -> Vec<String> {
                vec![]
            }
        }

        let selector = NoToolsSelector;
        let tools = vec![make_tool("search"), make_tool("calculator")];
        let selected = selector.select(&tools, &[]).await;
        assert!(selected.is_empty());
    }
}