tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Declarative providers (#5084)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Michael Neale <michael.neale@gmail.com>
This commit is contained in:
Douwe Osinga
2025-10-15 09:48:14 -04:00
committed by GitHub
parent 925a042cb1
commit 9251da4314
28 changed files with 1219 additions and 802 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/auth.rs
+3
View File
@@ -10,6 +10,9 @@ pub async fn check_token(
request: Request,
next: Next,
) -> Result<Response, StatusCode> {
if request.uri().path() == "/status" {
return Ok(next.run(request).await);
}
let secret_key = request
.headers()
.get("X-Secret-Key")
crates/goose-server/src/openapi.rs
+11 -2
View File
@@ -5,7 +5,7 @@ use goose::config::permission::PermissionLevel;
use goose::config::ExtensionEntry;
use goose::conversation::Conversation;
use goose::permission::permission_confirmation::PrincipalType;
use goose::providers::base::{ConfigKey, ModelInfo, ProviderMetadata};
use goose::providers::base::{ConfigKey, ModelInfo, ProviderMetadata, ProviderType};
use goose::session::{Session, SessionInsights};
use rmcp::model::{
Annotations, Content, EmbeddedResource, Icon, ImageContent, JsonObject, RawAudioContent,
@@ -14,6 +14,9 @@ use rmcp::model::{
};
use utoipa::{OpenApi, ToSchema};
use goose::config::declarative_providers::{
DeclarativeProviderConfig, LoadedProvider, ProviderEngine,
};
use goose::conversation::message::{
ContextLengthExceeded, FrontendToolRequest, Message, MessageContent, MessageMetadata,
RedactedThinkingContent, SummarizationRequested, ThinkingContent, ToolConfirmationRequest,
@@ -335,6 +338,8 @@ derive_utoipa!(Icon as IconSchema);
super::routes::config_management::get_provider_models,
super::routes::config_management::upsert_permissions,
super::routes::config_management::create_custom_provider,
super::routes::config_management::get_custom_provider,
super::routes::config_management::update_custom_provider,
super::routes::config_management::remove_custom_provider,
super::routes::agent::start_agent,
super::routes::agent::resume_agent,
@@ -386,7 +391,7 @@ derive_utoipa!(Icon as IconSchema);
super::routes::config_management::ExtensionQuery,
super::routes::config_management::ToolPermission,
super::routes::config_management::UpsertPermissionsQuery,
super::routes::config_management::CreateCustomProviderRequest,
super::routes::config_management::UpdateCustomProviderRequest,
super::routes::reply::PermissionConfirmationRequest,
super::routes::reply::ChatRequest,
super::routes::context::ContextManageRequest,
@@ -420,6 +425,10 @@ derive_utoipa!(Icon as IconSchema);
JsonObjectSchema,
RoleSchema,
ProviderMetadata,
ProviderType,
LoadedProvider,
ProviderEngine,
DeclarativeProviderConfig,
ExtensionEntry,
ExtensionConfig,
ConfigKey,
crates/goose-server/src/routes/config_management.rs
+92 -133
View File
@@ -1,15 +1,17 @@
use crate::routes::utils::check_provider_configured;
use crate::state::AppState;
use axum::routing::put;
use axum::{
extract::Path,
routing::{delete, get, post},
Json, Router,
};
use goose::config::declarative_providers::LoadedProvider;
use goose::config::paths::Paths;
use goose::config::ExtensionEntry;
use goose::config::{Config, ConfigError};
use goose::model::ModelConfig;
use goose::providers::base::ProviderMetadata;
use goose::providers::base::{ProviderMetadata, ProviderType};
use goose::providers::pricing::{
get_all_pricing, get_model_pricing, parse_model_id, refresh_pricing,
};
@@ -57,6 +59,7 @@ pub struct ProviderDetails {
pub name: String,
pub metadata: ProviderMetadata,
pub is_configured: bool,
pub provider_type: ProviderType,
}
#[derive(Serialize, ToSchema)]
@@ -76,8 +79,8 @@ pub struct UpsertPermissionsQuery {
}
#[derive(Deserialize, ToSchema)]
pub struct CreateCustomProviderRequest {
pub provider_type: String,
pub struct UpdateCustomProviderRequest {
pub engine: String,
pub display_name: String,
pub api_url: String,
pub api_key: String,
@@ -225,9 +228,7 @@ pub async fn add_extension(
(status = 500, description = "Internal server error")
)
)]
pub async fn remove_extension(
axum::extract::Path(name): axum::extract::Path<String>,
) -> Result<Json<String>, StatusCode> {
pub async fn remove_extension(Path(name): Path<String>) -> Result<Json<String>, StatusCode> {
let key = goose::config::extensions::name_to_key(&name);
goose::config::remove_extension(&key);
Ok(Json(format!("Removed extension {}", name)))
@@ -258,72 +259,17 @@ pub async fn read_all_config() -> Result<Json<ConfigResponse>, StatusCode> {
)
)]
pub async fn providers() -> Result<Json<Vec<ProviderDetails>>, StatusCode> {
let mut providers_metadata = get_providers().await;
let custom_providers_dir = goose::config::custom_providers::custom_providers_dir();
if custom_providers_dir.exists() {
if let Ok(entries) = std::fs::read_dir(&custom_providers_dir) {
for entry in entries.flatten() {
if let Some(extension) = entry.path().extension() {
if extension == "json" {
if let Ok(content) = std::fs::read_to_string(entry.path()) {
if let Ok(custom_provider) = serde_json::from_str::<
goose::config::custom_providers::CustomProviderConfig,
>(&content)
{
// CustomProviderConfig => ProviderMetadata
let default_model = custom_provider
.models
.first()
.map(|m| m.name.clone())
.unwrap_or_default();
let metadata = goose::providers::base::ProviderMetadata {
name: custom_provider.name.clone(),
display_name: custom_provider.display_name.clone(),
description: custom_provider
.description
.clone()
.unwrap_or_else(|| {
format!("{} (custom)", custom_provider.display_name)
}),
default_model,
known_models: custom_provider.models.clone(),
model_doc_link: "Custom provider".to_string(),
config_keys: vec![
goose::providers::base::ConfigKey::new(
&custom_provider.api_key_env,
true,
true,
None,
),
goose::providers::base::ConfigKey::new(
"CUSTOM_PROVIDER_BASE_URL",
true,
false,
Some(&custom_provider.base_url),
),
],
};
providers_metadata.push(metadata);
}
}
}
}
}
}
}
let providers_response: Vec<ProviderDetails> = providers_metadata
let providers = get_providers().await;
let providers_response: Vec<ProviderDetails> = providers
.into_iter()
.map(|metadata| {
let is_configured = check_provider_configured(&metadata);
.map(|(metadata, provider_type)| {
let is_configured = check_provider_configured(&metadata, provider_type);
ProviderDetails {
name: metadata.name.clone(),
metadata,
is_configured,
provider_type,
}
})
.collect();
@@ -347,11 +293,28 @@ pub async fn providers() -> Result<Json<Vec<ProviderDetails>>, StatusCode> {
pub async fn get_provider_models(
Path(name): Path<String>,
) -> Result<Json<Vec<String>>, StatusCode> {
let all = get_providers().await;
let Some(metadata) = all.into_iter().find(|m| m.name == name) else {
let loaded_provider = goose::config::declarative_providers::load_provider(name.as_str()).ok();
// TODO(Douwe): support a get models url for custom providers
if let Some(loaded_provider) = loaded_provider {
return Ok(Json(
loaded_provider
.config
.models
.into_iter()
.map(|m| m.name)
.collect::<Vec<_>>(),
));
}
let all = get_providers()
.await
.into_iter()
//.map(|(m, p)| m)
.collect::<Vec<_>>();
let Some((metadata, provider_type)) = all.into_iter().find(|(m, _)| m.name == name) else {
return Err(StatusCode::BAD_REQUEST);
};
if !check_provider_configured(&metadata) {
if !check_provider_configured(&metadata, provider_type) {
return Err(StatusCode::BAD_REQUEST);
}
@@ -449,12 +412,9 @@ pub async fn get_pricing(
}
}
} else {
// Get only configured providers' pricing
let providers_metadata = get_providers().await;
for metadata in providers_metadata {
for (metadata, provider_type) in get_providers().await {
// Skip unconfigured providers if filtering
if !check_provider_configured(&metadata) {
if !check_provider_configured(&metadata, provider_type) {
continue;
}
@@ -647,25 +607,10 @@ pub async fn validate_config() -> Result<Json<String>, StatusCode> {
}
}
#[utoipa::path(
get,
path = "/config/current-model",
responses(
(status = 200, description = "Current model retrieved successfully", body = String),
)
)]
pub async fn get_current_model() -> Result<Json<Value>, StatusCode> {
let current_model = goose::providers::base::get_current_model();
Ok(Json(serde_json::json!({
"model": current_model
})))
}
#[utoipa::path(
post,
path = "/config/custom-providers",
request_body = CreateCustomProviderRequest,
request_body = UpdateCustomProviderRequest,
responses(
(status = 200, description = "Custom provider created successfully", body = String),
(status = 400, description = "Invalid request"),
@@ -673,10 +618,10 @@ pub async fn get_current_model() -> Result<Json<Value>, StatusCode> {
)
)]
pub async fn create_custom_provider(
Json(request): Json<CreateCustomProviderRequest>,
Json(request): Json<UpdateCustomProviderRequest>,
) -> Result<Json<String>, StatusCode> {
let config = goose::config::custom_providers::CustomProviderConfig::create_and_save(
&request.provider_type,
let config = goose::config::declarative_providers::create_custom_provider(
&request.engine,
request.display_name,
request.api_url,
request.api_key,
@@ -692,6 +637,24 @@ pub async fn create_custom_provider(
Ok(Json(format!("Custom provider added - ID: {}", config.id())))
}
#[utoipa::path(
get,
path = "/config/custom-providers/{id}",
responses(
(status = 200, description = "Custom provider retrieved successfully", body = LoadedProvider),
(status = 404, description = "Provider not found"),
(status = 500, description = "Internal server error")
)
)]
pub async fn get_custom_provider(
Path(id): Path<String>,
) -> Result<Json<LoadedProvider>, StatusCode> {
let loaded_provider = goose::config::declarative_providers::load_provider(id.as_str())
.map_err(|_| StatusCode::NOT_FOUND)?;
Ok(Json(loaded_provider))
}
#[utoipa::path(
delete,
path = "/config/custom-providers/{id}",
@@ -701,10 +664,8 @@ pub async fn create_custom_provider(
(status = 500, description = "Internal server error")
)
)]
pub async fn remove_custom_provider(
axum::extract::Path(id): axum::extract::Path<String>,
) -> Result<Json<String>, StatusCode> {
goose::config::custom_providers::CustomProviderConfig::remove(&id)
pub async fn remove_custom_provider(Path(id): Path<String>) -> Result<Json<String>, StatusCode> {
goose::config::declarative_providers::remove_custom_provider(&id)
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
if let Err(e) = goose::providers::refresh_custom_providers().await {
@@ -714,6 +675,38 @@ pub async fn remove_custom_provider(
Ok(Json(format!("Removed custom provider: {}", id)))
}
#[utoipa::path(
put,
path = "/config/custom-providers/{id}",
request_body = UpdateCustomProviderRequest,
responses(
(status = 200, description = "Custom provider updated successfully", body = String),
(status = 404, description = "Provider not found"),
(status = 500, description = "Internal server error")
)
)]
pub async fn update_custom_provider(
Path(id): Path<String>,
Json(request): Json<UpdateCustomProviderRequest>,
) -> Result<Json<String>, StatusCode> {
goose::config::declarative_providers::update_custom_provider(
&id,
&request.engine,
request.display_name,
request.api_url,
request.api_key,
request.models,
request.supports_streaming,
)
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
if let Err(e) = goose::providers::refresh_custom_providers().await {
tracing::warn!("Failed to refresh custom providers after update: {}", e);
}
Ok(Json(format!("Updated custom provider: {}", id)))
}
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/config", get(read_all_config))
@@ -731,12 +724,13 @@ pub fn routes(state: Arc<AppState>) -> Router {
.route("/config/recover", post(recover_config))
.route("/config/validate", get(validate_config))
.route("/config/permissions", post(upsert_permissions))
.route("/config/current-model", get(get_current_model))
.route("/config/custom-providers", post(create_custom_provider))
.route(
"/config/custom-providers/{id}",
delete(remove_custom_provider),
)
.route("/config/custom-providers/{id}", put(update_custom_provider))
.route("/config/custom-providers/{id}", get(get_custom_provider))
.with_state(state)
}
@@ -768,39 +762,4 @@ mod tests {
assert!(gpt4_limit.is_some());
assert_eq!(gpt4_limit.unwrap().context_limit, 128_000);
}
#[tokio::test]
async fn test_get_provider_models_unknown_provider() {
let mut headers = HeaderMap::new();
headers.insert("X-Secret-Key", "test".parse().unwrap());
let result = get_provider_models(Path("unknown_provider".to_string())).await;
assert!(result.is_err());
assert_eq!(result.unwrap_err(), StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn test_get_provider_models_openai_configured() {
std::env::set_var("OPENAI_API_KEY", "test-key");
let mut headers = HeaderMap::new();
headers.insert("X-Secret-Key", "test".parse().unwrap());
let result = get_provider_models(Path("openai".to_string())).await;
// The response should be BAD_REQUEST since the API key is invalid (authentication error)
assert!(
result.is_err(),
"Expected error response from OpenAI provider with invalid key"
);
let status_code = result.unwrap_err();
assert!(status_code == StatusCode::BAD_REQUEST,
"Expected BAD_REQUEST (authentication error) or INTERNAL_SERVER_ERROR (other errors), got: {}",
status_code
);
std::env::remove_var("OPENAI_API_KEY");
}
}
crates/goose-server/src/routes/utils.rs
+13 -3
View File
@@ -1,5 +1,6 @@
use goose::config::declarative_providers::load_provider;
use goose::config::Config;
use goose::providers::base::{ConfigKey, ProviderMetadata};
use goose::providers::base::{ConfigKey, ProviderMetadata, ProviderType};
use serde::{Deserialize, Serialize};
use std::env;
use std::error::Error;
@@ -27,7 +28,7 @@ pub fn inspect_key(key_name: &str, is_secret: bool) -> Result<KeyInfo, Box<dyn E
let config = Config::global();
// Check environment variable first
let env_value = std::env::var(key_name).ok();
let env_value = env::var(key_name).ok();
if let Some(value) = env_value {
return Ok(KeyInfo {
@@ -90,9 +91,18 @@ pub fn inspect_keys(
Ok(results)
}
pub fn check_provider_configured(metadata: &ProviderMetadata) -> bool {
pub fn check_provider_configured(metadata: &ProviderMetadata, provider_type: ProviderType) -> bool {
let config = Config::global();
// TODO(Douwe): if the provider doesn't need an API key, it should be considered configured always
if provider_type == ProviderType::Custom || provider_type == ProviderType::Declarative {
if let Ok(loaded_provider) = load_provider(metadata.name.as_str()) {
return config
.get_secret::<String>(&loaded_provider.config.api_key_env)
.map(|s| !s.is_empty())
.unwrap_or(false);
}
}
// Special case: Zero-config providers (no config keys)
if metadata.config_keys.is_empty() {
// Check if the provider has been explicitly configured via the UI