提交qiming-mcp-proxy

qiming-mcp-proxy/fastembed/src/handlers/embeddings.rs

@@ -0,0 +1,167 @@
+use axum::{Json, extract::State, http::StatusCode};
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+use std::time::Instant;
+use utoipa::ToSchema;
+
+use crate::models::{ModelInfo, get_or_init_model, parse_model};
+use crate::server::AppState;
+
+/// 文本嵌入请求
+#[derive(Debug, Deserialize, ToSchema)]
+pub struct EmbedRequest {
+    /// 模型名称（变体名或模型代码）
+    #[schema(example = "BGELargeZHV15")]
+    pub model: Option<String>,
+
+    /// 待嵌入的文本列表
+    #[schema(example = json!(["query: 搜索文本", "passage: 文档内容"]))]
+    pub texts: Vec<String>,
+
+    /// 批处理大小
+    #[schema(example = 256)]
+    pub batch_size: Option<usize>,
+}
+
+/// 文本嵌入响应
+#[derive(Debug, Serialize, ToSchema)]
+pub struct EmbedResponse {
+    /// 模型信息
+    pub model: ModelInfo,
+
+    /// 嵌入向量数量
+    #[schema(example = 2)]
+    pub count: usize,
+
+    /// 嵌入向量列表
+    #[schema(example = json!([[0.00123, -0.00456], [0.00078, 0.00234]]))]
+    pub embeddings: Vec<Vec<f32>>,
+
+    /// 耗时（毫秒）
+    #[schema(example = 12)]
+    pub elapsed_ms: u128,
+}
+
+/// 错误响应
+#[derive(Debug, Serialize, ToSchema)]
+pub struct ErrorResponse {
+    /// 错误代码
+    #[schema(example = "INVALID_MODEL")]
+    pub error: String,
+
+    /// 错误消息
+    #[schema(example = "未知模型")]
+    pub message: String,
+
+    /// HTTP 状态码
+    #[schema(example = 400)]
+    pub status: u16,
+}
+
+/// 文本嵌入处理器
+#[utoipa::path(
+    post,
+    path = "/api/embeddings",
+    tag = "文本嵌入",
+    request_body = EmbedRequest,
+    responses(
+        (status = 200, description = "嵌入成功", body = EmbedResponse),
+        (status = 400, description = "请求参数错误", body = ErrorResponse),
+        (status = 413, description = "请求负载过大", body = ErrorResponse),
+        (status = 500, description = "服务器错误", body = ErrorResponse)
+    )
+)]
+pub async fn handle_embed(
+    State(state): State<Arc<AppState>>,
+    Json(req): Json<EmbedRequest>,
+) -> Result<Json<EmbedResponse>, (StatusCode, Json<ErrorResponse>)> {
+    let start = Instant::now();
+
+    // 参数验证
+    if req.texts.is_empty() {
+        return Err((
+            StatusCode::BAD_REQUEST,
+            Json(ErrorResponse {
+                error: "EMPTY_TEXTS".to_string(),
+                message: "texts 不能为空".to_string(),
+                status: 400,
+            }),
+        ));
+    }
+
+    // 检查文本数量限制（最大 1024）
+    if req.texts.len() > 1024 {
+        return Err((
+            StatusCode::PAYLOAD_TOO_LARGE,
+            Json(ErrorResponse {
+                error: "TOO_MANY_TEXTS".to_string(),
+                message: format!("texts 数量不能超过 1024，当前: {}", req.texts.len()),
+                status: 413,
+            }),
+        ));
+    }
+
+    // 解析模型
+    let model_name = req
+        .model
+        .as_deref()
+        .unwrap_or(&state.config.fastembed.default_model);
+    let embedding_model = parse_model(model_name).map_err(|e| {
+        (
+            StatusCode::BAD_REQUEST,
+            Json(ErrorResponse {
+                error: "INVALID_MODEL".to_string(),
+                message: format!("未知模型: {}, 错误: {}", model_name, e),
+                status: 400,
+            }),
+        )
+    })?;
+
+    // 获取或初始化模型
+    let model_arc = get_or_init_model(
+        embedding_model.clone(),
+        Some(state.config.fastembed.cache_dir.clone()),
+        None, // 使用模型默认的 max_length
+    )
+    .map_err(|e| {
+        tracing::error!("Model initialization failed: {}", e);
+        (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            Json(ErrorResponse {
+                error: "MODEL_INIT_ERROR".to_string(),
+                message: format!("模型初始化失败: {}", e),
+                status: 500,
+            }),
+        )
+    })?;
+
+    // 执行嵌入
+    let batch_size = req.batch_size.unwrap_or(state.config.fastembed.batch_size);
+
+    let mut model_guard = model_arc.lock().unwrap();
+    let embeddings = model_guard
+        .embed(req.texts.clone(), Some(batch_size))
+        .map_err(|e| {
+            tracing::error!("Embedding calculation failed: {}", e);
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                Json(ErrorResponse {
+                    error: "EMBED_ERROR".to_string(),
+                    message: format!("嵌入计算失败: {}", e),
+                    status: 500,
+                }),
+            )
+        })?;
+
+    // 转换为 Vec<Vec<f32>>
+    let embeddings_vec: Vec<Vec<f32>> = embeddings.into_iter().map(|e| e.to_vec()).collect();
+
+    let elapsed = start.elapsed();
+
+    Ok(Json(EmbedResponse {
+        model: ModelInfo::from_embedding_model(&embedding_model),
+        count: embeddings_vec.len(),
+        embeddings: embeddings_vec,
+        elapsed_ms: elapsed.as_millis(),
+    }))
+}

