Implement field-aware tool_search ranking

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<ToolSearchSourceInfo>,
 }
 
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct ToolSearchIdentity {
+    pub canonical_aliases: Vec<String>,
+    pub tool_aliases: Vec<String>,
+    pub source_aliases: Vec<String>,
+}
+
 impl ToolSearchInfo {
     pub fn from_tool_spec(
         spec: ToolSpec,
@@ -31,6 +39,20 @@ impl ToolSearchInfo {
         search_text: String,
         spec: ToolSpec,
         source_info: Option<ToolSearchSourceInfo>,
+    ) -> Option<Self> {
+        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<ToolSearchSourceInfo>,
+        identity: ToolSearchIdentity,
     ) -> Option<Self> {
         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<String>) {
-    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<String>) {
-    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<String>) {
     }
 }
 
+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<String>, 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<String>, 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<String>, part: String) {
     let part = part.trim();
     if !part.is_empty() {