public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add Rig native model providers

Browse Source
This commit is contained in:
rkw6086
2026-08-06 15:03:00 -05:00
parent 634ce7ba00
commit 3fda5d414b
34 changed files with 2134 additions and 452 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/blocklist/action_model/execute/shell_command.rs
+5 -3
View File
@@ -684,9 +684,11 @@ impl ShellCommandExecutor {
.force_refresh_senders
.keys()
.find(|selector| {
selector
.get_block(&terminal_model)
.is_some_and(|block| block.id() == block_id)
selector.get_block(&terminal_model).is_some_and(|block| {
block.id() == block_id
&& block.is_active_and_long_running()
&& !block.finished()
})
})
.cloned();
drop(terminal_model);
app/src/ai/blocklist/action_model/execute/shell_command_tests.rs
+15 -1
View File
@@ -103,11 +103,14 @@ fn force_refresh_block_reports_and_resolves_matching_poll() {
ActiveSession::new(sessions.clone(), model_event_dispatcher.clone(), ctx)
});
let terminal_model = Arc::new(FairMutex::new(TerminalModel::mock(None, None)));
terminal_model
.lock()
.simulate_long_running_block("sleep 120", "still running");
let block_id = terminal_model.lock().active_block_id().clone();
let executor = app.add_model(|ctx| {
ShellCommandExecutor::new(
active_session,
terminal_model,
terminal_model.clone(),
&model_event_dispatcher,
terminal_view_id,
ctx,
@@ -124,6 +127,17 @@ fn force_refresh_block_reports_and_resolves_matching_poll() {
});
assert!(matches!(rx.try_recv(), Ok(Some(()))));
let (tx, _rx) = oneshot::channel();
executor.update(&mut app, |executor, _| {
executor
.force_refresh_senders
.insert(BlockSelector::Id(block_id.clone()), tx);
});
terminal_model.lock().finish_block();
assert!(executor.update(&mut app, |executor, _| {
!executor.force_refresh_block(&block_id)
}));
});
}
app/src/ai/blocklist/block/cli_controller.rs
+119 -24
View File
@@ -265,6 +265,8 @@ impl CLISubagentController {
let block_id = block.id().clone();
let conversation_id = block.ai_conversation_id();
let requested_command_action_id = block.requested_command_action_id().cloned();
let should_skip_completion_assessment =
!should_request_completion_assessment(block.long_running_control_state());
let completion = match (&block_completed_event.block_type, conversation_id) {
(BlockType::User(completed), Some(conversation_id)) => {
let command = if completed.command_with_obfuscated_secrets.is_empty() {
@@ -310,17 +312,49 @@ impl CLISubagentController {
};
drop(terminal_model);
let Some(subagent_state) = me.active_subagents_by_block.get_mut(&block_id) else {
let Some(has_last_snapshot) = me
.active_subagents_by_block
.get(&block_id)
.map(|state| state.last_snapshot_at.is_some())
else {
return;
};
if subagent_state.last_snapshot_at.is_some() {
if has_last_snapshot {
ctx.emit(CLISubagentEvent::UpdatedLastSnapshot);
}
subagent_state.completion = completion;
if subagent_state.completion.is_none() {
// A Stop takeover intentionally cancels the subagent. The command may still
// finish later, but that completion must not start a new assessment turn. Also
// clean up the in-memory monitor state so the stopped subagent cannot linger in
// the UI or intercept later refreshes.
if should_skip_completion_assessment {
me.finish_subagent(
&block_id,
conversation_id,
requested_command_action_id,
ctx,
);
return;
}
let has_completion = {
let Some(subagent_state) = me.active_subagents_by_block.get_mut(&block_id)
else {
return;
};
subagent_state.completion = completion;
subagent_state.completion.is_some()
};
if !has_completion {
log::warn!(
"CLI monitor block {block_id:?} completed without final command metadata"
);
me.finish_subagent(
&block_id,
conversation_id,
requested_command_action_id,
ctx,
);
return;
}
me.advance_completed_subagent(&block_id, ctx);
@@ -380,7 +414,12 @@ impl CLISubagentController {
}
if completion.final_turn_started {
self.finish_completed_subagent(block_id, ctx);
self.finish_subagent(
block_id,
Some(completion.conversation_id),
completion.initial_requested_command_action_id,
ctx,
);
return;
}
@@ -404,38 +443,55 @@ impl CLISubagentController {
}
}
fn finish_completed_subagent(&mut self, block_id: &BlockId, ctx: &mut ModelContext<Self>) {
fn finish_subagent(
&mut self,
block_id: &BlockId,
conversation_id: Option<AIConversationId>,
initial_requested_command_action_id: Option<AIAgentActionId>,
ctx: &mut ModelContext<Self>,
) {
let Some(state) = self.active_subagents_by_block.remove(block_id) else {
return;
};
let Some(completion) = state.completion else {
return;
};
let conversation_id = conversation_id.or_else(|| {
state
.completion
.as_ref()
.map(|completion| completion.conversation_id)
});
let initial_requested_command_action_id = initial_requested_command_action_id
.or_else(|| {
state
.completion
.as_ref()
.and_then(|completion| completion.initial_requested_command_action_id.clone())
})
.or(state.initial_requested_command_action_id);
let deactivate_result =
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, _| {
history_model.deactivate_cli_subagent_task_for_conversation(
block_id,
completion.conversation_id,
)
});
if let Err(error) = deactivate_result {
log::error!(
"Failed to deactivate completed CLI monitor for block {block_id:?}: {error:?}"
);
if let Some(conversation_id) = conversation_id {
let deactivate_result =
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, _| {
history_model
.deactivate_cli_subagent_task_for_conversation(block_id, conversation_id)
});
if let Err(error) = deactivate_result {
log::error!("Failed to deactivate CLI monitor for block {block_id:?}: {error:?}");
}
}
ctx.emit(CLISubagentEvent::FinishedSubagent {
block_id: block_id.clone(),
conversation_id: Some(completion.conversation_id),
initial_requested_command_action_id: completion.initial_requested_command_action_id,
conversation_id,
initial_requested_command_action_id,
});
if let Some(agent_view_controller) = &self.agent_view_controller {
if let (Some(agent_view_controller), Some(conversation_id)) =
(&self.agent_view_controller, conversation_id)
{
agent_view_controller.update(ctx, |controller, ctx| {
let is_this_inline_conversation = controller.is_inline()
&& controller.agent_view_state().active_conversation_id()
== Some(completion.conversation_id);
== Some(conversation_id);
if is_this_inline_conversation {
controller.exit_agent_view(ctx);
}
@@ -919,3 +975,42 @@ fn command_finished_block_id(result: &AIAgentActionResultType) -> Option<&BlockI
| AIAgentActionResultType::WaitForEvents(_) => None,
}
}
fn should_request_completion_assessment(
control_state: Option<&LongRunningCommandControlState>,
) -> bool {
!control_state
.and_then(LongRunningCommandControlState::user_take_over_reason)
.is_some_and(UserTakeOverReason::is_stop)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn stop_takeover_does_not_request_a_completion_assessment() {
let state = LongRunningCommandControlState::User {
reason: UserTakeOverReason::Stop,
};
assert!(!should_request_completion_assessment(Some(&state)));
}
#[test]
fn non_stop_control_states_can_request_a_completion_assessment() {
let agent_state = LongRunningCommandControlState::Agent {
is_blocked: false,
should_hide_responses: false,
};
let transfer_state = LongRunningCommandControlState::User {
reason: UserTakeOverReason::TransferFromAgent {
reason: "needs user input".to_owned(),
},
};
assert!(should_request_completion_assessment(None));
assert!(should_request_completion_assessment(Some(&agent_state)));
assert!(should_request_completion_assessment(Some(&transfer_state)));
}
}
app/src/ai/blocklist/controller/response_stream.rs
+2
View File
@@ -230,6 +230,8 @@ impl ResponseStream {
kind: client_config.kind,
base_url: client_config.base_url.clone(),
api_key: client_config.api_key.clone(),
project_id: client_config.project_id.clone(),
location: client_config.location.clone(),
model: client_config
.model
.clone()
app/src/ai/crosscheck/reviewer.rs
+2
View File
@@ -167,6 +167,8 @@ impl CrosscheckReviewer {
kind: client_config.kind,
base_url: client_config.base_url.clone(),
api_key: client_config.api_key.clone(),
project_id: client_config.project_id.clone(),
location: client_config.location.clone(),
model: client_config
.model
.clone()
app/src/ai/llms.rs
+105 -8
View File
@@ -5,6 +5,11 @@ use std::sync::{Arc, OnceLock};
use ai::api_keys::ApiKeyManager;
pub use ai::LLMId;
#[cfg(not(target_family = "wasm"))]
use galaxy_agent_rig::{
discover_anthropic_models, discover_gemini_models, validate_vertex_ai_credentials,
vertex_ai_model_catalog, RigModelInfo,
};
use galaxy_core::features::FeatureFlag;
use galaxy_core::ui::icons::Icon;
use galaxy_core::user_preferences::GetUserPreferences;
@@ -1005,6 +1010,8 @@ impl LLMPreferences {
bool,
String,
Option<String>,
Option<String>,
Option<String>,
Vec<OpenAIModelConfig>,
);
let mut provider_entries: Vec<OpenAIProviderEntry> = Vec::new();
@@ -1037,6 +1044,8 @@ impl LLMPreferences {
true,
base_url,
api_key,
None,
None,
single_provider_models,
));
}
@@ -1047,11 +1056,19 @@ impl LLMPreferences {
.value()
.iter()
.filter_map(|provider| {
if !provider.enabled
|| (provider.kind == OpenAIProviderKind::OpenAICompatible
&& provider.base_url.trim().is_empty())
|| provider.models.is_empty()
{
let missing_credentials = match provider.kind {
OpenAIProviderKind::OpenAICompatible => provider.base_url.trim().is_empty(),
OpenAIProviderKind::Anthropic | OpenAIProviderKind::Gemini => provider
.api_key
.as_deref()
.is_none_or(|key| key.trim().is_empty()),
OpenAIProviderKind::VertexAI => provider
.project_id
.as_deref()
.is_none_or(|project| project.trim().is_empty()),
OpenAIProviderKind::ChatGPTSubscription => false,
};
if !provider.enabled || missing_credentials || provider.models.is_empty() {
return None;
}
Some((
@@ -1060,6 +1077,8 @@ impl LLMPreferences {
provider.enabled,
provider.base_url.clone(),
provider.api_key.clone(),
provider.project_id.clone(),
provider.location.clone(),
provider.models.clone(),
))
}),
@@ -1071,8 +1090,16 @@ impl LLMPreferences {
let mut total_injected = 0;
let mut seen_model_ids: HashSet<String> = HashSet::new();
for (provider_name, provider_kind, provider_enabled, base_url, api_key, models) in
provider_entries
for (
provider_name,
provider_kind,
provider_enabled,
base_url,
api_key,
provider_project_id,
provider_location,
models,
) in provider_entries
{
if !provider_enabled {
continue;
@@ -1109,12 +1136,14 @@ impl LLMPreferences {
kind: provider_kind,
base_url: base_url.clone(),
api_key: api_key.clone(),
project_id: provider_project_id.clone(),
location: provider_location.clone(),
model: Some(model.model_id.clone()),
reasoning_effort: reasoning_effort.clone(),
max_input_tokens: Some(openai_model_context_size(model)),
max_output_tokens: model.max_output_tokens,
use_rig: model.use_rig
|| provider_kind == OpenAIProviderKind::ChatGPTSubscription,
|| !matches!(provider_kind, OpenAIProviderKind::OpenAICompatible),
supports_system_messages: model.supports_system_messages(),
};
self.openai_provider_routing
@@ -1557,6 +1586,54 @@ impl LLMPreferences {
pub(crate) async fn discover_openai_provider_models(
provider: OpenAIProviderConfig,
) -> Result<Vec<OpenAIModelConfig>, String> {
let native_models = match provider.kind {
OpenAIProviderKind::Anthropic => {
let api_key = provider
.api_key
.as_deref()
.filter(|key| !key.trim().is_empty())
.ok_or_else(|| {
"Enter an Anthropic API key before testing the connection.".to_string()
})?;
Some(discover_anthropic_models(api_key).await?)
}
OpenAIProviderKind::Gemini => {
let api_key = provider
.api_key
.as_deref()
.filter(|key| !key.trim().is_empty())
.ok_or_else(|| {
"Enter a Gemini API key before testing the connection.".to_string()
})?;
Some(discover_gemini_models(api_key).await?)
}
OpenAIProviderKind::VertexAI => {
if provider
.project_id
.as_deref()
.is_none_or(|project| project.trim().is_empty())
{
return Err(
"Enter a Google Cloud project ID before testing the connection."
.to_string(),
);
}
validate_vertex_ai_credentials(
provider.project_id.as_deref().unwrap_or_default(),
provider.location.as_deref().unwrap_or("global"),
)?;
Some(vertex_ai_model_catalog())
}
OpenAIProviderKind::OpenAICompatible | OpenAIProviderKind::ChatGPTSubscription => None,
};
if let Some(models) = native_models {
if models.is_empty() {
return Err("The provider responded, but no models were found.".to_string());
}
return Ok(Self::rig_models_to_openai_models(models));
}
if provider.base_url.trim().is_empty() {
return Err("Enter a provider URL before testing the connection.".to_string());
}
@@ -1586,6 +1663,26 @@ impl LLMPreferences {
Ok(models)
}
#[cfg(not(target_family = "wasm"))]
fn rig_models_to_openai_models(models: Vec<RigModelInfo>) -> Vec<OpenAIModelConfig> {
models
.into_iter()
.map(|model| OpenAIModelConfig {
model_id: model.id,
display_name: model.display_name,
vision_supported: false,
context_size: model.context_size.unwrap_or(128_000),
max_input_tokens: model.context_size,
max_output_tokens: None,
provider: None,
use_rig: true,
supports_system_messages: Some(true),
reasoning_efforts: Vec::new(),
enabled: true,
})
.collect()
}
/// Returns the `LLMInfo` for the base LLM to be used for an Agent Mode request.
pub fn get_active_base_model<'a>(
&'a self,
app/src/ai/openai/client.rs
+2
View File
@@ -11,6 +11,8 @@ pub struct OpenAIClientConfig {
pub kind: OpenAIProviderKind,
pub base_url: String,
pub api_key: Option<String>,
pub project_id: Option<String>,
pub location: Option<String>,
pub model: Option<String>,
pub reasoning_effort: Option<String>,
pub max_input_tokens: Option<u32>,
app/src/ai/runtime/rig.rs
+49 -2
View File
@@ -8,8 +8,9 @@ use galaxy_agent_core::{
ToolCallDecision, ToolEvent, ToolPolicy, ToolResult, TurnCommand,
};
use galaxy_agent_rig::{
ChatGPTSubscriptionRuntime, ChatGPTSubscriptionRuntimeConfig, OpenAICompatibleRuntime,
OpenAICompatibleRuntimeConfig,
AnthropicRuntime, AnthropicRuntimeConfig, ChatGPTSubscriptionRuntime,
ChatGPTSubscriptionRuntimeConfig, GeminiRuntime, GeminiRuntimeConfig, OpenAICompatibleRuntime,
OpenAICompatibleRuntimeConfig, VertexAiRuntime, VertexAiRuntimeConfig,
};
use uuid::Uuid;
use warp_multi_agent_api::ToolType;
@@ -73,6 +74,52 @@ pub(crate) fn rig_openai_response_stream(
cancellation_rx,
)
}
OpenAIProviderKind::Anthropic => {
let runtime = AnthropicRuntime::new(AnthropicRuntimeConfig {
api_key: config.api_key.unwrap_or_default(),
model: model_id,
max_output_tokens: config.max_output_tokens.map(u64::from),
});
rig_response_stream(
runtime,
prepared,
skill_path_origin,
config.max_input_tokens,
"rig_anthropic",
cancellation_rx,
)
}
OpenAIProviderKind::Gemini => {
let runtime = GeminiRuntime::new(GeminiRuntimeConfig {
api_key: config.api_key.unwrap_or_default(),
model: model_id,
max_output_tokens: config.max_output_tokens.map(u64::from),
});
rig_response_stream(
runtime,
prepared,
skill_path_origin,
config.max_input_tokens,
"rig_gemini",
cancellation_rx,
)
}
OpenAIProviderKind::VertexAI => {
let runtime = VertexAiRuntime::new(VertexAiRuntimeConfig {
project_id: config.project_id.unwrap_or_default(),
location: config.location.unwrap_or_else(|| "global".to_string()),
model: model_id,
max_output_tokens: config.max_output_tokens.map(u64::from),
});
rig_response_stream(
runtime,
prepared,
skill_path_origin,
config.max_input_tokens,
"rig_vertex_ai",
cancellation_rx,
)
}
}
}
app/src/ai/runtime/rig_request_tests.rs
+2
View File
@@ -23,6 +23,8 @@ fn config() -> OpenAIClientConfig {
kind: crate::settings::OpenAIProviderKind::OpenAICompatible,
base_url: "http://localhost:4000/v1".to_string(),
api_key: None,
project_id: None,
location: None,
model: Some("provider-model".to_string()),
reasoning_effort: None,
max_input_tokens: Some(128_000),