use criterion::{Criterion, criterion_group, criterion_main};
use daimon::agent::Agent;
use daimon::error::Result;
use daimon::memory::{SlidingWindowMemory, TokenWindowMemory};
use daimon::model::Model;
use daimon::model::types::{
ChatRequest, ChatResponse, Message, Role, StopReason, Usage,
};
use daimon::orchestration::dag::{Dag, DagContext, FnDagNode, END, START};
use daimon::orchestration::chain::{Chain, TransformStep};
use daimon::stream::ResponseStream;
use daimon::tool::{Tool, ToolOutput, ToolRegistry};
struct InstantModel;
impl Model for InstantModel {
async fn generate(&self, request: &ChatRequest) -> Result<ChatResponse> {
let reply = if !request.tools.is_empty()
&& request.messages.last().is_some_and(|m| m.role == Role::User)
{
ChatResponse {
message: Message::assistant("done"),
stop_reason: StopReason::EndTurn,
usage: Some(Usage {
input_tokens: 100,
output_tokens: 50,
cached_tokens: 0,
}),
}
} else {
ChatResponse {
message: Message::assistant("response"),
stop_reason: StopReason::EndTurn,
usage: Some(Usage {
input_tokens: 100,
output_tokens: 50,
cached_tokens: 0,
}),
}
};
Ok(reply)
}
async fn generate_stream(&self, _request: &ChatRequest) -> Result<ResponseStream> {
Ok(Box::pin(futures::stream::empty()))
}
}
struct NoOpTool;
impl Tool for NoOpTool {
fn name(&self) -> &str {
"noop"
}
fn description(&self) -> &str {
"does nothing"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({ "type": "object", "properties": {} })
}
async fn execute(&self, _input: &serde_json::Value) -> Result<ToolOutput> {
Ok(ToolOutput::text("ok"))
}
}
fn bench_agent_prompt(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let agent = Agent::builder()
.model(InstantModel)
.system_prompt("You are helpful.")
.build()
.unwrap();
c.bench_function("agent_prompt_simple", |b| {
b.iter(|| {
rt.block_on(async {
let _ = agent.prompt("hello").await.unwrap();
})
})
});
}
fn bench_agent_with_tools(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let agent = Agent::builder()
.model(InstantModel)
.tool(NoOpTool)
.build()
.unwrap();
c.bench_function("agent_prompt_with_tools", |b| {
b.iter(|| {
rt.block_on(async {
let _ = agent.prompt("hello").await.unwrap();
})
})
});
}
fn bench_sliding_window_memory(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
c.bench_function("memory_sliding_window_50", |b| {
b.iter(|| {
rt.block_on(async {
use daimon::memory::Memory;
let mem = SlidingWindowMemory::new(50);
for i in 0..100 {
mem.add_message(Message::user(format!("msg {i}"))).await.unwrap();
}
let msgs = mem.get_messages().await.unwrap();
assert_eq!(msgs.len(), 50);
})
})
});
}
fn bench_token_window_memory(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
c.bench_function("memory_token_window_1000", |b| {
b.iter(|| {
rt.block_on(async {
use daimon::memory::Memory;
let mem = TokenWindowMemory::new(1000);
for i in 0..100 {
mem.add_message(Message::user(format!("message number {i} with some content")))
.await
.unwrap();
}
let msgs = mem.get_messages().await.unwrap();
assert!(msgs.len() <= 100);
})
})
});
}
fn bench_tool_registry(c: &mut Criterion) {
let mut registry = ToolRegistry::new();
for i in 0..50 {
registry.register(NamedTool(format!("tool_{i}"))).unwrap();
}
c.bench_function("tool_registry_lookup_50", |b| {
b.iter(|| {
let _ = registry.get("tool_25");
})
});
c.bench_function("tool_registry_specs_50_uncached", |b| {
b.iter(|| {
let _ = registry.tool_specs();
})
});
registry.warm_cache();
c.bench_function("tool_registry_specs_50_cached", |b| {
b.iter(|| {
let _ = registry.tool_specs();
})
});
c.bench_function("tool_registry_validate_cached", |b| {
let input = serde_json::json!({"key": "value"});
b.iter(|| {
let _ = registry.validate_input("tool_25", &input);
})
});
}
struct NamedTool(String);
impl Tool for NamedTool {
fn name(&self) -> &str {
&self.0
}
fn description(&self) -> &str {
"a tool"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({ "type": "object" })
}
async fn execute(&self, _input: &serde_json::Value) -> Result<ToolOutput> {
Ok(ToolOutput::text("ok"))
}
}
fn bench_dag_fan_out(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let dag = Dag::builder()
.node(
"a",
FnDagNode::new(|ctx| {
Box::pin(async move {
ctx.set("a", serde_json::json!(1));
Ok(())
})
}),
)
.node(
"b",
FnDagNode::new(|ctx| {
Box::pin(async move {
ctx.set("b", serde_json::json!(2));
Ok(())
})
}),
)
.node(
"c",
FnDagNode::new(|ctx| {
Box::pin(async move {
ctx.set("c", serde_json::json!(3));
Ok(())
})
}),
)
.node(
"merge",
FnDagNode::new(|ctx| {
Box::pin(async move {
let sum = ctx.get("a").and_then(|v| v.as_i64()).unwrap_or(0)
+ ctx.get("b").and_then(|v| v.as_i64()).unwrap_or(0)
+ ctx.get("c").and_then(|v| v.as_i64()).unwrap_or(0);
ctx.set("sum", serde_json::json!(sum));
Ok(())
})
}),
)
.edge(START, "a")
.edge(START, "b")
.edge(START, "c")
.edge("a", "merge")
.edge("b", "merge")
.edge("c", "merge")
.edge("merge", END)
.build()
.unwrap();
c.bench_function("dag_fan_out_3_merge", |b| {
b.iter(|| {
rt.block_on(async {
let _ = dag.run(DagContext::new()).await.unwrap();
})
})
});
}
fn bench_chain_3_steps(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let chain = Chain::builder()
.step(TransformStep::new(|mut ctx| {
Box::pin(async move {
ctx.metadata.insert("step1".to_string(), serde_json::json!(true));
Ok(ctx)
})
}))
.step(TransformStep::new(|mut ctx| {
Box::pin(async move {
ctx.metadata.insert("step2".to_string(), serde_json::json!(true));
Ok(ctx)
})
}))
.step(TransformStep::new(|mut ctx| {
Box::pin(async move {
ctx.metadata.insert("step3".to_string(), serde_json::json!(true));
Ok(ctx)
})
}))
.build()
.unwrap();
c.bench_function("chain_3_transforms", |b| {
b.iter(|| {
rt.block_on(async {
let _ = chain.run("input").await.unwrap();
})
})
});
}
fn bench_token_estimation(c: &mut Criterion) {
let text = "The quick brown fox jumps over the lazy dog. ".repeat(100);
c.bench_function("token_estimate_4500_chars", |b| {
b.iter(|| {
let _estimate = text.len().div_ceil(4);
})
});
}
fn bench_hot_swap_prompt(c: &mut Criterion) {
use daimon::agent::hot_swap::HotSwapAgent;
let rt = tokio::runtime::Runtime::new().unwrap();
let agent = Agent::builder()
.model(InstantModel)
.system_prompt("You are helpful.")
.build()
.unwrap();
let hot = HotSwapAgent::new(agent);
c.bench_function("hot_swap_prompt_simple", |b| {
b.iter(|| {
rt.block_on(async {
let _ = hot.prompt("hello").await.unwrap();
})
})
});
}
fn bench_hot_swap_swap_model(c: &mut Criterion) {
use daimon::agent::hot_swap::HotSwapAgent;
let rt = tokio::runtime::Runtime::new().unwrap();
let agent = Agent::builder()
.model(InstantModel)
.build()
.unwrap();
let hot = HotSwapAgent::new(agent);
c.bench_function("hot_swap_swap_model", |b| {
b.iter(|| {
rt.block_on(async {
hot.swap_model(InstantModel).await;
})
})
});
}
fn bench_in_process_broker(c: &mut Criterion) {
use daimon::distributed::{InProcessBroker, TaskBroker, AgentTask, TaskResult};
let rt = tokio::runtime::Runtime::new().unwrap();
c.bench_function("broker_submit_receive_complete", |b| {
b.iter(|| {
rt.block_on(async {
let broker = InProcessBroker::new(64);
let task = AgentTask::new("bench input");
let id = broker.submit(task).await.unwrap();
let received = broker.receive().await.unwrap().unwrap();
assert_eq!(received.task_id, id);
let result = TaskResult {
task_id: id,
output: "done".into(),
iterations: 1,
cost: 0.0,
error: None,
};
broker.complete(&received.task_id, result).await.unwrap();
})
})
});
}
fn bench_event_bus(c: &mut Criterion) {
use daimon::distributed::streaming::{
InProcessEventBus, SerializableStreamEvent, TaskEventBus, TaskStreamEvent,
};
let rt = tokio::runtime::Runtime::new().unwrap();
let bus = InProcessEventBus::new(1024);
let mut rx = bus.subscribe();
c.bench_function("event_bus_publish_receive", |b| {
b.iter(|| {
rt.block_on(async {
bus.publish(TaskStreamEvent {
task_id: "t1".into(),
event: SerializableStreamEvent::TextDelta("chunk".into()),
})
.await
.unwrap();
let _ = rx.recv().await.unwrap();
})
})
});
}
fn bench_checkpoint_memory(c: &mut Criterion) {
use daimon::checkpoint::{Checkpoint, CheckpointState, InMemoryCheckpoint};
use daimon::model::types::Message;
let rt = tokio::runtime::Runtime::new().unwrap();
c.bench_function("checkpoint_save_load_memory", |b| {
b.iter(|| {
rt.block_on(async {
let cp = InMemoryCheckpoint::new();
let state = CheckpointState::new(
"run1",
vec![Message::user("hello"), Message::assistant("hi")],
1,
);
cp.save(&state).await.unwrap();
let loaded = cp.load("run1").await.unwrap().unwrap();
assert_eq!(loaded.messages.len(), 2);
})
})
});
}
fn bench_serializable_stream_event(c: &mut Criterion) {
use daimon::distributed::streaming::SerializableStreamEvent;
let event = SerializableStreamEvent::Usage {
iteration: 3,
input_tokens: 512,
output_tokens: 128,
estimated_cost: 0.0023,
};
c.bench_function("stream_event_serialize", |b| {
b.iter(|| {
let _ = serde_json::to_string(&event).unwrap();
})
});
let json = serde_json::to_string(&event).unwrap();
c.bench_function("stream_event_deserialize", |b| {
b.iter(|| {
let _: SerializableStreamEvent = serde_json::from_str(&json).unwrap();
})
});
}
criterion_group!(
agent_benches,
bench_agent_prompt,
bench_agent_with_tools,
);
criterion_group!(
memory_benches,
bench_sliding_window_memory,
bench_token_window_memory,
);
criterion_group!(
tool_benches,
bench_tool_registry,
);
criterion_group!(
orchestration_benches,
bench_dag_fan_out,
bench_chain_3_steps,
);
criterion_group!(
misc_benches,
bench_token_estimation,
);
criterion_group!(
new_component_benches,
bench_hot_swap_prompt,
bench_hot_swap_swap_model,
bench_in_process_broker,
bench_event_bus,
bench_checkpoint_memory,
bench_serializable_stream_event,
);
criterion_main!(
agent_benches,
memory_benches,
tool_benches,
orchestration_benches,
misc_benches,
new_component_benches,
);