api_huggingface 0.5.1

//! Tests for Document Similarity & Search Engine Example
//!
//! This test suite verifies the functionality of a semantic search system that uses
//! `HuggingFace` embeddings to find relevant documents based on meaning rather than keywords.

#![allow(clippy::missing_inline_in_public_items)]

mod inc;

use api_huggingface::
{
  Client,
  environment::HuggingFaceEnvironmentImpl,
  components::embeddings::EmbeddingResponse,
};
use std::collections::HashMap;

#[ cfg( feature = "integration" ) ]
use api_huggingface::
{
  components::models::Models,
  secret::Secret,
};

#[ cfg( feature = "integration" ) ]
use std::time::Instant;

#[ allow( missing_docs ) ]
/// Represents a document in the search index
#[ derive( Debug, Clone ) ]
pub struct Document
{
  /// Unique document identifier
  pub id : String,
  /// Document title
  pub title : String,
  /// Document content
  pub content : String,
  /// Document metadata
  pub metadata : HashMap< String, String >,
  /// Document embedding vector (generated from content)
  pub embedding : Option< Vec< f32 > >,
}

/// Document categories for classification
#[ derive( Debug, Clone, Copy, PartialEq ) ]
pub enum DocumentCategory
{
  /// Technical documentation
  Technical,
  /// Academic research papers
  Academic,
  /// News articles
  News,
  /// Creative writing
  Creative,
  /// Legal documents
  Legal,
}

/// Search query with parameters
#[ derive( Debug, Clone ) ]
pub struct SearchQuery
{
  /// Query text
  pub text : String,
  /// Number of results to return
  pub limit : usize,
  /// Minimum similarity threshold (0.0 to 1.0)
  pub threshold : f32,
  /// Optional category filter
  pub category_filter : Option< DocumentCategory >,
  /// Optional metadata filters
  pub metadata_filters : HashMap< String, String >,
}

/// Search result with similarity score
#[ derive( Debug, Clone ) ]
pub struct SearchResult
{
  /// Document reference
  pub document : Document,
  /// Similarity score (0.0 to 1.0, higher is more similar)
  pub similarity : f32,
  /// Ranking position in results
  pub rank : usize,
}

/// High-performance document search engine with embedding-based similarity
#[ derive( Debug ) ]
pub struct DocumentSearchEngine
{
  client : Client< HuggingFaceEnvironmentImpl >,
  documents : HashMap< String, Document >,
  embedding_model : String,
}

impl DocumentSearchEngine
{
  /// Create new search engine with embedding model
  #[ must_use ]
  pub fn new( client : Client< HuggingFaceEnvironmentImpl >, embedding_model : String ) -> Self
  {
  Self
  {
      client,
      documents : HashMap::new(),
      embedding_model,
  }
  }

  /// Add document to search index with embedding generation
  ///
  /// # Errors
  /// Returns error if embedding generation fails or document processing encounters issues
  pub async fn add_document( &mut self, mut document : Document ) -> Result< String, Box< dyn std::error::Error > >
  {
  // Generate embedding for document content
  let response = self.client
      .embeddings()
      .create( document.content.clone(), &self.embedding_model )
      .await?;

  match response
  {
      EmbeddingResponse::Single( embedding_vectors ) =>
      {
  if let Some( embedding ) = embedding_vectors.first()
  {
          document.embedding = Some( embedding.clone() );
  }
  else
  {
          return Err( "No embedding generated".into() );
  }
      },
      EmbeddingResponse::Batch( batch_vectors ) =>
      {
  if let Some( first_doc_vectors ) = batch_vectors.first()
  {
          if let Some( embedding ) = first_doc_vectors.first()
          {
      document.embedding = Some( embedding.clone() );
          }
          else
          {
      return Err( "No embedding in batch result".into() );
          }
  }
  else
  {
          return Err( "No embedding generated".into() );
  }
      },
  }

  let doc_id = document.id.clone();
  self.documents.insert( doc_id.clone(), document );
  Ok( doc_id )
  }

