Sheikh Limon
d4a7bd39be
Signed-off-by: sheikhlimon <sheikhlimon404@gmail.com> Signed-off-by: Sheikh Limon <sheikhlimon404@gmail.com>
69 lines
2.2 KiB
Rust
69 lines
2.2 KiB
Rust
use crate::session::build_session;
|
|
use crate::session::SessionBuilderConfig;
|
|
use crate::{logging, CliSession};
|
|
use async_trait::async_trait;
|
|
use goose::conversation::Conversation;
|
|
use goose_bench::bench_session::{BenchAgent, BenchBaseSession};
|
|
use goose_bench::eval_suites::ExtensionRequirements;
|
|
use std::sync::Arc;
|
|
use tokio::sync::Mutex;
|
|
|
|
// allow session obj to be used in benchmarking
|
|
#[async_trait]
|
|
impl BenchBaseSession for CliSession {
|
|
async fn headless(&mut self, message: String) -> anyhow::Result<()> {
|
|
self.headless(message).await
|
|
}
|
|
fn message_history(&self) -> Conversation {
|
|
self.message_history()
|
|
}
|
|
fn get_total_token_usage(&self) -> anyhow::Result<Option<i32>> {
|
|
// Since the trait requires sync but the session method is async,
|
|
// we need to block on the async call
|
|
tokio::task::block_in_place(|| {
|
|
tokio::runtime::Handle::current().block_on(self.get_total_token_usage())
|
|
})
|
|
}
|
|
|
|
fn get_session_id(&self) -> anyhow::Result<String> {
|
|
Ok(self.session_id().to_string())
|
|
}
|
|
}
|
|
pub async fn agent_generator(
|
|
requirements: ExtensionRequirements,
|
|
session_id: String,
|
|
) -> BenchAgent {
|
|
let base_session = build_session(SessionBuilderConfig {
|
|
session_id: Some(session_id),
|
|
resume: false,
|
|
no_session: false,
|
|
extensions: requirements.external,
|
|
remote_extensions: requirements.remote,
|
|
streamable_http_extensions: Vec::new(),
|
|
builtins: requirements.builtin,
|
|
extensions_override: None,
|
|
additional_system_prompt: None,
|
|
settings: None,
|
|
provider: None,
|
|
model: None,
|
|
debug: false,
|
|
max_tool_repetitions: None,
|
|
interactive: false, // Benchmarking is non-interactive
|
|
scheduled_job_id: None,
|
|
max_turns: None,
|
|
quiet: false,
|
|
sub_recipes: None,
|
|
final_output_response: None,
|
|
retry_config: None,
|
|
output_format: "text".to_string(),
|
|
})
|
|
.await;
|
|
|
|
let bench_agent = BenchAgent::new(Box::new(base_session));
|
|
|
|
let errors = Some(Arc::new(Mutex::new(bench_agent.get_errors().await)));
|
|
logging::setup_logging(Some("bench"), errors).expect("Failed to initialize logging");
|
|
|
|
bench_agent
|
|
}
|