mirror of
https://github.com/openai/codex.git
synced 2026-10-28 18:58:46 +00:00
Add model-gated literature search to standalone web search (#57542)
## What changed Expose `literature_search_query` in standalone `web.run` when the active model lists it in `experimental_supported_tools`. Support query text, domain filters, and recency filters, and reject literature search calls for models without the capability. Pass the active model information to extensions for each step and report literature queries as search actions in web search events. ## Testing Add coverage for capability rejection, model-dependent schema exposure, app-server round trips with a turn-level model override, and a scenario that searches for a paper and opens the returned reference. GitOrigin-RevId: 87ef1fadd80ca7bfc2a7642c4e443f9b9f2e0e00
This commit is contained in:
rhan-oai
@@ -22,6 +22,7 @@ use codex_app_server_protocol::WebSearchAction;
|
||||
use codex_app_server_protocol::WebSearchItem;
|
||||
use codex_config::types::AuthCredentialsStoreMode;
|
||||
use codex_features::Feature;
|
||||
use codex_models_manager::model_info::model_info_from_slug;
|
||||
use core_test_support::responses;
|
||||
use core_test_support::responses::strip_response_item_ids_from_json;
|
||||
use pretty_assertions::assert_eq;
|
||||
@@ -44,12 +45,29 @@ const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
#[tokio::test]
|
||||
async fn standalone_web_search_round_trips_output() -> Result<()> {
|
||||
assert_standalone_web_search_round_trips_output(WebSearchProvider::ChatGpt).await
|
||||
assert_standalone_web_search_round_trips_output(
|
||||
WebSearchProvider::ChatGpt,
|
||||
/*literature_search*/ false,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standalone_web_search_round_trips_output_for_custom_provider() -> Result<()> {
|
||||
assert_standalone_web_search_round_trips_output(WebSearchProvider::CustomResponses).await
|
||||
assert_standalone_web_search_round_trips_output(
|
||||
WebSearchProvider::CustomResponses,
|
||||
/*literature_search*/ false,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standalone_literature_search_round_trips_for_an_opted_in_turn_model() -> Result<()> {
|
||||
assert_standalone_web_search_round_trips_output(
|
||||
WebSearchProvider::ChatGpt,
|
||||
/*literature_search*/ true,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
@@ -60,13 +78,27 @@ enum WebSearchProvider {
|
||||
|
||||
async fn assert_standalone_web_search_round_trips_output(
|
||||
provider: WebSearchProvider,
|
||||
literature_search: bool,
|
||||
) -> Result<()> {
|
||||
let call_id = "web-run-1";
|
||||
let (model, query, commands) = if literature_search {
|
||||
(
|
||||
"paper-model",
|
||||
"plasma biomarkers",
|
||||
json!({"literature_search_query": [{"q": "plasma biomarkers", "recency": 30, "domains": ["pubmed.ncbi.nlm.nih.gov"]}]}),
|
||||
)
|
||||
} else {
|
||||
(
|
||||
"mock-model",
|
||||
"standalone web search",
|
||||
json!({"search_query": [{"q": "standalone web search"}]}),
|
||||
)
|
||||
};
|
||||
let expected_model_id = "model-id-from-search-context";
|
||||
let search_context = json!({
|
||||
"telemetry_attributes": {
|
||||
"model_id": expected_model_id,
|
||||
"model_slug": "mock-model",
|
||||
"model_slug": model,
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
@@ -90,10 +122,7 @@ async fn assert_standalone_web_search_round_trips_output(
|
||||
call_id,
|
||||
"web",
|
||||
"run",
|
||||
&json!({
|
||||
"search_query": [{"q": "standalone web search"}],
|
||||
})
|
||||
.to_string(),
|
||||
&commands.to_string(),
|
||||
),
|
||||
responses::ev_completed("resp-1"),
|
||||
]),
|
||||
@@ -110,6 +139,23 @@ async fn assert_standalone_web_search_round_trips_output(
|
||||
.with_root_config(&format!("chatgpt_base_url = \"{}\"", server.uri()))
|
||||
.enable_feature(Feature::StandaloneWebSearch)
|
||||
.with_provider_config("supports_websockets = false");
|
||||
let config = if literature_search {
|
||||
let mut paper_model = model_info_from_slug(model);
|
||||
paper_model.experimental_supported_tools = vec!["literature_search_query".to_string()];
|
||||
let catalog_path = codex_home.path().join("models.json");
|
||||
std::fs::write(
|
||||
&catalog_path,
|
||||
serde_json::to_vec(
|
||||
&json!({"models": [model_info_from_slug("mock-model"), paper_model]}),
|
||||
)?,
|
||||
)?;
|
||||
config.with_root_config(&format!(
|
||||
"model_catalog_json = {}",
|
||||
serde_json::to_string(&catalog_path)?
|
||||
))
|
||||
} else {
|
||||
config
|
||||
};
|
||||
let config = match provider {
|
||||
WebSearchProvider::ChatGpt => config
|
||||
.with_model_provider("openai-custom")
|
||||
@@ -164,6 +210,7 @@ async fn assert_standalone_web_search_round_trips_output(
|
||||
.send_turn_start_request(TurnStartParams {
|
||||
thread_id: thread_id.clone(),
|
||||
client_user_message_id: None,
|
||||
model: literature_search.then(|| model.to_string()),
|
||||
input: vec![V2UserInput::Text {
|
||||
text: "Search the web".to_string(),
|
||||
text_elements: Vec::new(),
|
||||
@@ -199,6 +246,12 @@ async fn assert_standalone_web_search_round_trips_output(
|
||||
web_run.pointer("/parameters/properties/time/description"),
|
||||
Some(&json!("Get time for the given UTC offsets."))
|
||||
);
|
||||
assert_eq!(
|
||||
web_run
|
||||
.pointer("/parameters/properties/literature_search_query")
|
||||
.is_some(),
|
||||
literature_search
|
||||
);
|
||||
assert!(
|
||||
!has_hosted_web_search(&first_response),
|
||||
"standalone web search should replace hosted web search"
|
||||
@@ -242,13 +295,8 @@ async fn assert_standalone_web_search_round_trips_output(
|
||||
search_body.get("result_fields").is_none(),
|
||||
"standalone search should use the endpoint's default result projection"
|
||||
);
|
||||
assert_eq!(search_body["model"], json!("mock-model"));
|
||||
assert_eq!(
|
||||
search_body["commands"],
|
||||
json!({
|
||||
"search_query": [{"q": "standalone web search"}],
|
||||
})
|
||||
);
|
||||
assert_eq!(search_body["model"], json!(model));
|
||||
assert_eq!(search_body["commands"], commands);
|
||||
assert_eq!(
|
||||
search_body["settings"]["allowed_callers"],
|
||||
json!(["direct"])
|
||||
@@ -315,9 +363,9 @@ async fn assert_standalone_web_search_round_trips_output(
|
||||
);
|
||||
let expected_completed_item = ThreadItemPayload::WebSearch(WebSearchItem {
|
||||
id: call_id.to_string(),
|
||||
query: "standalone web search".to_string(),
|
||||
query: query.to_string(),
|
||||
action: Some(WebSearchAction::Search {
|
||||
query: Some("standalone web search".to_string()),
|
||||
query: Some(query.to_string()),
|
||||
queries: None,
|
||||
}),
|
||||
results: Some(vec![json!({
|
||||
|
||||
@@ -33,6 +33,9 @@ pub struct SearchCommands {
|
||||
/// Query the internet search engine for a given list of queries.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub search_query: Option<Vec<SearchQuery>>,
|
||||
/// Search for scientific papers. Open relevant papers and cite the primary sources you use.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub literature_search_query: Option<Vec<SearchQuery>>,
|
||||
/// Query the image search engine for a given list of queries.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub image_query: Option<Vec<SearchQuery>>,
|
||||
|
||||
@@ -4118,6 +4118,7 @@ impl Session {
|
||||
}
|
||||
let extension_data =
|
||||
codex_extension_api::ExtensionData::new(turn_context.sub_id.clone());
|
||||
extension_data.insert_arc(Arc::clone(&settings.model_info));
|
||||
if let Some(messages) = settings
|
||||
.model_info
|
||||
.model_messages
|
||||
|
||||
@@ -181,6 +181,9 @@ mod credential_masking;
|
||||
#[path = "scenarios_provider_capabilities.rs"]
|
||||
mod provider_capabilities;
|
||||
|
||||
#[path = "scenarios_web_search.rs"]
|
||||
mod web_search;
|
||||
|
||||
#[path = "scenarios_shared_instructions.rs"]
|
||||
mod shared_instructions;
|
||||
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
//! Model-visible context for standalone web search.
|
||||
|
||||
use super::*;
|
||||
use codex_protocol::config_types::WebSearchMode;
|
||||
use pretty_assertions::assert_eq;
|
||||
use wiremock::Mock;
|
||||
use wiremock::ResponseTemplate;
|
||||
use wiremock::matchers::body_partial_json;
|
||||
use wiremock::matchers::method;
|
||||
use wiremock::matchers::path;
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn literature_search_and_open() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
let server = start_mock_server().await;
|
||||
let search =
|
||||
json!({"literature_search_query": [{"q": "FlashAttention-2", "domains": ["arxiv.org"]}]});
|
||||
let open = json!({"open": [{"ref_id": "turn0academia0"}]});
|
||||
for (commands, output) in [
|
||||
(&search, "[turn0academia0] FlashAttention-2 by Tri Dao"),
|
||||
(
|
||||
&open,
|
||||
"FlashAttention-2 reports a 2x speedup over FlashAttention.",
|
||||
),
|
||||
] {
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/alpha/search"))
|
||||
.and(body_partial_json(json!({"commands": commands})))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"output": output})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
}
|
||||
let responses = mount_sse_sequence(
|
||||
&server,
|
||||
vec![
|
||||
sse(vec![
|
||||
ev_function_call_with_namespace("search-paper", "web", "run", &search.to_string()),
|
||||
ev_completed("resp-1"),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_function_call_with_namespace("open-paper", "web", "run", &open.to_string()),
|
||||
ev_completed("resp-2"),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_assistant_message("answer", "Tri Dao reports a 2x speedup."),
|
||||
ev_completed("resp-3"),
|
||||
]),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
|
||||
let auth = codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("dummy"));
|
||||
let mut extensions = ExtensionRegistryBuilder::new();
|
||||
codex_web_search_extension::install(&mut extensions, Arc::clone(&auth));
|
||||
let test = test_codex()
|
||||
.with_auth_manager(auth)
|
||||
.with_extensions(Arc::new(extensions.build()))
|
||||
.with_config(|config| {
|
||||
configure_scenario_catalog(config);
|
||||
config
|
||||
.features
|
||||
.enable(Feature::StandaloneWebSearch)
|
||||
.unwrap();
|
||||
config.web_search_mode.set(WebSearchMode::Live).unwrap();
|
||||
})
|
||||
.with_model_info_override("paper-model", |model| {
|
||||
model.tool_mode = Some(ToolMode::Direct);
|
||||
model
|
||||
.experimental_supported_tools
|
||||
.push("literature_search_query".to_string());
|
||||
})
|
||||
.build_with_auto_env(&server)
|
||||
.await?;
|
||||
test.submit_turn("Find the FlashAttention-2 paper and open it to check the reported speedup.")
|
||||
.await?;
|
||||
|
||||
let requests = responses.requests();
|
||||
assert_eq!(requests.len(), 3);
|
||||
let web_run = requests[0]
|
||||
.tool_by_name("web", "run")
|
||||
.expect("advertised web.run");
|
||||
let mut parameters = web_run
|
||||
.get("parameters")
|
||||
.expect("web.run parameters")
|
||||
.clone();
|
||||
parameters.sort_all_objects();
|
||||
insta::assert_snapshot!(
|
||||
"literature_search_web_run_schema",
|
||||
serde_json::to_string_pretty(¶meters)?
|
||||
);
|
||||
insta::assert_snapshot!(context_snapshot::format_request_history_snapshot(
|
||||
"A model opted into literature search finds a paper, opens the returned reference, and reads the result before answering.",
|
||||
&requests,
|
||||
&ContextSnapshotOptions::default()
|
||||
.rewrite_known_segments()
|
||||
.include_request_settings(),
|
||||
));
|
||||
Ok(())
|
||||
}
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
---
|
||||
source: core/tests/suite/scenarios_web_search.rs
|
||||
expression: "context_snapshot::format_request_history_snapshot(\"A model opted into literature search finds a paper, opens the returned reference, and reads the result before answering.\",\n&requests,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)"
|
||||
---
|
||||
Scenario: A model opted into literature search finds a paper, opens the returned reference, and reads the result before answering.
|
||||
|
||||
## Window 1
|
||||
Settings:
|
||||
include: ["reasoning.encrypted_content"]
|
||||
model: "paper-model"
|
||||
parallel_tool_calls: true
|
||||
prompt_cache_key: "<PROMPT_CACHE_KEY 1>"
|
||||
reasoning: {"effort":"medium"}
|
||||
store: false
|
||||
stream: true
|
||||
text: {"verbosity":"low"}
|
||||
tool_choice: "auto"
|
||||
tools (7; hash=b542554f4fc96824):
|
||||
- function/exec_command: Runs a command in a PTY, returning output or a session ID for ongoing interaction.; args=[cmd, justification, login, max_output_tokens, prefix_rule, sandbox_permissions, shell, tty, workdir, yield_time_ms]; hash=c04108e51a8548c5
|
||||
- function/write_stdin: Writes characters to an existing unified exec session and returns recent output.; args=[chars, max_output_tokens, session_id, yield_time_ms]; hash=7fb62d8a3e2df4e2
|
||||
- function/request_user_input: Request user input for one to three short questions and wait for the response. This tool is only ...; args=[questions]; hash=64a729b25c424fc6
|
||||
- custom/apply_patch: The `apply_patch` tool can be used to edit files. This is a FREEFORM tool, so do not wrap the pat...; hash=dea9f12b0111a4ba
|
||||
- function/view_image: View a local image file from the filesystem when visual inspection is needed. Use this for images...; args=[detail, path]; hash=3b766d5bdec5d50e
|
||||
- namespace/web: Tools in the web namespace.; hash=c60e4db83b9fd6b0
|
||||
- function/run
|
||||
- tool_search: # Tool discovery Searches over deferred tool metadata with BM25 and exposes matching tools for t...; args=[limit, query]; hash=400cd6ac61d524a9
|
||||
-- request 1 (turn) --
|
||||
00:message/developer:
|
||||
<MODEL_INSTRUCTIONS>
|
||||
01:message/developer:
|
||||
<PERMISSIONS_INSTRUCTIONS>
|
||||
02:message/user:
|
||||
<ENVIRONMENT_CONTEXT>
|
||||
03:message/user:
|
||||
Find the FlashAttention-2 paper and open it to check the reported speedup.
|
||||
-- request 2 (turn) --
|
||||
04:function_call/web.run:{"literature_search_query":[{"domains":["arxiv.org"],"q":"FlashAttention-2"}]}
|
||||
05:function_call_output:[turn0academia0] FlashAttention-2 by Tri Dao
|
||||
-- request 3 (turn) --
|
||||
06:function_call/web.run:{"open":[{"ref_id":"turn0academia0"}]}
|
||||
07:function_call_output:FlashAttention-2 reports a 2x speedup over FlashAttention.
|
||||
+320
@@ -0,0 +1,320 @@
|
||||
---
|
||||
source: core/tests/suite/scenarios_web_search.rs
|
||||
expression: "serde_json::to_string_pretty(¶meters)?"
|
||||
---
|
||||
{
|
||||
"properties": {
|
||||
"click": {
|
||||
"description": "Open links from previously opened pages.",
|
||||
"items": {
|
||||
"properties": {
|
||||
"id": {
|
||||
"description": "Numbered link id to open.",
|
||||
"type": "integer"
|
||||
},
|
||||
"ref_id": {
|
||||
"description": "Reference id containing the numbered link.",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"id",
|
||||
"ref_id"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"finance": {
|
||||
"description": "Look up prices for the given stock symbols.",
|
||||
"items": {
|
||||
"properties": {
|
||||
"market": {
|
||||
"description": "ISO 3166-1 alpha-3 country code, \"OTC\", or \"\" for cryptocurrency.",
|
||||
"type": "string"
|
||||
},
|
||||
"ticker": {
|
||||
"description": "Ticker symbol to look up.",
|
||||
"type": "string"
|
||||
},
|
||||
"type": {
|
||||
"description": "Asset type to look up.",
|
||||
"enum": [
|
||||
"equity",
|
||||
"fund",
|
||||
"crypto",
|
||||
"index"
|
||||
],
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"ticker",
|
||||
"type"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"find": {
|
||||
"description": "Find text patterns in pages.",
|
||||
"items": {
|
||||
"properties": {
|
||||
"pattern": {
|
||||
"description": "Text pattern to find.",
|
||||
"type": "string"
|
||||
},
|
||||
"ref_id": {
|
||||
"description": "Reference id or URL to search within.",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"pattern",
|
||||
"ref_id"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"image_query": {
|
||||
"description": "Query the image search engine for a given list of queries.",
|
||||
"items": {
|
||||
"properties": {
|
||||
"domains": {
|
||||
"description": "Whether to filter by a specific list of domains.",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"q": {
|
||||
"description": "Search query.",
|
||||
"type": "string"
|
||||
},
|
||||
"recency": {
|
||||
"description": "Whether to filter by recency, as a number of recent days.",
|
||||
"type": "integer"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"q"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"literature_search_query": {
|
||||
"description": "Search for scientific papers. Open relevant papers and cite the primary sources you use.",
|
||||
"items": {
|
||||
"properties": {
|
||||
"domains": {
|
||||
"description": "Whether to filter by a specific list of domains.",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"q": {
|
||||
"description": "Search query.",
|
||||
"type": "string"
|
||||
},
|
||||
"recency": {
|
||||
"description": "Whether to filter by recency, as a number of recent days.",
|
||||
"type": "integer"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"q"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"open": {
|
||||
"description": "Open pages by reference id or URL.",
|
||||
"items": {
|
||||
"properties": {
|
||||
"lineno": {
|
||||
"description": "Line number to position the page at.",
|
||||
"type": "integer"
|
||||
},
|
||||
"ref_id": {
|
||||
"description": "Reference id or URL to open.",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"ref_id"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"response_length": {
|
||||
"description": "Set the length of the response to be returned.",
|
||||
"enum": [
|
||||
"short",
|
||||
"medium",
|
||||
"long"
|
||||
],
|
||||
"type": "string"
|
||||
},
|
||||
"screenshot": {
|
||||
"description": "Take screenshots of PDF pages.",
|
||||
"items": {
|
||||
"properties": {
|
||||
"pageno": {
|
||||
"description": "Zero-indexed PDF page number.",
|
||||
"type": "integer"
|
||||
},
|
||||
"ref_id": {
|
||||
"description": "Reference id or URL to screenshot.",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"pageno",
|
||||
"ref_id"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"search_query": {
|
||||
"description": "Query the internet search engine for a given list of queries.",
|
||||
"items": {
|
||||
"properties": {
|
||||
"domains": {
|
||||
"description": "Whether to filter by a specific list of domains.",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"q": {
|
||||
"description": "Search query.",
|
||||
"type": "string"
|
||||
},
|
||||
"recency": {
|
||||
"description": "Whether to filter by recency, as a number of recent days.",
|
||||
"type": "integer"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"q"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"sports": {
|
||||
"description": "Look up sports schedules and standings.",
|
||||
"items": {
|
||||
"properties": {
|
||||
"date_from": {
|
||||
"description": "Start date in YYYY-MM-DD format.",
|
||||
"type": "string"
|
||||
},
|
||||
"date_to": {
|
||||
"description": "End date in YYYY-MM-DD format.",
|
||||
"type": "string"
|
||||
},
|
||||
"fn": {
|
||||
"description": "Sports function to call.",
|
||||
"enum": [
|
||||
"schedule",
|
||||
"standings"
|
||||
],
|
||||
"type": "string"
|
||||
},
|
||||
"league": {
|
||||
"description": "League to look up.",
|
||||
"enum": [
|
||||
"nba",
|
||||
"wnba",
|
||||
"nfl",
|
||||
"nhl",
|
||||
"mlb",
|
||||
"epl",
|
||||
"ncaamb",
|
||||
"ncaawb",
|
||||
"ipl"
|
||||
],
|
||||
"type": "string"
|
||||
},
|
||||
"locale": {
|
||||
"description": "Locale for the lookup.",
|
||||
"type": "string"
|
||||
},
|
||||
"num_games": {
|
||||
"description": "Number of games to return.",
|
||||
"type": "integer"
|
||||
},
|
||||
"opponent": {
|
||||
"description": "Opponent to use with `team` when narrowing the lookup.",
|
||||
"type": "string"
|
||||
},
|
||||
"team": {
|
||||
"description": "Team to look up, using the common 3 or 4 letter alias used in broadcasts.",
|
||||
"type": "string"
|
||||
},
|
||||
"tool": {
|
||||
"description": "Tool name for sports requests.",
|
||||
"enum": [
|
||||
"sports"
|
||||
],
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"fn",
|
||||
"league"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"time": {
|
||||
"description": "Get time for the given UTC offsets.",
|
||||
"items": {
|
||||
"properties": {
|
||||
"utc_offset": {
|
||||
"description": "UTC offset formatted like \"+03:00\".",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"utc_offset"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
"weather": {
|
||||
"description": "Look up weather forecasts.",
|
||||
"items": {
|
||||
"properties": {
|
||||
"duration": {
|
||||
"description": "Number of days to return. Defaults to 7.",
|
||||
"type": "integer"
|
||||
},
|
||||
"location": {
|
||||
"description": "Location in \"Country, Area, City\" format.",
|
||||
"type": "string"
|
||||
},
|
||||
"start": {
|
||||
"description": "Start date in YYYY-MM-DD format. Defaults to today.",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"location"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
"type": "array"
|
||||
}
|
||||
},
|
||||
"type": "object"
|
||||
}
|
||||
@@ -96,11 +96,19 @@ impl ExtensionData {
|
||||
|
||||
/// Stores `value` as the attachment of type `T`, returning any previous value.
|
||||
pub fn insert<T>(&self, value: T) -> Option<Arc<T>>
|
||||
where
|
||||
T: Any + Send + Sync,
|
||||
{
|
||||
self.insert_arc(Arc::new(value))
|
||||
}
|
||||
|
||||
/// Stores a shared attachment without cloning its contents.
|
||||
pub fn insert_arc<T>(&self, value: Arc<T>) -> Option<Arc<T>>
|
||||
where
|
||||
T: Any + Send + Sync,
|
||||
{
|
||||
self.entries()
|
||||
.insert(TypeId::of::<T>(), Arc::new(value))
|
||||
.insert(TypeId::of::<T>(), value)
|
||||
.map(downcast_data)
|
||||
}
|
||||
|
||||
|
||||
@@ -23,7 +23,9 @@ use codex_model_provider::create_model_provider;
|
||||
use codex_model_provider_info::ModelProviderInfo;
|
||||
use codex_protocol::config_types::WebSearchContextSize;
|
||||
use codex_protocol::config_types::WebSearchMode;
|
||||
use codex_protocol::openai_models::ModelInfo;
|
||||
|
||||
use crate::schema::LITERATURE_SEARCH_QUERY;
|
||||
use crate::tool::WebSearchTool;
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -120,10 +122,11 @@ impl ConfigContributor<Config> for WebSearchExtension {
|
||||
}
|
||||
|
||||
impl ToolContributor for WebSearchExtension {
|
||||
fn tools(
|
||||
fn tools_for_step(
|
||||
&self,
|
||||
session_store: &ExtensionData,
|
||||
thread_store: &ExtensionData,
|
||||
step_store: &ExtensionData,
|
||||
) -> Vec<
|
||||
Arc<dyn for<'call> codex_extension_api::ToolExecutor<codex_extension_api::ToolCall<'call>>>,
|
||||
> {
|
||||
@@ -142,6 +145,12 @@ impl ToolContributor for WebSearchExtension {
|
||||
Some(self.auth_manager.clone()),
|
||||
),
|
||||
settings: config.settings.clone(),
|
||||
supports_literature_search: step_store.get::<ModelInfo>().is_some_and(|model| {
|
||||
model
|
||||
.experimental_supported_tools
|
||||
.iter()
|
||||
.any(|tool| tool == LITERATURE_SEARCH_QUERY)
|
||||
}),
|
||||
originator: thread_store
|
||||
.get::<ThreadOriginator>()
|
||||
.map(|originator| originator.0.clone()),
|
||||
@@ -207,6 +216,7 @@ mod tests {
|
||||
let registry = builder.build();
|
||||
let session_store = ExtensionData::new("session");
|
||||
let thread_store = ExtensionData::new("11111111-1111-4111-8111-111111111111");
|
||||
let step_store = ExtensionData::new("step");
|
||||
thread_store.insert(WebSearchExtensionConfig {
|
||||
available: true,
|
||||
http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
|
||||
@@ -217,7 +227,9 @@ mod tests {
|
||||
let tool_names = registry
|
||||
.tool_contributors()
|
||||
.iter()
|
||||
.flat_map(|contributor| contributor.tools(&session_store, &thread_store))
|
||||
.flat_map(|contributor| {
|
||||
contributor.tools_for_step(&session_store, &thread_store, &step_store)
|
||||
})
|
||||
.map(|tool| (tool.tool_name(), tool.supports_parallel_tool_calls()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
|
||||
@@ -3,7 +3,9 @@ use schemars::r#gen::SchemaSettings;
|
||||
use serde_json::Map;
|
||||
use serde_json::Value;
|
||||
|
||||
pub(crate) fn commands_schema() -> Value {
|
||||
pub(crate) const LITERATURE_SEARCH_QUERY: &str = "literature_search_query";
|
||||
|
||||
pub(crate) fn commands_schema(supports_literature_search: bool) -> Value {
|
||||
let schema = SchemaSettings::draft2019_09()
|
||||
.with(|settings| {
|
||||
settings.inline_subschemas = true;
|
||||
@@ -18,6 +20,11 @@ pub(crate) fn commands_schema() -> Value {
|
||||
let Value::Object(mut schema) = schema else {
|
||||
unreachable!("search commands schema must be an object");
|
||||
};
|
||||
if !supports_literature_search
|
||||
&& let Some(Value::Object(properties)) = schema.get_mut("properties")
|
||||
{
|
||||
properties.remove(LITERATURE_SEARCH_QUERY);
|
||||
}
|
||||
|
||||
let mut tool_schema = Map::new();
|
||||
for key in [
|
||||
|
||||
@@ -48,6 +48,7 @@ pub(crate) struct WebSearchTool {
|
||||
pub(crate) http_client_factory: HttpClientFactory,
|
||||
pub(crate) provider: SharedModelProvider,
|
||||
pub(crate) settings: SearchSettings,
|
||||
pub(crate) supports_literature_search: bool,
|
||||
pub(crate) originator: Option<String>,
|
||||
}
|
||||
|
||||
@@ -58,7 +59,8 @@ impl<'call> ToolExecutor<ToolCall<'call>> for WebSearchTool {
|
||||
|
||||
fn spec(&self) -> ToolSpec {
|
||||
// parse schema without compaction that removes field metadata/descriptions to match hosted tool definition
|
||||
let parameters = match parse_tool_input_schema_without_compaction(&commands_schema()) {
|
||||
let schema = commands_schema(self.supports_literature_search);
|
||||
let parameters = match parse_tool_input_schema_without_compaction(&schema) {
|
||||
Ok(parameters) => parameters,
|
||||
Err(err) => panic!("search command schema should parse: {err}"),
|
||||
};
|
||||
@@ -98,7 +100,7 @@ impl WebSearchTool {
|
||||
&self,
|
||||
call: ToolCall<'_>,
|
||||
) -> Result<Box<dyn ToolOutput>, FunctionCallError> {
|
||||
let commands = parse_commands(&call)?;
|
||||
let commands = parse_commands(call.function_arguments()?, self.supports_literature_search)?;
|
||||
let command_action = command_action(&commands);
|
||||
let (provider, auth) = self
|
||||
.provider
|
||||
@@ -201,14 +203,22 @@ fn search_request_headers(originator: Option<&str>, turn_metadata: Option<&str>)
|
||||
headers
|
||||
}
|
||||
|
||||
fn parse_commands(call: &ToolCall<'_>) -> Result<SearchCommands, FunctionCallError> {
|
||||
let arguments = call.function_arguments()?;
|
||||
fn parse_commands(
|
||||
arguments: &str,
|
||||
supports_literature_search: bool,
|
||||
) -> Result<SearchCommands, FunctionCallError> {
|
||||
if arguments.trim().is_empty() {
|
||||
return Ok(SearchCommands::default());
|
||||
}
|
||||
|
||||
serde_json::from_str(arguments)
|
||||
.map_err(|err| FunctionCallError::RespondToModel(err.to_string()))
|
||||
let commands: SearchCommands = serde_json::from_str(arguments)
|
||||
.map_err(|err| FunctionCallError::RespondToModel(err.to_string()))?;
|
||||
if !supports_literature_search && commands.literature_search_query.is_some() {
|
||||
return Err(FunctionCallError::RespondToModel(
|
||||
"literature_search_query is unavailable for this model".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(commands)
|
||||
}
|
||||
|
||||
fn command_action(commands: &SearchCommands) -> WebSearchAction {
|
||||
@@ -216,6 +226,12 @@ fn command_action(commands: &SearchCommands) -> WebSearchAction {
|
||||
.search_query
|
||||
.as_deref()
|
||||
.and_then(query_action)
|
||||
.or_else(|| {
|
||||
commands
|
||||
.literature_search_query
|
||||
.as_deref()
|
||||
.and_then(query_action)
|
||||
})
|
||||
.or_else(|| commands.image_query.as_deref().and_then(query_action))
|
||||
.or_else(|| {
|
||||
commands
|
||||
@@ -272,8 +288,10 @@ mod tests {
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::command_action;
|
||||
use super::parse_commands;
|
||||
use super::search_request_headers;
|
||||
use codex_core::X_CODEX_TURN_METADATA_HEADER;
|
||||
use codex_extension_api::FunctionCallError;
|
||||
|
||||
#[test]
|
||||
fn search_request_headers_forward_thread_originator_and_turn_metadata() {
|
||||
@@ -292,6 +310,17 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_literature_search_without_the_model_capability() {
|
||||
let arguments =
|
||||
r#"{"search_query":[{"q":"regular"}],"literature_search_query":[{"q":"papers"}]}"#;
|
||||
assert!(matches!(
|
||||
parse_commands(arguments, /*supports_literature_search*/ false),
|
||||
Err(FunctionCallError::RespondToModel(message))
|
||||
if message == "literature_search_query is unavailable for this model"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn command_action_reports_queries_and_navigation_detail() {
|
||||
let cases = [
|
||||
|
||||
Reference in New Issue
Block a user