  /// Add multiple documents in batch for better performance
  ///
  /// # Errors
  /// Returns error if batch processing fails or any document cannot be processed
  pub async fn add_documents_batch( &mut self, documents : Vec< Document > ) -> Result< Vec< String >, Box< dyn std::error::Error > >
  {
  let mut added_ids = Vec::new();

  // Process in smaller batches to avoid API limits
  let batch_size = 5;
  for batch in documents.chunks( batch_size )
  {
      let batch_contents : Vec< String > = batch.iter().map( | doc | doc.content.clone() ).collect();

      let response = self.client
  .embeddings()
  .create_batch( batch_contents, &self.embedding_model )
  .await?;

      match response
      {
  EmbeddingResponse::Batch( batch_vectors ) =>
  {
          for ( i, doc_vectors ) in batch_vectors.iter().enumerate()
          {
      if let Some( mut doc ) = batch.get( i ).cloned()
      {
              if let Some( embedding ) = doc_vectors.first()
              {
        doc.embedding = Some( embedding.clone() );
        let doc_id = doc.id.clone();
        self.documents.insert( doc_id.clone(), doc );
        added_ids.push( doc_id );
              }
      }
          }
  },
  EmbeddingResponse::Single( embedding_vectors ) =>
  {
          // Handle single response for batch (shouldn't happen but be safe)
          if let Some( mut doc ) = batch.first().cloned()
          {
      if let Some( embedding ) = embedding_vectors.first()
      {
              doc.embedding = Some( embedding.clone() );
              let doc_id = doc.id.clone();
              self.documents.insert( doc_id.clone(), doc );
              added_ids.push( doc_id );
      }
          }
  },
      }
  }

  Ok( added_ids )
  }

  /// Search for documents similar to query
  ///
  /// # Errors
  /// Returns error if query embedding generation fails or search processing encounters issues
  pub async fn search( &self, query : SearchQuery ) -> Result< Vec< SearchResult >, Box< dyn std::error::Error > >
  {
  // Generate embedding for query
  let response = self.client
      .embeddings()
      .create( query.text.clone(), &self.embedding_model )
      .await?;

  let query_embedding = match response
  {
      EmbeddingResponse::Single( embedding_vectors ) => 
      {
  embedding_vectors.first()
          .ok_or( "No query embedding generated" )?
          .clone()
      },
      EmbeddingResponse::Batch( batch_vectors ) => 
      {
  batch_vectors.first()
          .and_then( | doc_vectors | doc_vectors.first() )
          .ok_or( "No query embedding generated" )?
          .clone()
      },
  };

  // Calculate similarities and filter results
  let mut results = Vec::new();

  for document in self.documents.values()
  {
      if let Some( doc_embedding ) = &document.embedding
      {
  let similarity = Self::cosine_similarity( &query_embedding, doc_embedding );

  // Apply threshold filter
  if similarity >= query.threshold
  {
          // Apply category filter if specified
          if let Some( category_filter ) = query.category_filter
          {
      if let Some( doc_category_str ) = document.metadata.get( "category" )
      {
              let doc_category = Self::parse_category( doc_category_str );
              if doc_category != Some( category_filter )
              {
        continue;
              }
      }
      else
      {
              continue; // Skip documents without category if filter specified
      }
          }

          // Apply metadata filters
          let mut passes_metadata_filters = true;
          for ( key, value ) in &query.metadata_filters
          {
      if document.metadata.get( key ) != Some( value )
      {
              passes_metadata_filters = false;
              break;
      }
          }

          if passes_metadata_filters
          {
      results.push( SearchResult
      {
              document : document.clone(),
              similarity,
              rank : 0, // Will be set after sorting
      } );
          }
  }
      }
  }

  // Sort by similarity (highest first) and apply limit
  results.sort_by( | a, b | b.similarity.partial_cmp( &a.similarity ).unwrap_or( core::cmp::Ordering::Equal ) );
  results.truncate( query.limit );

  // Set ranking positions
  for ( i, result ) in results.iter_mut().enumerate()
  {
      result.rank = i + 1;
  }

  Ok( results )
  }

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

