提交qiming-mcp-proxy

qiming-mcp-proxy/fastembed/src/cli/mod.rs

@@ -0,0 +1,106 @@
+use clap::{Parser, Subcommand};
+use std::path::PathBuf;
+
+pub mod models;
+
+/// FastEmbed - 文本向量化服务
+#[derive(Parser, Debug)]
+#[command(name = "fastembed")]
+#[command(about = "FastEmbed 文本向量化服务", long_about = None)]
+pub struct Cli {
+    #[command(subcommand)]
+    pub command: Commands,
+}
+
+#[derive(Subcommand, Debug)]
+pub enum Commands {
+    /// 启动 HTTP 服务
+    Server(ServerArgs),
+
+    /// 模型管理
+    Models(ModelsCmd),
+}
+
+/// HTTP 服务启动参数
+#[derive(Parser, Debug)]
+pub struct ServerArgs {
+    /// 监听端口
+    #[arg(short, long, default_value = "8080")]
+    pub port: u16,
+
+    /// 配置文件路径
+    #[arg(short, long)]
+    pub config: Option<PathBuf>,
+}
+
+/// 模型管理子命令
+#[derive(Parser, Debug)]
+pub struct ModelsCmd {
+    #[command(subcommand)]
+    pub command: ModelsSubcommand,
+}
+
+#[derive(Subcommand, Debug)]
+pub enum ModelsSubcommand {
+    /// 下载模型到本地缓存
+    Download(DownloadArgs),
+
+    /// 列出已下载的模型
+    List(ListArgs),
+}
+
+/// 模型下载参数
+#[derive(Parser, Debug)]
+pub struct DownloadArgs {
+    /// 模型类型: text | image | sparse
+    #[arg(long, default_value = "text")]
+    pub r#type: String,
+
+    /// 内置模型变体名，如 BGELargeZHV15
+    #[arg(long)]
+    pub model: Option<String>,
+
+    /// Hugging Face 模型代码，如 Xenova/bge-large-zh-v1.5
+    #[arg(long)]
+    pub code: Option<String>,
+
+    /// BYO 模式：ONNX 文件名
+    #[arg(long)]
+    pub onnx: Option<String>,
+
+    /// BYO 模式：Tokenizer 文件名
+    #[arg(long)]
+    pub tokenizer: Option<String>,
+
+    /// BYO 模式：Config 文件名
+    #[arg(long)]
+    pub config: Option<String>,
+
+    /// BYO 模式：Special tokens map 文件名
+    #[arg(long, alias = "special_tokens")]
+    pub special_tokens_map: Option<String>,
+
+    /// BYO 模式：Tokenizer config 文件名
+    #[arg(long)]
+    pub tokenizer_config: Option<String>,
+
+    /// 缓存目录
+    #[arg(long, default_value = ".fastembed_cache")]
+    pub cache_dir: PathBuf,
+
+    /// 显示下载进度
+    #[arg(long, default_value_t = true)]
+    pub progress: bool,
+}
+
+/// 模型列表参数
+#[derive(Parser, Debug)]
+pub struct ListArgs {
+    /// 模型类型筛选: text | image | sparse
+    #[arg(long, default_value = "text")]
+    pub r#type: String,
+
+    /// 缓存目录
+    #[arg(long, default_value = ".fastembed_cache")]
+    pub cache_dir: PathBuf,
+}

qiming-mcp-proxy/fastembed/src/cli/models.rs

@@ -0,0 +1,96 @@
+use super::{DownloadArgs, ListArgs};
+use crate::models::{ModelInfo, list_available_models};
+use anyhow::Result;
+
+/// 执行模型下载
+pub async fn download_model(args: DownloadArgs) -> Result<()> {
+    use fastembed::{InitOptions, TextEmbedding};
+
+    tracing::info!("Start downloading the model...");
+
+    // 解析模型
+    let model = if let Some(model_name) = args.model {
+        // 使用内置模型变体名
+        crate::models::parse_model(&model_name)?
+    } else if let Some(code) = args.code {
+        // 使用模型代码
+        crate::models::parse_model(&code)?
+    } else {
+        anyhow::bail!("必须指定 --model 或 --code 参数");
+    };
+
+    // 显示下载信息
+    let model_info = ModelInfo::from_embedding_model(&model);
+    println!("📦 Download model:");
+    println!("Variant name: {}", model_info.variant);
+    println!("Model code: {}", model_info.code);
+    println!("Vector dimensions: {}", model_info.dim);
+    println!("Cache directory: {}", args.cache_dir.display());
+    println!();
+
+    // 初始化模型（会自动下载）
+    let mut options = InitOptions::new(model.clone());
+    options = options.with_cache_dir(args.cache_dir.clone());
+    options = options.with_show_download_progress(args.progress);
+
+    println!("⬇️ Downloading model files...");
+    let start = std::time::Instant::now();
+
+    let _embedding = TextEmbedding::try_new(options)?;
+
+    let elapsed = start.elapsed();
+
+    println!();
+    println!("✅ Model download completed!");
+    println!("Time taken: {:?}", elapsed);
+    println!("Cache location: {}", args.cache_dir.display());
+
+    // 验证文件
+    println!();
+    println!("🔍 Verify model file...");
+    let available = list_available_models(args.cache_dir.to_str().unwrap())?;
+
+    if available.iter().any(|m| m.variant == model_info.variant) {
+        println!("✅ Model file verification successful!");
+    } else {
+        println!("⚠️ WARNING: Model file may be incomplete");
+    }
+
+    Ok(())
+}
+
+/// 列出已下载的模型
+pub async fn list_models(args: ListArgs) -> Result<()> {
+    use crate::models::list_available_models;
+
+    println!("📋 Query downloaded models...");
+    println!("Type: {}", args.r#type);
+    println!("Cache directory: {}", args.cache_dir.display());
+    println!();
+
+    // 检查缓存目录是否存在
+    if !args.cache_dir.exists() {
+        println!("⚠️ The cache directory does not exist: {}", args.cache_dir.display());
+        println!("Tip: Please download the model first");
+        return Ok(());
+    }
+
+    // 列出可用模型
+    let models = list_available_models(args.cache_dir.to_str().unwrap())?;
+
+    if models.is_empty() {
+        println!("📭 No downloaded model found");
+        println!("Tip: Use 'fastembed models download --model BGELargeZHV15' to download the model");
+    } else {
+        println!("✅ Found {} downloaded models:", models.len());
+        println!();
+        println!("{:<20} {:<40} {:<10}", "Variant", "Model Code", "Dim");
+        println!("{}", "─".repeat(72));
+
+        for model in models {
+            println!("{:<20} {:<40} {:<10}", model.variant, model.code, model.dim);
+        }
+    }
+
+    Ok(())
+}