use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use harmony_response::{
chat::{Conversation, DeveloperContent, Message, Role, SystemContent, ToolDescription, ToolNamespaceConfig},
load_harmony_encoding, HarmonyEncodingName, StreamableParser,
};
fn try_load_encoding() -> Option<harmony_response::HarmonyEncoding> {
load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).ok()
}
fn create_simple_conversation() -> Conversation {
Conversation::from_messages([
Message::from_role_and_content(Role::System, SystemContent::new()),
Message::from_role_and_content(Role::User, "Hello, world!"),
Message::from_role_and_content(Role::Assistant, "Hello! How can I help you?"),
])
}
fn create_complex_conversation() -> Conversation {
let tools = vec![
ToolDescription::new(
"calculate",
"Performs mathematical calculations",
Some(serde_json::json!({
"type": "object",
"properties": {
"expression": {"type": "string"},
"precision": {"type": "number", "default": 2}
},
"required": ["expression"]
}))
),
ToolDescription::new(
"search",
"Searches for information",
Some(serde_json::json!({
"type": "object",
"properties": {
"query": {"type": "string"},
"limit": {"type": "number", "default": 10}
},
"required": ["query"]
}))
)
];
let system_content = SystemContent::new()
.with_model_identity("Benchmark Assistant")
.with_required_channels(["analysis", "commentary", "final"])
.with_browser_tool()
.with_conversation_start_date("2024-01-01")
.with_knowledge_cutoff("2024-06");
let dev_content = DeveloperContent::new()
.with_instructions("You are a helpful benchmark assistant. Be thorough and accurate.")
.with_function_tools(tools);
Conversation::from_messages([
Message::from_role_and_content(Role::System, system_content),
Message::from_role_and_content(Role::Developer, dev_content),
Message::from_role_and_content(Role::User, "Can you help me calculate 15 * 23 and then search for information about the result?"),
Message::from_role_and_content(Role::Assistant, "I'll help you with both tasks. Let me start by calculating 15 * 23.")
.with_channel("analysis"),
Message::from_role_and_content(Role::Assistant, "First, I need to perform the calculation.")
.with_channel("commentary"),
Message::from_role_and_content(Role::Assistant, "I'll calculate this step by step.")
.with_recipient("functions")
.with_channel("commentary"),
Message::from_role_and_content(Role::Tool, r#"{"result": 345}"#)
.with_recipient("assistant"),
Message::from_role_and_content(Role::Assistant, "The calculation 15 * 23 equals 345. Now I'll search for information about this number.")
.with_channel("final"),
])
}
fn create_large_conversation() -> Conversation {
let mut messages = vec![
Message::from_role_and_content(
Role::System,
SystemContent::new()
.with_model_identity("Large Conversation Benchmark")
.with_required_channels(["analysis", "commentary", "final"])
)
];
for i in 0..50 {
messages.push(Message::from_role_and_content(
Role::User,
format!("This is user message number {} with some content to make it longer and more realistic for benchmarking purposes.", i + 1)
));
messages.push(Message::from_role_and_content(
Role::Assistant,
format!("I'm analyzing user message {}. This requires careful consideration.", i + 1)
).with_channel("analysis"));
messages.push(Message::from_role_and_content(
Role::Assistant,
format!("Let me provide some commentary on message {}.", i + 1)
).with_channel("commentary"));
messages.push(Message::from_role_and_content(
Role::Assistant,
format!("This is my response to message number {}. I understand your request and will help you with whatever you need.", i + 1)
).with_channel("final"));
}
Conversation::from_messages(messages)
}
fn bench_conversation_creation(c: &mut Criterion) {
let mut group = c.benchmark_group("conversation_creation");
group.bench_function("simple", |b| {
b.iter(|| {
black_box(create_simple_conversation());
});
});
group.bench_function("complex", |b| {
b.iter(|| {
black_box(create_complex_conversation());
});
});
group.bench_function("large", |b| {
b.iter(|| {
black_box(create_large_conversation());
});
});
group.finish();
}
fn bench_message_operations(c: &mut Criterion) {
let mut group = c.benchmark_group("message_operations");
group.bench_function("create_user_message", |b| {
b.iter(|| {
black_box(Message::from_role_and_content(Role::User, "Test message"));
});
});
group.bench_function("create_system_message", |b| {
b.iter(|| {
black_box(Message::from_role_and_content(Role::System, SystemContent::new()));
});
});
group.bench_function("create_message_with_channel", |b| {
b.iter(|| {
black_box(
Message::from_role_and_content(Role::Assistant, "Test response")
.with_channel("final")
.with_recipient("user")
);
});
});
group.finish();
}
fn bench_system_content_builder(c: &mut Criterion) {
let mut group = c.benchmark_group("system_content_builder");
group.bench_function("basic", |b| {
b.iter(|| {
black_box(SystemContent::new());
});
});
group.bench_function("with_tools", |b| {
b.iter(|| {
black_box(
SystemContent::new()
.with_browser_tool()
.with_python_tool()
.with_required_channels(["analysis", "commentary", "final"])
);
});
});
group.bench_function("full_configuration", |b| {
b.iter(|| {
black_box(
SystemContent::new()
.with_model_identity("Benchmark Assistant")
.with_required_channels(["analysis", "commentary", "final"])
.with_browser_tool()
.with_python_tool()
.with_conversation_start_date("2024-01-01")
.with_knowledge_cutoff("2024-06")
);
});
});
group.finish();
}
fn bench_encoding_operations(c: &mut Criterion) {
if let Some(encoding) = try_load_encoding() {
let mut group = c.benchmark_group("encoding_operations");
let simple_conv = create_simple_conversation();
let complex_conv = create_complex_conversation();
let large_conv = create_large_conversation();
group.bench_function("render_simple", |b| {
b.iter(|| {
black_box(encoding.render_conversation(&simple_conv, None).unwrap());
});
});
group.bench_function("render_complex", |b| {
b.iter(|| {
black_box(encoding.render_conversation(&complex_conv, None).unwrap());
});
});
group.bench_function("render_large", |b| {
b.iter(|| {
black_box(encoding.render_conversation(&large_conv, None).unwrap());
});
});
group.bench_function("render_for_completion_simple", |b| {
b.iter(|| {
black_box(
encoding.render_conversation_for_completion(&simple_conv, Role::Assistant, None)
.unwrap()
);
});
});
group.bench_function("render_for_training_simple", |b| {
b.iter(|| {
black_box(
encoding.render_conversation_for_training(&simple_conv, None)
.unwrap()
);
});
});
group.finish();
} else {
eprintln!("⚠️ Encoding benchmarks skipped (no network access)");
}
}
fn bench_streaming_parser(c: &mut Criterion) {
if let Some(encoding) = try_load_encoding() {
let mut group = c.benchmark_group("streaming_parser");
let simple_conv = create_simple_conversation();
let tokens = encoding.render_conversation_for_completion(&simple_conv, Role::Assistant, None)
.unwrap_or_else(|_| vec![200006, 1234, 5678, 200008, 9012, 200007]);
group.throughput(Throughput::Elements(tokens.len() as u64));
group.bench_function("create_parser", |b| {
b.iter(|| {
black_box(StreamableParser::new(encoding.clone(), Some(Role::Assistant)).unwrap());
});
});
group.bench_function("process_tokens", |b| {
b.iter(|| {
let mut parser = StreamableParser::new(encoding.clone(), Some(Role::Assistant)).unwrap();
for &token in &tokens {
black_box(parser.process(token).unwrap());
}
black_box(parser.process_eos().unwrap());
});
});
group.bench_function("process_single_token", |b| {
b.iter(|| {
let mut parser = StreamableParser::new(encoding.clone(), Some(Role::Assistant)).unwrap();
black_box(parser.process(200006).unwrap()); });
});
group.finish();
} else {
eprintln!("⚠️ Streaming parser benchmarks skipped (no network access)");
}
}
fn bench_role_operations(c: &mut Criterion) {
let mut group = c.benchmark_group("role_operations");
let role_strings = ["user", "assistant", "system", "developer", "tool"];
group.bench_function("role_from_str", |b| {
b.iter(|| {
for role_str in &role_strings {
black_box(Role::try_from(*role_str).unwrap());
}
});
});
group.bench_function("role_to_str", |b| {
let roles = [Role::User, Role::Assistant, Role::System, Role::Developer, Role::Tool];
b.iter(|| {
for role in &roles {
black_box(role.as_str());
}
});
});
group.finish();
}
fn bench_serialization(c: &mut Criterion) {
let mut group = c.benchmark_group("serialization");
let simple_conv = create_simple_conversation();
let complex_conv = create_complex_conversation();
group.bench_function("serialize_simple_conversation", |b| {
b.iter(|| {
black_box(serde_json::to_string(&simple_conv).unwrap());
});
});
group.bench_function("serialize_complex_conversation", |b| {
b.iter(|| {
black_box(serde_json::to_string(&complex_conv).unwrap());
});
});
let simple_json = serde_json::to_string(&simple_conv).unwrap();
let complex_json = serde_json::to_string(&complex_conv).unwrap();
group.bench_function("deserialize_simple_conversation", |b| {
b.iter(|| {
black_box(serde_json::from_str::<Conversation>(&simple_json).unwrap());
});
});
group.bench_function("deserialize_complex_conversation", |b| {
b.iter(|| {
black_box(serde_json::from_str::<Conversation>(&complex_json).unwrap());
});
});
group.finish();
}
criterion_group!(
benches,
bench_conversation_creation,
bench_message_operations,
bench_system_content_builder,
bench_encoding_operations,
bench_streaming_parser,
bench_role_operations,
bench_serialization
);
criterion_main!(benches);