  let dot_product : f32 = a.iter().zip( b.iter() ).map( | ( x, y ) | x * y ).sum();
  let norm_a : f32 = a.iter().map( | x | x * x ).sum::< f32 >().sqrt();
  let norm_b : f32 = b.iter().map( | x | x * x ).sum::< f32 >().sqrt();

  if norm_a == 0.0 || norm_b == 0.0
  {
      0.0
  }
  else
  {
      dot_product / ( norm_a * norm_b )
  }
  }

  /// Parse category from string
  fn parse_category( category_str : &str ) -> Option< DocumentCategory >
  {
  match category_str.to_lowercase().as_str()
  {
      "technical" => Some( DocumentCategory::Technical ),
      "academic" => Some( DocumentCategory::Academic ),
      "news" => Some( DocumentCategory::News ),
      "creative" => Some( DocumentCategory::Creative ),
      "legal" => Some( DocumentCategory::Legal ),
      _ => None,
  }
  }

  /// Get search index statistics
  #[ must_use ]
  pub fn get_stats( &self ) -> SearchIndexStats
  {
  let total_documents = self.documents.len();
  let indexed_documents = self.documents.values()
      .filter( | doc | doc.embedding.is_some() )
      .count();

  let mut category_counts = HashMap::new();
  for document in self.documents.values()
  {
      if let Some( category_str ) = document.metadata.get( "category" )
      {
  *category_counts.entry( category_str.clone() ).or_insert( 0 ) += 1;
      }
  }

  SearchIndexStats
  {
      total_documents,
      indexed_documents,
      category_counts,
      embedding_model : self.embedding_model.clone(),
  }
  }

  /// Remove document from index
  pub fn remove_document( &mut self, document_id : &str ) -> Option< Document >
  {
  self.documents.remove( document_id )
  }

  /// Clear all documents from index
  pub fn clear( &mut self )
  {
  self.documents.clear();
  }
}

/// Search index statistics
#[ derive( Debug, Clone ) ]
pub struct SearchIndexStats
{
  /// Total number of documents
  pub total_documents : usize,
  /// Number of documents with embeddings
  pub indexed_documents : usize,
  /// Document count by category
  pub category_counts : HashMap< String, usize >,
  /// Embedding model being used
  pub embedding_model : String,
}

#[ cfg( test ) ]
mod tests
{
  use super::*;

  #[ cfg( feature = "integration" ) ]
  fn create_integration_client() -> Client< HuggingFaceEnvironmentImpl >
  {
  let api_key = crate::inc::get_api_key_for_integration();
  let secret = Secret::new( api_key );
  let env = HuggingFaceEnvironmentImpl::build( secret, None )
      .expect( "Failed to build environment" );
  Client::build( env ).expect( "Failed to create client" )
  }

  fn create_sample_documents() -> Vec< Document >
  {
  vec!
  [
      Document
      {
  id : "doc1".to_string(),
  title : "Rust Programming Guide".to_string(),
  content : "Rust is a systems programming language focused on safety, speed, and concurrency.".to_string(),
  metadata : { let mut m = HashMap::new(); m.insert( "category".to_string(), "technical".to_string() ); m.insert( "author".to_string(), "tech_writer".to_string() ); m },
  embedding : None,
      },
      Document
      {
  id : "doc2".to_string(),
  title : "Machine Learning Fundamentals".to_string(),
  content : "Machine learning involves algorithms that can learn patterns from data without explicit programming.".to_string(),
  metadata : { let mut m = HashMap::new(); m.insert( "category".to_string(), "academic".to_string() ); m.insert( "author".to_string(), "researcher".to_string() ); m },
  embedding : None,
      },
      Document
      {
  id : "doc3".to_string(),
  title : "Climate Change Report".to_string(),
  content : "Global warming is causing significant changes to weather patterns and ecosystems worldwide.".to_string(),
  metadata : { let mut m = HashMap::new(); m.insert( "category".to_string(), "news".to_string() ); m.insert( "author".to_string(), "journalist".to_string() ); m },
  embedding : None,
      },
      Document
      {
  id : "doc4".to_string(),
  title : "Poetry Collection".to_string(),
  content : "In whispered dreams and moonlit nights, the soul finds peace in gentle lights.".to_string(),
  metadata : { let mut m = HashMap::new(); m.insert( "category".to_string(), "creative".to_string() ); m.insert( "author".to_string(), "poet".to_string() ); m },
  embedding : None,
      },
      Document
      {
  id : "doc5".to_string(),
  title : "Legal Contract Template".to_string(),
  content : "This agreement establishes the terms and conditions governing the relationship between parties.".to_string(),
  metadata : { let mut m = HashMap::new(); m.insert( "category".to_string(), "legal".to_string() ); m.insert( "author".to_string(), "lawyer".to_string() ); m },
  embedding : None,
      },
  ]
  }

