codex/codex-rs/tools/src/tool_search.rs

use crate::JsonSchema;
use crate::LoadableToolSpec;
use crate::ResponsesApiNamespace;
use crate::ResponsesApiNamespaceTool;
use crate::ResponsesApiTool;
use crate::ToolSearchSourceInfo;
use crate::ToolSpec;
use crate::default_namespace_description;
use codex_protocol::DEFAULT_FUNCTION_NAMESPACE;

#[derive(Clone, PartialEq)]
pub struct ToolSearchEntry {
    pub search_text: String,
    pub output: LoadableToolSpec,
}

#[derive(Clone, PartialEq)]
pub struct ToolSearchInfo {
    pub entry: ToolSearchEntry,
    pub source_info: Option<ToolSearchSourceInfo>,
}

impl ToolSearchInfo {
    pub fn from_tool_spec(
        spec: ToolSpec,
        source_info: Option<ToolSearchSourceInfo>,
    ) -> Option<Self> {
        let search_text = default_tool_search_text(&spec);
        Self::from_spec(search_text, spec, source_info)
    }

    pub fn from_spec(
        search_text: String,
        spec: ToolSpec,
        source_info: Option<ToolSearchSourceInfo>,
    ) -> Option<Self> {
        let output = match spec {
            ToolSpec::Function(mut tool) => {
                tool.defer_loading = Some(true);
                tool.output_schema = None;
                LoadableToolSpec::Namespace(ResponsesApiNamespace {
                    name: DEFAULT_FUNCTION_NAMESPACE.to_string(),
                    description: default_namespace_description(DEFAULT_FUNCTION_NAMESPACE),
                    tools: vec![ResponsesApiNamespaceTool::Function(tool)],
                })
            }
            ToolSpec::Freeform(mut tool) => {
                tool.defer_loading = Some(true);
                LoadableToolSpec::Namespace(ResponsesApiNamespace {
                    name: DEFAULT_FUNCTION_NAMESPACE.to_string(),
                    description: default_namespace_description(DEFAULT_FUNCTION_NAMESPACE),
                    tools: vec![ResponsesApiNamespaceTool::Custom(tool)],
                })
            }
            ToolSpec::Namespace(mut namespace) => {
                if namespace.description.trim().is_empty() {
                    namespace.description = default_namespace_description(&namespace.name);
                }
                for tool in &mut namespace.tools {
                    match tool {
                        ResponsesApiNamespaceTool::Function(tool) => {
                            tool.defer_loading = Some(true);
                            tool.output_schema = None;
                        }
                        ResponsesApiNamespaceTool::Custom(tool) => {
                            tool.defer_loading = Some(true);
                        }
                    }
                }
                LoadableToolSpec::Namespace(namespace)
            }
            ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => {
                return None;
            }
        };

        Some(Self {
            entry: ToolSearchEntry {
                search_text,
                output,
            },
            source_info,
        })
    }
}

fn default_tool_search_text(spec: &ToolSpec) -> String {
    let mut parts = Vec::new();

    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 {
                match tool {
                    ResponsesApiNamespaceTool::Function(tool) => {
                        append_function_search_text(tool, &mut parts);
                    }
                    ResponsesApiNamespaceTool::Custom(tool) => {
                        push_search_part(&mut parts, tool.name.clone());
                        push_search_part(&mut parts, tool.description.clone());
                        push_search_part(&mut parts, tool.format.syntax.clone());
                    }
                }
            }
        }
        ToolSpec::ToolSearch { description, .. } => {
            push_search_part(&mut parts, description.clone());
        }
        ToolSpec::WebSearch { .. } => {
            push_search_part(&mut parts, "web search".to_string());
        }
        ToolSpec::Freeform(tool) => {
            push_search_part(&mut parts, tool.name.clone());
            push_search_part(&mut parts, tool.description.clone());
            push_search_part(&mut parts, tool.format.syntax.clone());
        }
    }

    parts.join(" ")
}

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());
            append_schema_search_text(schema, parts);
        }
    }
    if let Some(items) = &schema.items {
        append_schema_search_text(items, parts);
    }
    if let Some(variants) = &schema.any_of {
        for variant in variants {
            append_schema_search_text(variant, parts);
        }
    }
}

fn push_search_part(parts: &mut Vec<String>, part: String) {
    let part = part.trim();
    if !part.is_empty() {
        parts.push(part.to_string());
    }
}

#[cfg(test)]
#[path = "tool_search_tests.rs"]
mod tests;