overhaul provider inventory and agent/model selection (#8652)
Signed-off-by: Bradley Axen <baxen@squareup.com>
This commit is contained in:
@@ -416,6 +416,11 @@ mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
let tmp_dir = tempfile::tempdir().unwrap();
|
||||
let temp_root = tmp_dir.path().display().to_string();
|
||||
let _guard = env_lock::lock_env([
|
||||
("HOME", Some(temp_root.as_str())),
|
||||
("GOOSE_PATH_ROOT", Some(temp_root.as_str())),
|
||||
]);
|
||||
let session_manager = Arc::new(SessionManager::new(tmp_dir.path().to_path_buf()));
|
||||
let session = session_manager
|
||||
.create_session(
|
||||
|
||||
@@ -2,6 +2,7 @@ use crate::config::paths::Paths;
|
||||
use crate::config::Config;
|
||||
use crate::providers::anthropic::AnthropicProvider;
|
||||
use crate::providers::base::{ModelInfo, ProviderType};
|
||||
use crate::providers::inventory::declarative_inventory_identity;
|
||||
use crate::providers::ollama::OllamaProvider;
|
||||
use crate::providers::openai::OpenAiProvider;
|
||||
use anyhow::Result;
|
||||
@@ -460,38 +461,59 @@ pub fn register_declarative_provider(
|
||||
match config.engine {
|
||||
ProviderEngine::OpenAI => {
|
||||
let captured = config.clone();
|
||||
registry.register_with_name::<OpenAiProvider, _>(
|
||||
let identity_config = config.clone();
|
||||
registry.register_with_name::<OpenAiProvider, _, _>(
|
||||
&config,
|
||||
provider_type,
|
||||
config.dynamic_models.unwrap_or(false),
|
||||
move |model| {
|
||||
let mut cfg = captured.clone();
|
||||
resolve_config(&mut cfg)?;
|
||||
OpenAiProvider::from_custom_config(model, cfg)
|
||||
},
|
||||
move || {
|
||||
let mut cfg = identity_config.clone();
|
||||
resolve_config(&mut cfg)?;
|
||||
declarative_inventory_identity(&cfg)
|
||||
},
|
||||
);
|
||||
}
|
||||
ProviderEngine::Ollama => {
|
||||
let captured = config.clone();
|
||||
registry.register_with_name::<OllamaProvider, _>(
|
||||
let identity_config = config.clone();
|
||||
registry.register_with_name::<OllamaProvider, _, _>(
|
||||
&config,
|
||||
provider_type,
|
||||
config.dynamic_models.unwrap_or(false),
|
||||
move |model| {
|
||||
let mut cfg = captured.clone();
|
||||
resolve_config(&mut cfg)?;
|
||||
OllamaProvider::from_custom_config(model, cfg)
|
||||
},
|
||||
move || {
|
||||
let mut cfg = identity_config.clone();
|
||||
resolve_config(&mut cfg)?;
|
||||
declarative_inventory_identity(&cfg)
|
||||
},
|
||||
);
|
||||
}
|
||||
ProviderEngine::Anthropic => {
|
||||
let captured = config.clone();
|
||||
registry.register_with_name::<AnthropicProvider, _>(
|
||||
let identity_config = config.clone();
|
||||
registry.register_with_name::<AnthropicProvider, _, _>(
|
||||
&config,
|
||||
provider_type,
|
||||
config.dynamic_models.unwrap_or(false),
|
||||
move |model| {
|
||||
let mut cfg = captured.clone();
|
||||
resolve_config(&mut cfg)?;
|
||||
AnthropicProvider::from_custom_config(model, cfg)
|
||||
},
|
||||
move || {
|
||||
let mut cfg = identity_config.clone();
|
||||
resolve_config(&mut cfg)?;
|
||||
declarative_inventory_identity(&cfg)
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
use crate::config::search_path::SearchPaths;
|
||||
use crate::providers::inventory::InventoryIdentityInput;
|
||||
use anyhow::Result;
|
||||
use std::path::PathBuf;
|
||||
|
||||
pub fn acp_adapter_installed(command: &str) -> bool {
|
||||
resolve_acp_command(command).is_ok()
|
||||
}
|
||||
|
||||
pub fn acp_inventory_identity(provider_id: &str, command: &str) -> Result<InventoryIdentityInput> {
|
||||
let resolved_command = resolve_acp_command(command)?;
|
||||
Ok(InventoryIdentityInput::new(provider_id, provider_id)
|
||||
.with_public("command", resolved_command.display().to_string()))
|
||||
}
|
||||
|
||||
fn resolve_acp_command(command: &str) -> Result<PathBuf> {
|
||||
SearchPaths::builder().with_npm().resolve(command)
|
||||
}
|
||||
@@ -9,7 +9,9 @@ use crate::acp::{
|
||||
use crate::config::search_path::SearchPaths;
|
||||
use crate::config::{Config, GooseMode};
|
||||
use crate::model::ModelConfig;
|
||||
use crate::providers::acp_tooling::{acp_adapter_installed, acp_inventory_identity};
|
||||
use crate::providers::base::{ProviderDef, ProviderMetadata};
|
||||
use crate::providers::inventory::InventoryIdentityInput;
|
||||
|
||||
const AMP_ACP_PROVIDER_NAME: &str = "amp-acp";
|
||||
const AMP_ACP_DOC_URL: &str = "https://ampcode.com";
|
||||
@@ -37,6 +39,7 @@ impl ProviderDef for AmpAcpProvider {
|
||||
"Set in your goose config file (`~/.config/goose/config.yaml` on macOS/Linux):\n GOOSE_PROVIDER: amp-acp\n GOOSE_MODEL: current",
|
||||
"Restart goose for changes to take effect",
|
||||
])
|
||||
.with_model_selection_hint("Use the Amp CLI to configure models")
|
||||
}
|
||||
|
||||
fn from_env(
|
||||
@@ -49,10 +52,12 @@ impl ProviderDef for AmpAcpProvider {
|
||||
let goose_mode = config.get_goose_mode().unwrap_or(GooseMode::Auto);
|
||||
|
||||
let mode_mapping = HashMap::from([
|
||||
(GooseMode::Auto, "auto".to_string()),
|
||||
(GooseMode::Approve, "approve".to_string()),
|
||||
(GooseMode::SmartApprove, "smart-approve".to_string()),
|
||||
(GooseMode::Chat, "chat".to_string()),
|
||||
// "bypass" skips confirmations, closest to autonomous mode.
|
||||
(GooseMode::Auto, "bypass".to_string()),
|
||||
// "default" prompts before risky actions.
|
||||
(GooseMode::Approve, "default".to_string()),
|
||||
(GooseMode::SmartApprove, "default".to_string()),
|
||||
(GooseMode::Chat, "default".to_string()),
|
||||
]);
|
||||
|
||||
let provider_config = AcpProviderConfig {
|
||||
@@ -71,4 +76,16 @@ impl ProviderDef for AmpAcpProvider {
|
||||
AcpProvider::connect(metadata.name, model, goose_mode, provider_config).await
|
||||
})
|
||||
}
|
||||
|
||||
fn supports_inventory_refresh() -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn inventory_identity() -> Result<InventoryIdentityInput> {
|
||||
acp_inventory_identity(AMP_ACP_PROVIDER_NAME, AMP_ACP_BINARY)
|
||||
}
|
||||
|
||||
fn inventory_configured() -> bool {
|
||||
acp_adapter_installed(AMP_ACP_BINARY)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@ use super::errors::ProviderError;
|
||||
use super::formats::anthropic::{
|
||||
create_request, response_to_streaming_message, thinking_type, ThinkingType,
|
||||
};
|
||||
use super::inventory::{config_secret_value, serialize_string_map, InventoryIdentityInput};
|
||||
use super::openai_compatible::handle_status_openai_compat;
|
||||
use super::openai_compatible::map_http_error_to_provider_error;
|
||||
use super::retry::ProviderRetry;
|
||||
@@ -235,6 +236,33 @@ impl ProviderDef for AnthropicProvider {
|
||||
) -> BoxFuture<'static, Result<Self::Provider>> {
|
||||
Box::pin(Self::from_env(model))
|
||||
}
|
||||
|
||||
fn supports_inventory_refresh() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn inventory_identity() -> Result<InventoryIdentityInput> {
|
||||
let config = crate::config::Config::global();
|
||||
let mut identity =
|
||||
InventoryIdentityInput::new(ANTHROPIC_PROVIDER_NAME, ANTHROPIC_PROVIDER_NAME)
|
||||
.with_public(
|
||||
"host",
|
||||
config
|
||||
.get_param::<String>("ANTHROPIC_HOST")
|
||||
.unwrap_or_else(|_| "https://api.anthropic.com".to_string()),
|
||||
);
|
||||
|
||||
if let Some(api_key) = config_secret_value(config, "ANTHROPIC_API_KEY") {
|
||||
identity = identity.with_secret("api_key", api_key);
|
||||
}
|
||||
if let Ok(headers) = config
|
||||
.get_secret::<std::collections::HashMap<String, String>>("ANTHROPIC_CUSTOM_HEADERS")
|
||||
{
|
||||
identity = identity.with_secret("headers", serialize_string_map(&headers)?);
|
||||
}
|
||||
|
||||
Ok(identity)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
|
||||
@@ -6,9 +6,10 @@ use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::canonical::{map_to_canonical_model, CanonicalModelRegistry};
|
||||
use super::errors::ProviderError;
|
||||
use super::inventory::{default_inventory_identity, InventoryIdentityInput};
|
||||
use super::retry::RetryConfig;
|
||||
use crate::config::base::ConfigValue;
|
||||
use crate::config::{ExtensionConfig, GooseMode};
|
||||
use crate::config::{Config, ExtensionConfig, GooseMode};
|
||||
use crate::conversation::message::{Message, MessageContent};
|
||||
use crate::conversation::Conversation;
|
||||
use crate::model::ModelConfig;
|
||||
@@ -179,6 +180,9 @@ pub struct ProviderMetadata {
|
||||
/// step-by-step instructions for set up providers eg: api key
|
||||
#[serde(default)]
|
||||
pub setup_steps: Vec<String>,
|
||||
/// Hint shown in the model picker when this provider manages its own model selection.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub model_selection_hint: Option<String>,
|
||||
}
|
||||
|
||||
impl ProviderMetadata {
|
||||
@@ -212,6 +216,7 @@ impl ProviderMetadata {
|
||||
model_doc_link: model_doc_link.to_string(),
|
||||
config_keys,
|
||||
setup_steps: vec![],
|
||||
model_selection_hint: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -233,6 +238,7 @@ impl ProviderMetadata {
|
||||
model_doc_link: model_doc_link.to_string(),
|
||||
config_keys,
|
||||
setup_steps: vec![],
|
||||
model_selection_hint: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -246,6 +252,7 @@ impl ProviderMetadata {
|
||||
model_doc_link: "".to_string(),
|
||||
config_keys: vec![],
|
||||
setup_steps: vec![],
|
||||
model_selection_hint: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -253,6 +260,11 @@ impl ProviderMetadata {
|
||||
self.setup_steps = steps.into_iter().map(|s| s.to_string()).collect();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_model_selection_hint(mut self, hint: &str) -> Self {
|
||||
self.model_selection_hint = Some(hint.to_string());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration key metadata for provider setup
|
||||
@@ -492,6 +504,34 @@ pub trait ProviderDef: Send + Sync {
|
||||
) -> BoxFuture<'static, Result<Self::Provider>>
|
||||
where
|
||||
Self: Sized;
|
||||
|
||||
fn supports_inventory_refresh() -> bool
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
false
|
||||
}
|
||||
|
||||
fn inventory_identity() -> Result<InventoryIdentityInput>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
let metadata = Self::metadata();
|
||||
Ok(default_inventory_identity(
|
||||
&metadata.name,
|
||||
&metadata.name,
|
||||
&metadata.config_keys,
|
||||
Config::global(),
|
||||
))
|
||||
}
|
||||
|
||||
fn inventory_configured() -> bool
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
let metadata = Self::metadata();
|
||||
super::inventory::default_inventory_configured(&metadata.config_keys, Config::global())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
@@ -588,7 +628,7 @@ pub trait Provider: Send + Sync {
|
||||
false
|
||||
}
|
||||
|
||||
/// Fetch models filtered by canonical registry and usability
|
||||
/// Fetch inventory models filtered by canonical registry and usability.
|
||||
async fn fetch_recommended_models(&self) -> Result<Vec<String>, ProviderError> {
|
||||
let all_models = self.fetch_supported_models().await?;
|
||||
|
||||
@@ -637,15 +677,15 @@ pub trait Provider: Send + Sync {
|
||||
(None, None) => a.0.cmp(&b.0),
|
||||
});
|
||||
|
||||
let recommended_models: Vec<String> = models_with_dates
|
||||
let inventory_models: Vec<String> = models_with_dates
|
||||
.into_iter()
|
||||
.map(|(name, _)| name)
|
||||
.collect();
|
||||
|
||||
if recommended_models.is_empty() {
|
||||
if inventory_models.is_empty() {
|
||||
Ok(all_models)
|
||||
} else {
|
||||
Ok(recommended_models)
|
||||
Ok(inventory_models)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -9,7 +9,9 @@ use crate::acp::{
|
||||
use crate::config::search_path::SearchPaths;
|
||||
use crate::config::{Config, GooseMode};
|
||||
use crate::model::ModelConfig;
|
||||
use crate::providers::acp_tooling::{acp_adapter_installed, acp_inventory_identity};
|
||||
use crate::providers::base::{ProviderDef, ProviderMetadata};
|
||||
use crate::providers::inventory::InventoryIdentityInput;
|
||||
|
||||
const CLAUDE_ACP_PROVIDER_NAME: &str = "claude-acp";
|
||||
const CLAUDE_ACP_DOC_URL: &str = "https://github.com/zed-industries/claude-agent-acp";
|
||||
@@ -78,4 +80,16 @@ impl ProviderDef for ClaudeAcpProvider {
|
||||
AcpProvider::connect(metadata.name, model, goose_mode, provider_config).await
|
||||
})
|
||||
}
|
||||
|
||||
fn supports_inventory_refresh() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn inventory_identity() -> Result<InventoryIdentityInput> {
|
||||
acp_inventory_identity(CLAUDE_ACP_PROVIDER_NAME, CLAUDE_ACP_BINARY)
|
||||
}
|
||||
|
||||
fn inventory_configured() -> bool {
|
||||
acp_adapter_installed(CLAUDE_ACP_BINARY)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,7 +9,9 @@ use crate::acp::{
|
||||
use crate::config::search_path::SearchPaths;
|
||||
use crate::config::{Config, GooseMode};
|
||||
use crate::model::ModelConfig;
|
||||
use crate::providers::acp_tooling::{acp_adapter_installed, acp_inventory_identity};
|
||||
use crate::providers::base::{ProviderDef, ProviderMetadata};
|
||||
use crate::providers::inventory::InventoryIdentityInput;
|
||||
|
||||
const CODEX_ACP_PROVIDER_NAME: &str = "codex-acp";
|
||||
const CODEX_ACP_DOC_URL: &str = "https://github.com/zed-industries/codex-acp";
|
||||
@@ -98,6 +100,18 @@ impl ProviderDef for CodexAcpProvider {
|
||||
AcpProvider::connect(metadata.name, model, goose_mode, provider_config).await
|
||||
})
|
||||
}
|
||||
|
||||
fn supports_inventory_refresh() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn inventory_identity() -> Result<InventoryIdentityInput> {
|
||||
acp_inventory_identity(CODEX_ACP_PROVIDER_NAME, CODEX_ACP_PROVIDER_NAME)
|
||||
}
|
||||
|
||||
fn inventory_configured() -> bool {
|
||||
acp_adapter_installed(CODEX_ACP_PROVIDER_NAME)
|
||||
}
|
||||
}
|
||||
|
||||
// Codex sandbox scope determines what needs approval: operations within the
|
||||
|
||||
@@ -9,7 +9,9 @@ use crate::acp::{
|
||||
use crate::config::search_path::SearchPaths;
|
||||
use crate::config::{Config, GooseMode};
|
||||
use crate::model::ModelConfig;
|
||||
use crate::providers::acp_tooling::{acp_adapter_installed, acp_inventory_identity};
|
||||
use crate::providers::base::{ProviderDef, ProviderMetadata};
|
||||
use crate::providers::inventory::InventoryIdentityInput;
|
||||
|
||||
const COPILOT_ACP_PROVIDER_NAME: &str = "copilot-acp";
|
||||
const COPILOT_ACP_DOC_URL: &str = "https://github.com/github/copilot-cli";
|
||||
@@ -84,4 +86,16 @@ impl ProviderDef for CopilotAcpProvider {
|
||||
AcpProvider::connect(metadata.name, model, goose_mode, provider_config).await
|
||||
})
|
||||
}
|
||||
|
||||
fn supports_inventory_refresh() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn inventory_identity() -> Result<InventoryIdentityInput> {
|
||||
acp_inventory_identity(COPILOT_ACP_PROVIDER_NAME, COPILOT_ACP_BINARY)
|
||||
}
|
||||
|
||||
fn inventory_configured() -> bool {
|
||||
acp_adapter_installed(COPILOT_ACP_BINARY)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -340,6 +340,10 @@ impl ProviderDef for DatabricksProvider {
|
||||
) -> BoxFuture<'static, Result<Self::Provider>> {
|
||||
Box::pin(Self::from_env(model))
|
||||
}
|
||||
|
||||
fn supports_inventory_refresh() -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
|
||||
@@ -1100,12 +1100,16 @@ mod tests {
|
||||
fn test_create_request_enabled_thinking_with_budget() -> anyhow::Result<()> {
|
||||
let _guard = env_lock::lock_env([
|
||||
("CLAUDE_THINKING_TYPE", None::<&str>),
|
||||
("CLAUDE_THINKING_ENABLED", Some("1")),
|
||||
("CLAUDE_THINKING_ENABLED", None::<&str>),
|
||||
("CLAUDE_THINKING_BUDGET", Some("10000")),
|
||||
]);
|
||||
|
||||
let mut model_config = ModelConfig::new_or_fail("databricks-claude-3-7-sonnet");
|
||||
model_config.max_tokens = Some(4096);
|
||||
model_config = model_config.with_request_params(Some(std::collections::HashMap::from([(
|
||||
"thinking_type".to_string(),
|
||||
json!("enabled"),
|
||||
)])));
|
||||
|
||||
let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?;
|
||||
|
||||
|
||||
@@ -149,6 +149,10 @@ pub async fn get_from_registry(name: &str) -> Result<ProviderEntry> {
|
||||
.cloned()
|
||||
}
|
||||
|
||||
pub async fn inventory_identity(name: &str) -> Result<super::inventory::InventoryIdentityInput> {
|
||||
get_from_registry(name).await?.inventory_identity()
|
||||
}
|
||||
|
||||
pub async fn create(
|
||||
name: &str,
|
||||
model: ModelConfig,
|
||||
|
||||
@@ -0,0 +1,970 @@
|
||||
use super::base::{ConfigKey, ModelInfo};
|
||||
use super::canonical::{map_provider_name, map_to_canonical_model, CanonicalModelRegistry};
|
||||
use crate::config::declarative_providers::{DeclarativeProviderConfig, ProviderEngine};
|
||||
use crate::config::Config;
|
||||
use crate::session::session_manager::SessionStorage;
|
||||
use anyhow::Result;
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use sqlx::{Pool, Row, Sqlite, Transaction};
|
||||
use std::collections::{BTreeMap, HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
const STALE_AFTER_HOURS: i64 = 24;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ProviderInventoryEntry {
|
||||
pub provider_id: String,
|
||||
pub provider_name: String,
|
||||
pub configured: bool,
|
||||
pub supports_refresh: bool,
|
||||
pub refreshing: bool,
|
||||
pub models: Vec<InventoryModel>,
|
||||
pub last_updated_at: Option<DateTime<Utc>>,
|
||||
pub last_refresh_attempt_at: Option<DateTime<Utc>>,
|
||||
pub last_refresh_error: Option<String>,
|
||||
pub model_selection_hint: Option<String>,
|
||||
}
|
||||
|
||||
/// Families whose latest model should be surfaced in the compact picker.
|
||||
/// Each entry is matched against the `family` field of enriched models.
|
||||
const RECOMMENDED_FAMILIES: &[&str] = &[
|
||||
"claude-opus",
|
||||
"claude-sonnet",
|
||||
"gpt",
|
||||
"gpt-mini",
|
||||
"glm",
|
||||
"gemini-pro",
|
||||
"gemini-flash",
|
||||
"gemma",
|
||||
];
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct InventoryModel {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub family: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub context_limit: Option<usize>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub reasoning: Option<bool>,
|
||||
/// Whether this model should appear in the compact recommended picker.
|
||||
pub recommended: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct InventoryIdentity {
|
||||
pub provider_id: String,
|
||||
pub provider_family: String,
|
||||
pub inventory_key: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct InventoryIdentityInput {
|
||||
pub provider_id: String,
|
||||
pub provider_family: String,
|
||||
pub public_inputs: BTreeMap<String, String>,
|
||||
pub secret_inputs: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
impl InventoryIdentityInput {
|
||||
pub fn new(
|
||||
provider_id: impl Into<String>,
|
||||
provider_family: impl Into<String>,
|
||||
) -> InventoryIdentityInput {
|
||||
InventoryIdentityInput {
|
||||
provider_id: provider_id.into(),
|
||||
provider_family: provider_family.into(),
|
||||
public_inputs: BTreeMap::new(),
|
||||
secret_inputs: BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_public(
|
||||
mut self,
|
||||
key: impl Into<String>,
|
||||
value: impl Into<String>,
|
||||
) -> InventoryIdentityInput {
|
||||
self.public_inputs.insert(key.into(), value.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_secret(
|
||||
mut self,
|
||||
key: impl Into<String>,
|
||||
value: impl Into<String>,
|
||||
) -> InventoryIdentityInput {
|
||||
self.secret_inputs.insert(key.into(), value.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn into_identity(self) -> Result<InventoryIdentity> {
|
||||
let InventoryIdentityInput {
|
||||
provider_id,
|
||||
provider_family,
|
||||
public_inputs,
|
||||
secret_inputs,
|
||||
} = self;
|
||||
let payload = serde_json::json!({
|
||||
"provider_family": provider_family,
|
||||
"public_inputs": public_inputs,
|
||||
"secret_inputs": secret_inputs,
|
||||
});
|
||||
let digest = Sha256::digest(serde_json::to_vec(&payload)?);
|
||||
Ok(InventoryIdentity {
|
||||
provider_id,
|
||||
provider_family,
|
||||
inventory_key: format!("{digest:x}"),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum RefreshSkipReason {
|
||||
UnknownProvider,
|
||||
NotConfigured,
|
||||
DoesNotSupportRefresh,
|
||||
AlreadyRefreshing,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RefreshSkip {
|
||||
pub provider_id: String,
|
||||
pub reason: RefreshSkipReason,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct RefreshPlan {
|
||||
pub started: Vec<String>,
|
||||
pub skipped: Vec<RefreshSkip>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ProviderInventoryService {
|
||||
storage: Arc<SessionStorage>,
|
||||
refreshing_keys: Arc<RwLock<HashSet<String>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct InventorySnapshot {
|
||||
models: Vec<InventoryModel>,
|
||||
last_updated_at: Option<DateTime<Utc>>,
|
||||
last_refresh_attempt_at: Option<DateTime<Utc>>,
|
||||
last_refresh_error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct ProviderDescriptor {
|
||||
provider_id: String,
|
||||
provider_name: String,
|
||||
identity: InventoryIdentity,
|
||||
configured: bool,
|
||||
supports_refresh: bool,
|
||||
static_models: Vec<ModelInfo>,
|
||||
model_selection_hint: Option<String>,
|
||||
}
|
||||
|
||||
impl ProviderInventoryService {
|
||||
pub fn new(storage: Arc<SessionStorage>) -> ProviderInventoryService {
|
||||
ProviderInventoryService {
|
||||
storage,
|
||||
refreshing_keys: Arc::new(RwLock::new(HashSet::new())),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn entry_for_provider(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
) -> Result<Option<ProviderInventoryEntry>> {
|
||||
let Some(descriptor) = self.describe_provider(provider_id).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let snapshot = self.read_snapshot(&descriptor.identity).await?;
|
||||
let refreshing = self
|
||||
.refreshing_keys
|
||||
.read()
|
||||
.await
|
||||
.contains(&descriptor.identity.inventory_key);
|
||||
let models = inventory_models_from_snapshot(
|
||||
snapshot.as_ref(),
|
||||
&descriptor.identity.provider_family,
|
||||
&descriptor.static_models,
|
||||
);
|
||||
|
||||
Ok(Some(ProviderInventoryEntry {
|
||||
provider_id: descriptor.provider_id,
|
||||
provider_name: descriptor.provider_name,
|
||||
configured: descriptor.configured,
|
||||
supports_refresh: descriptor.supports_refresh,
|
||||
refreshing,
|
||||
models,
|
||||
last_updated_at: snapshot
|
||||
.as_ref()
|
||||
.and_then(|snapshot| snapshot.last_updated_at),
|
||||
last_refresh_attempt_at: snapshot
|
||||
.as_ref()
|
||||
.and_then(|snapshot| snapshot.last_refresh_attempt_at),
|
||||
last_refresh_error: snapshot.and_then(|snapshot| snapshot.last_refresh_error),
|
||||
model_selection_hint: descriptor.model_selection_hint,
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn entries(&self, provider_ids: &[String]) -> Result<Vec<ProviderInventoryEntry>> {
|
||||
let ids = self.resolve_provider_ids(provider_ids).await;
|
||||
let mut entries = Vec::with_capacity(ids.len());
|
||||
for provider_id in ids {
|
||||
if let Some(entry) = self.entry_for_provider(&provider_id).await? {
|
||||
entries.push(entry);
|
||||
}
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
pub async fn plan_refresh(&self, provider_ids: &[String]) -> Result<RefreshPlan> {
|
||||
let ids = self.resolve_provider_ids(provider_ids).await;
|
||||
let mut plan = RefreshPlan::default();
|
||||
|
||||
for provider_id in ids {
|
||||
let Some(descriptor) = self.describe_provider(&provider_id).await? else {
|
||||
plan.skipped.push(RefreshSkip {
|
||||
provider_id,
|
||||
reason: RefreshSkipReason::UnknownProvider,
|
||||
});
|
||||
continue;
|
||||
};
|
||||
|
||||
if !descriptor.supports_refresh {
|
||||
plan.skipped.push(RefreshSkip {
|
||||
provider_id: descriptor.provider_id,
|
||||
reason: RefreshSkipReason::DoesNotSupportRefresh,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
if !descriptor.configured {
|
||||
plan.skipped.push(RefreshSkip {
|
||||
provider_id: descriptor.provider_id,
|
||||
reason: RefreshSkipReason::NotConfigured,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
let mut refreshing_keys = self.refreshing_keys.write().await;
|
||||
if refreshing_keys.contains(&descriptor.identity.inventory_key) {
|
||||
plan.skipped.push(RefreshSkip {
|
||||
provider_id: descriptor.provider_id,
|
||||
reason: RefreshSkipReason::AlreadyRefreshing,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
refreshing_keys.insert(descriptor.identity.inventory_key.clone());
|
||||
drop(refreshing_keys);
|
||||
|
||||
self.mark_refresh_started(&descriptor.identity).await?;
|
||||
plan.started.push(descriptor.provider_id);
|
||||
}
|
||||
|
||||
Ok(plan)
|
||||
}
|
||||
|
||||
pub async fn store_refreshed_models(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
model_ids: &[String],
|
||||
) -> Result<()> {
|
||||
let descriptor = self.require_provider(provider_id).await?;
|
||||
let models =
|
||||
enrich_model_ids_with_canonical(&descriptor.identity.provider_family, model_ids);
|
||||
let now = Utc::now();
|
||||
let pool = self.storage.pool().await?;
|
||||
let mut tx = pool.begin().await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO provider_inventory_entries (
|
||||
inventory_key,
|
||||
provider_id,
|
||||
provider_family,
|
||||
last_updated_at,
|
||||
last_refresh_attempt_at,
|
||||
last_refresh_error,
|
||||
updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, NULL, CURRENT_TIMESTAMP)
|
||||
ON CONFLICT(inventory_key) DO UPDATE SET
|
||||
provider_id = excluded.provider_id,
|
||||
provider_family = excluded.provider_family,
|
||||
last_updated_at = excluded.last_updated_at,
|
||||
last_refresh_attempt_at = excluded.last_refresh_attempt_at,
|
||||
last_refresh_error = NULL,
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
"#,
|
||||
)
|
||||
.bind(&descriptor.identity.inventory_key)
|
||||
.bind(&descriptor.identity.provider_id)
|
||||
.bind(&descriptor.identity.provider_family)
|
||||
.bind(now.to_rfc3339())
|
||||
.bind(now.to_rfc3339())
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
sqlx::query("DELETE FROM provider_inventory_models WHERE inventory_key = ?")
|
||||
.bind(&descriptor.identity.inventory_key)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
for (ordinal, model) in models.iter().enumerate() {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO provider_inventory_models (
|
||||
inventory_key,
|
||||
ordinal,
|
||||
model_id,
|
||||
name,
|
||||
family,
|
||||
context_limit,
|
||||
reasoning,
|
||||
recommended
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(&descriptor.identity.inventory_key)
|
||||
.bind(i64::try_from(ordinal)?)
|
||||
.bind(&model.id)
|
||||
.bind(&model.name)
|
||||
.bind(&model.family)
|
||||
.bind(model.context_limit.map(i64::try_from).transpose()?)
|
||||
.bind(model.reasoning)
|
||||
.bind(model.recommended)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
}
|
||||
|
||||
tx.commit().await?;
|
||||
self.refreshing_keys
|
||||
.write()
|
||||
.await
|
||||
.remove(&descriptor.identity.inventory_key);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn store_refresh_error(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
error: impl Into<String>,
|
||||
) -> Result<()> {
|
||||
let descriptor = self.require_provider(provider_id).await?;
|
||||
let error = error.into();
|
||||
let existing = self.read_snapshot(&descriptor.identity).await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO provider_inventory_entries (
|
||||
inventory_key,
|
||||
provider_id,
|
||||
provider_family,
|
||||
last_updated_at,
|
||||
last_refresh_attempt_at,
|
||||
last_refresh_error,
|
||||
updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
|
||||
ON CONFLICT(inventory_key) DO UPDATE SET
|
||||
provider_id = excluded.provider_id,
|
||||
provider_family = excluded.provider_family,
|
||||
last_updated_at = excluded.last_updated_at,
|
||||
last_refresh_attempt_at = excluded.last_refresh_attempt_at,
|
||||
last_refresh_error = excluded.last_refresh_error,
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
"#,
|
||||
)
|
||||
.bind(&descriptor.identity.inventory_key)
|
||||
.bind(&descriptor.identity.provider_id)
|
||||
.bind(&descriptor.identity.provider_family)
|
||||
.bind(existing.and_then(|snapshot| snapshot.last_updated_at.map(|time| time.to_rfc3339())))
|
||||
.bind(Utc::now().to_rfc3339())
|
||||
.bind(error)
|
||||
.execute(self.storage.pool().await?)
|
||||
.await?;
|
||||
|
||||
self.refreshing_keys
|
||||
.write()
|
||||
.await
|
||||
.remove(&descriptor.identity.inventory_key);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn is_stale(entry: &ProviderInventoryEntry) -> bool {
|
||||
let Some(last_updated_at) = entry.last_updated_at else {
|
||||
return false;
|
||||
};
|
||||
entry.supports_refresh && Utc::now() - last_updated_at > Duration::hours(STALE_AFTER_HOURS)
|
||||
}
|
||||
|
||||
async fn describe_provider(&self, provider_id: &str) -> Result<Option<ProviderDescriptor>> {
|
||||
let entry = match crate::providers::get_from_registry(provider_id).await {
|
||||
Ok(entry) => entry,
|
||||
Err(_) => return Ok(None),
|
||||
};
|
||||
let metadata = entry.metadata().clone();
|
||||
let identity = crate::providers::inventory_identity(provider_id)
|
||||
.await
|
||||
.unwrap_or_else(|_| fallback_inventory_identity(provider_id))
|
||||
.into_identity()?;
|
||||
|
||||
Ok(Some(ProviderDescriptor {
|
||||
provider_id: metadata.name.clone(),
|
||||
provider_name: metadata.display_name.clone(),
|
||||
identity,
|
||||
configured: entry.inventory_configured(),
|
||||
supports_refresh: entry.supports_inventory_refresh(),
|
||||
static_models: metadata.known_models,
|
||||
model_selection_hint: metadata.model_selection_hint,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn require_provider(&self, provider_id: &str) -> Result<ProviderDescriptor> {
|
||||
self.describe_provider(provider_id)
|
||||
.await?
|
||||
.ok_or_else(|| anyhow::anyhow!("Unknown provider: {}", provider_id))
|
||||
}
|
||||
|
||||
async fn mark_refresh_started(&self, identity: &InventoryIdentity) -> Result<()> {
|
||||
let existing = self.read_snapshot(identity).await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO provider_inventory_entries (
|
||||
inventory_key,
|
||||
provider_id,
|
||||
provider_family,
|
||||
last_updated_at,
|
||||
last_refresh_attempt_at,
|
||||
last_refresh_error,
|
||||
updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, NULL, CURRENT_TIMESTAMP)
|
||||
ON CONFLICT(inventory_key) DO UPDATE SET
|
||||
provider_id = excluded.provider_id,
|
||||
provider_family = excluded.provider_family,
|
||||
last_updated_at = excluded.last_updated_at,
|
||||
last_refresh_attempt_at = excluded.last_refresh_attempt_at,
|
||||
last_refresh_error = NULL,
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
"#,
|
||||
)
|
||||
.bind(&identity.inventory_key)
|
||||
.bind(&identity.provider_id)
|
||||
.bind(&identity.provider_family)
|
||||
.bind(existing.and_then(|snapshot| snapshot.last_updated_at.map(|time| time.to_rfc3339())))
|
||||
.bind(Utc::now().to_rfc3339())
|
||||
.execute(self.storage.pool().await?)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn read_snapshot(
|
||||
&self,
|
||||
identity: &InventoryIdentity,
|
||||
) -> Result<Option<InventorySnapshot>> {
|
||||
let pool = self.storage.pool().await?;
|
||||
let entry = sqlx::query(
|
||||
r#"
|
||||
SELECT last_updated_at, last_refresh_attempt_at, last_refresh_error
|
||||
FROM provider_inventory_entries
|
||||
WHERE inventory_key = ?
|
||||
"#,
|
||||
)
|
||||
.bind(&identity.inventory_key)
|
||||
.fetch_optional(pool)
|
||||
.await?;
|
||||
|
||||
let Some(entry) = entry else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let last_updated_at = parse_optional_datetime(entry.try_get("last_updated_at")?)?;
|
||||
let last_refresh_attempt_at =
|
||||
parse_optional_datetime(entry.try_get("last_refresh_attempt_at")?)?;
|
||||
let last_refresh_error = entry.try_get("last_refresh_error")?;
|
||||
|
||||
let rows = sqlx::query(
|
||||
r#"
|
||||
SELECT model_id, name, family, context_limit, reasoning, recommended
|
||||
FROM provider_inventory_models
|
||||
WHERE inventory_key = ?
|
||||
ORDER BY ordinal
|
||||
"#,
|
||||
)
|
||||
.bind(&identity.inventory_key)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
|
||||
let models = rows
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
Ok(InventoryModel {
|
||||
id: row.try_get("model_id")?,
|
||||
name: row.try_get("name")?,
|
||||
family: row.try_get("family")?,
|
||||
context_limit: row
|
||||
.try_get::<Option<i64>, _>("context_limit")?
|
||||
.map(usize::try_from)
|
||||
.transpose()?,
|
||||
reasoning: row.try_get("reasoning")?,
|
||||
recommended: row
|
||||
.try_get::<Option<bool>, _>("recommended")?
|
||||
.unwrap_or(false),
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, anyhow::Error>>()?;
|
||||
|
||||
Ok(Some(InventorySnapshot {
|
||||
models,
|
||||
last_updated_at,
|
||||
last_refresh_attempt_at,
|
||||
last_refresh_error,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn resolve_provider_ids(&self, provider_ids: &[String]) -> Vec<String> {
|
||||
let mut ids = if provider_ids.is_empty() {
|
||||
crate::providers::providers()
|
||||
.await
|
||||
.into_iter()
|
||||
.map(|(metadata, _)| metadata.name)
|
||||
.collect::<Vec<_>>()
|
||||
} else {
|
||||
provider_ids.to_vec()
|
||||
};
|
||||
ids.sort();
|
||||
ids.dedup();
|
||||
ids
|
||||
}
|
||||
}
|
||||
|
||||
pub fn default_inventory_identity(
|
||||
provider_id: &str,
|
||||
provider_family: &str,
|
||||
config_keys: &[ConfigKey],
|
||||
config: &Config,
|
||||
) -> InventoryIdentityInput {
|
||||
let mut identity = InventoryIdentityInput::new(provider_id, provider_family);
|
||||
|
||||
for key in config_keys {
|
||||
if key.secret {
|
||||
if let Some(value) = config_secret_value(config, &key.name) {
|
||||
identity.secret_inputs.insert(key.name.clone(), value);
|
||||
}
|
||||
} else if let Some(value) = config_param_value(config, &key.name) {
|
||||
identity.public_inputs.insert(key.name.clone(), value);
|
||||
}
|
||||
}
|
||||
|
||||
identity
|
||||
}
|
||||
|
||||
pub fn default_inventory_configured(config_keys: &[ConfigKey], config: &Config) -> bool {
|
||||
config_keys.iter().all(|key| {
|
||||
if !key.required {
|
||||
return true;
|
||||
}
|
||||
if key.default.is_some() {
|
||||
return true;
|
||||
}
|
||||
if key.secret {
|
||||
config.get_secret::<serde_json::Value>(&key.name).is_ok()
|
||||
} else {
|
||||
config.get_param::<serde_json::Value>(&key.name).is_ok()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
pub fn declarative_inventory_identity(
|
||||
config: &DeclarativeProviderConfig,
|
||||
) -> Result<InventoryIdentityInput> {
|
||||
let global = Config::global();
|
||||
let mut identity = InventoryIdentityInput::new(
|
||||
config.name.clone(),
|
||||
config
|
||||
.catalog_provider_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| match config.engine {
|
||||
ProviderEngine::OpenAI => "openai".to_string(),
|
||||
ProviderEngine::Anthropic => "anthropic".to_string(),
|
||||
ProviderEngine::Ollama => "ollama".to_string(),
|
||||
}),
|
||||
);
|
||||
|
||||
identity
|
||||
.public_inputs
|
||||
.insert("base_url".to_string(), config.base_url.clone());
|
||||
|
||||
if let Some(base_path) = &config.base_path {
|
||||
identity
|
||||
.public_inputs
|
||||
.insert("base_path".to_string(), base_path.clone());
|
||||
}
|
||||
if let Some(catalog_provider_id) = &config.catalog_provider_id {
|
||||
identity.public_inputs.insert(
|
||||
"catalog_provider_id".to_string(),
|
||||
catalog_provider_id.clone(),
|
||||
);
|
||||
}
|
||||
if let Some(dynamic_models) = config.dynamic_models {
|
||||
identity
|
||||
.public_inputs
|
||||
.insert("dynamic_models".to_string(), dynamic_models.to_string());
|
||||
}
|
||||
identity.public_inputs.insert(
|
||||
"skip_canonical_filtering".to_string(),
|
||||
config.skip_canonical_filtering.to_string(),
|
||||
);
|
||||
if !config.models.is_empty() {
|
||||
identity.public_inputs.insert(
|
||||
"models".to_string(),
|
||||
serde_json::to_string(
|
||||
&config
|
||||
.models
|
||||
.iter()
|
||||
.map(|model| &model.name)
|
||||
.collect::<Vec<_>>(),
|
||||
)?,
|
||||
);
|
||||
}
|
||||
if let Some(headers) = &config.headers {
|
||||
identity
|
||||
.public_inputs
|
||||
.insert("headers".to_string(), serialize_string_map(headers)?);
|
||||
}
|
||||
if config.requires_auth && !config.api_key_env.is_empty() {
|
||||
if let Some(value) = config_secret_value(global, &config.api_key_env) {
|
||||
identity
|
||||
.secret_inputs
|
||||
.insert(config.api_key_env.clone(), value);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(identity)
|
||||
}
|
||||
|
||||
pub fn config_param_value(config: &Config, key: &str) -> Option<String> {
|
||||
config
|
||||
.get_param::<serde_json::Value>(key)
|
||||
.ok()
|
||||
.and_then(|value| normalize_json_value(&value))
|
||||
}
|
||||
|
||||
pub fn config_secret_value(config: &Config, key: &str) -> Option<String> {
|
||||
config
|
||||
.get_secret::<serde_json::Value>(key)
|
||||
.ok()
|
||||
.and_then(|value| normalize_json_value(&value))
|
||||
}
|
||||
|
||||
pub fn serialize_string_map(map: &HashMap<String, String>) -> Result<String> {
|
||||
let ordered = map
|
||||
.iter()
|
||||
.map(|(key, value)| (key.clone(), value.clone()))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
Ok(serde_json::to_string(&ordered)?)
|
||||
}
|
||||
|
||||
fn parse_optional_datetime(value: Option<String>) -> Result<Option<DateTime<Utc>>> {
|
||||
value
|
||||
.map(|value| value.parse::<DateTime<Utc>>())
|
||||
.transpose()
|
||||
.map_err(Into::into)
|
||||
}
|
||||
|
||||
fn normalize_json_value(value: &serde_json::Value) -> Option<String> {
|
||||
match value {
|
||||
serde_json::Value::Null => None,
|
||||
serde_json::Value::String(value) if value.is_empty() => None,
|
||||
serde_json::Value::String(value) => Some(value.clone()),
|
||||
other => serde_json::to_string(other).ok(),
|
||||
}
|
||||
}
|
||||
|
||||
fn fallback_inventory_identity(provider_id: &str) -> InventoryIdentityInput {
|
||||
InventoryIdentityInput::new(
|
||||
provider_id.to_string(),
|
||||
map_provider_name(provider_id).to_string(),
|
||||
)
|
||||
}
|
||||
|
||||
fn enrich_model_ids_with_canonical(
|
||||
provider_family: &str,
|
||||
model_ids: &[String],
|
||||
) -> Vec<InventoryModel> {
|
||||
let mut models: Vec<InventoryModel> = Vec::new();
|
||||
let mut seen_names: HashSet<String> = HashSet::new();
|
||||
|
||||
for id in model_ids {
|
||||
let model = enriched_model(provider_family, id, None);
|
||||
if !seen_names.insert(model.name.clone()) {
|
||||
continue;
|
||||
}
|
||||
models.push(model);
|
||||
}
|
||||
|
||||
// For databricks, prefer goose- prefixed model_ids when there are duplicates.
|
||||
// Re-scan: if a later model_id with "goose-" prefix maps to the same display name,
|
||||
// swap it in.
|
||||
if provider_family == "databricks" {
|
||||
let mut name_to_idx: HashMap<String, usize> = HashMap::new();
|
||||
for (idx, model) in models.iter().enumerate() {
|
||||
name_to_idx.insert(model.name.clone(), idx);
|
||||
}
|
||||
for id in model_ids {
|
||||
if !id.starts_with("goose-") {
|
||||
continue;
|
||||
}
|
||||
let candidate = enriched_model(provider_family, id, None);
|
||||
if let Some(&idx) = name_to_idx.get(&candidate.name) {
|
||||
if !models[idx].id.starts_with("goose-") {
|
||||
models[idx].id = candidate.id;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Mark the latest model per recommended family.
|
||||
let mut seen_recommended_families: HashSet<String> = HashSet::new();
|
||||
for model in &mut models {
|
||||
if let Some(family) = &model.family {
|
||||
if RECOMMENDED_FAMILIES.contains(&family.as_str())
|
||||
&& seen_recommended_families.insert(family.clone())
|
||||
{
|
||||
model.recommended = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
models
|
||||
}
|
||||
|
||||
fn configured_models_to_inventory(
|
||||
provider_family: &str,
|
||||
models: &[ModelInfo],
|
||||
) -> Vec<InventoryModel> {
|
||||
let mut result: Vec<InventoryModel> = Vec::new();
|
||||
let mut seen_names: HashSet<String> = HashSet::new();
|
||||
for model in models {
|
||||
let enriched = enriched_model(provider_family, &model.name, Some(model.context_limit));
|
||||
if seen_names.insert(enriched.name.clone()) {
|
||||
result.push(enriched);
|
||||
}
|
||||
}
|
||||
|
||||
let mut seen_recommended_families: HashSet<String> = HashSet::new();
|
||||
for model in &mut result {
|
||||
if let Some(family) = &model.family {
|
||||
if RECOMMENDED_FAMILIES.contains(&family.as_str())
|
||||
&& seen_recommended_families.insert(family.clone())
|
||||
{
|
||||
model.recommended = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn inventory_models_from_snapshot(
|
||||
snapshot: Option<&InventorySnapshot>,
|
||||
provider_family: &str,
|
||||
configured_models: &[ModelInfo],
|
||||
) -> Vec<InventoryModel> {
|
||||
match snapshot {
|
||||
Some(snapshot) if !snapshot.models.is_empty() || snapshot.last_updated_at.is_some() => {
|
||||
snapshot.models.clone()
|
||||
}
|
||||
_ => configured_models_to_inventory(provider_family, configured_models),
|
||||
}
|
||||
}
|
||||
|
||||
fn enriched_model(
|
||||
provider_family: &str,
|
||||
model_id: &str,
|
||||
fallback_context_limit: Option<usize>,
|
||||
) -> InventoryModel {
|
||||
let registry = CanonicalModelRegistry::bundled().ok();
|
||||
let canonical = registry.as_ref().and_then(|registry| {
|
||||
let canonical_id = map_to_canonical_model(provider_family, model_id, registry)?;
|
||||
let (provider, model) = canonical_id.split_once('/')?;
|
||||
registry.get(provider, model).cloned()
|
||||
});
|
||||
|
||||
InventoryModel {
|
||||
id: model_id.to_string(),
|
||||
name: canonical
|
||||
.as_ref()
|
||||
.map(|model| model.name.clone())
|
||||
.unwrap_or_else(|| model_id.to_string()),
|
||||
family: canonical.as_ref().and_then(|model| model.family.clone()),
|
||||
context_limit: canonical
|
||||
.as_ref()
|
||||
.map(|model| model.limit.context)
|
||||
.or(fallback_context_limit),
|
||||
reasoning: canonical.as_ref().and_then(|model| model.reasoning),
|
||||
recommended: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn create_tables(pool: &Pool<Sqlite>) -> Result<()> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS provider_inventory_entries (
|
||||
inventory_key TEXT PRIMARY KEY,
|
||||
provider_id TEXT NOT NULL,
|
||||
provider_family TEXT NOT NULL,
|
||||
last_updated_at TEXT,
|
||||
last_refresh_attempt_at TEXT,
|
||||
last_refresh_error TEXT,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS provider_inventory_models (
|
||||
inventory_key TEXT NOT NULL REFERENCES provider_inventory_entries(inventory_key) ON DELETE CASCADE,
|
||||
ordinal INTEGER NOT NULL,
|
||||
model_id TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
family TEXT,
|
||||
context_limit INTEGER,
|
||||
reasoning BOOLEAN,
|
||||
recommended BOOLEAN,
|
||||
PRIMARY KEY (inventory_key, ordinal)
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS idx_provider_inventory_provider_id ON provider_inventory_entries(provider_id)",
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn create_tables_in_tx(tx: &mut Transaction<'_, Sqlite>) -> Result<()> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS provider_inventory_entries (
|
||||
inventory_key TEXT PRIMARY KEY,
|
||||
provider_id TEXT NOT NULL,
|
||||
provider_family TEXT NOT NULL,
|
||||
last_updated_at TEXT,
|
||||
last_refresh_attempt_at TEXT,
|
||||
last_refresh_error TEXT,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS provider_inventory_models (
|
||||
inventory_key TEXT NOT NULL REFERENCES provider_inventory_entries(inventory_key) ON DELETE CASCADE,
|
||||
ordinal INTEGER NOT NULL,
|
||||
model_id TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
family TEXT,
|
||||
context_limit INTEGER,
|
||||
reasoning BOOLEAN,
|
||||
recommended BOOLEAN,
|
||||
PRIMARY KEY (inventory_key, ordinal)
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS idx_provider_inventory_provider_id ON provider_inventory_entries(provider_id)",
|
||||
)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn inventory_identity_hash_changes_with_secret_inputs() {
|
||||
let left = InventoryIdentityInput::new("openai", "openai")
|
||||
.with_public("host", "https://api.openai.com")
|
||||
.with_secret("api_key", "secret-a")
|
||||
.into_identity()
|
||||
.unwrap();
|
||||
let right = InventoryIdentityInput::new("openai", "openai")
|
||||
.with_public("host", "https://api.openai.com")
|
||||
.with_secret("api_key", "secret-b")
|
||||
.into_identity()
|
||||
.unwrap();
|
||||
|
||||
assert_ne!(left.inventory_key, right.inventory_key);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn configured_models_use_canonical_enrichment() {
|
||||
let models =
|
||||
configured_models_to_inventory("anthropic", &[ModelInfo::new("claude-sonnet-4-5", 0)]);
|
||||
|
||||
assert_eq!(models.len(), 1);
|
||||
assert!(models[0].name.contains("Claude"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inventory_uses_configured_models_before_first_successful_refresh() {
|
||||
let configured_models = [ModelInfo::new("claude-sonnet-4-5", 0)];
|
||||
let snapshot = InventorySnapshot {
|
||||
models: vec![],
|
||||
last_updated_at: None,
|
||||
last_refresh_attempt_at: Some(Utc::now()),
|
||||
last_refresh_error: Some("auth failed".to_string()),
|
||||
};
|
||||
|
||||
let models =
|
||||
inventory_models_from_snapshot(Some(&snapshot), "anthropic", &configured_models);
|
||||
|
||||
assert_eq!(models.len(), 1);
|
||||
assert_eq!(models[0].id, "claude-sonnet-4-5");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inventory_preserves_empty_models_after_successful_refresh() {
|
||||
let configured_models = [ModelInfo::new("claude-sonnet-4-5", 0)];
|
||||
let snapshot = InventorySnapshot {
|
||||
models: vec![],
|
||||
last_updated_at: Some(Utc::now()),
|
||||
last_refresh_attempt_at: Some(Utc::now()),
|
||||
last_refresh_error: None,
|
||||
};
|
||||
|
||||
let models =
|
||||
inventory_models_from_snapshot(Some(&snapshot), "anthropic", &configured_models);
|
||||
|
||||
assert!(models.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
mod acp_tooling;
|
||||
pub mod amp_acp;
|
||||
pub mod anthropic;
|
||||
pub mod api_client;
|
||||
@@ -28,6 +29,7 @@ pub mod gemini_oauth;
|
||||
pub mod githubcopilot;
|
||||
pub mod google;
|
||||
mod init;
|
||||
pub mod inventory;
|
||||
pub mod kimicode;
|
||||
pub mod litellm;
|
||||
#[cfg(feature = "local-inference")]
|
||||
@@ -56,6 +58,6 @@ pub mod xai;
|
||||
|
||||
pub use init::{
|
||||
cleanup_provider, create, create_with_default_model, create_with_named_model,
|
||||
get_from_registry, providers, refresh_custom_providers,
|
||||
get_from_registry, inventory_identity, providers, refresh_custom_providers,
|
||||
};
|
||||
pub use retry::{retry_operation, RetryConfig};
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use super::api_client::{ApiClient, AuthMethod};
|
||||
use super::base::{ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata};
|
||||
use super::errors::ProviderError;
|
||||
use super::inventory::InventoryIdentityInput;
|
||||
use super::openai_compatible::handle_status_openai_compat;
|
||||
use super::retry::{ProviderRetry, RetryConfig};
|
||||
use super::utils::{ImageFormat, RequestLog};
|
||||
@@ -256,6 +257,22 @@ impl ProviderDef for OllamaProvider {
|
||||
) -> BoxFuture<'static, Result<Self::Provider>> {
|
||||
Box::pin(Self::from_env(model))
|
||||
}
|
||||
|
||||
fn supports_inventory_refresh() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn inventory_identity() -> Result<InventoryIdentityInput> {
|
||||
let config = crate::config::Config::global();
|
||||
Ok(
|
||||
InventoryIdentityInput::new(OLLAMA_PROVIDER_NAME, OLLAMA_PROVIDER_NAME).with_public(
|
||||
"host",
|
||||
config
|
||||
.get_param::<String>("OLLAMA_HOST")
|
||||
.unwrap_or_else(|_| OLLAMA_HOST.to_string()),
|
||||
),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
|
||||
@@ -7,6 +7,7 @@ use super::formats::openai_responses::{
|
||||
create_responses_request, get_responses_usage, responses_api_to_message,
|
||||
responses_api_to_streaming_message, ResponsesApiResponse,
|
||||
};
|
||||
use super::inventory::{config_secret_value, InventoryIdentityInput};
|
||||
use super::openai_compatible::{
|
||||
handle_response_openai_compat, handle_status_openai_compat, stream_openai_compat,
|
||||
};
|
||||
@@ -425,6 +426,58 @@ impl ProviderDef for OpenAiProvider {
|
||||
) -> BoxFuture<'static, Result<Self::Provider>> {
|
||||
Box::pin(Self::from_env(model))
|
||||
}
|
||||
|
||||
fn supports_inventory_refresh() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn inventory_configured() -> bool {
|
||||
let config = crate::config::Config::global();
|
||||
// If the host is explicitly set to something non-default, trust the user's
|
||||
// custom setup (e.g. a local server that doesn't require an API key).
|
||||
if let Ok(host) = config.get_param::<String>("OPENAI_HOST") {
|
||||
if host != "https://api.openai.com" {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
// Standard OpenAI endpoint requires an API key.
|
||||
config
|
||||
.get_secret::<serde_json::Value>("OPENAI_API_KEY")
|
||||
.is_ok()
|
||||
}
|
||||
|
||||
fn inventory_identity() -> Result<InventoryIdentityInput> {
|
||||
let config = crate::config::Config::global();
|
||||
let mut identity =
|
||||
InventoryIdentityInput::new(OPEN_AI_PROVIDER_NAME, OPEN_AI_PROVIDER_NAME)
|
||||
.with_public(
|
||||
"host",
|
||||
config
|
||||
.get_param::<String>("OPENAI_HOST")
|
||||
.unwrap_or_else(|_| "https://api.openai.com".to_string()),
|
||||
)
|
||||
.with_public(
|
||||
"base_path",
|
||||
config
|
||||
.get_param::<String>("OPENAI_BASE_PATH")
|
||||
.unwrap_or_else(|_| OPEN_AI_DEFAULT_BASE_PATH.to_string()),
|
||||
);
|
||||
|
||||
if let Ok(organization) = config.get_param::<String>("OPENAI_ORGANIZATION") {
|
||||
identity = identity.with_public("organization", organization);
|
||||
}
|
||||
if let Ok(project) = config.get_param::<String>("OPENAI_PROJECT") {
|
||||
identity = identity.with_public("project", project);
|
||||
}
|
||||
if let Some(api_key) = config_secret_value(config, "OPENAI_API_KEY") {
|
||||
identity = identity.with_secret("api_key", api_key);
|
||||
}
|
||||
if let Some(custom_headers) = config_secret_value(config, "OPENAI_CUSTOM_HEADERS") {
|
||||
identity = identity.with_secret("custom_headers", custom_headers);
|
||||
}
|
||||
|
||||
Ok(identity)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
|
||||
@@ -9,7 +9,9 @@ use crate::acp::{
|
||||
use crate::config::search_path::SearchPaths;
|
||||
use crate::config::{Config, GooseMode};
|
||||
use crate::model::ModelConfig;
|
||||
use crate::providers::acp_tooling::{acp_adapter_installed, acp_inventory_identity};
|
||||
use crate::providers::base::{ProviderDef, ProviderMetadata};
|
||||
use crate::providers::inventory::InventoryIdentityInput;
|
||||
|
||||
const PI_ACP_PROVIDER_NAME: &str = "pi-acp";
|
||||
const PI_ACP_DOC_URL: &str = "https://github.com/anthropics/pi";
|
||||
@@ -36,6 +38,7 @@ impl ProviderDef for PiAcpProvider {
|
||||
"Set in your goose config file (`~/.config/goose/config.yaml` on macOS/Linux):\n GOOSE_PROVIDER: pi-acp\n GOOSE_MODEL: current",
|
||||
"Restart goose for changes to take effect",
|
||||
])
|
||||
.with_model_selection_hint("Use the Pi CLI to configure models")
|
||||
}
|
||||
|
||||
fn from_env(
|
||||
@@ -70,4 +73,16 @@ impl ProviderDef for PiAcpProvider {
|
||||
AcpProvider::connect(metadata.name, model, goose_mode, provider_config).await
|
||||
})
|
||||
}
|
||||
|
||||
fn supports_inventory_refresh() -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn inventory_identity() -> Result<InventoryIdentityInput> {
|
||||
acp_inventory_identity(PI_ACP_PROVIDER_NAME, PI_ACP_BINARY)
|
||||
}
|
||||
|
||||
fn inventory_configured() -> bool {
|
||||
acp_adapter_installed(PI_ACP_BINARY)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use super::base::{ModelInfo, Provider, ProviderDef, ProviderMetadata, ProviderType};
|
||||
use super::inventory::InventoryIdentityInput;
|
||||
use crate::config::{DeclarativeProviderConfig, ExtensionConfig};
|
||||
use crate::model::ModelConfig;
|
||||
use anyhow::Result;
|
||||
@@ -14,12 +15,20 @@ pub type ProviderConstructor = Arc<
|
||||
|
||||
pub type ProviderCleanup = Arc<dyn Fn() -> BoxFuture<'static, Result<()>> + Send + Sync>;
|
||||
|
||||
pub type ProviderInventoryIdentityResolver =
|
||||
Arc<dyn Fn() -> Result<InventoryIdentityInput> + Send + Sync>;
|
||||
|
||||
pub type ProviderInventoryConfiguredResolver = Arc<dyn Fn() -> bool + Send + Sync>;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ProviderEntry {
|
||||
metadata: ProviderMetadata,
|
||||
pub(crate) constructor: ProviderConstructor,
|
||||
pub(crate) inventory_identity: ProviderInventoryIdentityResolver,
|
||||
pub(crate) inventory_configured: ProviderInventoryConfiguredResolver,
|
||||
pub(crate) cleanup: Option<ProviderCleanup>,
|
||||
provider_type: ProviderType,
|
||||
supports_inventory_refresh: bool,
|
||||
}
|
||||
|
||||
impl ProviderEntry {
|
||||
@@ -27,6 +36,22 @@ impl ProviderEntry {
|
||||
&self.metadata
|
||||
}
|
||||
|
||||
pub fn provider_type(&self) -> ProviderType {
|
||||
self.provider_type
|
||||
}
|
||||
|
||||
pub fn supports_inventory_refresh(&self) -> bool {
|
||||
self.supports_inventory_refresh
|
||||
}
|
||||
|
||||
pub fn inventory_identity(&self) -> Result<InventoryIdentityInput> {
|
||||
(self.inventory_identity)()
|
||||
}
|
||||
|
||||
pub fn inventory_configured(&self) -> bool {
|
||||
(self.inventory_configured)()
|
||||
}
|
||||
|
||||
fn normalize_model_config(&self, mut model: ModelConfig) -> ModelConfig {
|
||||
model = model.with_canonical_limits(&self.metadata.name);
|
||||
|
||||
@@ -92,24 +117,30 @@ impl ProviderRegistry {
|
||||
Ok(Arc::new(provider) as Arc<dyn Provider>)
|
||||
})
|
||||
}),
|
||||
inventory_identity: Arc::new(F::inventory_identity),
|
||||
inventory_configured: Arc::new(F::inventory_configured),
|
||||
cleanup: None,
|
||||
provider_type: if preferred {
|
||||
ProviderType::Preferred
|
||||
} else {
|
||||
ProviderType::Builtin
|
||||
},
|
||||
supports_inventory_refresh: F::supports_inventory_refresh(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
pub fn register_with_name<P, F>(
|
||||
pub fn register_with_name<P, F, G>(
|
||||
&mut self,
|
||||
config: &DeclarativeProviderConfig,
|
||||
provider_type: ProviderType,
|
||||
supports_inventory_refresh: bool,
|
||||
constructor: F,
|
||||
inventory_identity: G,
|
||||
) where
|
||||
P: ProviderDef + 'static,
|
||||
F: Fn(ModelConfig) -> Result<P::Provider> + Send + Sync + 'static,
|
||||
G: Fn() -> Result<InventoryIdentityInput> + Send + Sync + 'static,
|
||||
{
|
||||
let base_metadata = P::metadata();
|
||||
let description = config
|
||||
@@ -174,7 +205,9 @@ impl ProviderRegistry {
|
||||
model_doc_link: base_metadata.model_doc_link,
|
||||
config_keys,
|
||||
setup_steps: vec![],
|
||||
model_selection_hint: None,
|
||||
};
|
||||
let inventory_config_keys = custom_metadata.config_keys.clone();
|
||||
|
||||
self.entries.insert(
|
||||
config.name.clone(),
|
||||
@@ -187,8 +220,16 @@ impl ProviderRegistry {
|
||||
Ok(Arc::new(provider) as Arc<dyn Provider>)
|
||||
})
|
||||
}),
|
||||
inventory_identity: Arc::new(inventory_identity),
|
||||
inventory_configured: Arc::new(move || {
|
||||
super::inventory::default_inventory_configured(
|
||||
&inventory_config_keys,
|
||||
crate::config::Config::global(),
|
||||
)
|
||||
}),
|
||||
cleanup: None,
|
||||
provider_type,
|
||||
supports_inventory_refresh,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
@@ -19,7 +19,7 @@ use std::sync::{Arc, LazyLock};
|
||||
use tracing::{info, warn};
|
||||
use utoipa::ToSchema;
|
||||
|
||||
pub const CURRENT_SCHEMA_VERSION: i32 = 10;
|
||||
pub const CURRENT_SCHEMA_VERSION: i32 = 11;
|
||||
pub const SESSIONS_FOLDER: &str = "sessions";
|
||||
pub const DB_NAME: &str = "sessions.db";
|
||||
|
||||
@@ -717,6 +717,8 @@ impl SessionStorage {
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
crate::providers::inventory::create_tables(pool).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1060,6 +1062,9 @@ impl SessionStorage {
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
11 => {
|
||||
crate::providers::inventory::create_tables_in_tx(tx).await?;
|
||||
}
|
||||
_ => {
|
||||
anyhow::bail!("Unknown migration version: {}", version);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user