use super::AnyEngine;
use crate::config::model::{EngineType, ModelConfig, Precision};
use crate::error::VecboostError;
use std::path::Path;
pub fn should_use_quantized_engine(model_path: &Path, quantized: bool) -> bool {
if !quantized {
return false;
}
model_path
.extension()
.and_then(|e| e.to_str())
.map(|e| e.eq_ignore_ascii_case("gguf"))
.unwrap_or(false)
}
pub struct EngineFactory;
impl EngineFactory {
pub fn create(
engine_type: EngineType,
config: &ModelConfig,
) -> Result<AnyEngine, VecboostError> {
if should_use_quantized_engine(&config.model_path, config.quantized) {
#[cfg(feature = "quantized-gguf")]
{
let engine =
super::quantized_engine::QuantizedCandleEngine::load(&config.model_path)?;
return Ok(AnyEngine::Quantized(engine));
}
#[cfg(not(feature = "quantized-gguf"))]
{
return Err(VecboostError::ConfigError(
"检测到 GGUF 量化模型配置(.gguf + quantized=true),但本次构建未启用 \
`quantized-gguf` feature;请以 `--features quantized-gguf` 重新构建"
.to_string(),
));
}
}
let precision = Precision::Fp32;
AnyEngine::new(config, engine_type, precision)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::model::{DeviceType, ModelConfig};
use std::path::PathBuf;
fn test_config() -> ModelConfig {
ModelConfig {
name: "test-factory".to_string(),
engine_type: EngineType::Candle,
model_path: PathBuf::from("/nonexistent/model"),
tokenizer_path: None,
device: DeviceType::Cpu,
max_batch_size: 32,
pooling_mode: None,
expected_dimension: Some(1024),
memory_limit_bytes: None,
oom_fallback_enabled: true,
model_sha256: None,
quantized: false,
}
}
#[test]
fn test_create_candle_returns_error_for_missing_model() {
let config = test_config();
let result = EngineFactory::create(EngineType::Candle, &config);
assert!(result.is_err());
}
#[test]
fn test_quantized_routing_matrix() {
let cases = [
("model.gguf", true, true),
("model.GGUF", true, true),
("model.safetensors", true, false),
("model.gguf", false, false),
("model", true, false),
("model.safetensors", false, false),
];
for (p, q, expected) in cases {
assert_eq!(
should_use_quantized_engine(&PathBuf::from(p), q),
expected,
"路由矩阵: path={} quantized={}",
p,
q
);
}
}
#[test]
fn test_create_gguf_without_feature_returns_config_error() {
let mut config = test_config();
config.model_path = PathBuf::from("/tmp/model-q8_0.gguf");
config.quantized = true;
let result = EngineFactory::create(EngineType::Candle, &config);
assert!(result.is_err());
#[cfg(not(feature = "quantized-gguf"))]
match result {
Err(VecboostError::ConfigError(msg)) => {
assert!(
msg.contains("quantized-gguf"),
"应提示启用 feature: {}",
msg
);
}
Err(other) => panic!("期望 ConfigError,实际 {:?}", other),
Ok(_) => panic!("期望 GGUF 配置无 feature 时失败"),
}
}
#[cfg(feature = "onnx")]
#[test]
fn test_create_onnx_engine_missing_model() {
let mut config = test_config();
config.engine_type = EngineType::Onnx;
config.model_path = PathBuf::from("/nonexistent/onnx/model");
let result = EngineFactory::create(EngineType::Onnx, &config);
assert!(
result.is_err(),
"ONNX engine should return error for missing model path"
);
}
#[test]
fn test_engine_type_display() {
assert_eq!(EngineType::Candle.to_string(), "candle");
}
#[test]
fn test_removed_engine_types_return_error() {
let json = "\"tensorrt\"";
let result: Result<EngineType, _> = serde_json::from_str(json);
assert!(
result.is_err(),
"tensorrt should not deserialize after removal"
);
let json = "\"openvino\"";
let result: Result<EngineType, _> = serde_json::from_str(json);
assert!(
result.is_err(),
"openvino should not deserialize after removal"
);
}
}