qiming-mcp-proxy/fastembed/src/handlers/health.rs

@@ -0,0 +1,41 @@
+use axum::{Json, extract::State};
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+use utoipa::ToSchema;
+
+use crate::server::AppState;
+
+/// 健康检查响应
+#[derive(Debug, Serialize, Deserialize, ToSchema)]
+pub struct HealthResponse {
+    /// 服务状态
+    #[schema(example = "ok")]
+    pub status: String,
+
+    /// 服务运行时长（毫秒）
+    #[schema(example = 123456)]
+    pub uptime_ms: u128,
+
+    /// 模型缓存是否就绪
+    #[schema(example = true)]
+    pub model_cache_ready: bool,
+}
+
+/// 健康检查处理器
+#[utoipa::path(
+    get,
+    path = "/health",
+    tag = "健康检查",
+    responses(
+        (status = 200, description = "服务健康", body = HealthResponse)
+    )
+)]
+pub async fn handle_health(State(state): State<Arc<AppState>>) -> Json<HealthResponse> {
+    let uptime = state.start_time.elapsed();
+
+    Json(HealthResponse {
+        status: "ok".to_string(),
+        uptime_ms: uptime.as_millis(),
+        model_cache_ready: *state.model_cache_ready.lock().unwrap(),
+    })
+}

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

@@ -0,0 +1,3 @@
+pub mod embeddings;
+pub mod health;
+pub mod models;

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

@@ -0,0 +1,87 @@
+use axum::{
+    Json,
+    extract::{Query, State},
+    http::StatusCode,
+};
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+use utoipa::{IntoParams, ToSchema};
+
+use crate::handlers::embeddings::ErrorResponse;
+use crate::models::{ModelInfo, list_available_models};
+use crate::server::AppState;
+
+/// 查询参数
+#[derive(Debug, Deserialize, IntoParams)]
+pub struct ModelsQuery {
+    /// 模型类型: text | image | sparse
+    #[serde(rename = "type")]
+    #[param(example = "text")]
+    pub model_type: Option<String>,
+}
+
+/// 模型列表响应
+#[derive(Debug, Serialize, ToSchema)]
+pub struct ModelsResponse {
+    /// 模型类型
+    #[schema(example = "text")]
+    pub r#type: String,
+
+    /// 模型数量
+    #[schema(example = 2)]
+    pub count: usize,
+
+    /// 模型列表
+    pub models: Vec<ModelInfo>,
+}
+
+/// 列出可用模型处理器
+#[utoipa::path(
+    get,
+    path = "/api/models/available",
+    tag = "模型管理",
+    params(ModelsQuery),
+    responses(
+        (status = 200, description = "模型列表", body = ModelsResponse),
+        (status = 400, description = "请求参数错误", body = ErrorResponse),
+        (status = 500, description = "服务器错误", body = ErrorResponse)
+    )
+)]
+pub async fn handle_list_models(
+    State(state): State<Arc<AppState>>,
+    Query(query): Query<ModelsQuery>,
+) -> Result<Json<ModelsResponse>, (StatusCode, Json<ErrorResponse>)> {
+    // 验证类型参数
+    let model_type = query.model_type.as_deref().unwrap_or("text");
+
+    // 目前仅支持 text 类型
+    if model_type != "text" {
+        return Err((
+            StatusCode::BAD_REQUEST,
+            Json(ErrorResponse {
+                error: "INVALID_TYPE".to_string(),
+                message: format!("不支持的模型类型: {}，当前仅支持 text", model_type),
+                status: 400,
+            }),
+        ));
+    }
+
+    // 列出可用模型
+    let models = list_available_models(&state.config.fastembed.cache_dir).map_err(|e| {
+        tracing::error!("Failed to list available models: {}", e);
+        (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            Json(ErrorResponse {
+                error: "LIST_ERROR".to_string(),
+                message: format!("列出可用模型失败: {}", e),
+                status: 500,
+            }),
+        )
+    })?;
+
+    Ok(Json(ModelsResponse {
+        r#type: model_type.to_string(),
+        count: models.len(),
+        models,
+    }))
+}