mirror of
https://github.com/block/goose.git
synced 2026-07-17 12:56:20 +02:00
rig provider PoC
This commit is contained in:
Jack Amadeo
Generated
+77
@@ -187,6 +187,12 @@ version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
|
||||
|
||||
[[package]]
|
||||
name = "as-any"
|
||||
version = "0.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b0f477b951e452a0b6b4a10b53ccd569042d1d01729b519e02074a9c0958a063"
|
||||
|
||||
[[package]]
|
||||
name = "ascii"
|
||||
version = "1.1.0"
|
||||
@@ -3328,6 +3334,17 @@ dependencies = [
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "eventsource-stream"
|
||||
version = "0.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "74fef4569247a5f429d9156b9d0a2599914385dd189c539334c625d8099d90ab"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"nom 7.1.3",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "exr"
|
||||
version = "1.74.0"
|
||||
@@ -3713,6 +3730,12 @@ version = "0.3.31"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988"
|
||||
|
||||
[[package]]
|
||||
name = "futures-timer"
|
||||
version = "3.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24"
|
||||
|
||||
[[package]]
|
||||
name = "futures-util"
|
||||
version = "0.3.31"
|
||||
@@ -4138,6 +4161,7 @@ dependencies = [
|
||||
"rand 0.8.5",
|
||||
"regex",
|
||||
"reqwest 0.12.28",
|
||||
"rig-core",
|
||||
"rmcp 0.14.0",
|
||||
"rubato",
|
||||
"schemars 1.2.1",
|
||||
@@ -6767,6 +6791,15 @@ version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d"
|
||||
|
||||
[[package]]
|
||||
name = "ordered-float"
|
||||
version = "5.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7f4779c6901a562440c3786d08192c6fbda7c1c2060edd10006b05ee35d10f2d"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ordered-multimap"
|
||||
version = "0.7.3"
|
||||
@@ -7966,6 +7999,38 @@ dependencies = [
|
||||
"bytemuck",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rig-core"
|
||||
version = "0.30.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a8f7a3f0c7c00eaced15a68ee16e1bd6bb709ff598d11b9aedac8b628217dc09"
|
||||
dependencies = [
|
||||
"as-any",
|
||||
"async-stream",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"eventsource-stream",
|
||||
"fastrand",
|
||||
"futures",
|
||||
"futures-timer",
|
||||
"glob",
|
||||
"http 1.4.0",
|
||||
"mime",
|
||||
"mime_guess",
|
||||
"nanoid",
|
||||
"ordered-float",
|
||||
"pin-project-lite",
|
||||
"reqwest 0.12.28",
|
||||
"schemars 1.2.1",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"tracing-futures",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ring"
|
||||
version = "0.17.9"
|
||||
@@ -10780,6 +10845,18 @@ dependencies = [
|
||||
"valuable",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tracing-futures"
|
||||
version = "0.2.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2"
|
||||
dependencies = [
|
||||
"futures",
|
||||
"futures-task",
|
||||
"pin-project",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tracing-log"
|
||||
version = "0.2.0"
|
||||
|
||||
@@ -132,6 +132,7 @@ ignore = { workspace = true }
|
||||
which = { workspace = true }
|
||||
pctx_code_mode = "^0.2.2"
|
||||
unbinder = "0.1.7"
|
||||
rig-core = { version = "0.30.0", default-features = false, features = ["reqwest-rustls"] }
|
||||
|
||||
[target.'cfg(target_os = "windows")'.dependencies]
|
||||
winapi = { version = "0.3", features = ["wincred"] }
|
||||
|
||||
@@ -27,12 +27,13 @@ use super::{
|
||||
xai::XaiProvider,
|
||||
};
|
||||
use crate::model::ModelConfig;
|
||||
use crate::providers::base::ProviderType;
|
||||
use crate::providers::{base::ProviderType, rig::RigProvider};
|
||||
use crate::{
|
||||
config::declarative_providers::register_declarative_providers,
|
||||
providers::provider_registry::ProviderEntry,
|
||||
};
|
||||
use anyhow::Result;
|
||||
use rig::providers::anthropic;
|
||||
use tokio::sync::OnceCell;
|
||||
|
||||
const DEFAULT_LEAD_TURNS: usize = 3;
|
||||
@@ -64,6 +65,7 @@ async fn init_registry() -> RwLock<ProviderRegistry> {
|
||||
registry.register::<TetrateProvider>(true);
|
||||
registry.register::<VeniceProvider>(false);
|
||||
registry.register::<XaiProvider>(false);
|
||||
registry.register::<RigProvider<anthropic::completion::CompletionModel>>(false);
|
||||
});
|
||||
if let Err(e) = load_custom_providers_into_registry(&mut registry) {
|
||||
tracing::warn!("Failed to load custom providers: {}", e);
|
||||
|
||||
@@ -30,6 +30,7 @@ pub mod openrouter;
|
||||
pub mod provider_registry;
|
||||
pub mod provider_test;
|
||||
mod retry;
|
||||
pub mod rig;
|
||||
pub mod sagemaker_tgi;
|
||||
pub mod snowflake;
|
||||
pub mod testprovider;
|
||||
|
||||
@@ -0,0 +1,761 @@
|
||||
//! Adapters to use rig CompletionModels as goose Providers and vice-versa.
|
||||
//!
|
||||
//! This module provides bidirectional integration between goose's Provider trait
|
||||
//! and rig's CompletionModel trait, allowing:
|
||||
//! - Using any rig CompletionModel as a goose Provider via `RigProvider`
|
||||
//! - Using any goose Provider as a rig CompletionModel via `GooseProvider`
|
||||
|
||||
use std::borrow::Cow;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::future::BoxFuture;
|
||||
use futures::stream;
|
||||
use rig::client::{CompletionClient, ProviderClient};
|
||||
use rig::completion::{
|
||||
CompletionError, CompletionModel, CompletionRequest, CompletionRequestBuilder,
|
||||
CompletionResponse, GetTokenUsage, ToolDefinition,
|
||||
};
|
||||
use rig::message::{
|
||||
AssistantContent, DocumentSourceKind, Image, ImageMediaType, Message as RigMessage, Reasoning,
|
||||
Text, ToolCall as RigToolCall, ToolFunction, ToolResult as RigToolResult, ToolResultContent,
|
||||
UserContent,
|
||||
};
|
||||
use rig::providers::anthropic::completion::CLAUDE_4_OPUS;
|
||||
use rig::streaming::{RawStreamingChoice, StreamingCompletionResponse, StreamingResult};
|
||||
use rig::OneOrMany;
|
||||
use rmcp::model::{CallToolRequestParams, CallToolResult, Content, Role, Tool};
|
||||
use serde_json::Value;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::conversation::message::{
|
||||
Message, MessageContent, ThinkingContent, ToolRequest, ToolResponse,
|
||||
};
|
||||
use crate::model::ModelConfig;
|
||||
use crate::providers::base::{Provider, ProviderDef, ProviderUsage, Usage};
|
||||
use crate::providers::errors::ProviderError;
|
||||
use rig::providers::anthropic;
|
||||
|
||||
pub fn claude_opus() -> anthropic::completion::CompletionModel {
|
||||
let config = crate::config::Config::global();
|
||||
let api_key: String = config.get_secret("ANTHROPIC_API_KEY").unwrap();
|
||||
|
||||
let client = anthropic::Client::from_val(api_key);
|
||||
|
||||
client.completion_model(CLAUDE_4_OPUS)
|
||||
}
|
||||
|
||||
/// Wraps a rig `CompletionModel` to implement goose's `Provider` trait.
|
||||
///
|
||||
/// This allows any model from rig (OpenAI, Anthropic, etc.) to be used
|
||||
/// as a goose provider.
|
||||
pub struct RigProvider<M>
|
||||
where
|
||||
M: CompletionModel + Send + Sync,
|
||||
{
|
||||
model: M,
|
||||
model_config: ModelConfig,
|
||||
name: String,
|
||||
}
|
||||
|
||||
impl<M> RigProvider<M>
|
||||
where
|
||||
M: CompletionModel + Send + Sync,
|
||||
{
|
||||
pub fn new(model: M, model_config: ModelConfig, name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
model,
|
||||
model_config,
|
||||
name: name.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert goose messages to rig messages
|
||||
fn goose_messages_to_rig(messages: &[Message]) -> Vec<RigMessage> {
|
||||
let mut rig_messages = Vec::new();
|
||||
|
||||
for msg in messages {
|
||||
match msg.role {
|
||||
Role::User => {
|
||||
let mut user_content = Vec::new();
|
||||
|
||||
for content in &msg.content {
|
||||
match content {
|
||||
MessageContent::Text(text) => {
|
||||
user_content.push(UserContent::Text(Text {
|
||||
text: text.text.clone(),
|
||||
}));
|
||||
}
|
||||
MessageContent::ToolResponse(tool_response) => {
|
||||
let result_content = match &tool_response.tool_result {
|
||||
Ok(call_result) => {
|
||||
let text = call_result
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|c: &Content| {
|
||||
c.as_text().map(|s| s.text.clone())
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
OneOrMany::one(ToolResultContent::text(text))
|
||||
}
|
||||
Err(error_data) => OneOrMany::one(ToolResultContent::text(
|
||||
format!("Error: {}", error_data.message),
|
||||
)),
|
||||
};
|
||||
user_content.push(UserContent::ToolResult(RigToolResult {
|
||||
id: tool_response.id.clone(),
|
||||
call_id: None,
|
||||
content: result_content,
|
||||
}));
|
||||
}
|
||||
MessageContent::Image(image) => {
|
||||
let media_type = mime_to_image_media_type(&image.mime_type);
|
||||
let rig_image = Image {
|
||||
data: DocumentSourceKind::Base64(image.data.clone()),
|
||||
media_type: Some(media_type),
|
||||
detail: None,
|
||||
additional_params: None,
|
||||
};
|
||||
user_content.push(UserContent::Image(rig_image));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if !user_content.is_empty() {
|
||||
rig_messages.push(RigMessage::User {
|
||||
content: OneOrMany::many(user_content)
|
||||
.expect("user_content should not be empty"),
|
||||
});
|
||||
}
|
||||
}
|
||||
Role::Assistant => {
|
||||
let mut assistant_content = Vec::new();
|
||||
|
||||
for content in &msg.content {
|
||||
match content {
|
||||
MessageContent::Text(text) => {
|
||||
assistant_content.push(AssistantContent::Text(Text {
|
||||
text: text.text.clone(),
|
||||
}));
|
||||
}
|
||||
MessageContent::ToolRequest(tool_request) => {
|
||||
if let Ok(tool_call) = &tool_request.tool_call {
|
||||
assistant_content.push(AssistantContent::ToolCall(RigToolCall {
|
||||
id: tool_request.id.clone(),
|
||||
call_id: None,
|
||||
function: ToolFunction {
|
||||
name: tool_call.name.to_string(),
|
||||
arguments: tool_call
|
||||
.arguments
|
||||
.clone()
|
||||
.map(Value::Object)
|
||||
.unwrap_or(Value::Null),
|
||||
},
|
||||
signature: None,
|
||||
additional_params: None,
|
||||
}));
|
||||
}
|
||||
}
|
||||
MessageContent::Thinking(thinking) => {
|
||||
let signature = if thinking.signature.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(thinking.signature.clone())
|
||||
};
|
||||
let reasoning =
|
||||
Reasoning::new(&thinking.thinking).with_signature(signature);
|
||||
assistant_content.push(AssistantContent::Reasoning(reasoning));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if !assistant_content.is_empty() {
|
||||
rig_messages.push(RigMessage::Assistant {
|
||||
id: None,
|
||||
content: OneOrMany::many(assistant_content)
|
||||
.expect("assistant_content should not be empty"),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
rig_messages
|
||||
}
|
||||
|
||||
/// Convert a MIME type string to rig's ImageMediaType
|
||||
fn mime_to_image_media_type(mime: &str) -> ImageMediaType {
|
||||
match mime.to_lowercase().as_str() {
|
||||
"image/png" => ImageMediaType::PNG,
|
||||
"image/gif" => ImageMediaType::GIF,
|
||||
"image/webp" => ImageMediaType::WEBP,
|
||||
_ => ImageMediaType::JPEG,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert goose tools to rig tool definitions
|
||||
fn goose_tools_to_rig(tools: &[Tool]) -> Vec<ToolDefinition> {
|
||||
tools
|
||||
.iter()
|
||||
.map(|tool| ToolDefinition {
|
||||
name: tool.name.to_string(),
|
||||
description: tool.description.as_deref().unwrap_or("").to_string(),
|
||||
parameters: serde_json::to_value(&*tool.input_schema).unwrap_or(Value::Null),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Convert rig AssistantContent to goose MessageContent
|
||||
fn rig_content_to_goose(content: &AssistantContent) -> Option<MessageContent> {
|
||||
match content {
|
||||
AssistantContent::Text(text) => Some(MessageContent::text(&text.text)),
|
||||
AssistantContent::ToolCall(tool_call) => {
|
||||
let arguments = match &tool_call.function.arguments {
|
||||
Value::Object(obj) => Some(obj.clone()),
|
||||
Value::Null => None,
|
||||
other => {
|
||||
let mut map = serde_json::Map::new();
|
||||
map.insert("value".to_string(), other.clone());
|
||||
Some(map)
|
||||
}
|
||||
};
|
||||
Some(MessageContent::ToolRequest(ToolRequest {
|
||||
id: tool_call.id.clone(),
|
||||
tool_call: Ok(CallToolRequestParams {
|
||||
meta: None,
|
||||
name: Cow::Owned(tool_call.function.name.clone()),
|
||||
arguments,
|
||||
task: None,
|
||||
}),
|
||||
metadata: None,
|
||||
tool_meta: None,
|
||||
}))
|
||||
}
|
||||
AssistantContent::Reasoning(reasoning) => Some(MessageContent::Thinking(ThinkingContent {
|
||||
thinking: reasoning.reasoning.join("\n"),
|
||||
signature: reasoning.signature.clone().unwrap_or_default(),
|
||||
})),
|
||||
AssistantContent::Image(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert rig response to goose Message
|
||||
fn rig_response_to_goose<T>(response: &CompletionResponse<T>) -> Message {
|
||||
let content: Vec<MessageContent> = response
|
||||
.choice
|
||||
.iter()
|
||||
.filter_map(rig_content_to_goose)
|
||||
.collect();
|
||||
|
||||
let mut message = Message::assistant();
|
||||
for c in content {
|
||||
message = message.with_content(c);
|
||||
}
|
||||
message
|
||||
}
|
||||
|
||||
/// Convert rig Usage to goose Usage
|
||||
fn rig_usage_to_goose(usage: &rig::completion::Usage) -> Usage {
|
||||
Usage {
|
||||
input_tokens: Some(usage.input_tokens as i32),
|
||||
output_tokens: Some(usage.output_tokens as i32),
|
||||
total_tokens: Some(usage.total_tokens as i32),
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<M> Provider for RigProvider<M>
|
||||
where
|
||||
M: CompletionModel + Send + Sync + 'static,
|
||||
M::Response: Unpin + Send,
|
||||
{
|
||||
fn get_name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
fn get_model_config(&self) -> ModelConfig {
|
||||
self.model_config.clone()
|
||||
}
|
||||
|
||||
async fn complete_with_model(
|
||||
&self,
|
||||
_session_id: Option<&str>,
|
||||
model_config: &ModelConfig,
|
||||
system: &str,
|
||||
messages: &[Message],
|
||||
tools: &[Tool],
|
||||
) -> Result<(Message, ProviderUsage), ProviderError> {
|
||||
let rig_messages = goose_messages_to_rig(messages);
|
||||
let rig_tools = goose_tools_to_rig(tools);
|
||||
|
||||
let chat_history = if rig_messages.is_empty() {
|
||||
OneOrMany::one(RigMessage::user(""))
|
||||
} else {
|
||||
OneOrMany::many(rig_messages).expect("rig_messages should not be empty")
|
||||
};
|
||||
|
||||
let request = CompletionRequest {
|
||||
preamble: Some(system.to_string()),
|
||||
chat_history,
|
||||
documents: vec![],
|
||||
tools: rig_tools,
|
||||
temperature: model_config.temperature.map(|t| t as f64),
|
||||
max_tokens: model_config.max_tokens.map(|t| t as u64),
|
||||
tool_choice: None,
|
||||
additional_params: None,
|
||||
};
|
||||
|
||||
let response = self
|
||||
.model
|
||||
.completion(request)
|
||||
.await
|
||||
.map_err(|e| ProviderError::ExecutionError(e.to_string()))?;
|
||||
|
||||
let message = rig_response_to_goose(&response);
|
||||
let usage = rig_usage_to_goose(&response.usage);
|
||||
let provider_usage = ProviderUsage::new(model_config.model_name.clone(), usage);
|
||||
|
||||
Ok((message, provider_usage))
|
||||
}
|
||||
}
|
||||
|
||||
/// Wraps a goose `Provider` to implement rig's `CompletionModel` trait.
|
||||
///
|
||||
/// This allows any goose provider to be used with rig's agent framework.
|
||||
#[derive(Clone)]
|
||||
pub struct GooseProvider {
|
||||
provider: Arc<dyn Provider>,
|
||||
system_prompt: Arc<Mutex<String>>,
|
||||
}
|
||||
|
||||
impl GooseProvider {
|
||||
pub fn new(provider: Arc<dyn Provider>) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
system_prompt: Arc::new(Mutex::new(String::new())),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_system_prompt(
|
||||
provider: Arc<dyn Provider>,
|
||||
system_prompt: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
system_prompt: Arc::new(Mutex::new(system_prompt.into())),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn set_system_prompt(&self, prompt: impl Into<String>) {
|
||||
let mut sp = self.system_prompt.lock().await;
|
||||
*sp = prompt.into();
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert rig messages to goose messages
|
||||
fn rig_messages_to_goose(messages: &OneOrMany<RigMessage>) -> Vec<Message> {
|
||||
let mut goose_messages = Vec::new();
|
||||
|
||||
for msg in messages.iter() {
|
||||
match msg {
|
||||
RigMessage::User { content } => {
|
||||
let mut message = Message::user();
|
||||
|
||||
for item in content.iter() {
|
||||
match item {
|
||||
UserContent::Text(text) => {
|
||||
message = message.with_text(&text.text);
|
||||
}
|
||||
UserContent::ToolResult(result) => {
|
||||
let content_vec: Vec<Content> = result
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
ToolResultContent::Text(t) => Some(Content::text(&t.text)),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
message
|
||||
.content
|
||||
.push(MessageContent::ToolResponse(ToolResponse {
|
||||
id: result.id.clone(),
|
||||
tool_result: Ok(CallToolResult {
|
||||
content: content_vec,
|
||||
structured_content: None,
|
||||
is_error: None,
|
||||
meta: None,
|
||||
}),
|
||||
metadata: None,
|
||||
}));
|
||||
}
|
||||
UserContent::Image(_)
|
||||
| UserContent::Audio(_)
|
||||
| UserContent::Document(_)
|
||||
| UserContent::Video(_) => {}
|
||||
}
|
||||
}
|
||||
|
||||
if !message.content.is_empty() {
|
||||
goose_messages.push(message);
|
||||
}
|
||||
}
|
||||
RigMessage::Assistant { content, .. } => {
|
||||
let mut message = Message::assistant();
|
||||
|
||||
for item in content.iter() {
|
||||
match item {
|
||||
AssistantContent::Text(text) => {
|
||||
message = message.with_text(&text.text);
|
||||
}
|
||||
AssistantContent::ToolCall(tool_call) => {
|
||||
let arguments = match &tool_call.function.arguments {
|
||||
Value::Object(obj) => Some(obj.clone()),
|
||||
Value::Null => None,
|
||||
other => {
|
||||
let mut map = serde_json::Map::new();
|
||||
map.insert("value".to_string(), other.clone());
|
||||
Some(map)
|
||||
}
|
||||
};
|
||||
message
|
||||
.content
|
||||
.push(MessageContent::ToolRequest(ToolRequest {
|
||||
id: tool_call.id.clone(),
|
||||
tool_call: Ok(CallToolRequestParams {
|
||||
meta: None,
|
||||
name: Cow::Owned(tool_call.function.name.clone()),
|
||||
arguments,
|
||||
task: None,
|
||||
}),
|
||||
metadata: None,
|
||||
tool_meta: None,
|
||||
}));
|
||||
}
|
||||
AssistantContent::Reasoning(reasoning) => {
|
||||
message
|
||||
.content
|
||||
.push(MessageContent::Thinking(ThinkingContent {
|
||||
thinking: reasoning.reasoning.join("\n"),
|
||||
signature: reasoning.signature.clone().unwrap_or_default(),
|
||||
}));
|
||||
}
|
||||
AssistantContent::Image(_) => {}
|
||||
}
|
||||
}
|
||||
|
||||
if !message.content.is_empty() {
|
||||
goose_messages.push(message);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
goose_messages
|
||||
}
|
||||
|
||||
/// Convert rig tool definitions to goose tools
|
||||
fn rig_tools_to_goose(tools: &[ToolDefinition]) -> Vec<Tool> {
|
||||
tools
|
||||
.iter()
|
||||
.map(|tool| {
|
||||
let input_schema = match &tool.parameters {
|
||||
Value::Object(obj) => obj.clone(),
|
||||
_ => serde_json::Map::new(),
|
||||
};
|
||||
Tool {
|
||||
name: tool.name.clone().into(),
|
||||
title: None,
|
||||
description: Some(tool.description.clone().into()),
|
||||
input_schema: Arc::new(input_schema),
|
||||
output_schema: None,
|
||||
annotations: None,
|
||||
icons: None,
|
||||
meta: None,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Convert goose message to rig AssistantContent
|
||||
fn goose_message_to_rig_content(message: &Message) -> OneOrMany<AssistantContent> {
|
||||
let content: Vec<AssistantContent> = message
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
MessageContent::Text(text) => Some(AssistantContent::Text(Text {
|
||||
text: text.text.clone(),
|
||||
})),
|
||||
MessageContent::ToolRequest(req) => {
|
||||
if let Ok(tool_call) = &req.tool_call {
|
||||
Some(AssistantContent::ToolCall(RigToolCall {
|
||||
id: req.id.clone(),
|
||||
call_id: None,
|
||||
function: ToolFunction {
|
||||
name: tool_call.name.to_string(),
|
||||
arguments: tool_call
|
||||
.arguments
|
||||
.clone()
|
||||
.map(Value::Object)
|
||||
.unwrap_or(Value::Null),
|
||||
},
|
||||
signature: None,
|
||||
additional_params: None,
|
||||
}))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
MessageContent::Thinking(thinking) => {
|
||||
let signature = if thinking.signature.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(thinking.signature.clone())
|
||||
};
|
||||
let reasoning = Reasoning::new(&thinking.thinking).with_signature(signature);
|
||||
Some(AssistantContent::Reasoning(reasoning))
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
if content.is_empty() {
|
||||
OneOrMany::one(AssistantContent::text(""))
|
||||
} else {
|
||||
OneOrMany::many(content).unwrap_or_else(|_| OneOrMany::one(AssistantContent::text("")))
|
||||
}
|
||||
}
|
||||
|
||||
/// Response type for GooseProvider that implements GetTokenUsage
|
||||
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||
pub struct GooseResponse {
|
||||
pub usage: Option<Usage>,
|
||||
}
|
||||
|
||||
impl GetTokenUsage for GooseResponse {
|
||||
fn token_usage(&self) -> Option<rig::completion::Usage> {
|
||||
self.usage.map(|u| rig::completion::Usage {
|
||||
input_tokens: u.input_tokens.unwrap_or(0) as u64,
|
||||
output_tokens: u.output_tokens.unwrap_or(0) as u64,
|
||||
total_tokens: u.total_tokens.unwrap_or(0) as u64,
|
||||
cached_input_tokens: 0,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Dummy client type for GooseProvider
|
||||
#[derive(Clone)]
|
||||
pub struct GooseProviderClient;
|
||||
|
||||
impl CompletionModel for GooseProvider {
|
||||
type Response = GooseResponse;
|
||||
type StreamingResponse = GooseResponse;
|
||||
type Client = GooseProviderClient;
|
||||
|
||||
fn make(_client: &Self::Client, _model: impl Into<String>) -> Self {
|
||||
panic!("GooseProvider cannot be created via make(). Use GooseProvider::new() instead.")
|
||||
}
|
||||
|
||||
async fn completion(
|
||||
&self,
|
||||
request: CompletionRequest,
|
||||
) -> Result<CompletionResponse<Self::Response>, CompletionError> {
|
||||
let system_prompt = self.system_prompt.lock().await.clone();
|
||||
let system = request.preamble.as_deref().unwrap_or(&system_prompt);
|
||||
|
||||
let messages = rig_messages_to_goose(&request.chat_history);
|
||||
let tools = rig_tools_to_goose(&request.tools);
|
||||
|
||||
let (message, provider_usage) = self
|
||||
.provider
|
||||
.complete("", system, &messages, &tools)
|
||||
.await
|
||||
.map_err(|e| CompletionError::ProviderError(e.to_string()))?;
|
||||
|
||||
let choice = goose_message_to_rig_content(&message);
|
||||
let usage = provider_usage.usage;
|
||||
|
||||
Ok(CompletionResponse {
|
||||
choice,
|
||||
usage: rig::completion::Usage {
|
||||
input_tokens: usage.input_tokens.unwrap_or(0) as u64,
|
||||
output_tokens: usage.output_tokens.unwrap_or(0) as u64,
|
||||
total_tokens: usage.total_tokens.unwrap_or(0) as u64,
|
||||
cached_input_tokens: 0,
|
||||
},
|
||||
raw_response: GooseResponse { usage: Some(usage) },
|
||||
})
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
request: CompletionRequest,
|
||||
) -> Result<StreamingCompletionResponse<Self::StreamingResponse>, CompletionError> {
|
||||
// Implement streaming by calling completion and yielding the result.
|
||||
// A proper implementation would use the provider's native streaming if available.
|
||||
let response = self.completion(request).await?;
|
||||
|
||||
let text_content: String = response
|
||||
.choice
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
AssistantContent::Text(t) => Some(t.text.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("");
|
||||
|
||||
let tool_calls: Vec<_> = response
|
||||
.choice
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
AssistantContent::ToolCall(tc) => Some(tc.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let final_response = response.raw_response.clone();
|
||||
|
||||
let mut items: Vec<RawStreamingChoice<GooseResponse>> = Vec::new();
|
||||
|
||||
if !text_content.is_empty() {
|
||||
items.push(RawStreamingChoice::Message(text_content));
|
||||
}
|
||||
|
||||
for tool_call in tool_calls {
|
||||
items.push(RawStreamingChoice::ToolCall(
|
||||
rig::streaming::RawStreamingToolCall {
|
||||
id: tool_call.id.clone(),
|
||||
internal_call_id: nanoid::nanoid!(),
|
||||
call_id: tool_call.call_id.clone(),
|
||||
name: tool_call.function.name.clone(),
|
||||
arguments: tool_call.function.arguments.clone(),
|
||||
signature: tool_call.signature.clone(),
|
||||
additional_params: tool_call.additional_params.clone(),
|
||||
},
|
||||
));
|
||||
}
|
||||
|
||||
items.push(RawStreamingChoice::FinalResponse(final_response));
|
||||
|
||||
let stream: StreamingResult<GooseResponse> =
|
||||
Box::pin(stream::iter(items.into_iter().map(Ok)));
|
||||
|
||||
Ok(StreamingCompletionResponse::stream(stream))
|
||||
}
|
||||
|
||||
fn completion_request(&self, prompt: impl Into<RigMessage>) -> CompletionRequestBuilder<Self> {
|
||||
CompletionRequestBuilder::new(self.clone(), prompt)
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderDef for RigProvider<anthropic::completion::CompletionModel> {
|
||||
type Provider = RigProvider<anthropic::completion::CompletionModel>;
|
||||
|
||||
fn metadata() -> super::base::ProviderMetadata
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
super::base::ProviderMetadata {
|
||||
name: "anthropic-rig".to_owned(),
|
||||
display_name: "Anthropic (rig)".to_owned(),
|
||||
description: "".to_owned(),
|
||||
default_model: "claude-4-opus".to_owned(),
|
||||
known_models: vec![],
|
||||
model_doc_link: "".to_owned(),
|
||||
config_keys: vec![],
|
||||
allows_unlisted_models: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn from_env(_model: ModelConfig) -> BoxFuture<'static, anyhow::Result<Self::Provider>>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
let model = claude_opus();
|
||||
Box::pin(async move {
|
||||
Ok(RigProvider::new(
|
||||
model,
|
||||
ModelConfig::new("claude opus (rig)").unwrap(),
|
||||
"claude opus (rig)",
|
||||
))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_goose_messages_to_rig_user() {
|
||||
let messages = vec![Message::user().with_text("Hello, world!")];
|
||||
let rig_messages = goose_messages_to_rig(&messages);
|
||||
|
||||
assert_eq!(rig_messages.len(), 1);
|
||||
match &rig_messages[0] {
|
||||
RigMessage::User { content } => {
|
||||
assert_eq!(content.iter().count(), 1);
|
||||
}
|
||||
_ => panic!("Expected User message"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_goose_messages_to_rig_assistant() {
|
||||
let messages = vec![Message::assistant().with_text("Hello!")];
|
||||
let rig_messages = goose_messages_to_rig(&messages);
|
||||
|
||||
assert_eq!(rig_messages.len(), 1);
|
||||
match &rig_messages[0] {
|
||||
RigMessage::Assistant { content, .. } => {
|
||||
assert_eq!(content.iter().count(), 1);
|
||||
}
|
||||
_ => panic!("Expected Assistant message"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rig_messages_to_goose_user() {
|
||||
let messages = OneOrMany::one(RigMessage::user("Hello!"));
|
||||
let goose_messages = rig_messages_to_goose(&messages);
|
||||
|
||||
assert_eq!(goose_messages.len(), 1);
|
||||
assert_eq!(goose_messages[0].role, Role::User);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rig_messages_to_goose_assistant() {
|
||||
let messages = OneOrMany::one(RigMessage::assistant("Hello!"));
|
||||
let goose_messages = rig_messages_to_goose(&messages);
|
||||
|
||||
assert_eq!(goose_messages.len(), 1);
|
||||
assert_eq!(goose_messages[0].role, Role::Assistant);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_conversion_roundtrip() {
|
||||
let goose_tool = Tool {
|
||||
name: "test_tool".into(),
|
||||
title: None,
|
||||
description: Some("A test tool".into()),
|
||||
input_schema: Arc::new(serde_json::Map::new()),
|
||||
output_schema: None,
|
||||
annotations: None,
|
||||
icons: None,
|
||||
meta: None,
|
||||
};
|
||||
|
||||
let rig_tools = goose_tools_to_rig(std::slice::from_ref(&goose_tool));
|
||||
assert_eq!(rig_tools.len(), 1);
|
||||
assert_eq!(rig_tools[0].name, "test_tool");
|
||||
assert_eq!(rig_tools[0].description, "A test tool");
|
||||
|
||||
let back_to_goose = rig_tools_to_goose(&rig_tools);
|
||||
assert_eq!(back_to_goose.len(), 1);
|
||||
assert_eq!(back_to_goose[0].name, "test_tool");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user