embedding_demo/

embedding_demo.rs

use llm_brain::{ChunkingStrategy, LLMBrain, Result};
use serde_json::json;

/// Embedding functionality demonstration
#[tokio::main]
async fn main() -> Result<()> {
    // Initialize LLMBrain instance
    println!("Initializing LLMBrain...");
    let llm_brain = LLMBrain::launch().await?;

    // Demo 1: Basic embedding generation and query
    println!("\n--- Basic Embedding Demo ---");

    // Add several memories
    let memories = vec![
        (
            "Cats are small mammals with fur, whiskers, and tails. They eat mice and birds, and are related to tigers and lions.",
            json!({
                "name": "Cat",
                "type": "Animal",
                "properties": {
                    "characteristics": ["fur", "whiskers", "tail"],
                    "diet": "carnivorous"
                }
            }),
        ),
        (
            "Dogs are one of the earliest domesticated animals, loyal companions to humans, with various breeds and purposes.",
            json!({
                "name": "Dog",
                "type": "Animal",
                "properties": {
                    "characteristics": ["fur", "loyalty", "keen sense of smell"],
                    "uses": ["pet", "working dog", "guide dog"]
                }
            }),
        ),
        (
            "Birds are warm-blooded vertebrates, covered with feathers, with forelimbs evolved into wings, and most can fly.",
            json!({
                "name": "Bird",
                "type": "Animal",
                "properties": {
                    "characteristics": ["feathers", "wings", "beak"],
                    "abilities": ["flying", "nest building", "singing"]
                }
            }),
        ),
    ];

    // Add memories to database
    for (content, metadata) in memories {
        let memory_id = llm_brain.add_memory(content.to_owned(), metadata).await?;
        println!("Added memory: ID={memory_id}");
    }

    // Wait for embedding processing to complete
    tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;

    // Query related memories using semantic search
    println!("\nQuerying memories about 'pets':");
    let results = llm_brain.recall("pets and human relationships", 2).await?;
    for (i, (memory, score)) in results.iter().enumerate() {
        println!("\nResult {}: Similarity {:.4}", i + 1, score);
        println!("Content: {}", memory.content);
        println!(
            "Metadata: {}",
            serde_json::to_string_pretty(&memory.metadata)?
        );
    }

    // Demo 2: Long text processing
    println!("\n--- Long Text Processing Demo ---");

    // Prepare a longer text
    let long_text = r#"
        Artificial Intelligence (AI) is a branch of computer science aimed at creating systems capable of simulating human intelligence behaviors.
        AI research includes multiple sub-fields such as machine learning, deep learning, natural language processing, computer vision, expert systems, etc.
        Machine learning is one of the core technologies of AI, using statistical techniques to enable computer systems to "learn" (i.e., progressively improve performance) without explicit programming.
        Deep learning is a subset of machine learning that uses multi-layered neural networks to process data. These networks extract features from data, with each layer building on the output of the previous one.
        Large Language Models (LLMs) are a significant breakthrough in AI in recent years, generating human-like text, understanding context, and performing various language tasks by learning from vast amounts of text data.
        Vector embeddings are techniques for representing text, images, or other data as points in a multi-dimensional vector space, playing an important role in information retrieval, recommendation systems, and semantic search.
    "#;

    println!("Long text content:\n{long_text}");

    // Process long text using different strategies
    println!("\nUsing default strategy (no chunking):");
    let embedding1 = llm_brain.process_long_text(long_text, None).await?;
    println!("Generated embedding vector length: {}", embedding1.len());

    println!("\nUsing chunk and average strategy:");
    let chunk_strategy = ChunkingStrategy::ChunkAndAverage {
        chunk_size: 100,
        chunk_overlap: 20,
    };
    let embedding2 = llm_brain
        .process_long_text(long_text, Some(chunk_strategy))
        .await?;
    println!("Generated embedding vector length: {}", embedding2.len());

    // Calculate similarity between embeddings generated by different strategies
    let similarity = cosine_similarity(&embedding1, &embedding2);
    println!("\nCosine similarity between embeddings from two strategies: {similarity:.4}");

    // Demo 3: Adding long text to memory
    println!("\n--- Adding Long Text Memory Demo ---");

    let metadata = json!({
        "name": "Artificial Intelligence Overview",
        "type": "Knowledge",
        "tags": ["AI", "Machine Learning", "Deep Learning"]
    });

    let memory_id = llm_brain.add_memory(long_text.to_owned(), metadata).await?;
    println!("Added long text memory: ID={memory_id}");

    // Wait for embedding processing to complete
    tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;

    // Query related content
    println!("\nQuerying memories about 'large language models':");
    let results = llm_brain
        .recall("large language models and vector embeddings", 1)
        .await?;
    for (i, (memory, score)) in results.iter().enumerate() {
        println!("\nResult {}: Similarity {:.4}", i + 1, score);
        println!("Content summary: {}", truncate_text(&memory.content, 100));
        println!(
            "Metadata: {}",
            serde_json::to_string_pretty(&memory.metadata)?
        );
    }

    println!("\nDemo completed.");
    Ok(())
}

/// Calculate cosine similarity between two vectors
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    if a.len() != b.len() || a.is_empty() {
        return 0.0;
    }

    let mut dot_product = 0.0;
    let mut norm_a = 0.0;
    let mut norm_b = 0.0;

    for i in 0..a.len() {
        dot_product += a[i] * b[i];
        norm_a += a[i] * a[i];
        norm_b += b[i] * b[i];
    }

    if norm_a <= 0.0 || norm_b <= 0.0 {
        return 0.0;
    }

    dot_product / (norm_a.sqrt() * norm_b.sqrt())
}

/// Truncate text to specified length, adding ellipsis
fn truncate_text(text: &str, max_length: usize) -> String {
    if text.len() <= max_length {
        return text.to_owned();
    }

    let truncated = &text[0..max_length];
    format!("{truncated}...")
}