Adrian Cole
44 lines
1.2 KiB
Rust
44 lines
1.2 KiB
Rust
use anyhow::Result;
|
|
use dotenvy::dotenv;
|
|
use goose::conversation::message::Message;
|
|
use goose::providers::create_with_named_model;
|
|
use goose::providers::databricks::DATABRICKS_DEFAULT_MODEL;
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
dotenv().ok();
|
|
|
|
// Clear any token to force OAuth
|
|
std::env::remove_var("DATABRICKS_TOKEN");
|
|
|
|
// Create the provider
|
|
let provider =
|
|
create_with_named_model("databricks", DATABRICKS_DEFAULT_MODEL, Vec::new()).await?;
|
|
|
|
// Create a simple message
|
|
let message = Message::user().with_text("Tell me a short joke about programming.");
|
|
|
|
// Get a response
|
|
let (response, usage) = provider
|
|
.complete_with_model(
|
|
None,
|
|
&provider.get_model_config(),
|
|
"You are a helpful assistant.",
|
|
&[message],
|
|
&[],
|
|
)
|
|
.await?;
|
|
|
|
println!("\nResponse from AI:");
|
|
println!("---------------");
|
|
println!("{:?}", response);
|
|
|
|
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(())
|
|
}
|