mirror of
https://github.com/openai/codex.git
synced 2026-11-01 19:33:35 +00:00
Use request originator for model catalog requests
This commit is contained in:
@@ -1141,7 +1141,9 @@ impl MessageProcessor {
|
||||
.await
|
||||
}
|
||||
ClientRequest::ModelList { params, .. } => {
|
||||
self.catalog_processor.model_list(params).await
|
||||
self.catalog_processor
|
||||
.model_list(params, &request_context)
|
||||
.await
|
||||
}
|
||||
ClientRequest::ExperimentalFeatureList { params, .. } => {
|
||||
self.catalog_processor
|
||||
|
||||
@@ -5,6 +5,7 @@ use codex_app_server_protocol::ModelServiceTier;
|
||||
use codex_app_server_protocol::ModelUpgradeInfo;
|
||||
use codex_app_server_protocol::ReasoningEffortOption;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_login::default_client::Originator;
|
||||
use codex_models_manager::manager::RefreshStrategy;
|
||||
use codex_protocol::openai_models::ModelPreset;
|
||||
use codex_protocol::openai_models::ReasoningEffortPreset;
|
||||
@@ -12,9 +13,10 @@ use codex_protocol::openai_models::ReasoningEffortPreset;
|
||||
pub async fn supported_models(
|
||||
thread_manager: Arc<ThreadManager>,
|
||||
include_hidden: bool,
|
||||
originator: Option<Originator>,
|
||||
) -> Vec<Model> {
|
||||
thread_manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached)
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, originator)
|
||||
.await
|
||||
.into_iter()
|
||||
.filter(|preset| include_hidden || preset.show_in_picker)
|
||||
|
||||
@@ -140,10 +140,15 @@ impl CatalogRequestProcessor {
|
||||
pub(crate) async fn model_list(
|
||||
&self,
|
||||
params: ModelListParams,
|
||||
request_context: &RequestContext,
|
||||
) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
|
||||
Self::list_models(self.thread_manager.clone(), params)
|
||||
.await
|
||||
.map(|response| Some(response.into()))
|
||||
Self::list_models(
|
||||
self.thread_manager.clone(),
|
||||
params,
|
||||
Some(request_context.originator().clone()),
|
||||
)
|
||||
.await
|
||||
.map(|response| Some(response.into()))
|
||||
}
|
||||
|
||||
pub(crate) async fn experimental_feature_list(
|
||||
@@ -223,13 +228,15 @@ impl CatalogRequestProcessor {
|
||||
async fn list_models(
|
||||
thread_manager: Arc<ThreadManager>,
|
||||
params: ModelListParams,
|
||||
originator: Option<Originator>,
|
||||
) -> Result<ModelListResponse, JSONRPCErrorError> {
|
||||
let ModelListParams {
|
||||
limit,
|
||||
cursor,
|
||||
include_hidden,
|
||||
} = params;
|
||||
let models = supported_models(thread_manager, include_hidden.unwrap_or(false)).await;
|
||||
let models =
|
||||
supported_models(thread_manager, include_hidden.unwrap_or(false), originator).await;
|
||||
let total = models.len();
|
||||
|
||||
if total == 0 {
|
||||
|
||||
@@ -1750,7 +1750,7 @@ async fn run_debug_models_command(
|
||||
AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true).await;
|
||||
let models_manager = build_models_manager(&config, auth_manager);
|
||||
models_manager
|
||||
.raw_model_catalog(RefreshStrategy::OnlineIfUncached)
|
||||
.raw_model_catalog(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await
|
||||
};
|
||||
|
||||
|
||||
@@ -650,7 +650,10 @@ pub(super) async fn run_guardian_review_session(
|
||||
let available_models = session
|
||||
.services
|
||||
.models_manager
|
||||
.list_models(codex_models_manager::manager::RefreshStrategy::Offline)
|
||||
.list_models(
|
||||
codex_models_manager::manager::RefreshStrategy::Offline,
|
||||
/*originator*/ None,
|
||||
)
|
||||
.await;
|
||||
let preferred_reasoning_effort = |supports_low: bool, fallback| {
|
||||
if supports_low {
|
||||
|
||||
@@ -537,7 +537,9 @@ impl Codex {
|
||||
codex_models_manager::manager::RefreshStrategy::Offline
|
||||
)
|
||||
{
|
||||
let _ = models_manager.list_models(refresh_strategy).await;
|
||||
let _ = models_manager
|
||||
.list_models(refresh_strategy, Some(originator.clone()))
|
||||
.await;
|
||||
}
|
||||
let model = models_manager
|
||||
.get_default_model(&config.model, refresh_strategy)
|
||||
|
||||
@@ -28,7 +28,10 @@ pub(super) async fn spawn_review_thread(
|
||||
let available_models = sess
|
||||
.services
|
||||
.models_manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached)
|
||||
.list_models(
|
||||
RefreshStrategy::OnlineIfUncached,
|
||||
Some(parent_turn_context.originator.clone()),
|
||||
)
|
||||
.await;
|
||||
let shell_command_backend = shell_command_backend_for_features(review_features.get());
|
||||
let unified_exec_shell_mode = UnifiedExecShellMode::for_session(
|
||||
|
||||
@@ -204,7 +204,10 @@ impl TurnContext {
|
||||
);
|
||||
let features = self.features.clone();
|
||||
let available_models = models_manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached)
|
||||
.list_models(
|
||||
RefreshStrategy::OnlineIfUncached,
|
||||
Some(self.originator.clone()),
|
||||
)
|
||||
.await;
|
||||
|
||||
Self {
|
||||
|
||||
@@ -437,10 +437,14 @@ impl ThreadManager {
|
||||
self.state.models_manager.clone()
|
||||
}
|
||||
|
||||
pub async fn list_models(&self, refresh_strategy: RefreshStrategy) -> Vec<ModelPreset> {
|
||||
pub async fn list_models(
|
||||
&self,
|
||||
refresh_strategy: RefreshStrategy,
|
||||
originator: Option<Originator>,
|
||||
) -> Vec<ModelPreset> {
|
||||
self.state
|
||||
.models_manager
|
||||
.list_models(refresh_strategy)
|
||||
.list_models(refresh_strategy, originator)
|
||||
.await
|
||||
}
|
||||
|
||||
|
||||
@@ -980,7 +980,9 @@ async fn new_uses_active_provider_for_model_refresh() {
|
||||
/*attestation_provider*/ None,
|
||||
);
|
||||
|
||||
let _ = manager.list_models(RefreshStrategy::Online).await;
|
||||
let _ = manager
|
||||
.list_models(RefreshStrategy::Online, /*originator*/ None)
|
||||
.await;
|
||||
assert_eq!(models_mock.requests().len(), 1);
|
||||
}
|
||||
|
||||
|
||||
@@ -302,7 +302,7 @@ pub(crate) async fn apply_requested_spawn_agent_model_overrides(
|
||||
let available_models = session
|
||||
.services
|
||||
.models_manager
|
||||
.list_models(RefreshStrategy::Offline)
|
||||
.list_models(RefreshStrategy::Offline, /*originator*/ None)
|
||||
.await;
|
||||
let selected_model_name = find_spawn_agent_model_name(&available_models, requested_model)?;
|
||||
let selected_model_info = session
|
||||
|
||||
@@ -418,7 +418,7 @@ async fn model_change_from_image_to_text_strips_prior_image_content() -> Result<
|
||||
let test = builder.build(&server).await?;
|
||||
let models_manager = test.thread_manager.get_models_manager();
|
||||
let _ = models_manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached)
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
let image_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR4nGNgYAAAAAMAASsJTYQAAAAASUVORK5CYII="
|
||||
.to_string();
|
||||
@@ -537,7 +537,7 @@ async fn generated_image_is_replayed_for_image_capable_models() -> Result<()> {
|
||||
let _ = std::fs::remove_file(&saved_path);
|
||||
let models_manager = test.thread_manager.get_models_manager();
|
||||
let _ = models_manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached)
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
|
||||
test.codex
|
||||
@@ -651,7 +651,7 @@ async fn model_change_from_generated_image_to_text_preserves_prior_generated_ima
|
||||
let _ = std::fs::remove_file(&saved_path);
|
||||
let models_manager = test.thread_manager.get_models_manager();
|
||||
let _ = models_manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached)
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
|
||||
test.codex
|
||||
@@ -767,7 +767,7 @@ async fn thread_rollback_after_generated_image_drops_entire_image_turn_history()
|
||||
let _ = std::fs::remove_file(&saved_path);
|
||||
let models_manager = test.thread_manager.get_models_manager();
|
||||
let _ = models_manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached)
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
|
||||
test.codex
|
||||
@@ -920,7 +920,9 @@ async fn model_switch_to_smaller_model_updates_token_context_window() -> Result<
|
||||
let test = builder.build(&server).await?;
|
||||
|
||||
let models_manager = test.thread_manager.get_models_manager();
|
||||
let available_models = models_manager.list_models(RefreshStrategy::Online).await;
|
||||
let available_models = models_manager
|
||||
.list_models(RefreshStrategy::Online, /*originator*/ None)
|
||||
.await;
|
||||
assert!(
|
||||
available_models
|
||||
.iter()
|
||||
|
||||
@@ -70,7 +70,7 @@ async fn renews_cache_ttl_on_matching_models_etag() -> Result<()> {
|
||||
// Populate cache via initial refresh.
|
||||
let models_manager = test.thread_manager.get_models_manager();
|
||||
let _ = models_manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached)
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
|
||||
let cache_path = config.codex_home.join(CACHE_FILE);
|
||||
@@ -134,7 +134,7 @@ async fn renews_cache_ttl_on_matching_models_etag() -> Result<()> {
|
||||
// Cached models remain usable offline.
|
||||
let offline_models = test
|
||||
.thread_manager
|
||||
.list_models(RefreshStrategy::Offline)
|
||||
.list_models(RefreshStrategy::Offline, /*originator*/ None)
|
||||
.await;
|
||||
assert!(
|
||||
offline_models
|
||||
@@ -177,7 +177,7 @@ async fn uses_cache_when_version_matches() -> Result<()> {
|
||||
let test = builder.build(&server).await?;
|
||||
let models_manager = test.thread_manager.get_models_manager();
|
||||
let models = models_manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached)
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
@@ -224,7 +224,7 @@ async fn refreshes_when_cache_version_missing() -> Result<()> {
|
||||
let test = builder.build(&server).await?;
|
||||
let models_manager = test.thread_manager.get_models_manager();
|
||||
let models = models_manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached)
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
@@ -272,7 +272,7 @@ async fn refreshes_when_cache_version_differs() -> Result<()> {
|
||||
let test = builder.build(&server).await?;
|
||||
let models_manager = test.thread_manager.get_models_manager();
|
||||
let models = models_manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached)
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
|
||||
@@ -785,7 +785,9 @@ async fn user_turn_personality_remote_model_template_includes_update_message() -
|
||||
async fn wait_for_model_available(manager: &SharedModelsManager, slug: &str) {
|
||||
let deadline = Instant::now() + Duration::from_secs(2);
|
||||
loop {
|
||||
let models = manager.list_models(RefreshStrategy::OnlineIfUncached).await;
|
||||
let models = manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
if models.iter().any(|model| model.model == slug) {
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -102,7 +102,9 @@ async fn remote_models_get_model_info_uses_longest_matching_prefix() -> Result<(
|
||||
provider,
|
||||
);
|
||||
|
||||
manager.list_models(RefreshStrategy::OnlineIfUncached).await;
|
||||
manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
|
||||
let model_info = manager
|
||||
.get_model_info("gpt-5.3-codex-test", &config.to_models_manager_config())
|
||||
@@ -848,7 +850,9 @@ async fn remote_models_do_not_append_removed_builtin_presets() -> Result<()> {
|
||||
provider,
|
||||
);
|
||||
|
||||
let available = manager.list_models(RefreshStrategy::OnlineIfUncached).await;
|
||||
let available = manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
let remote = available
|
||||
.iter()
|
||||
.find(|model| model.model == "remote-alpha")
|
||||
@@ -909,7 +913,9 @@ async fn remote_models_merge_adds_new_high_priority_first() -> Result<()> {
|
||||
provider,
|
||||
);
|
||||
|
||||
let available = manager.list_models(RefreshStrategy::OnlineIfUncached).await;
|
||||
let available = manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
assert_eq!(
|
||||
available.first().map(|model| model.model.as_str()),
|
||||
Some("remote-top")
|
||||
@@ -956,7 +962,9 @@ async fn remote_models_merge_replaces_overlapping_model() -> Result<()> {
|
||||
provider,
|
||||
);
|
||||
|
||||
let available = manager.list_models(RefreshStrategy::OnlineIfUncached).await;
|
||||
let available = manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
let overridden = available
|
||||
.iter()
|
||||
.find(|model| model.model == slug)
|
||||
@@ -1000,7 +1008,9 @@ async fn remote_models_merge_preserves_bundled_models_on_empty_response() -> Res
|
||||
provider,
|
||||
);
|
||||
|
||||
let available = manager.list_models(RefreshStrategy::OnlineIfUncached).await;
|
||||
let available = manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
let bundled_slug = bundled_model_slug();
|
||||
assert!(
|
||||
available.iter().any(|model| model.model == bundled_slug),
|
||||
@@ -1117,7 +1127,9 @@ async fn remote_models_hide_picker_only_models() -> Result<()> {
|
||||
.await;
|
||||
assert_eq!(selected, bundled_default_model_slug());
|
||||
|
||||
let available = manager.list_models(RefreshStrategy::OnlineIfUncached).await;
|
||||
let available = manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
let hidden = available
|
||||
.iter()
|
||||
.find(|model| model.model == "codex-auto-balanced")
|
||||
@@ -1138,7 +1150,9 @@ async fn wait_for_model_available(manager: &SharedModelsManager, slug: &str) ->
|
||||
let deadline = Instant::now() + Duration::from_secs(2);
|
||||
loop {
|
||||
if let Some(model) = {
|
||||
let guard = manager.list_models(RefreshStrategy::OnlineIfUncached).await;
|
||||
let guard = manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
guard.iter().find(|model| model.model == slug).cloned()
|
||||
} {
|
||||
return model;
|
||||
|
||||
@@ -1349,7 +1349,7 @@ async fn stdio_image_responses_are_sanitized_for_text_only_model() -> anyhow::Re
|
||||
fixture
|
||||
.thread_manager
|
||||
.get_models_manager()
|
||||
.list_models(RefreshStrategy::Online)
|
||||
.list_models(RefreshStrategy::Online, /*originator*/ None)
|
||||
.await;
|
||||
assert_eq!(models_mock.requests().len(), 1);
|
||||
|
||||
|
||||
@@ -87,7 +87,9 @@ fn test_model_info(
|
||||
async fn wait_for_model_available(manager: &SharedModelsManager, slug: &str) {
|
||||
let deadline = Instant::now() + Duration::from_secs(2);
|
||||
loop {
|
||||
let available_models = manager.list_models(RefreshStrategy::Online).await;
|
||||
let available_models = manager
|
||||
.list_models(RefreshStrategy::Online, /*originator*/ None)
|
||||
.await;
|
||||
if available_models.iter().any(|model| model.model == slug) {
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -82,6 +82,7 @@ impl ModelsEndpointClient for OpenAiModelsEndpoint {
|
||||
async fn list_models(
|
||||
&self,
|
||||
client_version: &str,
|
||||
originator: Option<Originator>,
|
||||
) -> CoreResult<(Vec<ModelInfo>, Option<String>)> {
|
||||
let _timer =
|
||||
codex_otel::start_global_timer("codex.remote_models.fetch_update.duration_ms", &[]);
|
||||
@@ -89,8 +90,9 @@ impl ModelsEndpointClient for OpenAiModelsEndpoint {
|
||||
let auth_mode = auth.as_ref().map(CodexAuth::auth_mode);
|
||||
let api_provider = self.provider_info.to_api_provider(auth_mode)?;
|
||||
let api_auth = resolve_provider_auth(auth.as_ref(), &self.provider_info)?;
|
||||
let originator = Originator::process_default();
|
||||
let transport = ReqwestTransport::new(build_reqwest_client(&originator));
|
||||
let originator = originator.unwrap_or_else(Originator::process_default);
|
||||
let reqwest_client = build_reqwest_client(&originator);
|
||||
let transport = ReqwestTransport::new(reqwest_client);
|
||||
let auth_telemetry = auth_header_telemetry(api_auth.as_ref());
|
||||
let request_telemetry: Arc<dyn RequestTelemetry> = Arc::new(ModelsRequestTelemetry {
|
||||
auth_mode: auth_mode.map(|mode| TelemetryAuthMode::from(mode).to_string()),
|
||||
|
||||
@@ -500,7 +500,9 @@ mod tests {
|
||||
let manager =
|
||||
provider.models_manager(test_codex_home(), /*config_model_catalog*/ None);
|
||||
|
||||
let catalog = manager.raw_model_catalog(RefreshStrategy::Online).await;
|
||||
let catalog = manager
|
||||
.raw_model_catalog(RefreshStrategy::Online, /*originator*/ None)
|
||||
.await;
|
||||
let model_ids = catalog
|
||||
.models
|
||||
.iter()
|
||||
@@ -517,7 +519,7 @@ mod tests {
|
||||
);
|
||||
|
||||
let default_model = manager
|
||||
.list_models(RefreshStrategy::Online)
|
||||
.list_models(RefreshStrategy::Online, /*originator*/ None)
|
||||
.await
|
||||
.into_iter()
|
||||
.find(|preset| preset.is_default)
|
||||
@@ -542,7 +544,9 @@ mod tests {
|
||||
}),
|
||||
);
|
||||
|
||||
let catalog = manager.raw_model_catalog(RefreshStrategy::Online).await;
|
||||
let catalog = manager
|
||||
.raw_model_catalog(RefreshStrategy::Online, /*originator*/ None)
|
||||
.await;
|
||||
|
||||
assert_eq!(catalog.models.len(), 1);
|
||||
assert_eq!(catalog.models[0].slug, "custom-bedrock-model");
|
||||
@@ -578,7 +582,9 @@ mod tests {
|
||||
|
||||
let manager =
|
||||
provider.models_manager(test_codex_home(), /*config_model_catalog*/ None);
|
||||
let catalog = manager.raw_model_catalog(RefreshStrategy::Online).await;
|
||||
let catalog = manager
|
||||
.raw_model_catalog(RefreshStrategy::Online, /*originator*/ None)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
catalog
|
||||
|
||||
@@ -5,6 +5,7 @@ use crate::model_info;
|
||||
use async_trait::async_trait;
|
||||
use codex_app_server_protocol::AuthMode;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::default_client::Originator;
|
||||
use codex_protocol::config_types::CollaborationModeMask;
|
||||
use codex_protocol::error::Result as CoreResult;
|
||||
use codex_protocol::openai_models::ModelInfo;
|
||||
@@ -41,6 +42,7 @@ pub trait ModelsEndpointClient: fmt::Debug + Send + Sync {
|
||||
async fn list_models(
|
||||
&self,
|
||||
client_version: &str,
|
||||
originator: Option<Originator>,
|
||||
) -> CoreResult<(Vec<ModelInfo>, Option<String>)>;
|
||||
}
|
||||
|
||||
@@ -79,9 +81,13 @@ pub trait ModelsManager: fmt::Debug + Send + Sync {
|
||||
/// List all available models, refreshing according to the specified strategy.
|
||||
///
|
||||
/// Returns model presets sorted by priority and filtered by auth mode and visibility.
|
||||
async fn list_models(&self, refresh_strategy: RefreshStrategy) -> Vec<ModelPreset> {
|
||||
async fn list_models(
|
||||
&self,
|
||||
refresh_strategy: RefreshStrategy,
|
||||
originator: Option<Originator>,
|
||||
) -> Vec<ModelPreset> {
|
||||
async move {
|
||||
let catalog = self.raw_model_catalog(refresh_strategy).await;
|
||||
let catalog = self.raw_model_catalog(refresh_strategy, originator).await;
|
||||
self.build_available_models(catalog.models)
|
||||
}
|
||||
.instrument(tracing::info_span!(
|
||||
@@ -92,7 +98,11 @@ pub trait ModelsManager: fmt::Debug + Send + Sync {
|
||||
}
|
||||
|
||||
/// Return the active raw model catalog, refreshing according to the specified strategy.
|
||||
async fn raw_model_catalog(&self, refresh_strategy: RefreshStrategy) -> ModelsResponse;
|
||||
async fn raw_model_catalog(
|
||||
&self,
|
||||
refresh_strategy: RefreshStrategy,
|
||||
originator: Option<Originator>,
|
||||
) -> ModelsResponse;
|
||||
|
||||
/// Return the current in-memory remote model catalog without refreshing or loading cache state.
|
||||
async fn get_remote_models(&self) -> Vec<ModelInfo>;
|
||||
@@ -147,7 +157,10 @@ pub trait ModelsManager: fmt::Debug + Send + Sync {
|
||||
if let Some(model) = model.as_ref() {
|
||||
return model.to_string();
|
||||
}
|
||||
default_model_from_available(self.list_models(refresh_strategy).await)
|
||||
default_model_from_available(
|
||||
self.list_models(refresh_strategy, /*originator*/ None)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
.instrument(tracing::info_span!(
|
||||
"get_default_model",
|
||||
@@ -226,8 +239,15 @@ impl StaticModelsManager {
|
||||
|
||||
#[async_trait]
|
||||
impl ModelsManager for OpenAiModelsManager {
|
||||
async fn raw_model_catalog(&self, refresh_strategy: RefreshStrategy) -> ModelsResponse {
|
||||
if let Err(err) = self.refresh_available_models(refresh_strategy).await {
|
||||
async fn raw_model_catalog(
|
||||
&self,
|
||||
refresh_strategy: RefreshStrategy,
|
||||
originator: Option<Originator>,
|
||||
) -> ModelsResponse {
|
||||
if let Err(err) = self
|
||||
.refresh_available_models_with_originator(refresh_strategy, originator)
|
||||
.await
|
||||
{
|
||||
error!("failed to refresh available models: {err}");
|
||||
}
|
||||
ModelsResponse {
|
||||
@@ -268,6 +288,15 @@ impl ModelsManager for OpenAiModelsManager {
|
||||
impl OpenAiModelsManager {
|
||||
/// Refresh available models according to the specified strategy.
|
||||
async fn refresh_available_models(&self, refresh_strategy: RefreshStrategy) -> CoreResult<()> {
|
||||
self.refresh_available_models_with_originator(refresh_strategy, /*originator*/ None)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn refresh_available_models_with_originator(
|
||||
&self,
|
||||
refresh_strategy: RefreshStrategy,
|
||||
originator: Option<Originator>,
|
||||
) -> CoreResult<()> {
|
||||
if !self.should_refresh_models().await {
|
||||
if matches!(
|
||||
refresh_strategy,
|
||||
@@ -291,18 +320,21 @@ impl OpenAiModelsManager {
|
||||
return Ok(());
|
||||
}
|
||||
info!("models cache: cache miss, fetching remote models");
|
||||
self.fetch_and_update_models().await
|
||||
self.fetch_and_update_models(originator).await
|
||||
}
|
||||
RefreshStrategy::Online => {
|
||||
// Always fetch from network
|
||||
self.fetch_and_update_models().await
|
||||
self.fetch_and_update_models(originator).await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_and_update_models(&self) -> CoreResult<()> {
|
||||
async fn fetch_and_update_models(&self, originator: Option<Originator>) -> CoreResult<()> {
|
||||
let client_version = crate::client_version_to_whole();
|
||||
let (models, etag) = self.endpoint_client.list_models(&client_version).await?;
|
||||
let (models, etag) = self
|
||||
.endpoint_client
|
||||
.list_models(&client_version, originator)
|
||||
.await?;
|
||||
self.apply_remote_models(models.clone()).await;
|
||||
*self.etag.write().await = etag.clone();
|
||||
self.cache_manager
|
||||
@@ -381,7 +413,11 @@ impl OpenAiModelsManager {
|
||||
|
||||
#[async_trait]
|
||||
impl ModelsManager for StaticModelsManager {
|
||||
async fn raw_model_catalog(&self, _refresh_strategy: RefreshStrategy) -> ModelsResponse {
|
||||
async fn raw_model_catalog(
|
||||
&self,
|
||||
_refresh_strategy: RefreshStrategy,
|
||||
_originator: Option<Originator>,
|
||||
) -> ModelsResponse {
|
||||
ModelsResponse {
|
||||
models: self.get_remote_models().await,
|
||||
}
|
||||
|
||||
@@ -157,6 +157,7 @@ impl ModelsEndpointClient for TestModelsEndpoint {
|
||||
async fn list_models(
|
||||
&self,
|
||||
_client_version: &str,
|
||||
_originator: Option<Originator>,
|
||||
) -> CoreResult<(Vec<ModelInfo>, Option<String>)> {
|
||||
self.fetch_count.fetch_add(1, Ordering::SeqCst);
|
||||
let models = self
|
||||
@@ -350,7 +351,9 @@ async fn refresh_available_models_sorts_by_priority() {
|
||||
let cached_remote = manager.get_remote_models().await;
|
||||
assert_models_contain(&cached_remote, &remote_models);
|
||||
|
||||
let available = manager.list_models(RefreshStrategy::OnlineIfUncached).await;
|
||||
let available = manager
|
||||
.list_models(RefreshStrategy::OnlineIfUncached, /*originator*/ None)
|
||||
.await;
|
||||
let high_idx = available
|
||||
.iter()
|
||||
.position(|model| model.model == "priority-high")
|
||||
@@ -733,6 +736,7 @@ impl ModelsEndpointClient for TestAuthAwareModelsEndpoint {
|
||||
async fn list_models(
|
||||
&self,
|
||||
_client_version: &str,
|
||||
_originator: Option<Originator>,
|
||||
) -> CoreResult<(Vec<ModelInfo>, Option<String>)> {
|
||||
self.fetch_count.fetch_add(1, Ordering::SeqCst);
|
||||
let models = self
|
||||
@@ -897,7 +901,9 @@ async fn static_manager_reads_latest_auth_mode() {
|
||||
},
|
||||
);
|
||||
|
||||
let chatgpt_models = manager.list_models(RefreshStrategy::Online).await;
|
||||
let chatgpt_models = manager
|
||||
.list_models(RefreshStrategy::Online, /*originator*/ None)
|
||||
.await;
|
||||
assert_eq!(
|
||||
chatgpt_models
|
||||
.iter()
|
||||
@@ -907,7 +913,9 @@ async fn static_manager_reads_latest_auth_mode() {
|
||||
);
|
||||
|
||||
auth_manager.set_external_auth(Arc::new(TestExternalApiKeyAuth));
|
||||
let api_models = manager.list_models(RefreshStrategy::Online).await;
|
||||
let api_models = manager
|
||||
.list_models(RefreshStrategy::Online, /*originator*/ None)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
api_models
|
||||
|
||||
Reference in New Issue
Block a user