aimo-core 0.4.0

AiMo Network core protocol Rust specs
Documentation
use aimo_core::receipt::RequestReceipt;
use solana_sdk::{signature::Keypair, signer::Signer};

#[test]
fn test_request_receipt_builder_all_fields() {
    let sender_keypair = Keypair::new();
    let request_id_keypair = Keypair::new();

    let receipt = RequestReceipt::new()
        .sender(sender_keypair.pubkey())
        .request_id(request_id_keypair.pubkey())
        .service_id("test-service-123")
        .max_est_cost(1000)
        .timestamp(1684512000)
        .finish_reason("stop")
        .prompt_tokens(150)
        .completion_tokens(75)
        .amount(500)
        .token_mint("SOL")
        .model_name("gpt-4")
        .cache_discount(50)
        .streamed(true)
        .cancelled(false)
        .latency(250)
        .generation_time(1200)
        .moderation_latency(50)
        .num_media_prompt(2)
        .num_media_completion(1)
        .num_search_results(5)
        .app_id("my-app")
        .origin("https://example.com");

    // Verify all required fields
    assert_eq!(receipt.sender, sender_keypair.pubkey().to_string());
    assert_eq!(receipt.request_id, request_id_keypair.pubkey().to_string());
    assert_eq!(receipt.service_id, "test-service-123");
    assert_eq!(receipt.max_est_cost, 1000);
    assert_eq!(receipt.timestamp, 1684512000);
    assert_eq!(receipt.finish_reason, "stop");
    assert_eq!(receipt.prompt_tokens, 150);
    assert_eq!(receipt.completion_tokens, 75);
    assert_eq!(receipt.amount, 500);
    assert_eq!(receipt.token_mint, "SOL");
    assert_eq!(receipt.model_name, "gpt-4");
    assert_eq!(receipt.cache_discount, 50);

    // Verify optional model statistics
    assert_eq!(receipt.streamed, Some(true));
    assert_eq!(receipt.cancelled, Some(false));
    assert_eq!(receipt.latency, Some(250));
    assert_eq!(receipt.generation_time, Some(1200));
    assert_eq!(receipt.moderation_latency, Some(50));
    assert_eq!(receipt.num_media_prompt, Some(2));
    assert_eq!(receipt.num_media_completion, Some(1));
    assert_eq!(receipt.num_search_results, Some(5));

    // Verify optional fields
    assert_eq!(receipt.app_id, Some("my-app".to_string()));
    assert_eq!(receipt.origin, Some("https://example.com".to_string()));
}

#[test]
fn test_request_receipt_builder_partial() {
    let sender_keypair = Keypair::new();
    let request_id_keypair = Keypair::new();

    let receipt = RequestReceipt::new()
        .sender(sender_keypair.pubkey())
        .request_id(request_id_keypair.pubkey())
        .service_id("minimal-service")
        .amount(100)
        .token_mint("USDC")
        .model_name("gpt-3.5-turbo");

    // Verify set fields
    assert_eq!(receipt.sender, sender_keypair.pubkey().to_string());
    assert_eq!(receipt.request_id, request_id_keypair.pubkey().to_string());
    assert_eq!(receipt.service_id, "minimal-service");
    assert_eq!(receipt.amount, 100);
    assert_eq!(receipt.token_mint, "USDC");
    assert_eq!(receipt.model_name, "gpt-3.5-turbo");

    // Verify default/unset fields
    assert_eq!(receipt.max_est_cost, 0);
    assert_eq!(receipt.timestamp, 0);
    assert_eq!(receipt.finish_reason, "");
    assert_eq!(receipt.prompt_tokens, 0);
    assert_eq!(receipt.completion_tokens, 0);
    assert_eq!(receipt.cache_discount, 0);

    // Verify optional fields are None
    assert_eq!(receipt.streamed, None);
    assert_eq!(receipt.cancelled, None);
    assert_eq!(receipt.latency, None);
    assert_eq!(receipt.generation_time, None);
    assert_eq!(receipt.moderation_latency, None);
    assert_eq!(receipt.num_media_prompt, None);
    assert_eq!(receipt.num_media_completion, None);
    assert_eq!(receipt.num_search_results, None);
    assert_eq!(receipt.app_id, None);
    assert_eq!(receipt.origin, None);
}

#[test]
fn test_request_receipt_builder_chaining() {
    let sender_keypair = Keypair::new();
    let request_id_keypair = Keypair::new();

    // Test that builder methods can be chained in any order
    let receipt1 = RequestReceipt::new()
        .amount(100)
        .sender(sender_keypair.pubkey())
        .model_name("test-model")
        .request_id(request_id_keypair.pubkey())
        .service_id("test-service");

    let receipt2 = RequestReceipt::new()
        .sender(sender_keypair.pubkey())
        .request_id(request_id_keypair.pubkey())
        .service_id("test-service")
        .model_name("test-model")
        .amount(100);

    // Both should have the same values
    assert_eq!(receipt1.sender, receipt2.sender);
    assert_eq!(receipt1.request_id, receipt2.request_id);
    assert_eq!(receipt1.service_id, receipt2.service_id);
    assert_eq!(receipt1.model_name, receipt2.model_name);
    assert_eq!(receipt1.amount, receipt2.amount);
}

#[test]
fn test_request_receipt_builder_into_string_types() {
    let sender_keypair = Keypair::new();

    // Test that string fields accept both &str and String
    let receipt = RequestReceipt::new()
        .sender(sender_keypair.pubkey())
        .service_id("string_literal")
        .finish_reason(String::from("owned_string"))
        .token_mint("SOL")
        .model_name(format!("formatted_{}", "string"))
        .app_id("app_literal")
        .origin(String::from("https://test.com"));

    assert_eq!(receipt.service_id, "string_literal");
    assert_eq!(receipt.finish_reason, "owned_string");
    assert_eq!(receipt.token_mint, "SOL");
    assert_eq!(receipt.model_name, "formatted_string");
    assert_eq!(receipt.app_id, Some("app_literal".to_string()));
    assert_eq!(receipt.origin, Some("https://test.com".to_string()));
}

#[test]
fn test_request_receipt_serialization() {
    let sender_keypair = Keypair::new();
    let request_id_keypair = Keypair::new();

    let receipt = RequestReceipt::new()
        .sender(sender_keypair.pubkey())
        .request_id(request_id_keypair.pubkey())
        .service_id("test-service")
        .amount(100)
        .prompt_tokens(50)
        .completion_tokens(25)
        .streamed(true)
        .app_id("test-app");

    // Test that the receipt can be serialized to bytes and hashed
    let bytes_result = receipt.clone().to_bytes();
    assert!(bytes_result.is_ok());

    let hash_result = receipt.to_hash();
    assert!(hash_result.is_ok());

    let hash = hash_result.unwrap();
    assert_eq!(hash.len(), 32); // SHA256 hash length
}