Vapora/crates/vapora-knowledge-graph/benches/kg_benchmarks.rs

use std::hint::black_box;

use chrono::Utc;
use criterion::{criterion_group, criterion_main, Criterion};
use vapora_knowledge_graph::{ExecutionRecord, TemporalKG};

async fn setup_kg_with_records(count: usize) -> TemporalKG {
    let kg = TemporalKG::new("ws://localhost:8000", "root", "root")
        .await
        .unwrap();

    for i in 0..count {
        let record = ExecutionRecord {
            id: format!("exec-{}", i),
            task_id: format!("task-{}", i),
            agent_id: format!("agent-{}", i % 10),
            agent_role: None,
            task_type: match i % 3 {
                0 => "coding".to_string(),
                1 => "analysis".to_string(),
                _ => "review".to_string(),
            },
            description: format!("Execute task {} with description", i),
            duration_ms: 1000 + (i as u64 * 100) % 5000,
            input_tokens: 100 + (i as u64 * 10),
            output_tokens: 200 + (i as u64 * 20),
            cost_cents: (i as u32 % 100 + 1) * 10,
            provider: "claude".to_string(),
            success: i % 10 != 0,
            error: if i % 10 == 0 {
                Some("timeout".to_string())
            } else {
                None
            },
            solution: Some(format!("Solution for task {}", i)),
            root_cause: None,
            timestamp: Utc::now(),
        };
        kg.record_execution(record).await.unwrap();
    }

    kg
}

fn kg_record_execution(c: &mut Criterion) {
    c.bench_function("record_single_execution", |b| {
        b.to_async(tokio::runtime::Runtime::new().unwrap())
            .iter(|| async {
                let kg = TemporalKG::new("ws://localhost:8000", "root", "root")
                    .await
                    .unwrap();
                let record = ExecutionRecord {
                    id: "test-exec".to_string(),
                    task_id: "test-task".to_string(),
                    agent_id: "test-agent".to_string(),
                    agent_role: None,
                    task_type: "coding".to_string(),
                    description: "Test execution".to_string(),
                    duration_ms: 1000,
                    input_tokens: 100,
                    output_tokens: 200,
                    cost_cents: 50,
                    provider: "claude".to_string(),
                    success: true,
                    error: None,
                    solution: None,
                    root_cause: None,
                    timestamp: Utc::now(),
                };
                black_box(kg.record_execution(black_box(record)).await)
            });
    });
}

fn kg_query_similar(c: &mut Criterion) {
    c.bench_function("query_similar_tasks_100_records", |b| {
        b.to_async(tokio::runtime::Runtime::new().unwrap())
            .iter_batched(
                || {
                    let rt = tokio::runtime::Runtime::new().unwrap();
                    rt.block_on(setup_kg_with_records(100))
                },
                |kg| async move {
                    black_box(
                        kg.query_similar_tasks("coding", "Write a function for processing data")
                            .await,
                    )
                },
                criterion::BatchSize::SmallInput,
            );
    });
}

fn kg_get_statistics(c: &mut Criterion) {
    c.bench_function("get_statistics_1000_records", |b| {
        b.to_async(tokio::runtime::Runtime::new().unwrap())
            .iter_batched(
                || {
                    let rt = tokio::runtime::Runtime::new().unwrap();
                    rt.block_on(setup_kg_with_records(1000))
                },
                |kg| async move { black_box(kg.get_statistics().await) },
                criterion::BatchSize::SmallInput,
            );
    });
}

fn kg_get_agent_profile(c: &mut Criterion) {
    c.bench_function("get_agent_profile_500_records", |b| {
        b.to_async(tokio::runtime::Runtime::new().unwrap())
            .iter_batched(
                || {
                    let rt = tokio::runtime::Runtime::new().unwrap();
                    rt.block_on(setup_kg_with_records(500))
                },
                |kg| async move { black_box(kg.get_agent_profile("agent-1").await) },
                criterion::BatchSize::SmallInput,
            );
    });
}

criterion_group!(
    benches,
    kg_record_execution,
    kg_query_similar,
    kg_get_statistics,
    kg_get_agent_profile
);
criterion_main!(benches);