325bf396af
Co-authored-by: Alex Hancock <alexhancock@block.xyz>
85 lines
3.1 KiB
Rust
85 lines
3.1 KiB
Rust
use anyhow::Result;
|
|
use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
|
|
use dotenvy::dotenv;
|
|
use goose::conversation::message::Message;
|
|
use goose::providers::anthropic::ANTHROPIC_DEFAULT_MODEL;
|
|
use goose::providers::create_with_named_model;
|
|
use goose::providers::databricks::DATABRICKS_DEFAULT_MODEL;
|
|
use goose::providers::openai::OPEN_AI_DEFAULT_MODEL;
|
|
use rmcp::model::{CallToolRequestParams, Content, Tool};
|
|
use rmcp::object;
|
|
use std::fs;
|
|
use std::sync::Arc;
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
// Load environment variables from .env file
|
|
dotenv().ok();
|
|
|
|
// Create providers
|
|
let providers: Vec<Arc<dyn goose::providers::base::Provider>> = vec![
|
|
create_with_named_model("databricks", DATABRICKS_DEFAULT_MODEL, Vec::new()).await?,
|
|
create_with_named_model("openai", OPEN_AI_DEFAULT_MODEL, Vec::new()).await?,
|
|
create_with_named_model("anthropic", ANTHROPIC_DEFAULT_MODEL, Vec::new()).await?,
|
|
];
|
|
for provider in providers {
|
|
// Read and encode test image
|
|
let image_data = fs::read("crates/goose/examples/test_assets/test_image.png")?;
|
|
let base64_image = BASE64.encode(image_data);
|
|
|
|
// Create a message sequence that includes a tool response with both text and image
|
|
let messages = vec![
|
|
Message::user().with_text("Read the image at ./test_image.png please"),
|
|
Message::assistant().with_tool_request(
|
|
"000",
|
|
Ok(CallToolRequestParams::new("view_image")
|
|
.with_arguments(object!({"path": "./test_image.png"}))),
|
|
),
|
|
Message::user().with_tool_response(
|
|
"000",
|
|
Ok(rmcp::model::CallToolResult::success(vec![Content::image(
|
|
base64_image,
|
|
"image/png",
|
|
)])),
|
|
),
|
|
];
|
|
|
|
// Get a response from the model about the image
|
|
let input_schema = object!({
|
|
"type": "object",
|
|
"required": ["path"],
|
|
"properties": {
|
|
"path": {
|
|
"type": "string",
|
|
"default": null,
|
|
"description": "The path to the image"
|
|
},
|
|
}
|
|
});
|
|
let model_config = provider.get_model_config();
|
|
let (response, usage) = provider
|
|
.complete(
|
|
&model_config,
|
|
"",
|
|
"You are a helpful assistant. Please describe any text you see in the image.",
|
|
&messages,
|
|
&[Tool::new("view_image", "View an image", input_schema)],
|
|
)
|
|
.await?;
|
|
|
|
// Print the response and usage statistics
|
|
println!("\nResponse from AI:");
|
|
println!("---------------");
|
|
for content in response.content {
|
|
println!("{:?}", content);
|
|
}
|
|
println!("\nToken Usage:");
|
|
println!("------------");
|
|
println!("Input tokens: {:?}", usage.usage.input_tokens);
|
|
println!("Output tokens: {:?}", usage.usage.output_tokens);
|
|
println!("Total tokens: {:?}", usage.usage.total_tokens);
|
|
}
|
|
|
|
Ok(())
|
|
}
|