use crate::errors::{Result, TrustformersError};
use semver::{Version, VersionReq};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PluginInfo {
name: String,
version: Version,
description: String,
author: Option<String>,
homepage: Option<String>,
license: Option<String>,
dependencies: Vec<Dependency>,
capabilities: Vec<String>,
tags: Vec<String>,
requirements: SystemRequirements,
entry_point: String,
metadata: HashMap<String, serde_json::Value>,
}
impl PluginInfo {
pub fn new(name: &str, version: &str, description: &str, dependencies: &[&str]) -> Self {
let version = Version::parse(version).expect("Invalid version string");
let deps = dependencies.iter().filter_map(|dep| Dependency::parse(dep).ok()).collect();
Self {
name: name.to_string(),
version,
description: description.to_string(),
author: None,
homepage: None,
license: None,
dependencies: deps,
capabilities: Vec::new(),
tags: Vec::new(),
requirements: SystemRequirements::default(),
entry_point: String::new(),
metadata: HashMap::new(),
}
}
pub fn name(&self) -> &str {
&self.name
}
pub fn version(&self) -> &Version {
&self.version
}
pub fn description(&self) -> &str {
&self.description
}
pub fn author(&self) -> Option<&str> {
self.author.as_deref()
}
pub fn set_author(&mut self, author: String) {
self.author = Some(author);
}
pub fn homepage(&self) -> Option<&str> {
self.homepage.as_deref()
}
pub fn set_homepage(&mut self, homepage: String) {
self.homepage = Some(homepage);
}
pub fn license(&self) -> Option<&str> {
self.license.as_deref()
}
pub fn set_license(&mut self, license: String) {
self.license = Some(license);
}
pub fn dependencies(&self) -> &[Dependency] {
&self.dependencies
}
pub fn add_dependency(&mut self, dependency: Dependency) {
self.dependencies.push(dependency);
}
pub fn capabilities(&self) -> &[String] {
&self.capabilities
}
pub fn add_capability(&mut self, capability: String) {
self.capabilities.push(capability);
}
pub fn tags(&self) -> &[String] {
&self.tags
}
pub fn add_tag(&mut self, tag: String) {
self.tags.push(tag);
}
pub fn requirements(&self) -> &SystemRequirements {
&self.requirements
}
pub fn set_requirements(&mut self, requirements: SystemRequirements) {
self.requirements = requirements;
}
pub fn entry_point(&self) -> &str {
&self.entry_point
}
pub fn set_entry_point(&mut self, entry_point: String) {
self.entry_point = entry_point;
}
pub fn metadata(&self) -> &HashMap<String, serde_json::Value> {
&self.metadata
}
pub fn add_metadata(&mut self, key: String, value: serde_json::Value) {
self.metadata.insert(key, value);
}
pub fn is_compatible_with(&self, name: &str, version: &str) -> bool {
if let Ok(dep_version) = Version::parse(version) {
for dep in &self.dependencies {
if dep.name == name {
return dep.requirement.matches(&dep_version);
}
}
}
true }
pub fn validate(&self) -> Result<()> {
if self.name.is_empty() {
return Err(TrustformersError::invalid_config(
"Plugin name cannot be empty".to_string(),
));
}
if self.description.is_empty() {
return Err(TrustformersError::invalid_config(
"Plugin description cannot be empty".to_string(),
));
}
if self.entry_point.is_empty() {
return Err(TrustformersError::invalid_config(
"Plugin entry point cannot be empty".to_string(),
));
}
for dep in &self.dependencies {
dep.validate()?;
}
self.requirements.validate()?;
Ok(())
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Dependency {
pub name: String,
pub requirement: VersionReq,
pub optional: bool,
pub features: Vec<String>,
}
impl Dependency {
pub fn new(name: &str, requirement: &str) -> Result<Self> {
let req = VersionReq::parse(requirement).map_err(|e| {
TrustformersError::invalid_config(format!(
"Invalid version requirement '{}': {}",
requirement, e
))
})?;
Ok(Self {
name: name.to_string(),
requirement: req,
optional: false,
features: Vec::new(),
})
}
pub fn parse(spec: &str) -> Result<Self> {
let parts: Vec<&str> = spec.splitn(2, ' ').collect();
let name = parts[0].to_string();
let requirement = if parts.len() > 1 {
VersionReq::parse(parts[1].trim()).map_err(|e| {
TrustformersError::invalid_config(format!(
"Invalid dependency spec '{}': {}",
spec, e
))
})?
} else {
VersionReq::STAR
};
Ok(Self {
name,
requirement,
optional: false,
features: Vec::new(),
})
}
pub fn optional(mut self) -> Self {
self.optional = true;
self
}
pub fn with_features(mut self, features: Vec<String>) -> Self {
self.features = features;
self
}
pub fn validate(&self) -> Result<()> {
if self.name.is_empty() {
return Err(TrustformersError::invalid_config(
"Dependency name cannot be empty".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct SystemRequirements {
pub min_memory_mb: Option<u64>,
pub min_disk_mb: Option<u64>,
pub cpu_features: Vec<String>,
pub gpu: Option<GpuRequirements>,
pub os: Vec<String>,
pub arch: Vec<String>,
}
impl SystemRequirements {
pub fn new() -> Self {
Self::default()
}
pub fn min_memory_mb(mut self, memory_mb: u64) -> Self {
self.min_memory_mb = Some(memory_mb);
self
}
pub fn min_disk_mb(mut self, disk_mb: u64) -> Self {
self.min_disk_mb = Some(disk_mb);
self
}
pub fn cpu_feature(mut self, feature: String) -> Self {
self.cpu_features.push(feature);
self
}
pub fn gpu(mut self, gpu: GpuRequirements) -> Self {
self.gpu = Some(gpu);
self
}
pub fn os(mut self, os: String) -> Self {
self.os.push(os);
self
}
pub fn arch(mut self, arch: String) -> Self {
self.arch.push(arch);
self
}
pub fn validate(&self) -> Result<()> {
if let Some(gpu) = &self.gpu {
gpu.validate()?;
}
Ok(())
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuRequirements {
pub min_memory_mb: u64,
pub compute_capability: Option<String>,
pub vendors: Vec<String>,
pub apis: Vec<String>,
}
impl GpuRequirements {
pub fn new(min_memory_mb: u64) -> Self {
Self {
min_memory_mb,
compute_capability: None,
vendors: Vec::new(),
apis: Vec::new(),
}
}
pub fn compute_capability(mut self, capability: String) -> Self {
self.compute_capability = Some(capability);
self
}
pub fn vendor(mut self, vendor: String) -> Self {
self.vendors.push(vendor);
self
}
pub fn api(mut self, api: String) -> Self {
self.apis.push(api);
self
}
pub fn validate(&self) -> Result<()> {
if self.min_memory_mb == 0 {
return Err(TrustformersError::invalid_config(
"GPU memory requirement must be greater than 0".to_string(),
));
}
Ok(())
}
}