  #[ test ]
  fn test_document_structure()
  {
  let doc = Document
  {
      id : "test-id".to_string(),
      title : "Test Document".to_string(),
      content : "Test content for document".to_string(),
      metadata : HashMap::new(),
      embedding : None,
  };

  assert_eq!( doc.id, "test-id" );
  assert_eq!( doc.title, "Test Document" );
  assert_eq!( doc.content, "Test content for document" );
  assert!( doc.metadata.is_empty() );
  assert!( doc.embedding.is_none() );
  }

  #[ test ]
  fn test_search_query_construction()
  {
  let query = SearchQuery
  {
      text : "machine learning".to_string(),
      limit : 10,
      threshold : 0.7,
      category_filter : Some( DocumentCategory::Academic ),
      metadata_filters : { let mut m = HashMap::new(); m.insert( "author".to_string(), "researcher".to_string() ); m },
  };

  assert_eq!( query.text, "machine learning" );
  assert_eq!( query.limit, 10 );
  assert!( ( query.threshold - 0.7 ).abs() < f32::EPSILON );
  assert_eq!( query.category_filter, Some( DocumentCategory::Academic ) );
  assert_eq!( query.metadata_filters.get( "author" ), Some( &"researcher".to_string() ) );
  }

  #[ test ]
  fn test_document_category_parsing()
  {
  assert_eq!( DocumentSearchEngine::parse_category( "technical" ), Some( DocumentCategory::Technical ) );
  assert_eq!( DocumentSearchEngine::parse_category( "Academic" ), Some( DocumentCategory::Academic ) );
  assert_eq!( DocumentSearchEngine::parse_category( "NEWS" ), Some( DocumentCategory::News ) );
  assert_eq!( DocumentSearchEngine::parse_category( "Creative" ), Some( DocumentCategory::Creative ) );
  assert_eq!( DocumentSearchEngine::parse_category( "legal" ), Some( DocumentCategory::Legal ) );
  assert_eq!( DocumentSearchEngine::parse_category( "unknown" ), None );
  }

  #[ test ]
  fn test_cosine_similarity_calculation()
  {
  let vec_a = vec![ 1.0, 0.0, 0.0 ];
  let vec_b = vec![ 1.0, 0.0, 0.0 ];
  let vec_c = vec![ 0.0, 1.0, 0.0 ];

  // Identical vectors should have similarity 1.0
  let sim_identical = DocumentSearchEngine::cosine_similarity( &vec_a, &vec_b );
  assert!( ( sim_identical - 1.0 ).abs() < f32::EPSILON );

  // Orthogonal vectors should have similarity 0.0
  let sim_orthogonal = DocumentSearchEngine::cosine_similarity( &vec_a, &vec_c );
  assert!( sim_orthogonal.abs() < f32::EPSILON );

  // Different length vectors should return 0.0
  let vec_d = vec![ 1.0, 0.0 ];
  let sim_different_length = DocumentSearchEngine::cosine_similarity( &vec_a, &vec_d );
  assert!( sim_different_length.abs() < f32::EPSILON );
  }

