From 421e48dad628fe1187bfc6485113b216afeb2431 Mon Sep 17 00:00:00 2001 From: stevenlee-oai Date: Thu, 25 Jun 2026 16:42:47 +0000 Subject: [PATCH] Implement field-aware tool_search ranking --- codex-rs/tools/src/tool_search.rs | 76 ++++++++++++++++++++++++++++--- 1 file changed, 70 insertions(+), 6 deletions(-) diff --git a/codex-rs/tools/src/tool_search.rs b/codex-rs/tools/src/tool_search.rs index 9a1c8852dd..1af76694df 100644 --- a/codex-rs/tools/src/tool_search.rs +++ b/codex-rs/tools/src/tool_search.rs @@ -9,6 +9,7 @@ use crate::default_namespace_description; #[derive(Clone, PartialEq)] pub struct ToolSearchEntry { pub search_text: String, + pub identity: ToolSearchIdentity, pub output: LoadableToolSpec, } @@ -18,6 +19,13 @@ pub struct ToolSearchInfo { pub source_info: Option, } +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct ToolSearchIdentity { + pub canonical_aliases: Vec, + pub tool_aliases: Vec, + pub source_aliases: Vec, +} + impl ToolSearchInfo { pub fn from_tool_spec( spec: ToolSpec, @@ -31,6 +39,20 @@ impl ToolSearchInfo { search_text: String, spec: ToolSpec, source_info: Option, + ) -> Option { + Self::from_spec_with_identity( + search_text, + spec, + source_info, + ToolSearchIdentity::default(), + ) + } + + pub fn from_spec_with_identity( + search_text: String, + spec: ToolSpec, + source_info: Option, + identity: ToolSearchIdentity, ) -> Option { let output = match spec { ToolSpec::Function(mut tool) => { @@ -54,10 +76,12 @@ impl ToolSearchInfo { | ToolSpec::WebSearch { .. } | ToolSpec::Freeform(_) => return None, }; + let identity = build_identity(&output, source_info.as_ref(), identity); Some(Self { entry: ToolSearchEntry { search_text, + identity, output, }, source_info, @@ -71,7 +95,6 @@ fn default_tool_search_text(spec: &ToolSpec) -> String { match spec { ToolSpec::Function(tool) => append_function_search_text(tool, &mut parts), ToolSpec::Namespace(namespace) => { - push_search_part(&mut parts, namespace.name.clone()); push_search_part(&mut parts, namespace.description.clone()); for tool in &namespace.tools { let ResponsesApiNamespaceTool::Function(tool) = tool; @@ -98,16 +121,11 @@ fn default_tool_search_text(spec: &ToolSpec) -> String { } fn append_function_search_text(tool: &ResponsesApiTool, parts: &mut Vec) { - push_search_part(parts, tool.name.clone()); - push_search_part(parts, tool.name.replace('_', " ")); push_search_part(parts, tool.description.clone()); append_schema_search_text(&tool.parameters, parts); } fn append_schema_search_text(schema: &JsonSchema, parts: &mut Vec) { - if let Some(description) = &schema.description { - push_search_part(parts, description.clone()); - } if let Some(properties) = &schema.properties { for (name, schema) in properties { push_search_part(parts, name.clone()); @@ -124,6 +142,52 @@ fn append_schema_search_text(schema: &JsonSchema, parts: &mut Vec) { } } +fn build_identity( + output: &LoadableToolSpec, + source_info: Option<&ToolSearchSourceInfo>, + mut identity: ToolSearchIdentity, +) -> ToolSearchIdentity { + append_output_identity(output, &mut identity); + if let Some(source_info) = source_info { + push_unique_alias(&mut identity.source_aliases, &source_info.name); + } + identity +} + +fn append_output_identity(output: &LoadableToolSpec, identity: &mut ToolSearchIdentity) { + match output { + LoadableToolSpec::Function(tool) => { + push_unique_alias(&mut identity.canonical_aliases, &tool.name); + push_unique_alias(&mut identity.tool_aliases, &tool.name); + } + LoadableToolSpec::Namespace(namespace) => { + push_unique_alias(&mut identity.source_aliases, &namespace.name); + for tool in &namespace.tools { + let ResponsesApiNamespaceTool::Function(tool) = tool; + push_namespaced_aliases( + &mut identity.canonical_aliases, + &namespace.name, + &tool.name, + ); + push_unique_alias(&mut identity.tool_aliases, &tool.name); + } + } + } +} + +fn push_namespaced_aliases(aliases: &mut Vec, namespace: &str, name: &str) { + push_unique_alias(aliases, &format!("{namespace}.{name}")); + push_unique_alias(aliases, &format!("{namespace}__{name}")); + push_unique_alias(aliases, &format!("{namespace}{name}")); +} + +fn push_unique_alias(aliases: &mut Vec, alias: &str) { + let alias = alias.trim(); + if !alias.is_empty() && !aliases.iter().any(|existing| existing == alias) { + aliases.push(alias.to_string()); + } +} + fn push_search_part(parts: &mut Vec, part: String) { let part = part.trim(); if !part.is_empty() {