9ae9045584
* Centralize session files to goose::session module * Write session metadata and messages in jsonl * Refactor CLI build_session to use goose::session functions * Track session's token usage by adding optional session_id in agent.reply(...) * NOTE: Only sessions saved through the updates goose::session functions will show up in GUI Co-authored-by: Bradley Axen <baxen@squareup.com>
42 lines
1.2 KiB
Rust
42 lines
1.2 KiB
Rust
use dotenv::dotenv;
|
|
use futures::StreamExt;
|
|
use goose::agents::{AgentFactory, ExtensionConfig};
|
|
use goose::config::DEFAULT_EXTENSION_TIMEOUT;
|
|
use goose::message::Message;
|
|
use goose::providers::databricks::DatabricksProvider;
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
// Setup a model provider from env vars
|
|
let _ = dotenv();
|
|
|
|
let provider = Box::new(DatabricksProvider::default());
|
|
|
|
// Setup an agent with the developer extension
|
|
let mut agent = AgentFactory::create("reference", provider).expect("default should exist");
|
|
|
|
let config = ExtensionConfig::stdio(
|
|
"developer",
|
|
"./target/debug/developer",
|
|
DEFAULT_EXTENSION_TIMEOUT,
|
|
);
|
|
agent.add_extension(config).await.unwrap();
|
|
|
|
println!("Extensions:");
|
|
for extension in agent.list_extensions().await {
|
|
println!(" {}", extension);
|
|
}
|
|
|
|
let messages = vec![Message::user()
|
|
.with_text("can you summarize the readme.md in this dir using just a haiku?")];
|
|
|
|
let mut stream = agent.reply(&messages, None).await.unwrap();
|
|
while let Some(message) = stream.next().await {
|
|
println!(
|
|
"{}",
|
|
serde_json::to_string_pretty(&message.unwrap()).unwrap()
|
|
);
|
|
println!("\n");
|
|
}
|
|
}
|