  #[ cfg( feature = "integration" ) ]
  #[ tokio::test ]
  async fn test_search_engine_creation()
  {
  let client = create_integration_client();

  let engine = DocumentSearchEngine::new( client, Models::all_minilm_l6_v2().to_string() );
  assert!( engine.documents.is_empty() );
  assert_eq!( engine.embedding_model, Models::all_minilm_l6_v2() );
  }

  #[ cfg( feature = "integration" ) ]
  #[ tokio::test ]
  async fn test_search_index_stats()
  {
  let client = create_integration_client();

  let mut engine = DocumentSearchEngine::new( client, Models::all_minilm_l6_v2().to_string() );
  let docs = create_sample_documents();

  // Add documents without embeddings first
  for doc in docs
  {
      engine.documents.insert( doc.id.clone(), doc );
  }

  let stats = engine.get_stats();
  assert_eq!( stats.total_documents, 5 );
  assert_eq!( stats.indexed_documents, 0 ); // No embeddings yet
  assert_eq!( stats.embedding_model, Models::all_minilm_l6_v2() );
  
  // Check category counts
  assert_eq!( stats.category_counts.get( "technical" ), Some( &1 ) );
  assert_eq!( stats.category_counts.get( "academic" ), Some( &1 ) );
  assert_eq!( stats.category_counts.get( "news" ), Some( &1 ) );
  }

  #[ cfg( feature = "integration" ) ]
  #[ tokio::test ]
  async fn test_document_management()
  {
  let client = create_integration_client();

  let mut engine = DocumentSearchEngine::new( client, Models::all_minilm_l6_v2().to_string() );
  let doc = create_sample_documents().into_iter().next().expect( "[test_document_management] create_sample_documents() should return at least 1 document - check create_sample_documents() implementation" );
  let doc_id = doc.id.clone();

  // Test document removal on empty index
  assert!( engine.remove_document( &doc_id ).is_none() );

  // Add document manually for testing removal
  engine.documents.insert( doc_id.clone(), doc );
  assert_eq!( engine.documents.len(), 1 );

  // Test document removal
  let removed_doc = engine.remove_document( &doc_id );
  assert!( removed_doc.is_some() );
  assert!( engine.documents.is_empty() );

  // Test clear functionality
  let docs = create_sample_documents();
  for doc in docs
  {
      engine.documents.insert( doc.id.clone(), doc );
  }
  assert_eq!( engine.documents.len(), 5 );

  engine.clear();
  assert!( engine.documents.is_empty() );
  }

  #[ tokio::test ]
  async fn test_search_result_ranking()
  {
  // Test SearchResult structure and ranking logic
  let doc = create_sample_documents().into_iter().next().expect( "[test_search_result_ranking] create_sample_documents() should return at least 1 document - check create_sample_documents() implementation" );
  let mut results =
  [
      SearchResult { document : doc.clone(), similarity : 0.8, rank : 0 },
      SearchResult { document : doc.clone(), similarity : 0.9, rank : 0 },
      SearchResult { document : doc, similarity : 0.7, rank : 0 },
  ].to_vec();

  // Sort by similarity (mimicking search engine behavior)
  results.sort_by( | a, b | b.similarity.partial_cmp( &a.similarity ).unwrap_or( core::cmp::Ordering::Equal ) );

  // Set rankings
  for ( i, result ) in results.iter_mut().enumerate()
  {
      result.rank = i + 1;
  }

  assert_eq!( results[ 0 ].rank, 1 );
  assert!( ( results[ 0 ].similarity - 0.9 ).abs() < f32::EPSILON );
  assert_eq!( results[ 1 ].rank, 2 );
  assert!( ( results[ 1 ].similarity - 0.8 ).abs() < f32::EPSILON );
  assert_eq!( results[ 2 ].rank, 3 );
  assert!( ( results[ 2 ].similarity - 0.7 ).abs() < f32::EPSILON );
  }

