github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-15 06:51:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(ai): unify model discovery and effort handling

Browse Source
This commit is contained in:
Guoyu Su authored and GitHub committed 2026-07-28 21:37:55 +08:00
1 parent c54b95e274
commit c7e7b7c9d9
41 files changed
+3291 -634
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/dbx-core/src/agent_events.rs
+4
View File
@@ -36,6 +36,10 @@ pub struct ToolCall {
pub id: String,
pub name: String,
pub arguments: serde_json::Value,
/// Opaque provider response data required to replay this tool call in a
/// follow-up request (for example, Gemini thought signatures).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_payload: Option<serde_json::Value>,
}
/// Result of executing a tool.
crates/dbx-core/src/agent_loop.rs
+127 -15
View File
@@ -84,14 +84,33 @@ pub struct AgentLoopContext {
pub max_agent_turns: u32,
}
/// Check if the provider supports function calling / tool use.
/// Returns false for providers that are known to lack reliable tool support.
fn provider_supports_function_calling(config: &AiConfig) -> bool {
match config.provider {
// Ollama function calling support varies by model/version; conservative default is false.
// Users with capable models can override via openai-compatible with an Ollama endpoint.
AiProvider::Ollama => false,
_ => true,
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum FunctionCallingSupport {
Supported,
Unsupported,
Unknown,
}
/// Resolve function calling support for the selected provider/model.
///
/// Native Ollama capabilities are model-specific. Missing metadata is kept as
/// unknown so older servers and compatible proxies can still attempt tools.
async fn provider_function_calling_support(config: &AiConfig) -> FunctionCallingSupport {
if !matches!(config.provider, AiProvider::Ollama) {
return FunctionCallingSupport::Supported;
}
match ai::ollama_selected_model_tool_support(config).await {
Ok(Some(true)) => FunctionCallingSupport::Supported,
Ok(Some(false)) => FunctionCallingSupport::Unsupported,
Ok(None) => FunctionCallingSupport::Unknown,
Err(error) => {
log::debug!(
"[agent][ollama] tool capability unavailable for model {}: {error}; attempting tool call",
config.model
);
FunctionCallingSupport::Unknown
}
}
}
@@ -100,8 +119,9 @@ fn provider_supports_function_calling(config: &AiConfig) -> bool {
/// The `on_event` callback receives streaming events for the frontend.
/// Returns the final accumulated assistant text.
///
/// If the provider does not support function calling (e.g., Ollama), automatically
/// degrades to a text-only completion with schema context injected into the system prompt.
/// If the selected model explicitly does not support function calling,
/// automatically degrades to a text-only completion with schema context
/// injected into the system prompt.
#[allow(clippy::too_many_arguments)]
pub async fn run_agent_loop(
config: &AiConfig,
@@ -149,8 +169,19 @@ pub async fn run_agent_loop(
return crate::ai_codex_cli::run_codex_agent(config, &prompt, options, cancelled, on_event).await;
}
// Auto-degrade: providers without function calling fall back to text-only completion.
if !provider_supports_function_calling(config) {
// Auto-degrade only models that explicitly advertise no tool support.
// Unknown Ollama capabilities still get a chance to use tools so older
// servers and compatible proxies are not disabled provider-wide.
let function_calling_support = tokio::select! {
support = provider_function_calling_support(config) => support,
_ = cancelled.notified() => {
let message = "Agent run was cancelled before producing output.".to_string();
on_event(AgentEvent::TextDelta { delta: message.clone() });
on_event(AgentEvent::AgentEnd { input_tokens: None, output_tokens: None });
return Ok(message);
}
};
if function_calling_support == FunctionCallingSupport::Unsupported {
return run_agent_loop_text_only(
config,
system_prompt,
@@ -248,6 +279,29 @@ pub async fn run_agent_loop(
stream_result = Some((tool_calls, usage, accumulated_text));
break;
}
Err(err)
if turn == 0
&& matches!(config.provider, AiProvider::Ollama)
&& is_tool_unsupported_error(&err)
&& !emitted_any_chunk.load(Ordering::Relaxed) =>
{
log::debug!(
"[agent][ollama] model {} rejected tools before producing output; falling back to text-only mode",
config.model
);
on_event(AgentEvent::TurnEnd { turn });
return run_agent_loop_text_only(
config,
system_prompt,
messages,
agent_ctx,
on_event,
cancelled,
max_tokens,
task_contract.as_ref(),
)
.await;
}
Err(err)
if attempt == 0 && is_context_length_error(&err) && !emitted_any_chunk.load(Ordering::Relaxed) =>
{
@@ -315,7 +369,12 @@ pub async fn run_agent_loop(
tool_call_id: None,
tool_calls: collected_tool_calls
.iter()
.map(|tc| ai::ToolCallRef { id: tc.id.clone(), name: tc.name.clone(), arguments: tc.arguments.clone() })
.map(|tc| ai::ToolCallRef {
id: tc.id.clone(),
name: tc.name.clone(),
arguments: tc.arguments.clone(),
provider_payload: tc.provider_payload.clone(),
})
.collect(),
});
@@ -377,8 +436,12 @@ pub async fn run_agent_loop(
let (parallel_indices, sequential_indices): (Vec<usize>, Vec<usize>) = (0..collected_tool_calls.len())
.partition(|&i| *tool_parallel_map.get(collected_tool_calls[i].name.as_str()).unwrap_or(&false));
let make_tc =
|tc: &ToolCall| ToolCall { id: tc.id.clone(), name: tc.name.clone(), arguments: tc.arguments.clone() };
let make_tc = |tc: &ToolCall| ToolCall {
id: tc.id.clone(),
name: tc.name.clone(),
arguments: tc.arguments.clone(),
provider_payload: tc.provider_payload.clone(),
};
// Run parallel group
let parallel_futures: Vec<_> = parallel_indices
@@ -692,6 +755,24 @@ fn is_context_length_error(error: &str) -> bool {
.any(|marker| lower.contains(marker))
}
fn is_tool_unsupported_error(error: &str) -> bool {
let lower = error.to_lowercase();
let mentions_tool_use = lower.contains("tool") || lower.contains("function call");
let rejects_capability = [
"does not support",
"doesn't support",
"not supported",
"unsupported",
"unknown field",
"unknown parameter",
"unrecognized field",
"unrecognized parameter",
]
.iter()
.any(|marker| lower.contains(marker));
mentions_tool_use && rejects_capability
}
/// Text-only fallback for providers that don't support function calling.
///
/// Injects database schema context into the system prompt so the LLM can still
@@ -836,6 +917,11 @@ fn estimate_message_tokens(message: &AiMessage) -> u32 {
if let Ok(args) = serde_json::to_string(&tool_call.arguments) {
tokens += estimate_text_tokens(&args);
}
if let Some(provider_payload) = &tool_call.provider_payload {
if let Ok(payload) = serde_json::to_string(provider_payload) {
tokens += estimate_text_tokens(&payload);
}
}
}
tokens
@@ -1201,6 +1287,32 @@ mod tests {
assert_eq!(clamp_max_agent_turns(u32::MAX), MAX_MAX_AGENT_TURNS);
}
#[test]
fn identifies_explicit_tool_rejection_errors() {
for error in [
"model does not support tools",
"tool use is not supported by this model",
"unsupported parameter: tools",
"unknown field `tools`",
"function calling is not supported",
] {
assert!(is_tool_unsupported_error(error), "{error}");
}
}
#[test]
fn does_not_hide_unrelated_ollama_errors_as_tool_rejections() {
for error in [
"connection refused",
"model not found",
"context length exceeded",
"invalid tool arguments returned by model",
"request timed out",
] {
assert!(!is_tool_unsupported_error(error), "{error}");
}
}
fn generate_contract(user_request: &str, mode: &str) -> AiTaskContract {
AiTaskContract {
action: Some("generate".to_string()),
crates/dbx-core/src/agent_tools.rs
+1
View File
@@ -662,6 +662,7 @@ async fn execute_get_sample_data(
id: tool_call.id.clone(),
name: "execute_query".to_string(),
arguments: serde_json::json!({ "sql": sql, "limit": limit }),
provider_payload: None,
};
execute_execute_query(&synthetic_call, state, connection_id, database, db_type, AgentSqlPermissions::default())
.await
crates/dbx-core/src/ai.rs
+921 -78
View File
File diff suppressed because it is too large. Load diff
crates/dbx-core/src/ai_claude_code_cli.rs
+58 -23
View File
@@ -1,5 +1,5 @@
use crate::agent_events::AgentEvent;
use crate::ai::{AiConfig, AiEffortLevel, AiModelInfo, AiTestConnectionResult};
use crate::ai::{AiCapabilitySource, AiConfig, AiEffortCapability, AiEffortLevel, AiModelInfo, AiTestConnectionResult};
use crate::ai_cli_agent::{
build_cli_agent_prompt, cli_command, dbx_mcp_enabled_tools, dbx_mcp_scope_env, model_infos, parse_cli_jsonl_event,
run_cli_jsonl_agent, CliAgentCommandSpec, CliAgentJsonlDialect, CliAgentProcessSpec, CliAgentRunOptions,
@@ -284,9 +284,13 @@ fn build_claude_code_command_with_mcp_arg(
args.push("--model".to_string());
args.push(model.to_string());
}
if let Some(effort) = config.reasoning_level.as_claude_code_effort() {
let effort = match config.runtime_effort.as_ref() {
Some(effort) => effort.cli_value(),
None => config.reasoning_level.as_claude_code_effort().map(ToString::to_string),
};
if let Some(effort) = effort {
args.push("--effort".to_string());
args.push(effort.to_string());
args.push(effort);
}
ClaudeCodeCommandSpec { program: claude_code_program(config), args }
@@ -387,7 +391,17 @@ fn parse_claude_code_models(stdout: &str) -> Option<Vec<AiModelInfo>> {
.filter(|name| !name.is_empty())
.map(ToString::to_string);
let mut info = AiModelInfo::new(id, display_name);
info.supported_effort_levels = parse_claude_code_effort_levels(model);
let effort_levels = parse_claude_code_effort_level_strings(model);
info.supported_effort_levels =
effort_levels.iter().filter_map(|level| level.parse::<AiEffortLevel>().ok()).collect();
info.effort_capability =
if model.get("supportsEffort").or_else(|| model.get("supports_effort")).and_then(Value::as_bool)
== Some(false)
{
Some(AiEffortCapability::Unsupported)
} else {
crate::ai_effort::dynamic_enum_capability(effort_levels, AiCapabilitySource::LocalCli)
};
result.push(info);
}
@@ -403,7 +417,7 @@ fn parse_claude_code_models(stdout: &str) -> Option<Vec<AiModelInfo>> {
None
}
fn parse_claude_code_effort_levels(model: &Value) -> Vec<AiEffortLevel> {
fn parse_claude_code_effort_level_strings(model: &Value) -> Vec<String> {
if model.get("supportsEffort").or_else(|| model.get("supports_effort")).and_then(Value::as_bool) == Some(false) {
return Vec::new();
}
@@ -417,8 +431,10 @@ fn parse_claude_code_effort_levels(model: &Value) -> Vec<AiEffortLevel> {
levels
.iter()
.filter_map(Value::as_str)
.filter_map(|level| level.parse::<AiEffortLevel>().ok())
.filter(|level| seen.insert(*level))
.map(str::trim)
.filter(|level| !level.is_empty())
.filter(|level| seen.insert((*level).to_string()))
.map(ToString::to_string)
.collect()
}
@@ -532,7 +548,10 @@ mod tests {
#[cfg(unix)]
use super::{list_claude_code_models, run_claude_code_agent};
use crate::agent_events::AgentEvent;
use crate::ai::{AiApiStyle, AiAuthMethod, AiConfig, AiEffortLevel, AiModelInfo, AiProvider, AiReasoningLevel};
use crate::ai::{
AiApiStyle, AiAuthMethod, AiConfig, AiEffortCapability, AiEffortLevel, AiEffortSelection, AiProvider,
AiReasoningLevel,
};
use crate::ai_cli_agent::{model_infos, CliAgentCommandSpec};
#[cfg(unix)]
use std::os::unix::fs::PermissionsExt;
@@ -552,6 +571,7 @@ mod tests {
proxy_url: String::new(),
enable_thinking: true,
reasoning_level: AiReasoningLevel::Default,
runtime_effort: None,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
@@ -736,6 +756,21 @@ esac
assert!(!spec.args.contains(&"--effort".to_string()));
}
#[test]
fn runtime_effort_takes_priority_over_legacy_reasoning_level() {
let mut config = claude_code_config("sonnet");
config.reasoning_level = AiReasoningLevel::Xhigh;
config.runtime_effort = Some(AiEffortSelection::ProviderDefault);
let spec = build_claude_code_command(&config, "hello", &run_options());
assert!(!spec.args.contains(&"--effort".to_string()));
config.runtime_effort = Some(AiEffortSelection::Enum("future".to_string()));
let spec = build_claude_code_command(&config, "hello", &run_options());
let effort_pos = spec.args.iter().position(|arg| arg == "--effort").unwrap();
assert_eq!(spec.args[effort_pos + 1], "future");
}
#[test]
fn default_model_list_matches_supported_aliases() {
assert_eq!(model_infos(DEFAULT_CLAUDE_CODE_MODELS), model_infos(&["default", "sonnet", "opus", "fable"]));
@@ -755,21 +790,21 @@ esac
let models = parse_claude_code_models(stdout).unwrap();
assert_eq!(
models,
vec![
AiModelInfo::new("default", Some("Default".to_string())),
AiModelInfo {
id: "claude-sonnet-4-6".to_string(),
display_name: Some("Sonnet 4.6".to_string()),
supported_effort_levels: vec![
AiEffortLevel::Low,
AiEffortLevel::Medium,
AiEffortLevel::High,
AiEffortLevel::Max
],
},
AiModelInfo::new("claude-opus-4-8", Some("Opus 4.8".to_string())),
]
models.iter().map(|model| model.id.as_str()).collect::<Vec<_>>(),
["default", "claude-sonnet-4-6", "claude-opus-4-8"]
);
let sonnet = &models[1];
assert_eq!(sonnet.display_name.as_deref(), Some("Sonnet 4.6"));
assert_eq!(
sonnet.supported_effort_levels,
[AiEffortLevel::Low, AiEffortLevel::Medium, AiEffortLevel::High, AiEffortLevel::Max]
);
let AiEffortCapability::Enum { options, .. } = sonnet.effort_capability.as_ref().unwrap() else {
panic!("expected effort enum");
};
assert_eq!(
options.iter().map(|option| option.id.as_str()).collect::<Vec<_>>(),
["low", "medium", "high", "max", "future"]
);
}
crates/dbx-core/src/ai_codex_cli.rs
+39 -5
View File
@@ -1,5 +1,5 @@
use crate::agent_events::AgentEvent;
use crate::ai::{AiConfig, AiModelInfo, AiTestConnectionResult};
use crate::ai::{AiCapabilitySource, AiConfig, AiModelInfo, AiTestConnectionResult};
use crate::ai_cli_agent::{
append_config_overrides, build_cli_agent_prompt, cli_command, dbx_mcp_enabled_tools, dbx_mcp_scope_env,
model_infos, parse_cli_jsonl_event, run_cli_jsonl_agent, toml_string, toml_string_array, CliAgentCommandSpec,
@@ -417,8 +417,12 @@ pub fn build_codex_exec_command(config: &AiConfig, _prompt: &str, options: &Code
"read-only".to_string(),
];
let mut config_overrides = vec!["features.shell_tool=false".to_string(), "web_search=\"disabled\"".to_string()];
if let Some(reasoning_effort) = config.reasoning_level.as_codex_effort() {
config_overrides.push(format!("model_reasoning_effort={}", toml_string(reasoning_effort)));
let reasoning_effort = match config.runtime_effort.as_ref() {
Some(effort) => effort.cli_value(),
None => config.reasoning_level.as_codex_effort().map(ToString::to_string),
};
if let Some(reasoning_effort) = reasoning_effort {
config_overrides.push(format!("model_reasoning_effort={}", toml_string(&reasoning_effort)));
}
append_config_overrides(&mut args, config_overrides.into_iter().chain(codex_mcp_config_overrides(options)));
@@ -483,7 +487,21 @@ fn parse_codex_models(stdout: &str) -> Option<Vec<AiModelInfo>> {
.map(str::trim)
.filter(|name| !name.is_empty())
.map(ToString::to_string);
result.push(AiModelInfo::new(id, display_name));
let mut info = AiModelInfo::new(id, display_name);
let levels = model
.get("supported_reasoning_levels")
.or_else(|| model.get("supportedReasoningLevels"))
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(|level| {
level
.as_str()
.or_else(|| level.get("effort").and_then(Value::as_str))
.or_else(|| level.get("value").and_then(Value::as_str))
});
info.effort_capability = crate::ai_effort::dynamic_enum_capability(levels, AiCapabilitySource::LocalCli);
result.push(info);
}
(result.len() > 1).then_some(result)
@@ -596,7 +614,7 @@ mod tests {
windows_npm_codex_shim_command,
};
use crate::agent_events::AgentEvent;
use crate::ai::{AiApiStyle, AiAuthMethod, AiConfig, AiProvider, AiReasoningLevel};
use crate::ai::{AiApiStyle, AiAuthMethod, AiConfig, AiEffortSelection, AiProvider, AiReasoningLevel};
use crate::ai_cli_agent::{model_infos, CliAgentCommandSpec};
fn codex_config(model: &str) -> AiConfig {
@@ -612,6 +630,7 @@ mod tests {
proxy_url: String::new(),
enable_thinking: true,
reasoning_level: AiReasoningLevel::Default,
runtime_effort: None,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
@@ -687,6 +706,21 @@ mod tests {
assert!(spec.args.contains(&"model_reasoning_effort=\"high\"".to_string()));
}
#[test]
fn runtime_effort_takes_priority_over_legacy_reasoning_level() {
let mut config = codex_config("default");
config.reasoning_level = AiReasoningLevel::High;
config.runtime_effort = Some(AiEffortSelection::ProviderDefault);
let spec = build_codex_exec_command(&config, "hello", &run_options());
assert!(!spec.args.iter().any(|arg| arg.starts_with("model_reasoning_effort=")));
config.runtime_effort = Some(AiEffortSelection::Enum("xhigh".to_string()));
let spec = build_codex_exec_command(&config, "hello", &run_options());
assert!(spec.args.contains(&"model_reasoning_effort=\"xhigh\"".to_string()));
}
#[test]
fn path_like_codex_programs_are_detected() {
assert!(is_path_like_program("/opt/homebrew/bin/codex"));
crates/dbx-core/src/ai_effort.rs
+606
View File
@@ -0,0 +1,606 @@
use crate::ai::{
AiApiStyle, AiCapabilitySource, AiConfig, AiEffortCapability, AiEffortOption, AiEffortSelection, AiProvider,
};
use serde_json::{json, Map, Value};
const OPENAI_REASONING_DOCS: &str = "https://platform.openai.com/docs/guides/reasoning";
const GEMINI_THINKING_DOCS: &str = "https://ai.google.dev/gemini-api/docs/thinking";
const DEEPSEEK_THINKING_DOCS: &str = "https://api-docs.deepseek.com/guides/thinking_mode";
const QWEN_THINKING_DOCS: &str = "https://help.aliyun.com/en/model-studio/deep-thinking";
const OLLAMA_THINKING_DOCS: &str = "https://docs.ollama.com/capabilities/thinking";
pub const EFFORT_REGISTRY_LAST_VERIFIED: &str = "2026-07-26";
fn option(id: &str, label: &str, selection: AiEffortSelection) -> AiEffortOption {
AiEffortOption { id: id.to_string(), label: label.to_string(), description: None, selection }
}
fn enum_capability(values: &[&str], source: AiCapabilitySource) -> AiEffortCapability {
let options = values
.iter()
.map(|value| option(value, &title_case_effort(value), AiEffortSelection::Enum((*value).to_string())))
.collect::<Vec<_>>();
let default = options.first().map(|option| option.selection.clone()).unwrap_or(AiEffortSelection::ProviderDefault);
AiEffortCapability::Enum { options, default, source }
}
pub fn dynamic_enum_capability<I, S>(values: I, source: AiCapabilitySource) -> Option<AiEffortCapability>
where
I: IntoIterator<Item = S>,
S: AsRef<str>,
{
let mut seen = Vec::<String>::new();
for value in values {
let value = value.as_ref().trim();
if value.is_empty() || seen.iter().any(|existing| existing == value) {
continue;
}
seen.push(value.to_string());
}
(!seen.is_empty()).then(|| {
let refs = seen.iter().map(String::as_str).collect::<Vec<_>>();
enum_capability(&refs, source)
})
}
fn integer_capability(min: i64, max: i64, allow_disabled: bool, source: AiCapabilitySource) -> AiEffortCapability {
let mut special_values = vec![option("auto", "Auto", AiEffortSelection::Integer(-1))];
if allow_disabled {
special_values.push(option("off", "Off", AiEffortSelection::Disabled));
}
AiEffortCapability::Integer { min, max, step: 1, default: AiEffortSelection::Integer(-1), special_values, source }
}
fn boolean_capability(source: AiCapabilitySource) -> AiEffortCapability {
AiEffortCapability::Boolean { default: AiEffortSelection::ProviderDefault, source }
}
fn normalized_model_id(model_id: &str) -> String {
model_id.trim().trim_start_matches("models/").to_ascii_lowercase()
}
fn matches_family(model: &str, family: &str) -> bool {
model == family
|| model
.strip_prefix(family)
.is_some_and(|suffix| suffix.starts_with('-') || suffix.starts_with(':') || suffix.starts_with('@'))
}
pub fn static_effort_capability(config: &AiConfig, model_id: &str) -> Option<AiEffortCapability> {
let model = normalized_model_id(model_id);
let source = AiCapabilitySource::OfficialRegistry;
match config.provider {
AiProvider::Openai => openai_capability(&model, source),
AiProvider::Gemini => gemini_capability(&model, source),
AiProvider::Deepseek => deepseek_capability(&model, source),
AiProvider::Qwen => qwen_capability(&model, source),
AiProvider::Ollama => ollama_capability(&model, source),
AiProvider::OpenaiCompatible | AiProvider::Custom => {
Some(AiEffortCapability::FreeText { placeholder: None, source: AiCapabilitySource::Custom })
}
AiProvider::Claude | AiProvider::CodexCli | AiProvider::ClaudeCodeCli => None,
}
}
fn openai_capability(model: &str, source: AiCapabilitySource) -> Option<AiEffortCapability> {
if matches_family(model, "gpt-5-pro") {
return Some(enum_capability(&["high"], source));
}
if matches_family(model, "gpt-5.6") {
return Some(enum_capability(&["none", "low", "medium", "high", "xhigh", "max"], source));
}
if matches_family(model, "gpt-5.1") {
return Some(enum_capability(&["none", "low", "medium", "high"], source));
}
if matches_family(model, "gpt-5") || matches_family(model, "gpt-5.2") || matches_family(model, "gpt-5.4") {
return Some(enum_capability(&["minimal", "low", "medium", "high", "xhigh"], source));
}
if ["o1", "o3", "o3-mini", "o4-mini"].iter().any(|family| matches_family(model, family)) {
return Some(enum_capability(&["low", "medium", "high"], source));
}
None
}
fn gemini_capability(model: &str, source: AiCapabilitySource) -> Option<AiEffortCapability> {
if matches_family(model, "gemini-2.5-pro") {
return Some(integer_capability(128, 32_768, false, source));
}
if matches_family(model, "gemini-2.5-flash-lite") {
return Some(integer_capability(512, 24_576, true, source));
}
if matches_family(model, "gemini-2.5-flash")
|| matches_family(model, "robotics-er-1.6-preview")
|| model.starts_with("gemini-2.5-flash-live")
{
return Some(integer_capability(0, 24_576, true, source));
}
if model.starts_with("gemini-3.6-flash")
|| model.starts_with("gemini-3.5-flash")
|| model.starts_with("gemini-3-flash")
{
return Some(enum_capability(&["minimal", "low", "medium", "high"], source));
}
if model.starts_with("gemini-3.1-pro") {
return Some(enum_capability(&["low", "medium", "high"], source));
}
if model.starts_with("gemini-3.1-flash-lite-image") {
return Some(enum_capability(&["minimal", "high"], source));
}
if model.starts_with("gemini-3-pro") {
return Some(enum_capability(&["low", "high"], source));
}
None
}
fn deepseek_capability(model: &str, source: AiCapabilitySource) -> Option<AiEffortCapability> {
if matches_family(model, "deepseek-v4-flash") || matches_family(model, "deepseek-v4-pro") {
let mut capability = enum_capability(&["high", "max"], source);
if let AiEffortCapability::Enum { options, default, .. } = &mut capability {
options.insert(0, option("off", "Off", AiEffortSelection::Disabled));
*default = AiEffortSelection::Disabled;
}
return Some(capability);
}
None
}
fn qwen_capability(model: &str, source: AiCapabilitySource) -> Option<AiEffortCapability> {
if matches_family(model, "qwen3.8-max-preview") {
return Some(enum_capability(&["low", "medium", "xhigh"], source));
}
if model.starts_with("qwen3") || model.starts_with("qwq") {
return Some(boolean_capability(source));
}
None
}
fn ollama_capability(model: &str, source: AiCapabilitySource) -> Option<AiEffortCapability> {
if matches_family(model, "gpt-oss") {
return Some(enum_capability(&["low", "medium", "high"], source));
}
if model.starts_with("deepseek-r1")
|| model.starts_with("qwen3")
|| model.starts_with("qwq")
|| model.starts_with("nemotron")
|| model.starts_with("glm-4.7")
{
return Some(boolean_capability(source));
}
None
}
pub fn registry_source_url(provider: &AiProvider) -> Option<&'static str> {
match provider {
AiProvider::Openai => Some(OPENAI_REASONING_DOCS),
AiProvider::Gemini => Some(GEMINI_THINKING_DOCS),
AiProvider::Deepseek => Some(DEEPSEEK_THINKING_DOCS),
AiProvider::Qwen => Some(QWEN_THINKING_DOCS),
AiProvider::Ollama => Some(OLLAMA_THINKING_DOCS),
AiProvider::Claude
| AiProvider::OpenaiCompatible
| AiProvider::CodexCli
| AiProvider::ClaudeCodeCli
| AiProvider::Custom => None,
}
}
pub fn validate_runtime_effort(config: &AiConfig) -> Result<(), String> {
let Some(selection) = config.runtime_effort.as_ref() else {
return Ok(());
};
if matches!(selection, AiEffortSelection::ProviderDefault) {
return Ok(());
}
if matches!(config.provider, AiProvider::Claude | AiProvider::CodexCli | AiProvider::ClaudeCodeCli) {
return match selection {
AiEffortSelection::Enum(value) if !value.trim().is_empty() => Ok(()),
_ => Err("Invalid effort selection for dynamic provider".to_string()),
};
}
let capability = static_effort_capability(config, &config.model).unwrap_or(AiEffortCapability::Unsupported);
let valid = match capability {
AiEffortCapability::Enum { options, .. } => options.iter().any(|option| option.selection == *selection),
AiEffortCapability::Integer { min, max, step, special_values, .. } => {
special_values.iter().any(|option| option.selection == *selection)
|| matches!(selection, AiEffortSelection::Integer(value) if *value >= min && *value <= max && (*value - min) % step == 0)
}
AiEffortCapability::Boolean { .. } => {
matches!(selection, AiEffortSelection::Boolean(_) | AiEffortSelection::Disabled)
}
AiEffortCapability::FreeText { .. } => {
matches!(selection, AiEffortSelection::Text(value) if valid_custom_effort(value))
}
AiEffortCapability::Unsupported => false,
};
valid.then_some(()).ok_or_else(|| format!("Invalid effort selection for model '{}'", config.model))
}
pub fn apply_runtime_effort(body: &mut Value, config: &AiConfig) {
let Some(selection) = config.runtime_effort.as_ref() else {
return;
};
if matches!(selection, AiEffortSelection::ProviderDefault) {
return;
}
let Some(object) = body.as_object_mut() else {
return;
};
match config.provider {
AiProvider::Claude => apply_claude_effort(object, selection),
AiProvider::Gemini => apply_gemini_effort(object, &config.model, selection),
AiProvider::Deepseek => apply_deepseek_effort(object, selection),
AiProvider::Qwen => apply_qwen_effort(object, selection),
AiProvider::Ollama => apply_openai_effort(object, &config.api_style, selection),
AiProvider::Openai | AiProvider::OpenaiCompatible => apply_openai_effort(object, &config.api_style, selection),
AiProvider::Custom => {
if config.api_style == AiApiStyle::AnthropicMessages {
apply_claude_effort(object, selection);
} else {
apply_openai_effort(object, &config.api_style, selection);
}
}
AiProvider::CodexCli | AiProvider::ClaudeCodeCli => {}
}
}
fn apply_claude_effort(object: &mut Map<String, Value>, selection: &AiEffortSelection) {
if let Some(value) = effort_string(selection) {
object.insert("output_config".to_string(), json!({ "effort": value }));
}
}
fn apply_openai_effort(object: &mut Map<String, Value>, api_style: &AiApiStyle, selection: &AiEffortSelection) {
if let Some(value) = effort_string(selection) {
if *api_style == AiApiStyle::Responses {
object.insert("reasoning".to_string(), json!({ "effort": value }));
} else {
object.insert("reasoning_effort".to_string(), Value::String(value));
}
}
}
fn apply_gemini_effort(object: &mut Map<String, Value>, model_id: &str, selection: &AiEffortSelection) {
let generation_config = object.entry("generationConfig").or_insert_with(|| Value::Object(Map::new()));
let Some(generation_config) = generation_config.as_object_mut() else {
return;
};
let thinking_config = generation_config.entry("thinkingConfig").or_insert_with(|| Value::Object(Map::new()));
let Some(thinking_config) = thinking_config.as_object_mut() else {
return;
};
if normalized_model_id(model_id).starts_with("gemini-2.5")
|| normalized_model_id(model_id).starts_with("robotics-er")
{
let value = match selection {
AiEffortSelection::Disabled => Some(0),
AiEffortSelection::Integer(value) => Some(*value),
_ => None,
};
if let Some(value) = value {
thinking_config.insert("thinkingBudget".to_string(), Value::Number(value.into()));
}
} else if let Some(value) = effort_string(selection) {
thinking_config.insert("thinkingLevel".to_string(), Value::String(value));
}
}
fn apply_deepseek_effort(object: &mut Map<String, Value>, selection: &AiEffortSelection) {
match selection {
AiEffortSelection::Disabled | AiEffortSelection::Boolean(false) => {
object.insert("thinking".to_string(), json!({ "type": "disabled" }));
}
AiEffortSelection::Boolean(true) => {
object.insert("thinking".to_string(), json!({ "type": "enabled" }));
}
_ => {
if let Some(value) = effort_string(selection) {
object.insert("thinking".to_string(), json!({ "type": "enabled" }));
object.insert("reasoning_effort".to_string(), Value::String(value));
}
}
}
}
fn apply_qwen_effort(object: &mut Map<String, Value>, selection: &AiEffortSelection) {
object.remove("reasoning_effort");
object.remove("thinking_budget");
object.remove("enable_thinking");
match selection {
AiEffortSelection::Disabled | AiEffortSelection::Boolean(false) => {
object.insert("enable_thinking".to_string(), Value::Bool(false));
}
AiEffortSelection::Boolean(true) => {
object.insert("enable_thinking".to_string(), Value::Bool(true));
}
AiEffortSelection::Integer(value) => {
object.insert("thinking_budget".to_string(), Value::Number((*value).into()));
}
_ => {
if let Some(value) = effort_string(selection) {
object.insert("reasoning_effort".to_string(), Value::String(value));
}
}
}
}
fn valid_custom_effort(value: &str) -> bool {
let value = value.trim();
!value.is_empty() && value.chars().count() <= 64 && !value.chars().any(char::is_control)
}
fn effort_string(selection: &AiEffortSelection) -> Option<String> {
match selection {
AiEffortSelection::Enum(value) | AiEffortSelection::Text(value) => {
let value = value.trim();
(!value.is_empty()).then(|| value.to_string())
}
AiEffortSelection::Disabled => Some("none".to_string()),
AiEffortSelection::Boolean(value) => Some(if *value { "high" } else { "none" }.to_string()),
AiEffortSelection::ProviderDefault | AiEffortSelection::Integer(_) => None,
}
}
fn title_case_effort(value: &str) -> String {
let mut chars = value.chars();
match chars.next() {
Some(first) => first.to_uppercase().collect::<String>() + chars.as_str(),
None => String::new(),
}
}
#[cfg(test)]
mod tests {
use super::{apply_runtime_effort, dynamic_enum_capability, static_effort_capability, validate_runtime_effort};
use crate::ai::{
AiApiStyle, AiAuthMethod, AiCapabilitySource, AiConfig, AiEffortCapability, AiEffortSelection, AiProvider,
AiReasoningLevel,
};
use serde_json::json;
use std::collections::HashMap;
fn config(provider: AiProvider, model: &str) -> AiConfig {
AiConfig {
provider,
api_key: String::new(),
auth_method: AiAuthMethod::ApiKey,
endpoint: String::new(),
model: model.to_string(),
models: Vec::new(),
api_style: AiApiStyle::Completions,
proxy_enabled: false,
proxy_url: String::new(),
enable_thinking: true,
reasoning_level: AiReasoningLevel::Default,
runtime_effort: None,
context_window: None,
codex_cli_path: None,
codex_cli_env: HashMap::new(),
claude_code_cli_path: None,
claude_code_cli_env: HashMap::new(),
}
}
#[test]
fn dynamic_levels_preserve_unknown_values_and_order() {
let capability = dynamic_enum_capability(["low", "ultra", "low", ""], AiCapabilitySource::LocalCli).unwrap();
let AiEffortCapability::Enum { options, .. } = capability else {
panic!("expected enum capability");
};
assert_eq!(options.iter().map(|option| option.id.as_str()).collect::<Vec<_>>(), ["low", "ultra"]);
}
#[test]
fn free_text_effort_uses_the_translated_frontend_placeholder() {
for provider in [AiProvider::OpenaiCompatible, AiProvider::Custom] {
let capability = static_effort_capability(&config(provider, "custom-model"), "custom-model").unwrap();
let AiEffortCapability::FreeText { placeholder, .. } = capability else {
panic!("expected free-text capability");
};
assert_eq!(placeholder, None);
}
}
#[test]
fn gemini_25_pro_uses_numeric_budget_without_off() {
let capability =
static_effort_capability(&config(AiProvider::Gemini, "gemini-2.5-pro"), "gemini-2.5-pro").unwrap();
let AiEffortCapability::Integer { min, max, special_values, .. } = capability else {
panic!("expected integer capability");
};
assert_eq!((min, max), (128, 32_768));
assert_eq!(special_values.len(), 1);
assert_eq!(special_values[0].id, "auto");
}
#[test]
fn qwen_effort_fields_are_mutually_exclusive() {
let mut config = config(AiProvider::Qwen, "qwen3.8-max-preview");
config.runtime_effort = Some(AiEffortSelection::Enum("xhigh".to_string()));
let mut body = json!({
"reasoning_effort": "low",
"thinking_budget": 1024,
"enable_thinking": false
});
apply_runtime_effort(&mut body, &config);
assert_eq!(body["reasoning_effort"], "xhigh");
assert!(body.get("thinking_budget").is_none());
assert!(body.get("enable_thinking").is_none());
}
#[test]
fn gemini_25_never_sends_thinking_level() {
let mut config = config(AiProvider::Gemini, "gemini-2.5-flash");
config.runtime_effort = Some(AiEffortSelection::Integer(8192));
let mut body = json!({ "generationConfig": {} });
apply_runtime_effort(&mut body, &config);
assert_eq!(body["generationConfig"]["thinkingConfig"]["thinkingBudget"], 8192);
assert!(body["generationConfig"]["thinkingConfig"].get("thinkingLevel").is_none());
}
#[test]
fn openai_responses_and_chat_use_different_fields() {
let mut responses = config(AiProvider::Openai, "gpt-5.6");
responses.api_style = AiApiStyle::Responses;
responses.runtime_effort = Some(AiEffortSelection::Enum("high".to_string()));
let mut responses_body = json!({});
apply_runtime_effort(&mut responses_body, &responses);
assert_eq!(responses_body["reasoning"]["effort"], "high");
responses.api_style = AiApiStyle::Completions;
let mut chat_body = json!({});
apply_runtime_effort(&mut chat_body, &responses);
assert_eq!(chat_body["reasoning_effort"], "high");
}
#[test]
fn enum_capability_defaults_to_lowest_registered_level() {
let capability = static_effort_capability(&config(AiProvider::Openai, "gpt-5.6"), "gpt-5.6").unwrap();
let AiEffortCapability::Enum { default, .. } = capability else {
panic!("expected enum capability");
};
assert_eq!(default, AiEffortSelection::Enum("none".to_string()));
}
#[test]
fn deepseek_capability_defaults_to_off_and_maps_thinking_fields() {
let mut config = config(AiProvider::Deepseek, "deepseek-v4-pro");
let capability = static_effort_capability(&config, "deepseek-v4-pro").unwrap();
let AiEffortCapability::Enum { options, default, .. } = capability else {
panic!("expected enum capability");
};
assert_eq!(options[0].selection, AiEffortSelection::Disabled);
assert_eq!(default, AiEffortSelection::Disabled);
config.runtime_effort = Some(AiEffortSelection::Disabled);
let mut body = json!({});
apply_runtime_effort(&mut body, &config);
assert_eq!(body["thinking"]["type"], "disabled");
config.runtime_effort = Some(AiEffortSelection::Enum("max".to_string()));
let mut body = json!({});
apply_runtime_effort(&mut body, &config);
assert_eq!(body["thinking"]["type"], "enabled");
assert_eq!(body["reasoning_effort"], "max");
}
#[test]
fn qwen_boolean_effort_maps_to_enable_thinking() {
let mut config = config(AiProvider::Qwen, "qwen3");
config.runtime_effort = Some(AiEffortSelection::Boolean(true));
let mut body = json!({});
apply_runtime_effort(&mut body, &config);
assert_eq!(body["enable_thinking"], true);
config.runtime_effort = Some(AiEffortSelection::Disabled);
let mut body = json!({});
apply_runtime_effort(&mut body, &config);
assert_eq!(body["enable_thinking"], false);
}
#[test]
fn gemini_3_uses_thinking_level() {
let mut config = config(AiProvider::Gemini, "models/gemini-3-flash");
config.runtime_effort = Some(AiEffortSelection::Enum("medium".to_string()));
let mut body = json!({ "generationConfig": {} });
apply_runtime_effort(&mut body, &config);
assert_eq!(body["generationConfig"]["thinkingConfig"]["thinkingLevel"], "medium");
assert!(body["generationConfig"]["thinkingConfig"].get("thinkingBudget").is_none());
}
#[test]
fn rejects_out_of_range_gemini_budget() {
let mut config = config(AiProvider::Gemini, "gemini-2.5-pro");
config.runtime_effort = Some(AiEffortSelection::Integer(64));
assert!(validate_runtime_effort(&config).is_err());
config.runtime_effort = Some(AiEffortSelection::Integer(-1));
assert!(validate_runtime_effort(&config).is_ok());
}
#[test]
fn accepts_arbitrary_dynamic_cli_effort() {
let mut config = config(AiProvider::ClaudeCodeCli, "sonnet");
config.runtime_effort = Some(AiEffortSelection::Enum("future-effort".to_string()));
assert!(validate_runtime_effort(&config).is_ok());
}
#[test]
fn ollama_openai_compatibility_uses_reasoning_effort() {
let mut config = config(AiProvider::Ollama, "gpt-oss");
config.runtime_effort = Some(AiEffortSelection::Enum("medium".to_string()));
let mut body = json!({});
apply_runtime_effort(&mut body, &config);
assert_eq!(body["reasoning_effort"], "medium");
assert!(body.get("think").is_none());
}
#[test]
fn ollama_boolean_effort_maps_to_openai_compatible_values() {
let mut config = config(AiProvider::Ollama, "qwen3");
config.runtime_effort = Some(AiEffortSelection::Boolean(true));
let mut body = json!({});
apply_runtime_effort(&mut body, &config);
assert_eq!(body["reasoning_effort"], "high");
config.runtime_effort = Some(AiEffortSelection::Disabled);
let mut body = json!({});
apply_runtime_effort(&mut body, &config);
assert_eq!(body["reasoning_effort"], "none");
}
#[test]
fn unsupported_static_model_rejects_explicit_effort() {
let mut config = config(AiProvider::Openai, "gpt-4o");
config.runtime_effort = Some(AiEffortSelection::Enum("high".to_string()));
assert!(validate_runtime_effort(&config).is_err());
}
#[test]
fn provider_default_does_not_change_existing_request_fields() {
let mut config = config(AiProvider::Openai, "gpt-5.6");
config.runtime_effort = Some(AiEffortSelection::ProviderDefault);
let mut body = json!({ "reasoning_effort": "existing" });
apply_runtime_effort(&mut body, &config);
assert_eq!(body["reasoning_effort"], "existing");
}
#[test]
fn claude_effort_maps_to_output_config() {
let mut config = config(AiProvider::Claude, "claude-opus-4-6");
config.runtime_effort = Some(AiEffortSelection::Enum("high".to_string()));
let mut body = json!({});
apply_runtime_effort(&mut body, &config);
assert_eq!(body["output_config"]["effort"], "high");
}
#[test]
fn validates_custom_effort_text_bounds() {
let mut config = config(AiProvider::Custom, "custom-model");
config.runtime_effort = Some(AiEffortSelection::Text("provider-level".to_string()));
assert!(validate_runtime_effort(&config).is_ok());
config.runtime_effort = Some(AiEffortSelection::Text("x".repeat(65)));
assert!(validate_runtime_effort(&config).is_err());
config.runtime_effort = Some(AiEffortSelection::Text("invalid\nvalue".to_string()));
assert!(validate_runtime_effort(&config).is_err());
}
}
crates/dbx-core/src/ai_model_filter.rs
+269
View File
@@ -0,0 +1,269 @@
use crate::ai::{AiModelInfo, AiProvider};
use serde_json::Value;
// Provider model-family exclusions last checked against official catalogs on 2026-07-27.
fn normalized_model_id(model_id: &str) -> String {
let model_id = model_id.trim().trim_start_matches("models/").to_ascii_lowercase();
let model_id = model_id.rsplit('/').next().unwrap_or(&model_id);
model_id.strip_prefix("ft:").and_then(|model| model.split(':').next()).unwrap_or(model_id).to_string()
}
fn is_openai_non_assistant_model(model: &str) -> bool {
const PREFIXES: &[&str] = &[
"babbage-002",
"chatgpt-image",
"computer-use",
"dall-e",
"davinci-002",
"gpt-image",
"omni-moderation",
"sora",
"text-embedding",
"text-moderation",
"tts",
"whisper",
];
PREFIXES.iter().any(|prefix| model.starts_with(prefix))
|| model.contains("-audio")
|| model.contains("-realtime")
|| model.contains("-transcribe")
|| model.contains("-tts")
}
fn is_gemini_non_assistant_model(model: &str) -> bool {
const PREFIXES: &[&str] = &["antigravity", "chirp", "imagen", "lyria", "nano-banana", "veo"];
PREFIXES.iter().any(|prefix| model.starts_with(prefix))
|| model.contains("computer-use")
|| model.contains("deep-research")
|| model.contains("embedding")
|| model.contains("-image")
|| model.contains("gemini-omni")
|| model.contains("image-generation")
|| model.contains("-live")
|| model.contains("live-")
|| model.contains("native-audio")
|| model.contains("robotics")
|| model.contains("-tts")
}
fn is_qwen_non_assistant_model(model: &str) -> bool {
const PREFIXES: &[&str] = &[
"ccai-",
"cosyvoice",
"flux-",
"fun-asr",
"paraformer",
"qwen-image",
"qwen-mt-",
"qwen-tts",
"qwen3-tts",
"sambert",
"sensevoice",
"speech-",
"stable-diffusion",
"tongyi-tingwu",
"wan",
"wanx",
"z-image",
];
PREFIXES.iter().any(|prefix| model.starts_with(prefix))
|| model.contains("-asr")
|| model.contains("-captioner")
|| model.contains("embedding")
|| model.contains("livetranslate")
|| model.contains("moderation")
|| model.contains("-ocr")
|| model.contains("realtime")
|| model.contains("rerank")
|| model.contains("-s2s")
}
pub(crate) fn model_is_assistant_compatible(provider: &AiProvider, model_id: &str) -> bool {
let model = normalized_model_id(model_id);
match provider {
AiProvider::Openai => !is_openai_non_assistant_model(&model),
AiProvider::Gemini => !is_gemini_non_assistant_model(&model),
AiProvider::Qwen => !is_qwen_non_assistant_model(&model),
AiProvider::Claude
| AiProvider::Deepseek
| AiProvider::Ollama
| AiProvider::OpenaiCompatible
| AiProvider::CodexCli
| AiProvider::ClaudeCodeCli
| AiProvider::Custom => true,
}
}
pub(crate) fn retain_known_assistant_models(provider: &AiProvider, models: &mut Vec<AiModelInfo>) {
models.retain(|model| model_is_assistant_compatible(provider, &model.id));
}
pub(crate) fn gemini_item_is_assistant_compatible(item: &Value) -> bool {
let Some(model_id) = item["name"].as_str().or_else(|| item["id"].as_str()) else {
return false;
};
if !model_is_assistant_compatible(&AiProvider::Gemini, model_id) {
return false;
}
let Some(methods) = item["supportedGenerationMethods"].as_array() else {
return true;
};
if methods.is_empty() {
return true;
}
methods.iter().filter_map(Value::as_str).any(|method| {
method.eq_ignore_ascii_case("generateContent") || method.eq_ignore_ascii_case("streamGenerateContent")
})
}
fn ollama_capability(data: &Value, expected: &str) -> Option<bool> {
let capabilities = data["capabilities"].as_array()?;
if capabilities.is_empty() {
return None;
}
Some(capabilities.iter().filter_map(Value::as_str).any(|capability| capability == expected))
}
pub(crate) fn ollama_completion_capability(data: &Value) -> Option<bool> {
ollama_capability(data, "completion")
}
pub(crate) fn ollama_tool_capability(data: &Value) -> Option<bool> {
ollama_capability(data, "tools")
}
#[cfg(test)]
mod tests {
use super::{
gemini_item_is_assistant_compatible, model_is_assistant_compatible, ollama_completion_capability,
ollama_tool_capability,
};
use crate::ai::AiProvider;
#[test]
fn openai_filter_keeps_assistant_models_and_hides_specialized_endpoints() {
for model in ["gpt-4o", "gpt-5.6", "o4-mini", "ft:gpt-4o-mini:org:name:id"] {
assert!(model_is_assistant_compatible(&AiProvider::Openai, model), "{model}");
}
for model in [
"text-embedding-3-large",
"gpt-image-1",
"sora-2",
"omni-moderation-latest",
"gpt-4o-mini-transcribe",
"gpt-4o-realtime-preview",
"tts-1-hd",
] {
assert!(!model_is_assistant_compatible(&AiProvider::Openai, model), "{model}");
}
}
#[test]
fn qwen_filter_keeps_multimodal_assistants_and_hides_non_chat_families() {
for model in
["qwen-plus", "qwen3-max", "qwen-vl-max", "qwen-omni-turbo", "qwen3.5-omni-plus", "qwen3.5-omni-flash"]
{
assert!(model_is_assistant_compatible(&AiProvider::Qwen, model), "{model}");
}
for model in [
"ccai-pro",
"qwen-mt-plus",
"qwen-vl-ocr-2025-11-20",
"qwen3.5-livetranslate-flash",
"qwen3.5-omni-plus-realtime",
"qwen3-omni-30b-a3b-captioner",
"qwen3-s2s-flash-realtime",
"text-embedding-v4",
"tongyi-tingwu-slp",
"z-image-turbo",
"qwen3-vl-embedding",
"qwen3-reranker",
"qwen3-asr-flash",
"qwen-moderation-latest",
"wan2.1-t2v-turbo",
"qwen-image-plus",
"cosyvoice-v3-flash",
] {
assert!(!model_is_assistant_compatible(&AiProvider::Qwen, model), "{model}");
}
}
#[test]
fn generic_providers_do_not_filter_unknown_model_taxonomies() {
for provider in [AiProvider::OpenaiCompatible, AiProvider::Custom] {
for model in ["text-embedding-private-chat", "company/image-reasoner", "future-model"] {
assert!(model_is_assistant_compatible(&provider, model), "{}:{model}", provider.as_str());
}
}
}
#[test]
fn gemini_filter_uses_generation_methods_and_excludes_media_output_models() {
assert!(gemini_item_is_assistant_compatible(&serde_json::json!({
"name": "models/gemini-2.5-pro",
"supportedGenerationMethods": ["generateContent", "countTokens"]
})));
assert!(gemini_item_is_assistant_compatible(&serde_json::json!({
"name": "models/gemma-3-27b-it",
"supportedGenerationMethods": ["generateContent"]
})));
assert!(!gemini_item_is_assistant_compatible(&serde_json::json!({
"name": "models/gemini-embedding-001",
"supportedGenerationMethods": ["embedContent"]
})));
assert!(!gemini_item_is_assistant_compatible(&serde_json::json!({
"name": "models/gemini-3-pro-image-preview",
"supportedGenerationMethods": ["generateContent"]
})));
for model in [
"models/antigravity-preview-05-2026",
"models/deep-research-preview-04-2026",
"models/gemini-2.5-computer-use-preview-10-2025",
"models/gemini-omni-flash-preview",
"models/gemini-robotics-er-1.6-preview",
"models/nano-banana-pro-preview",
] {
assert!(
!gemini_item_is_assistant_compatible(&serde_json::json!({
"name": model,
"supportedGenerationMethods": ["generateContent"]
})),
"{model}"
);
}
assert!(gemini_item_is_assistant_compatible(&serde_json::json!({
"name": "models/future-chat-model"
})));
}
#[test]
fn ollama_capabilities_only_exclude_explicit_non_completion_models() {
assert_eq!(
ollama_completion_capability(&serde_json::json!({ "capabilities": ["completion", "vision"] })),
Some(true)
);
assert_eq!(ollama_completion_capability(&serde_json::json!({ "capabilities": ["embedding"] })), Some(false));
assert_eq!(ollama_completion_capability(&serde_json::json!({ "capabilities": [] })), None);
assert_eq!(ollama_completion_capability(&serde_json::json!({})), None);
}
#[test]
fn ollama_tool_capability_uses_explicit_model_metadata() {
assert_eq!(
ollama_tool_capability(&serde_json::json!({ "capabilities": ["completion", "tools", "thinking"] })),
Some(true)
);
assert_eq!(
ollama_tool_capability(&serde_json::json!({ "capabilities": ["completion", "thinking"] })),
Some(false)
);
assert_eq!(ollama_tool_capability(&serde_json::json!({ "capabilities": [] })), None);
assert_eq!(ollama_tool_capability(&serde_json::json!({})), None);
}
}
crates/dbx-core/src/cloud_sync.rs
+1
View File
@@ -1091,6 +1091,7 @@ mod tests {
proxy_url: String::new(),
enable_thinking: true,
reasoning_level: crate::ai::AiReasoningLevel::Default,
runtime_effort: None,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
crates/dbx-core/src/lib.rs
+2
View File
@@ -12,6 +12,8 @@ pub mod ai;
pub mod ai_claude_code_cli;
pub mod ai_cli_agent;
pub mod ai_codex_cli;
pub mod ai_effort;
mod ai_model_filter;
pub mod changelog;
pub mod cloud_sync;
pub mod connection;
crates/dbx-core/src/storage.rs
+36 -1
View File
@@ -7,7 +7,7 @@ use rusqlite::{params, params_from_iter, types::Value, Connection, DatabaseName,
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use crate::ai::{AiChatMessage, AiConfig, AiConfigItem, AiConversation, AiProvider};
use crate::ai::{AiChatMessage, AiChatSelectionState, AiConfig, AiConfigItem, AiConversation, AiProvider};
use crate::connection_secrets::{
MQ_AUTH_API_KEY_VALUE_KEY, MQ_AUTH_CLIENT_SECRET_KEY, MQ_AUTH_PASSWORD_KEY, MQ_AUTH_SECRET_PREFIX,
MQ_AUTH_TOKEN_KEY, MQ_TOKEN_SIGNING_KEY, MQ_TOKEN_SIGNING_SECRET_PREFIX, NACOS_AUTH_PASSWORD_KEY,
@@ -30,6 +30,7 @@ const APP_STATE_OPEN_TABS_KEY: &str = "open_tabs";
const APP_STATE_SAVED_SQL_EDITOR_POSITIONS_KEY: &str = "saved_sql_editor_positions";
const MCP_GLOBAL_POLICY_KEY: &str = "mcp_global_policy";
const APP_STATE_AI_GLOBAL_INSTRUCTIONS_KEY: &str = "ai_global_custom_instructions";
const APP_STATE_AI_CHAT_SELECTION_KEY: &str = "ai_chat_selection_v1";
const USER_DATA_TABLES: &[&str] = &[
"connections",
"connection_secrets",
@@ -1684,6 +1685,18 @@ impl Storage {
})
}
pub async fn save_ai_chat_selection(&self, selection: &AiChatSelectionState) -> Result<(), String> {
let value = serde_json::to_value(selection).map_err(|e| e.to_string())?;
self.save_app_state_value(APP_STATE_AI_CHAT_SELECTION_KEY, &value).await
}
pub async fn load_ai_chat_selection(&self) -> Result<Option<AiChatSelectionState>, String> {
self.load_app_state_value(APP_STATE_AI_CHAT_SELECTION_KEY)
.await?
.map(|value| serde_json::from_value(value).map_err(|e| e.to_string()))
.transpose()
}
pub async fn load_or_create_local_device_secret(&self) -> Result<String, String> {
let mut settings = self.load_app_settings_json().await?;
if let Some(secret) = settings.get("local_device_secret").and_then(|value| value.as_str()) {
@@ -3490,6 +3503,7 @@ mod tests {
maybe_import_user_data_db, DataDbImportResult, DesktopIconTheme, DesktopSettings, McpGlobalPolicy,
McpGlobalPolicyState, Storage, MCP_GLOBAL_POLICY_KEY,
};
use crate::ai::{AiActiveModelSelection, AiChatSelectionState, AiEffortSelection, AiModelEffortPreference};
use crate::connection_secrets::NACOS_RNACOS_CONSOLE_PASSWORD_KEY;
use crate::connection_secrets::{
MQ_AUTH_PASSWORD_KEY, MQ_AUTH_TOKEN_KEY, MQ_TOKEN_SIGNING_KEY, NACOS_AUTH_PASSWORD_KEY,
@@ -4693,6 +4707,26 @@ mod tests {
assert_eq!(storage.load_app_settings_json().await.unwrap().get("open_tabs"), None);
}
#[tokio::test]
async fn ai_chat_selection_roundtrips_in_local_app_state() {
let path = temp_db_path("ai-chat-selection");
let storage = Storage::open(&path).await.unwrap();
let selection = AiChatSelectionState {
version: 1,
active: Some(AiActiveModelSelection { config_id: "config-1".to_string(), model_id: "model-1".to_string() }),
effort_preferences: vec![AiModelEffortPreference {
config_id: "config-1".to_string(),
model_id: "model-1".to_string(),
selection: AiEffortSelection::Enum("high".to_string()),
}],
};
storage.save_ai_chat_selection(&selection).await.unwrap();
assert_eq!(storage.load_ai_chat_selection().await.unwrap(), Some(selection));
assert_eq!(storage.load_app_settings_json().await.unwrap().get("ai_chat_selection_v1"), None);
}
#[tokio::test]
async fn tab_runtime_cache_roundtrips_binary_payloads() {
let path = temp_db_path("tab-runtime-cache");
@@ -4877,6 +4911,7 @@ mod tests {
proxy_url: String::new(),
enable_thinking: true,
reasoning_level: AiReasoningLevel::Default,
runtime_effort: None,
context_window: None,
codex_cli_path: None,
codex_cli_env: std::collections::HashMap::new(),