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

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

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),
            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),
            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(),
                    task_type: "coding".to_string(),
                    description: "Test execution".to_string(),
                    duration_ms: 1000,
                    input_tokens: 100,
                    output_tokens: 200,
                    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);