  #[ tokio::test ]
  async fn test_search_query_filtering()
  {
  // Test various search query configurations
  let base_query = SearchQuery
  {
      text : "programming".to_string(),
      limit : 5,
      threshold : 0.5,
      category_filter : None,
      metadata_filters : HashMap::new(),
  };

  assert_eq!( base_query.limit, 5 );
  assert!( ( base_query.threshold - 0.5 ).abs() < f32::EPSILON );
  assert_eq!( base_query.category_filter, None );

  let filtered_query = SearchQuery
  {
      text : "programming".to_string(),
      limit : 3,
      threshold : 0.8,
      category_filter : Some( DocumentCategory::Technical ),
      metadata_filters : { let mut m = HashMap::new(); m.insert( "author".to_string(), "tech_writer".to_string() ); m },
  };

  assert_eq!( filtered_query.category_filter, Some( DocumentCategory::Technical ) );
  assert_eq!( filtered_query.metadata_filters.get( "author" ), Some( &"tech_writer".to_string() ) );
  }

  #[ cfg( feature = "integration" ) ]
  #[ tokio::test ]
  async fn test_embedding_model_variations()
  {
  let client = create_integration_client();

  // Test different embedding models
  let models = vec!
  [
      Models::all_minilm_l6_v2(),
      Models::all_minilm_l12_v2(),
      Models::bge_large_en_v1_5(),
  ];

  for model in models
  {
      let engine = DocumentSearchEngine::new( client.clone(), model.to_string() );
      assert_eq!( engine.embedding_model, model );
  }
  }

  #[ cfg( feature = "integration" ) ]
  #[ tokio::test ]
  async fn test_performance_characteristics()
  {
  let client = create_integration_client();

  let mut engine = DocumentSearchEngine::new( client, Models::all_minilm_l6_v2().to_string() );

  // Test performance with larger document sets
  let mut large_doc_set = Vec::new();
  for i in 0..20
  {
      large_doc_set.push( Document
      {
  id : format!( "perf_doc_{i}" ),
  title : format!( "Performance Test Document {i}" ),
  content : format!( "This is test document {i} for performance evaluation with some meaningful content about topic {i}." ),
  metadata : { let mut m = HashMap::new(); m.insert( "category".to_string(), "technical".to_string() ); m.insert( "batch".to_string(), format!( "{i}" ) ); m },
  embedding : None,
      } );
  }

  // Measure time for manual document insertion (simulating what add_documents_batch would do)
  let start_time = Instant::now();
  for doc in large_doc_set
  {
      engine.documents.insert( doc.id.clone(), doc );
  }
  let insertion_time = start_time.elapsed();

  println!( "Document insertion time for 20 documents : {insertion_time:?}" );
  assert_eq!( engine.documents.len(), 20 );

  // Test search performance simulation (without actual API calls)
  let stats = engine.get_stats();
  assert_eq!( stats.total_documents, 20 );
  assert_eq!( stats.indexed_documents, 0 ); // No embeddings in this test
  }

  #[ cfg( feature = "integration" ) ]
  #[ tokio::test ]
  async fn test_error_handling_scenarios()
  {
  let client = create_integration_client();

  let engine = DocumentSearchEngine::new( client, Models::bge_large_en_v1_5().to_string() );

  // Test search with empty index
  let query = SearchQuery
  {
      text : "test query".to_string(),
      limit : 5,
      threshold : 0.5,
      category_filter : None,
      metadata_filters : HashMap::new(),
  };

  // Search on empty index: query embedding is computed but no documents match
  let results = engine.search( query ).await
      .expect( "search on empty index should succeed (returns empty results)" );
  assert!( results.is_empty(), "Search on empty index should return no results, got {}", results.len() );
  }

  #[ cfg( feature = "integration" ) ]
  #[ tokio::test ]
  async fn test_batch_processing_structure()
  {
  let client = create_integration_client();

  let mut engine = DocumentSearchEngine::new( client, Models::bge_large_en_v1_5().to_string() );
  let docs = create_sample_documents();

  // Test that batch processing succeeds with valid credentials
  let ids = engine.add_documents_batch( docs ).await
      .expect( "add_documents_batch should succeed with valid credentials" );
  assert!( !ids.is_empty(), "Batch processing should return non-empty ID list" );
  }
}