xz-rerank 0.1.1

检索结果重排序 — 本地多信号融合 + 远程 Rerank API 适配
Documentation
use async_trait::async_trait;
use std::collections::HashMap;

use crate::error::RerankError;
use crate::traits::SignalPlugin;
use crate::types::RerankCandidate;

/// 元数据匹配信号
#[derive(Debug)]
pub struct MetadataMatchSignal {
    /// 需要匹配的元数据字段及权重
    field_weights: HashMap<String, f32>,
}

impl MetadataMatchSignal {
    /// 创建元数据匹配信号
    ///
    /// `field_weights` 指定需要匹配的元数据字段名及其权重。
    /// 例如:`{"author": 0.7, "category": 0.3}`。
    pub fn new(field_weights: HashMap<String, f32>) -> Self {
        Self { field_weights }
    }
}

impl Default for MetadataMatchSignal {
    /// 返回默认的元数据匹配信号(空的字段权重,所有候选项得分均为 1.0)
    fn default() -> Self {
        Self { field_weights: HashMap::new() }
    }
}

#[async_trait]
impl SignalPlugin for MetadataMatchSignal {
    fn name(&self) -> &str {
        "metadata_match"
    }

    fn weight_key(&self) -> &'static str {
        "metadata_match"
    }

    async fn score(&self, _query: &str, candidate: &RerankCandidate) -> Result<f32, RerankError> {
        if self.field_weights.is_empty() {
            return Ok(1.0);
        }

        let mut total_weight = 0.0f32;
        let mut matched_weight = 0.0f32;

        for (field, weight) in &self.field_weights {
            total_weight += weight;
            if candidate.metadata.contains_key(field.as_str()) {
                matched_weight += weight;
            }
        }

        if total_weight <= f32::EPSILON {
            return Ok(1.0);
        }
        Ok(matched_weight / total_weight)
    }
}