агенты: добавить поиск инструментов по каталогу

2026-07-21 13:12:46 +03:00
parent 63f8ee333f
commit 99bd05c145
35 changed files with 2088 additions and 155 deletions
@@ -1,5 +1,6 @@
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
+use thiserror::Error;
 use time::OffsetDateTime;

 use crate::ids::{AgentId, OperationId, WorkspaceId};
@@ -36,11 +37,149 @@ pub struct AgentVersion {
    pub version: u32,
    pub status: AgentStatus,
    pub instructions: Value,
-    pub tool_selection_policy: Value,
+    #[serde(default)]
+    pub tool_selection_policy: ToolSelectionPolicy,
    #[serde(with = "time::serde::rfc3339")]
    pub created_at: OffsetDateTime,
 }

+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum ToolAccessMode {
+    #[default]
+    #[serde(alias = "allow_list")]
+    Direct,
+    Search,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ToolGroup {
+    pub id: String,
+    pub name: String,
+    pub description: String,
+    #[serde(default)]
+    pub tool_names: Vec<String>,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ToolSearchSettings {
+    #[serde(default = "default_search_result_limit")]
+    pub max_results: usize,
+}
+
+impl Default for ToolSearchSettings {
+    fn default() -> Self {
+        Self {
+            max_results: default_search_result_limit(),
+        }
+    }
+}
+
+#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ToolSelectionPolicy {
+    #[serde(default)]
+    pub mode: ToolAccessMode,
+    #[serde(default)]
+    pub groups: Vec<ToolGroup>,
+    #[serde(default)]
+    pub search: ToolSearchSettings,
+}
+
+#[derive(Clone, Debug, Error, PartialEq, Eq)]
+pub enum ToolSelectionPolicyError {
+    #[error("search result limit must be between 1 and 20")]
+    InvalidSearchResultLimit,
+    #[error("tool groups are only allowed in search mode")]
+    GroupsRequireSearchMode,
+    #[error("tool group id {group_id} is invalid")]
+    InvalidGroupId { group_id: String },
+    #[error("tool group id {group_id} is duplicated")]
+    DuplicateGroupId { group_id: String },
+    #[error("tool group {group_id} must have a name and description")]
+    IncompleteGroup { group_id: String },
+    #[error("tool group {group_id} references unknown tool {tool_name}")]
+    UnknownTool { group_id: String, tool_name: String },
+    #[error("tool group {group_id} contains duplicate tool {tool_name}")]
+    DuplicateTool { group_id: String, tool_name: String },
+    #[error("tool name {tool_name} is reserved by search mode")]
+    ReservedToolName { tool_name: String },
+}
+
+impl ToolSelectionPolicy {
+    pub fn validate_for_tools<'a>(
+        &self,
+        tool_names: impl IntoIterator<Item = &'a str>,
+    ) -> Result<(), ToolSelectionPolicyError> {
+        use std::collections::BTreeSet;
+
+        if !(1..=20).contains(&self.search.max_results) {
+            return Err(ToolSelectionPolicyError::InvalidSearchResultLimit);
+        }
+        if self.mode == ToolAccessMode::Direct && !self.groups.is_empty() {
+            return Err(ToolSelectionPolicyError::GroupsRequireSearchMode);
+        }
+
+        let known_tools = tool_names.into_iter().collect::<BTreeSet<_>>();
+        if self.mode == ToolAccessMode::Search
+            && let Some(tool_name) = known_tools
+                .iter()
+                .find(|tool_name| matches!(**tool_name, "search_tools" | "call_tool"))
+        {
+            return Err(ToolSelectionPolicyError::ReservedToolName {
+                tool_name: (*tool_name).to_owned(),
+            });
+        }
+        let mut group_ids = BTreeSet::new();
+        for group in &self.groups {
+            if !valid_group_id(&group.id) {
+                return Err(ToolSelectionPolicyError::InvalidGroupId {
+                    group_id: group.id.clone(),
+                });
+            }
+            if !group_ids.insert(group.id.as_str()) {
+                return Err(ToolSelectionPolicyError::DuplicateGroupId {
+                    group_id: group.id.clone(),
+                });
+            }
+            if group.name.trim().is_empty() || group.description.trim().is_empty() {
+                return Err(ToolSelectionPolicyError::IncompleteGroup {
+                    group_id: group.id.clone(),
+                });
+            }
+
+            let mut grouped_tools = BTreeSet::new();
+            for tool_name in &group.tool_names {
+                if !known_tools.contains(tool_name.as_str()) {
+                    return Err(ToolSelectionPolicyError::UnknownTool {
+                        group_id: group.id.clone(),
+                        tool_name: tool_name.clone(),
+                    });
+                }
+                if !grouped_tools.insert(tool_name.as_str()) {
+                    return Err(ToolSelectionPolicyError::DuplicateTool {
+                        group_id: group.id.clone(),
+                        tool_name: tool_name.clone(),
+                    });
+                }
+            }
+        }
+
+        Ok(())
+    }
+}
+
+const fn default_search_result_limit() -> usize {
+    8
+}
+
+fn valid_group_id(value: &str) -> bool {
+    !value.is_empty()
+        && value.len() <= 64
+        && value
+            .bytes()
+            .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-')
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct AgentOperationBinding {
    pub operation_id: OperationId,
@@ -56,7 +195,10 @@ mod tests {
    use serde_json::json;
    use time::{OffsetDateTime, format_description::well_known::Rfc3339};

-    use super::{Agent, AgentStatus, AgentVersion};
+    use super::{
+        Agent, AgentStatus, AgentVersion, ToolAccessMode, ToolGroup, ToolSelectionPolicy,
+        ToolSelectionPolicyError,
+    };
    use crate::ids::{AgentId, WorkspaceId};

    #[test]
@@ -98,5 +240,28 @@ mod tests {
            version.created_at,
            OffsetDateTime::parse("2026-03-25T12:00:00Z", &Rfc3339).unwrap()
        );
+        assert_eq!(version.tool_selection_policy.mode, ToolAccessMode::Direct);
+    }
+
+    #[test]
+    fn search_policy_validates_groups_against_published_tools() {
+        let policy = ToolSelectionPolicy {
+            mode: ToolAccessMode::Search,
+            groups: vec![ToolGroup {
+                id: "finance".to_owned(),
+                name: "Finance".to_owned(),
+                description: "Invoices and payments".to_owned(),
+                tool_names: vec!["create_invoice".to_owned()],
+            }],
+            ..ToolSelectionPolicy::default()
+        };
+
+        assert_eq!(
+            policy.validate_for_tools(["list_invoices"]),
+            Err(ToolSelectionPolicyError::UnknownTool {
+                group_id: "finance".to_owned(),
+                tool_name: "create_invoice".to_owned(),
+            })
+        );
    }
 }
@@ -12,6 +12,7 @@ pub mod protocol;
 pub mod secret;
 pub mod tool_catalog;
 pub mod tool_quality;
+pub mod tool_search;
 pub mod workspace;

 pub mod domain {
@@ -19,7 +20,10 @@ pub mod domain {
        InvitationStatus, InvitationToken, Membership, MembershipRole, PlatformApiKey,
        PlatformApiKeyKind, PlatformApiKeyScope, PlatformApiKeyStatus, User, UserStatus,
    };
-    pub use crate::agent::{Agent, AgentOperationBinding, AgentStatus, AgentVersion};
+    pub use crate::agent::{
+        Agent, AgentOperationBinding, AgentStatus, AgentVersion, ToolAccessMode, ToolGroup,
+        ToolSearchSettings, ToolSelectionPolicy, ToolSelectionPolicyError,
+    };
    pub use crate::approval::{ApprovalRequest, ApprovalRequestStatus};
    pub use crate::auth::{
        ApiKeyHeaderAuthConfig, ApiKeyQueryAuthConfig, AuthConfig, AuthProfile, BasicAuthConfig,
@@ -56,6 +60,7 @@ pub mod domain {
        ToolQualityCatalogTool, ToolQualityFinding, ToolQualityMappingRule, ToolQualityMappingSet,
        ToolQualityReport, ToolQualitySchemaKind, ToolQualitySchemaNode, ToolQualitySeverity,
    };
+    pub use crate::tool_search::{SearchableTool, ToolSearchMatch, search_tool_catalog};
    pub use crate::workspace::{Workspace, WorkspaceStatus};
 }

@@ -92,7 +97,10 @@ pub use access::{
    InvitationStatus, InvitationToken, Membership, MembershipRole, PlatformApiKey,
    PlatformApiKeyKind, PlatformApiKeyScope, PlatformApiKeyStatus, User, UserStatus,
 };
-pub use agent::{Agent, AgentOperationBinding, AgentStatus, AgentVersion};
+pub use agent::{
+    Agent, AgentOperationBinding, AgentStatus, AgentVersion, ToolAccessMode, ToolGroup,
+    ToolSearchSettings, ToolSelectionPolicy, ToolSelectionPolicyError,
+};
 pub use approval::{ApprovalRequest, ApprovalRequestStatus};
 pub use auth::{
    ApiKeyHeaderAuthConfig, ApiKeyQueryAuthConfig, AuthConfig, AuthProfile, BasicAuthConfig,
@@ -153,4 +161,5 @@ pub use tool_quality::{
    analyze_agent_tool_catalog_quality, analyze_tool_identity_quality,
    analyze_tool_response_projection_quality, analyze_tool_schema_quality,
 };
+pub use tool_search::{SearchableTool, ToolSearchMatch, search_tool_catalog};
 pub use workspace::{Workspace, WorkspaceStatus};
@@ -0,0 +1,237 @@
+use std::collections::{BTreeMap, BTreeSet};
+
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct SearchableTool {
+    pub name: String,
+    pub title: String,
+    pub description: String,
+    pub input_schema: Value,
+    #[serde(default)]
+    pub group_ids: Vec<String>,
+    #[serde(default)]
+    pub group_context: String,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct ToolSearchMatch {
+    pub tool: SearchableTool,
+    pub score: f64,
+}
+
+pub fn search_tool_catalog(
+    tools: &[SearchableTool],
+    query: &str,
+    group_ids: &[String],
+    max_results: usize,
+) -> Vec<ToolSearchMatch> {
+    let query_tokens = tokenize(query);
+    if query_tokens.is_empty() || max_results == 0 {
+        return Vec::new();
+    }
+
+    let requested_groups = group_ids
+        .iter()
+        .map(String::as_str)
+        .collect::<BTreeSet<_>>();
+    let candidates = tools
+        .iter()
+        .filter(|tool| {
+            requested_groups.is_empty()
+                || tool
+                    .group_ids
+                    .iter()
+                    .any(|group_id| requested_groups.contains(group_id.as_str()))
+        })
+        .collect::<Vec<_>>();
+    if candidates.is_empty() {
+        return Vec::new();
+    }
+
+    let documents = candidates
+        .iter()
+        .map(|tool| weighted_tokens(tool))
+        .collect::<Vec<_>>();
+    let average_length =
+        documents.iter().map(Vec::len).sum::<usize>() as f64 / documents.len() as f64;
+    let candidate_count = documents.len();
+    let document_frequency = document_frequency(&documents, &query_tokens);
+    let normalized_query = query.trim().to_lowercase();
+
+    let mut matches = candidates
+        .into_iter()
+        .zip(documents)
+        .filter_map(|(tool, document)| {
+            let score = bm25_score(
+                &document,
+                &query_tokens,
+                &document_frequency,
+                candidate_count,
+                average_length,
+            ) + exact_match_bonus(tool, &normalized_query);
+            (score > 0.0).then(|| ToolSearchMatch {
+                tool: tool.clone(),
+                score,
+            })
+        })
+        .collect::<Vec<_>>();
+
+    matches.sort_by(|left, right| {
+        right
+            .score
+            .total_cmp(&left.score)
+            .then_with(|| left.tool.name.cmp(&right.tool.name))
+    });
+    matches.truncate(max_results);
+    matches
+}
+
+fn weighted_tokens(tool: &SearchableTool) -> Vec<String> {
+    let mut tokens = Vec::new();
+    for _ in 0..3 {
+        tokens.extend(tokenize(&tool.name));
+        tokens.extend(tokenize(&tool.title));
+    }
+    tokens.extend(tokenize(&tool.description));
+    for _ in 0..2 {
+        tokens.extend(tokenize(&tool.group_context));
+    }
+    tokens
+}
+
+fn document_frequency(
+    documents: &[Vec<String>],
+    query_tokens: &[String],
+) -> BTreeMap<String, usize> {
+    let mut frequencies = BTreeMap::new();
+    for query_token in query_tokens {
+        let count = documents
+            .iter()
+            .filter(|document| document.iter().any(|token| token == query_token))
+            .count();
+        frequencies.insert(query_token.clone(), count);
+    }
+    frequencies
+}
+
+fn bm25_score(
+    document: &[String],
+    query_tokens: &[String],
+    document_frequency: &BTreeMap<String, usize>,
+    document_count: usize,
+    average_length: f64,
+) -> f64 {
+    const K1: f64 = 1.2;
+    const B: f64 = 0.75;
+
+    query_tokens.iter().fold(0.0, |score, token| {
+        let term_frequency = document.iter().filter(|term| *term == token).count() as f64;
+        if term_frequency == 0.0 {
+            return score;
+        }
+        let frequency = document_frequency.get(token).copied().unwrap_or_default() as f64;
+        let inverse_document_frequency =
+            ((document_count as f64 - frequency + 0.5) / (frequency + 0.5) + 1.0).ln();
+        let length_ratio = if average_length > 0.0 {
+            document.len() as f64 / average_length
+        } else {
+            1.0
+        };
+        let normalized_frequency =
+            term_frequency * (K1 + 1.0) / (term_frequency + K1 * (1.0 - B + B * length_ratio));
+        score + inverse_document_frequency * normalized_frequency
+    })
+}
+
+fn exact_match_bonus(tool: &SearchableTool, query: &str) -> f64 {
+    if query.is_empty() {
+        return 0.0;
+    }
+    let name = tool.name.to_lowercase();
+    let title = tool.title.to_lowercase();
+    if name == query || title == query {
+        8.0
+    } else if name.contains(query) || title.contains(query) {
+        3.0
+    } else {
+        0.0
+    }
+}
+
+fn tokenize(value: &str) -> Vec<String> {
+    value
+        .to_lowercase()
+        .split(|character: char| !character.is_alphanumeric())
+        .filter(|token| token.chars().count() >= 2)
+        .map(ToOwned::to_owned)
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use serde_json::json;
+
+    use super::{SearchableTool, search_tool_catalog};
+
+    fn tool(name: &str, title: &str, description: &str, group_ids: &[&str]) -> SearchableTool {
+        SearchableTool {
+            name: name.to_owned(),
+            title: title.to_owned(),
+            description: description.to_owned(),
+            input_schema: json!({"type": "object"}),
+            group_ids: group_ids.iter().map(|value| (*value).to_owned()).collect(),
+            group_context: if group_ids.contains(&"finance") {
+                "Finance invoices and payments".to_owned()
+            } else {
+                "Customer support tickets".to_owned()
+            },
+        }
+    }
+
+    #[test]
+    fn ranks_title_and_name_matches_above_description_matches() {
+        let tools = vec![
+            tool(
+                "create_invoice",
+                "Create invoice",
+                "Issue a bill",
+                &["finance"],
+            ),
+            tool(
+                "list_customers",
+                "List customers",
+                "Customers with invoices",
+                &["support"],
+            ),
+        ];
+
+        let matches = search_tool_catalog(&tools, "create invoice", &[], 8);
+
+        assert_eq!(matches[0].tool.name, "create_invoice");
+    }
+
+    #[test]
+    fn filters_candidates_by_group() {
+        let tools = vec![
+            tool(
+                "create_invoice",
+                "Create invoice",
+                "Issue a bill",
+                &["finance"],
+            ),
+            tool(
+                "create_ticket",
+                "Create ticket",
+                "Open a support case",
+                &["support"],
+            ),
+        ];
+
+        let matches = search_tool_catalog(&tools, "create", &["support".to_owned()], 8);
+
+        assert_eq!(matches.len(), 1);
+        assert_eq!(matches[0].tool.name, "create_ticket");
+    }
+}