Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,975 @@
|
||||
//! Model search and discovery functionality.
|
||||
//!
|
||||
//! This module provides comprehensive search and discovery capabilities for models
|
||||
//! in the registry, including filtering, sorting, trending analysis, and similarity search.
|
||||
|
||||
use crate::{HubResult, ModelId, ModelInfo, ModelMetadata, ModelRegistry, QueryOptions, Registry};
|
||||
use chrono::{DateTime, Utc};
|
||||
use dashmap::DashMap;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tracing::debug;
|
||||
|
||||
/// Sort options for search results.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
|
||||
pub enum SortBy {
|
||||
/// Sort by total downloads
|
||||
Downloads,
|
||||
/// Sort by number of likes
|
||||
Likes,
|
||||
/// Sort by last modified date (most recent first)
|
||||
Recent,
|
||||
/// Sort by trending score
|
||||
Trending,
|
||||
/// Sort by relevance (for text search)
|
||||
Relevance,
|
||||
}
|
||||
|
||||
/// Search query for model discovery.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SearchQuery {
|
||||
/// Text search query (searches title, description, tags)
|
||||
pub query: Option<String>,
|
||||
/// Filter by task type (e.g., "text-classification", "image-generation")
|
||||
pub filter_by_task: Option<String>,
|
||||
/// Filter by library/framework
|
||||
pub filter_by_library: Option<String>,
|
||||
/// Filter by supported languages
|
||||
pub filter_by_language: Option<Vec<String>>,
|
||||
/// Filter by license
|
||||
pub filter_by_license: Option<Vec<String>>,
|
||||
/// Filter by author/organization
|
||||
pub filter_by_author: Option<String>,
|
||||
/// Sort by field
|
||||
pub sort_by: SortBy,
|
||||
/// Maximum number of results
|
||||
pub limit: usize,
|
||||
/// Offset for pagination
|
||||
pub offset: usize,
|
||||
}
|
||||
|
||||
impl Default for SearchQuery {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
query: None,
|
||||
filter_by_task: None,
|
||||
filter_by_library: None,
|
||||
filter_by_language: None,
|
||||
filter_by_license: None,
|
||||
filter_by_author: None,
|
||||
sort_by: SortBy::Relevance,
|
||||
limit: 20,
|
||||
offset: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SearchQuery {
|
||||
/// Create a new search query.
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Set text query.
|
||||
pub fn with_query(mut self, query: impl Into<String>) -> Self {
|
||||
self.query = Some(query.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Filter by task.
|
||||
pub fn with_task(mut self, task: impl Into<String>) -> Self {
|
||||
self.filter_by_task = Some(task.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Filter by library.
|
||||
pub fn with_library(mut self, library: impl Into<String>) -> Self {
|
||||
self.filter_by_library = Some(library.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Filter by languages.
|
||||
pub fn with_languages(mut self, languages: Vec<String>) -> Self {
|
||||
self.filter_by_language = Some(languages);
|
||||
self
|
||||
}
|
||||
|
||||
/// Filter by licenses.
|
||||
pub fn with_licenses(mut self, licenses: Vec<String>) -> Self {
|
||||
self.filter_by_license = Some(licenses);
|
||||
self
|
||||
}
|
||||
|
||||
/// Filter by author.
|
||||
pub fn with_author(mut self, author: impl Into<String>) -> Self {
|
||||
self.filter_by_author = Some(author.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set sort order.
|
||||
pub fn with_sort_by(mut self, sort_by: SortBy) -> Self {
|
||||
self.sort_by = sort_by;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set result limit.
|
||||
pub fn with_limit(mut self, limit: usize) -> Self {
|
||||
self.limit = limit;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set offset.
|
||||
pub fn with_offset(mut self, offset: usize) -> Self {
|
||||
self.offset = offset;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Model summary for search results.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ModelSummary {
|
||||
/// Model ID
|
||||
pub model_id: ModelId,
|
||||
/// Model title
|
||||
pub title: String,
|
||||
/// Short description
|
||||
pub description: String,
|
||||
/// Model author
|
||||
pub author: String,
|
||||
/// Model tags
|
||||
pub tags: Vec<String>,
|
||||
/// License
|
||||
pub license: Option<String>,
|
||||
/// Last modified
|
||||
pub last_modified: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl ModelSummary {
|
||||
/// Create a model summary from metadata.
|
||||
pub fn from_metadata(metadata: &ModelMetadata) -> Self {
|
||||
Self {
|
||||
model_id: metadata.id.clone(),
|
||||
title: metadata.title.clone(),
|
||||
description: metadata.description.clone(),
|
||||
author: metadata.author.clone(),
|
||||
tags: metadata.tags.clone(),
|
||||
license: metadata.license.clone(),
|
||||
last_modified: metadata.updated_at,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Search result with ranking information.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SearchResult {
|
||||
/// Model identifier
|
||||
pub model_id: ModelId,
|
||||
/// Model summary
|
||||
pub model_summary: ModelSummary,
|
||||
/// Total downloads
|
||||
pub downloads: u64,
|
||||
/// Number of likes
|
||||
pub likes: u64,
|
||||
/// Trending score (0.0 to 1.0)
|
||||
pub trending_score: f64,
|
||||
/// Last modified timestamp
|
||||
pub last_modified: DateTime<Utc>,
|
||||
/// Relevance score (0.0 to 1.0) for text search
|
||||
pub relevance_score: Option<f64>,
|
||||
}
|
||||
|
||||
/// Model discovery and search functionality.
|
||||
pub struct ModelDiscovery {
|
||||
/// Model registry
|
||||
registry: Arc<ModelRegistry>,
|
||||
/// Model metrics tracking
|
||||
metrics: Arc<DashMap<String, ModelMetrics>>,
|
||||
}
|
||||
|
||||
impl ModelDiscovery {
|
||||
/// Create a new model discovery instance.
|
||||
pub fn new(registry: ModelRegistry) -> Self {
|
||||
Self {
|
||||
registry: Arc::new(registry),
|
||||
metrics: Arc::new(DashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute a search query.
|
||||
pub async fn search(&self, query: SearchQuery) -> HubResult<Vec<SearchResult>> {
|
||||
debug!("Executing search query: {:?}", query);
|
||||
|
||||
// Build registry query options
|
||||
let mut options = QueryOptions::default();
|
||||
|
||||
if let Some(ref library) = query.filter_by_library {
|
||||
options.framework = Some(library.clone());
|
||||
}
|
||||
|
||||
if let Some(ref task) = query.filter_by_task {
|
||||
options.tags = vec![task.clone()];
|
||||
}
|
||||
|
||||
// Get models from registry
|
||||
let models = if let Some(ref text_query) = query.query {
|
||||
self.registry.search_models(text_query, options).await?
|
||||
} else {
|
||||
self.registry.list_models(options).await?
|
||||
};
|
||||
|
||||
// Convert to search results
|
||||
let mut results: Vec<SearchResult> = models
|
||||
.iter()
|
||||
.filter(|m| self.apply_filters(m, &query))
|
||||
.map(|m| self.to_search_result(m, &query))
|
||||
.collect();
|
||||
|
||||
// Sort results
|
||||
self.sort_results(&mut results, query.sort_by);
|
||||
|
||||
// Apply pagination
|
||||
let start = query.offset.min(results.len());
|
||||
let end = (query.offset + query.limit).min(results.len());
|
||||
results = results[start..end].to_vec();
|
||||
|
||||
debug!("Found {} search results", results.len());
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Get trending models.
|
||||
pub async fn get_trending(&self, limit: usize) -> HubResult<Vec<SearchResult>> {
|
||||
let query = SearchQuery::new()
|
||||
.with_sort_by(SortBy::Trending)
|
||||
.with_limit(limit);
|
||||
|
||||
self.search(query).await
|
||||
}
|
||||
|
||||
/// Get most downloaded models.
|
||||
pub async fn get_popular(&self, limit: usize) -> HubResult<Vec<SearchResult>> {
|
||||
let query = SearchQuery::new()
|
||||
.with_sort_by(SortBy::Downloads)
|
||||
.with_limit(limit);
|
||||
|
||||
self.search(query).await
|
||||
}
|
||||
|
||||
/// Get recently updated models.
|
||||
pub async fn get_recent(&self, limit: usize) -> HubResult<Vec<SearchResult>> {
|
||||
let query = SearchQuery::new()
|
||||
.with_sort_by(SortBy::Recent)
|
||||
.with_limit(limit);
|
||||
|
||||
self.search(query).await
|
||||
}
|
||||
|
||||
/// Get models for a specific task.
|
||||
pub async fn get_by_task(&self, task: &str, limit: usize) -> HubResult<Vec<SearchResult>> {
|
||||
let query = SearchQuery::new().with_task(task).with_limit(limit);
|
||||
|
||||
self.search(query).await
|
||||
}
|
||||
|
||||
/// Get similar models based on tags and task.
|
||||
pub async fn get_similar(
|
||||
&self,
|
||||
model_id: &ModelId,
|
||||
limit: usize,
|
||||
) -> HubResult<Vec<SearchResult>> {
|
||||
// Get the reference model
|
||||
let reference = self.registry.get_model(model_id, None).await?;
|
||||
|
||||
// Build query based on model tags
|
||||
let mut query = SearchQuery::new().with_limit(limit);
|
||||
|
||||
if !reference.metadata.tags.is_empty() {
|
||||
// Use first tag as task filter
|
||||
query = query.with_task(reference.metadata.tags[0].clone());
|
||||
}
|
||||
|
||||
// Search and filter out the reference model
|
||||
let mut results = self.search(query).await?;
|
||||
results.retain(|r| r.model_id != *model_id);
|
||||
|
||||
// Calculate similarity scores
|
||||
for result in &mut results {
|
||||
let similarity = self.calculate_similarity(&reference.metadata, &result.model_summary);
|
||||
result.relevance_score = Some(similarity);
|
||||
}
|
||||
|
||||
// Sort by similarity
|
||||
results.sort_by(|a, b| {
|
||||
b.relevance_score
|
||||
.unwrap_or(0.0)
|
||||
.partial_cmp(&a.relevance_score.unwrap_or(0.0))
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Record a download for metrics.
|
||||
pub async fn record_download(&self, model_id: &ModelId) -> HubResult<()> {
|
||||
let key = model_id.to_string();
|
||||
|
||||
self.metrics
|
||||
.entry(key)
|
||||
.and_modify(|m| {
|
||||
m.downloads += 1;
|
||||
m.last_downloaded = Some(Utc::now());
|
||||
})
|
||||
.or_insert_with(|| ModelMetrics {
|
||||
model_id: model_id.clone(),
|
||||
downloads: 1,
|
||||
likes: 0,
|
||||
last_downloaded: Some(Utc::now()),
|
||||
last_liked: None,
|
||||
daily_downloads: HashMap::new(),
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Record a like for metrics.
|
||||
pub async fn record_like(&self, model_id: &ModelId) -> HubResult<()> {
|
||||
let key = model_id.to_string();
|
||||
|
||||
self.metrics
|
||||
.entry(key)
|
||||
.and_modify(|m| {
|
||||
m.likes += 1;
|
||||
m.last_liked = Some(Utc::now());
|
||||
})
|
||||
.or_insert_with(|| ModelMetrics {
|
||||
model_id: model_id.clone(),
|
||||
downloads: 0,
|
||||
likes: 1,
|
||||
last_downloaded: None,
|
||||
last_liked: Some(Utc::now()),
|
||||
daily_downloads: HashMap::new(),
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Apply additional filters to model info.
|
||||
fn apply_filters(&self, model: &ModelInfo, query: &SearchQuery) -> bool {
|
||||
// Filter by author
|
||||
if let Some(ref author) = query.filter_by_author {
|
||||
if !model
|
||||
.metadata
|
||||
.author
|
||||
.to_lowercase()
|
||||
.contains(&author.to_lowercase())
|
||||
{
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Filter by license
|
||||
if let Some(ref licenses) = query.filter_by_license {
|
||||
if let Some(ref model_license) = model.metadata.license {
|
||||
if !licenses
|
||||
.iter()
|
||||
.any(|l| model_license.to_lowercase().contains(&l.to_lowercase()))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Filter by language (using tags as proxy)
|
||||
if let Some(ref languages) = query.filter_by_language {
|
||||
let has_language = languages.iter().any(|lang| {
|
||||
model
|
||||
.metadata
|
||||
.tags
|
||||
.iter()
|
||||
.any(|tag| tag.to_lowercase().contains(&lang.to_lowercase()))
|
||||
});
|
||||
|
||||
if !has_language {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// Convert ModelInfo to SearchResult.
|
||||
fn to_search_result(&self, model: &ModelInfo, query: &SearchQuery) -> SearchResult {
|
||||
let metrics = self.get_metrics(&model.metadata.id);
|
||||
|
||||
let trending_score = self.calculate_trending_score(&metrics, &model.metadata);
|
||||
|
||||
let relevance_score =
|
||||
if query.query.is_some() {
|
||||
Some(self.calculate_text_relevance(
|
||||
&model.metadata,
|
||||
query.query.as_deref().unwrap_or(""),
|
||||
))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
SearchResult {
|
||||
model_id: model.metadata.id.clone(),
|
||||
model_summary: ModelSummary::from_metadata(&model.metadata),
|
||||
downloads: metrics.downloads,
|
||||
likes: metrics.likes,
|
||||
trending_score,
|
||||
last_modified: model.metadata.updated_at,
|
||||
relevance_score,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get metrics for a model.
|
||||
fn get_metrics(&self, model_id: &ModelId) -> ModelMetrics {
|
||||
self.metrics
|
||||
.get(&model_id.to_string())
|
||||
.map(|m| m.clone())
|
||||
.unwrap_or_else(|| ModelMetrics {
|
||||
model_id: model_id.clone(),
|
||||
downloads: 0,
|
||||
likes: 0,
|
||||
last_downloaded: None,
|
||||
last_liked: None,
|
||||
daily_downloads: HashMap::new(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Calculate trending score based on recent activity.
|
||||
fn calculate_trending_score(&self, metrics: &ModelMetrics, metadata: &ModelMetadata) -> f64 {
|
||||
let now = Utc::now();
|
||||
let age_days = (now - metadata.created_at).num_days() as f64;
|
||||
|
||||
// Avoid division by zero
|
||||
let age_factor = if age_days > 0.0 {
|
||||
1.0 / (1.0 + age_days.ln())
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
|
||||
// Recent activity factor
|
||||
let recent_factor = if let Some(last_downloaded) = metrics.last_downloaded {
|
||||
let days_since = (now - last_downloaded).num_days() as f64;
|
||||
(-days_since / 7.0).exp() // Exponential decay over 7 days
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
// Combine downloads, likes, and recency
|
||||
let popularity = (metrics.downloads as f64).ln_1p() + (metrics.likes as f64 * 2.0).ln_1p();
|
||||
|
||||
let score = (popularity * age_factor * 0.4) + (recent_factor * 0.6);
|
||||
|
||||
// Normalize to 0.0-1.0 range
|
||||
score.min(1.0).max(0.0)
|
||||
}
|
||||
|
||||
/// Calculate text relevance score.
|
||||
fn calculate_text_relevance(&self, metadata: &ModelMetadata, query: &str) -> f64 {
|
||||
let query_lower = query.to_lowercase();
|
||||
let mut score = 0.0;
|
||||
|
||||
// Title match (highest weight)
|
||||
if metadata.title.to_lowercase().contains(&query_lower) {
|
||||
score += 0.5;
|
||||
}
|
||||
|
||||
// Description match
|
||||
if metadata.description.to_lowercase().contains(&query_lower) {
|
||||
score += 0.3;
|
||||
}
|
||||
|
||||
// Tag matches
|
||||
let tag_matches = metadata
|
||||
.tags
|
||||
.iter()
|
||||
.filter(|tag| tag.to_lowercase().contains(&query_lower))
|
||||
.count();
|
||||
score += (tag_matches as f64 * 0.1).min(0.2);
|
||||
|
||||
score.min(1.0)
|
||||
}
|
||||
|
||||
/// Calculate similarity between models.
|
||||
fn calculate_similarity(&self, reference: &ModelMetadata, candidate: &ModelSummary) -> f64 {
|
||||
let mut score = 0.0;
|
||||
|
||||
// Tag overlap
|
||||
let common_tags = reference
|
||||
.tags
|
||||
.iter()
|
||||
.filter(|tag| candidate.tags.contains(tag))
|
||||
.count();
|
||||
|
||||
if !reference.tags.is_empty() {
|
||||
score += (common_tags as f64 / reference.tags.len() as f64) * 0.6;
|
||||
}
|
||||
|
||||
// Same author
|
||||
if reference.author == candidate.author {
|
||||
score += 0.2;
|
||||
}
|
||||
|
||||
// Same license
|
||||
if reference.license == candidate.license {
|
||||
score += 0.2;
|
||||
}
|
||||
|
||||
score.min(1.0)
|
||||
}
|
||||
|
||||
/// Sort search results.
|
||||
fn sort_results(&self, results: &mut [SearchResult], sort_by: SortBy) {
|
||||
match sort_by {
|
||||
SortBy::Downloads => {
|
||||
results.sort_by(|a, b| b.downloads.cmp(&a.downloads));
|
||||
}
|
||||
SortBy::Likes => {
|
||||
results.sort_by(|a, b| b.likes.cmp(&a.likes));
|
||||
}
|
||||
SortBy::Recent => {
|
||||
results.sort_by(|a, b| b.last_modified.cmp(&a.last_modified));
|
||||
}
|
||||
SortBy::Trending => {
|
||||
results.sort_by(|a, b| {
|
||||
b.trending_score
|
||||
.partial_cmp(&a.trending_score)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
}
|
||||
SortBy::Relevance => {
|
||||
results.sort_by(|a, b| {
|
||||
let score_a = a.relevance_score.unwrap_or(0.0);
|
||||
let score_b = b.relevance_score.unwrap_or(0.0);
|
||||
score_b
|
||||
.partial_cmp(&score_a)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Model metrics for tracking popularity and trends.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ModelMetrics {
|
||||
/// Model ID
|
||||
pub model_id: ModelId,
|
||||
/// Total downloads
|
||||
pub downloads: u64,
|
||||
/// Total likes
|
||||
pub likes: u64,
|
||||
/// Last download timestamp
|
||||
pub last_downloaded: Option<DateTime<Utc>>,
|
||||
/// Last like timestamp
|
||||
pub last_liked: Option<DateTime<Utc>>,
|
||||
/// Daily download counts (date -> count)
|
||||
pub daily_downloads: HashMap<String, u64>,
|
||||
}
|
||||
|
||||
impl ModelMetrics {
|
||||
/// Create new metrics for a model.
|
||||
pub fn new(model_id: ModelId) -> Self {
|
||||
Self {
|
||||
model_id,
|
||||
downloads: 0,
|
||||
likes: 0,
|
||||
last_downloaded: None,
|
||||
last_liked: None,
|
||||
daily_downloads: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Increment download count.
|
||||
pub fn record_download(&mut self) {
|
||||
self.downloads += 1;
|
||||
self.last_downloaded = Some(Utc::now());
|
||||
|
||||
// Record daily download
|
||||
let date = Utc::now().format("%Y-%m-%d").to_string();
|
||||
*self.daily_downloads.entry(date).or_insert(0) += 1;
|
||||
}
|
||||
|
||||
/// Increment like count.
|
||||
pub fn record_like(&mut self) {
|
||||
self.likes += 1;
|
||||
self.last_liked = Some(Utc::now());
|
||||
}
|
||||
|
||||
/// Get downloads for a specific date.
|
||||
pub fn get_downloads_for_date(&self, date: &str) -> u64 {
|
||||
self.daily_downloads.get(date).copied().unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Get total downloads in the last N days.
|
||||
pub fn get_recent_downloads(&self, days: i64) -> u64 {
|
||||
let cutoff = Utc::now() - chrono::Duration::days(days);
|
||||
|
||||
self.daily_downloads
|
||||
.iter()
|
||||
.filter(|(date_str, _)| {
|
||||
if let Ok(date) = chrono::NaiveDate::parse_from_str(date_str, "%Y-%m-%d") {
|
||||
let datetime = date.and_hms_opt(0, 0, 0).unwrap();
|
||||
datetime.and_utc() >= cutoff
|
||||
} else {
|
||||
false
|
||||
}
|
||||
})
|
||||
.map(|(_, count)| count)
|
||||
.sum()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::model::ModelSchema;
|
||||
use crate::{ModelStatus, ModelVersion, RegistryConfig, StorageConfig};
|
||||
use semver::Version;
|
||||
use tempfile::TempDir;
|
||||
|
||||
async fn create_test_registry() -> (ModelRegistry, TempDir) {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let config = RegistryConfig {
|
||||
storage: StorageConfig::Local {
|
||||
base_path: temp_dir.path().to_path_buf(),
|
||||
},
|
||||
database_url: "sqlite::memory:".to_string(),
|
||||
enable_validation: false,
|
||||
enable_compression: false,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let registry = ModelRegistry::new(config).await.unwrap();
|
||||
(registry, temp_dir)
|
||||
}
|
||||
|
||||
fn create_test_metadata(id: ModelId, title: &str, tags: Vec<String>) -> ModelMetadata {
|
||||
ModelMetadata {
|
||||
id,
|
||||
version: ModelVersion::new(Version::parse("1.0.0").unwrap()),
|
||||
title: title.to_string(),
|
||||
description: format!("Description for {}", title),
|
||||
architecture: "transformer".to_string(),
|
||||
framework: "rustytorch".to_string(),
|
||||
framework_version: "1.0.0".to_string(),
|
||||
tags,
|
||||
author: "Test Author".to_string(),
|
||||
license: Some("MIT".to_string()),
|
||||
created_at: Utc::now(),
|
||||
updated_at: Utc::now(),
|
||||
status: ModelStatus::Available,
|
||||
size: 1024,
|
||||
content_hash: "test-hash".to_string(),
|
||||
dependencies: vec![],
|
||||
schema: ModelSchema {
|
||||
inputs: vec![],
|
||||
outputs: vec![],
|
||||
config: None,
|
||||
},
|
||||
metrics: HashMap::new(),
|
||||
metadata: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_search_query_builder() {
|
||||
let query = SearchQuery::new()
|
||||
.with_query("BERT")
|
||||
.with_task("text-classification")
|
||||
.with_library("rustytorch")
|
||||
.with_limit(10)
|
||||
.with_offset(5);
|
||||
|
||||
assert_eq!(query.query, Some("BERT".to_string()));
|
||||
assert_eq!(
|
||||
query.filter_by_task,
|
||||
Some("text-classification".to_string())
|
||||
);
|
||||
assert_eq!(query.filter_by_library, Some("rustytorch".to_string()));
|
||||
assert_eq!(query.limit, 10);
|
||||
assert_eq!(query.offset, 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_summary_from_metadata() {
|
||||
let model_id = ModelId::new("test", "model");
|
||||
let metadata =
|
||||
create_test_metadata(model_id.clone(), "Test Model", vec!["test".to_string()]);
|
||||
|
||||
let summary = ModelSummary::from_metadata(&metadata);
|
||||
|
||||
assert_eq!(summary.model_id, model_id);
|
||||
assert_eq!(summary.title, "Test Model");
|
||||
assert_eq!(summary.tags, vec!["test".to_string()]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_model_discovery_search() {
|
||||
let (registry, _temp_dir) = create_test_registry().await;
|
||||
|
||||
// Register test models
|
||||
let model1 = create_test_metadata(
|
||||
ModelId::new("test", "bert-base"),
|
||||
"BERT Base Model",
|
||||
vec!["nlp".to_string(), "text-classification".to_string()],
|
||||
);
|
||||
|
||||
let model2 = create_test_metadata(
|
||||
ModelId::new("test", "gpt2"),
|
||||
"GPT-2 Model",
|
||||
vec!["nlp".to_string(), "text-generation".to_string()],
|
||||
);
|
||||
|
||||
registry.register_model(model1).await.unwrap();
|
||||
registry.register_model(model2).await.unwrap();
|
||||
|
||||
let discovery = ModelDiscovery::new(registry);
|
||||
|
||||
// Search for BERT
|
||||
let query = SearchQuery::new().with_query("BERT");
|
||||
let results = discovery.search(query).await.unwrap();
|
||||
|
||||
assert_eq!(results.len(), 1);
|
||||
assert!(results[0].model_summary.title.contains("BERT"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_model_discovery_filter_by_task() {
|
||||
let (registry, _temp_dir) = create_test_registry().await;
|
||||
|
||||
let model1 = create_test_metadata(
|
||||
ModelId::new("test", "classifier"),
|
||||
"Classifier Model",
|
||||
vec!["text-classification".to_string()],
|
||||
);
|
||||
|
||||
let model2 = create_test_metadata(
|
||||
ModelId::new("test", "generator"),
|
||||
"Generator Model",
|
||||
vec!["text-generation".to_string()],
|
||||
);
|
||||
|
||||
registry.register_model(model1).await.unwrap();
|
||||
registry.register_model(model2).await.unwrap();
|
||||
|
||||
let discovery = ModelDiscovery::new(registry);
|
||||
|
||||
// Search by task
|
||||
let results = discovery
|
||||
.get_by_task("text-classification", 10)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(results.len(), 1);
|
||||
assert!(
|
||||
results[0]
|
||||
.model_summary
|
||||
.tags
|
||||
.contains(&"text-classification".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_model_discovery_trending() {
|
||||
let (registry, _temp_dir) = create_test_registry().await;
|
||||
|
||||
let model = create_test_metadata(
|
||||
ModelId::new("test", "model"),
|
||||
"Test Model",
|
||||
vec!["test".to_string()],
|
||||
);
|
||||
|
||||
registry.register_model(model).await.unwrap();
|
||||
|
||||
let discovery = ModelDiscovery::new(registry);
|
||||
|
||||
// Record some activity
|
||||
let model_id = ModelId::new("test", "model");
|
||||
discovery.record_download(&model_id).await.unwrap();
|
||||
discovery.record_like(&model_id).await.unwrap();
|
||||
|
||||
// Get trending
|
||||
let results = discovery.get_trending(10).await.unwrap();
|
||||
|
||||
assert!(!results.is_empty());
|
||||
assert!(results[0].trending_score > 0.0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_model_discovery_popular() {
|
||||
let (registry, _temp_dir) = create_test_registry().await;
|
||||
|
||||
let model = create_test_metadata(
|
||||
ModelId::new("test", "model"),
|
||||
"Test Model",
|
||||
vec!["test".to_string()],
|
||||
);
|
||||
|
||||
registry.register_model(model).await.unwrap();
|
||||
|
||||
let discovery = ModelDiscovery::new(registry);
|
||||
|
||||
// Record downloads
|
||||
let model_id = ModelId::new("test", "model");
|
||||
for _ in 0..5 {
|
||||
discovery.record_download(&model_id).await.unwrap();
|
||||
}
|
||||
|
||||
// Get popular
|
||||
let results = discovery.get_popular(10).await.unwrap();
|
||||
|
||||
assert!(!results.is_empty());
|
||||
assert_eq!(results[0].downloads, 5);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_model_discovery_recent() {
|
||||
let (registry, _temp_dir) = create_test_registry().await;
|
||||
|
||||
let model = create_test_metadata(
|
||||
ModelId::new("test", "model"),
|
||||
"Test Model",
|
||||
vec!["test".to_string()],
|
||||
);
|
||||
|
||||
registry.register_model(model).await.unwrap();
|
||||
|
||||
let discovery = ModelDiscovery::new(registry);
|
||||
|
||||
// Get recent
|
||||
let results = discovery.get_recent(10).await.unwrap();
|
||||
|
||||
assert!(!results.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_model_discovery_similar() {
|
||||
let (registry, _temp_dir) = create_test_registry().await;
|
||||
|
||||
let model1 = create_test_metadata(
|
||||
ModelId::new("test", "bert-base"),
|
||||
"BERT Base",
|
||||
vec!["nlp".to_string(), "encoder".to_string()],
|
||||
);
|
||||
|
||||
let model2 = create_test_metadata(
|
||||
ModelId::new("test", "roberta"),
|
||||
"RoBERTa",
|
||||
vec!["nlp".to_string(), "encoder".to_string()],
|
||||
);
|
||||
|
||||
let model3 = create_test_metadata(
|
||||
ModelId::new("test", "gpt2"),
|
||||
"GPT-2",
|
||||
vec!["nlp".to_string(), "decoder".to_string()],
|
||||
);
|
||||
|
||||
registry.register_model(model1).await.unwrap();
|
||||
registry.register_model(model2).await.unwrap();
|
||||
registry.register_model(model3).await.unwrap();
|
||||
|
||||
let discovery = ModelDiscovery::new(registry);
|
||||
|
||||
// Find similar to BERT
|
||||
let similar = discovery
|
||||
.get_similar(&ModelId::new("test", "bert-base"), 10)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!similar.is_empty());
|
||||
// RoBERTa should be more similar than GPT-2
|
||||
assert_eq!(similar[0].model_id.name, "roberta");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_model_discovery_record_metrics() {
|
||||
let (registry, _temp_dir) = create_test_registry().await;
|
||||
let discovery = ModelDiscovery::new(registry);
|
||||
|
||||
let model_id = ModelId::new("test", "model");
|
||||
|
||||
// Record downloads and likes
|
||||
discovery.record_download(&model_id).await.unwrap();
|
||||
discovery.record_download(&model_id).await.unwrap();
|
||||
discovery.record_like(&model_id).await.unwrap();
|
||||
|
||||
// Check metrics
|
||||
let metrics = discovery.get_metrics(&model_id);
|
||||
assert_eq!(metrics.downloads, 2);
|
||||
assert_eq!(metrics.likes, 1);
|
||||
assert!(metrics.last_downloaded.is_some());
|
||||
assert!(metrics.last_liked.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_metrics_record_download() {
|
||||
let model_id = ModelId::new("test", "model");
|
||||
let mut metrics = ModelMetrics::new(model_id);
|
||||
|
||||
metrics.record_download();
|
||||
|
||||
assert_eq!(metrics.downloads, 1);
|
||||
assert!(metrics.last_downloaded.is_some());
|
||||
assert!(!metrics.daily_downloads.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_metrics_record_like() {
|
||||
let model_id = ModelId::new("test", "model");
|
||||
let mut metrics = ModelMetrics::new(model_id);
|
||||
|
||||
metrics.record_like();
|
||||
|
||||
assert_eq!(metrics.likes, 1);
|
||||
assert!(metrics.last_liked.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_metrics_daily_downloads() {
|
||||
let model_id = ModelId::new("test", "model");
|
||||
let mut metrics = ModelMetrics::new(model_id);
|
||||
|
||||
// Record multiple downloads
|
||||
for _ in 0..3 {
|
||||
metrics.record_download();
|
||||
}
|
||||
|
||||
let today = Utc::now().format("%Y-%m-%d").to_string();
|
||||
assert_eq!(metrics.get_downloads_for_date(&today), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_metrics_recent_downloads() {
|
||||
let model_id = ModelId::new("test", "model");
|
||||
let mut metrics = ModelMetrics::new(model_id);
|
||||
|
||||
// Record downloads
|
||||
metrics.record_download();
|
||||
metrics.record_download();
|
||||
|
||||
let recent = metrics.get_recent_downloads(7);
|
||||
assert_eq!(recent, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_search_query_pagination() {
|
||||
let query = SearchQuery::new().with_limit(20).with_offset(10);
|
||||
|
||||
assert_eq!(query.limit, 20);
|
||||
assert_eq!(query.offset, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sort_by_variants() {
|
||||
assert_eq!(SortBy::Downloads, SortBy::Downloads);
|
||||
assert_ne!(SortBy::Downloads, SortBy::Likes);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user