diff --git a/.sqlx-sqlite/query-cc5af75f667fffc113e4369c75b0010c49269abff921e35e8bb52f08a3fb2871.json b/.sqlx-sqlite/query-36144402d6722221debd0ffdaf93b8585dbe86b44abc1f901e411f36ac84f530.json similarity index 50% rename from .sqlx-sqlite/query-cc5af75f667fffc113e4369c75b0010c49269abff921e35e8bb52f08a3fb2871.json rename to .sqlx-sqlite/query-36144402d6722221debd0ffdaf93b8585dbe86b44abc1f901e411f36ac84f530.json index c48206b6..0a25a3f3 100644 --- a/.sqlx-sqlite/query-cc5af75f667fffc113e4369c75b0010c49269abff921e35e8bb52f08a3fb2871.json +++ b/.sqlx-sqlite/query-36144402d6722221debd0ffdaf93b8585dbe86b44abc1f901e411f36ac84f530.json @@ -1,12 +1,12 @@ { "db_name": "SQLite", - "query": "INSERT INTO kb_knowledge_bases (id, name, description, slug, kind, indexing_strategy, embedding_model, embedding_dim, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)", + "query": "INSERT INTO kb_knowledge_bases (id, name, description, slug, kind, indexing_strategy, embedding_model, embedding_dim, rerank_model, rerank_window, rerank_threshold, created_at, tenant_id) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)", "describe": { "columns": [], "parameters": { - "Right": 9 + "Right": 13 }, "nullable": [] }, - "hash": "cc5af75f667fffc113e4369c75b0010c49269abff921e35e8bb52f08a3fb2871" + "hash": "36144402d6722221debd0ffdaf93b8585dbe86b44abc1f901e411f36ac84f530" } diff --git a/.sqlx-sqlite/query-5163b808b5a6443f0d2315e16c3a61bff034b5a0a94ca64aca46b53e2bc84f81.json b/.sqlx-sqlite/query-893ed416697c7606697ea438fa9bd868b91e3d824f0abf66d8bb111f8c431d22.json similarity index 52% rename from .sqlx-sqlite/query-5163b808b5a6443f0d2315e16c3a61bff034b5a0a94ca64aca46b53e2bc84f81.json rename to .sqlx-sqlite/query-893ed416697c7606697ea438fa9bd868b91e3d824f0abf66d8bb111f8c431d22.json index 8d4b0fff..f33ef123 100644 --- a/.sqlx-sqlite/query-5163b808b5a6443f0d2315e16c3a61bff034b5a0a94ca64aca46b53e2bc84f81.json +++ b/.sqlx-sqlite/query-893ed416697c7606697ea438fa9bd868b91e3d824f0abf66d8bb111f8c431d22.json @@ -1,12 +1,12 @@ { "db_name": "SQLite", - "query": "INSERT INTO kb_knowledge_bases (id, name, description, slug, kind, indexing_strategy, embedding_model, embedding_dim, created_at, tenant_id) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)", + "query": "INSERT INTO kb_knowledge_bases (id, name, description, slug, kind, indexing_strategy, embedding_model, embedding_dim, rerank_model, rerank_window, rerank_threshold, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)", "describe": { "columns": [], "parameters": { - "Right": 10 + "Right": 12 }, "nullable": [] }, - "hash": "5163b808b5a6443f0d2315e16c3a61bff034b5a0a94ca64aca46b53e2bc84f81" + "hash": "893ed416697c7606697ea438fa9bd868b91e3d824f0abf66d8bb111f8c431d22" } diff --git a/crates/core/src/agent/tools/kb.rs b/crates/core/src/agent/tools/kb.rs index 8c2c6952..d51f6a44 100644 --- a/crates/core/src/agent/tools/kb.rs +++ b/crates/core/src/agent/tools/kb.rs @@ -498,6 +498,7 @@ mod tests { vector: Arc::new(crate::kb::vectors::BruteForceIndex::new()), kbsearch: Arc::new(crate::kb::kbsearch::KbSearchEngine::open_in_memory().unwrap()), embedder: Arc::new(MockEmbedder), + reranker: None, router, emitter: crate::event::EventEmitter::eventbus_only(bus), } @@ -514,6 +515,9 @@ mod tests { indexing_strategy: None, embedding_model: Some("m".into()), embedding_dim: Some(4), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, tenant, ) @@ -578,6 +582,9 @@ mod tests { indexing_strategy: None, embedding_model: Some("m".into()), embedding_dim: Some(4), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "default", ) diff --git a/crates/core/src/config/app.rs b/crates/core/src/config/app.rs index b2fbc93c..0d1c46fc 100644 --- a/crates/core/src/config/app.rs +++ b/crates/core/src/config/app.rs @@ -737,6 +737,10 @@ impl AiConfig { /// | `RAISFAST_KB_FALLBACK_THRESHOLD` | f32 | `0.3` | Below this top score → "not covered" (no generation) | /// | `RAISFAST_KB_CONTEXT_TOKEN_BUDGET` | u32 | `4000` | S8 context assembly budget (chars/4 estimate) | /// | `RAISFAST_KB_EMBED_BATCH_SIZE` | usize | `32` | Texts per `/embeddings` request (WK `BATCH_EMBED_SIZE` analog; WK default 5 + goroutine pool, ours serializes so 32) | +/// | `RAISFAST_KB_RERANK_MODEL` | string | — | Global **default** rerank model via the llm 底座 (S5, kb-technical-design §6.1 revised 2026-09-17); each KB row's `rerank_model` overrides it; neither set = passthrough | +/// | `RAISFAST_KB_RERANK_WINDOW` | u32 | `30` | Global default rerank window (per-KB `rerank_window` overrides); S4 cuts to this window, S5 reranks then cuts back to `top_k` (§6.1.4 窗口重排) | +/// | `RAISFAST_KB_RERANK_THRESHOLD` | f32 | `0` | Global default rerank score floor (per-KB `rerank_threshold` overrides); below is dropped after S5 (0 = keep all) [抄WK:RerankThreshold 语义] | +/// | `RAISFAST_KB_RERANK_BATCH_SIZE` | usize | `64` | Docs per `/rerank` request (transport concern, mirrors `RAISFAST_KB_EMBED_BATCH_SIZE`) | #[derive(Debug, Clone, Serialize, Deserialize)] pub struct KbConfig { #[serde(default)] @@ -781,6 +785,29 @@ pub struct KbConfig { /// T3 cleanup sweeper. Env `RAISFAST_KB_TRACE_RETENTION_DAYS`. #[serde(default = "default_kb_trace_retention_days")] pub trace_retention_days: i64, + /// Global **default** rerank model routed via the llm 底座 (S5). Each + /// KB row's `rerank_model` overrides it; neither set = S5 passthrough + /// (kb-technical-design §6.1 revised 2026-09-17). Env + /// `RAISFAST_KB_RERANK_MODEL`. + #[serde(default)] + pub rerank_model: Option, + /// Global default rerank window; per-KB `rerank_window` overrides. + /// S4 cuts to this window (≥ `top_k`) so the reranker sees more than + /// the final keep set, S5 reranks then cuts back to `top_k` (§6.1.4 + /// 窗口重排). Env `RAISFAST_KB_RERANK_WINDOW` (default 30). + #[serde(default = "default_kb_rerank_window")] + pub rerank_window: u32, + /// Global default rerank score floor; per-KB `rerank_threshold` + /// overrides. Candidates scoring below are dropped after S5 + /// (0 = keep all) [抄WK:RerankThreshold 语义]. Not refilled when the + /// filter leaves fewer than `top_k` (§6.1.4). Env + /// `RAISFAST_KB_RERANK_THRESHOLD` (default 0). + #[serde(default)] + pub rerank_threshold: f32, + /// Docs per `/rerank` request (global transport concern, mirrors + /// `embed_batch_size`). Env `RAISFAST_KB_RERANK_BATCH_SIZE` (default 64). + #[serde(default = "default_kb_rerank_batch_size")] + pub rerank_batch_size: usize, } fn default_kb_top_k() -> u32 { @@ -819,6 +846,14 @@ fn default_kb_trace_retention_days() -> i64 { 14 } +fn default_kb_rerank_window() -> u32 { + 30 +} + +fn default_kb_rerank_batch_size() -> usize { + 64 +} + impl Default for KbConfig { fn default() -> Self { Self { @@ -834,6 +869,10 @@ impl Default for KbConfig { embed_batch_size: default_kb_embed_batch_size(), trace_mode: default_kb_trace_mode(), trace_retention_days: default_kb_trace_retention_days(), + rerank_model: None, + rerank_window: default_kb_rerank_window(), + rerank_threshold: 0.0, + rerank_batch_size: default_kb_rerank_batch_size(), } } } @@ -890,6 +929,24 @@ impl KbConfig { .and_then(|v| v.parse().ok()) .filter(|d| *d >= 0) .unwrap_or(defaults.trace_retention_days), + rerank_model: env::var("RAISFAST_KB_RERANK_MODEL") + .ok() + .filter(|v| !v.is_empty()), + rerank_window: env::var("RAISFAST_KB_RERANK_WINDOW") + .ok() + .and_then(|v| v.parse().ok()) + .filter(|w| *w > 0) + .unwrap_or(defaults.rerank_window), + rerank_threshold: env::var("RAISFAST_KB_RERANK_THRESHOLD") + .ok() + .and_then(|v| v.parse().ok()) + .filter(|t| (0.0..=1.0).contains(t)) + .unwrap_or(defaults.rerank_threshold), + rerank_batch_size: env::var("RAISFAST_KB_RERANK_BATCH_SIZE") + .ok() + .and_then(|v| v.parse().ok()) + .filter(|n| *n > 0) + .unwrap_or(defaults.rerank_batch_size), } } } diff --git a/crates/core/src/kb/distill.rs b/crates/core/src/kb/distill.rs index b9524d0e..966093c3 100644 --- a/crates/core/src/kb/distill.rs +++ b/crates/core/src/kb/distill.rs @@ -470,6 +470,7 @@ mod tests { vector: Arc::new(BruteForceIndex::new()), kbsearch: Arc::new(crate::kb::kbsearch::KbSearchEngine::open_in_memory().unwrap()), embedder: Arc::new(MockEmbedder), + reranker: None, router, emitter: crate::event::EventEmitter::eventbus_only(bus), } @@ -486,6 +487,9 @@ mod tests { indexing_strategy: None, embedding_model: Some("m".into()), embedding_dim: Some(4), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "default", ) diff --git a/crates/core/src/kb/eval.rs b/crates/core/src/kb/eval.rs index fab1d65a..3faddcc8 100644 --- a/crates/core/src/kb/eval.rs +++ b/crates/core/src/kb/eval.rs @@ -263,6 +263,7 @@ mod tests { vector: Arc::new(BruteForceIndex::new()), kbsearch: Arc::new(crate::kb::kbsearch::KbSearchEngine::open_in_memory().unwrap()), embedder: Arc::new(SumEmbedder), + reranker: None, router, emitter: crate::event::EventEmitter::eventbus_only(crate::eventbus::EventBus::new(16)), }; @@ -276,6 +277,9 @@ mod tests { indexing_strategy: None, embedding_model: Some("m".into()), embedding_dim: Some(4), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "default", ) diff --git a/crates/core/src/kb/handler.rs b/crates/core/src/kb/handler.rs index 1e2e7fc2..ded1e344 100644 --- a/crates/core/src/kb/handler.rs +++ b/crates/core/src/kb/handler.rs @@ -486,6 +486,7 @@ impl AppState { vector: kb.vector.clone(), kbsearch: kb.kbsearch.clone(), embedder: kb.embedder.clone(), + reranker: kb.reranker.clone(), router: self.llm_router.clone(), emitter: self.emitter.clone(), }) @@ -509,6 +510,17 @@ struct CreateKbRequest { /// Dimension pinned with the model (or directory params.dimension). #[serde(default)] embedding_dim: Option, + /// S5 rerank model for this KB (mutable afterwards — rerank is + /// query-time behavior with nothing baked). Empty = global + /// `RAISFAST_KB_RERANK_MODEL` default; neither = passthrough. + #[serde(default)] + rerank_model: Option, + /// Per-KB rerank window override (§6.1.4); null = global default. + #[serde(default)] + rerank_window: Option, + /// Per-KB rerank score floor override (0..=1); null = global default. + #[serde(default)] + rerank_threshold: Option, } fn default_kb_kind() -> String { @@ -528,6 +540,7 @@ async fn admin_create_kb( .into(), )); } + validate_rerank_params(req.rerank_window, req.rerank_threshold)?; let tenant = auth.tenant_id(); let model = resolve_kb_model(&state, tenant, &req).await?; let dim = resolve_kb_dim(&state, tenant, &model, &req).await?; @@ -541,6 +554,9 @@ async fn admin_create_kb( indexing_strategy: None, embedding_model: Some(model), embedding_dim: Some(i64::from(dim)), + rerank_model: req.rerank_model.filter(|m| !m.is_empty()), + rerank_window: req.rerank_window.map(i64::from), + rerank_threshold: req.rerank_threshold, }, &tenant_of(&auth), ) @@ -557,10 +573,18 @@ struct UpdateKbRequest { description: Option, slug: String, status: String, + /// Rerank trio (freely editable — query-time behavior, §6.1). + #[serde(default)] + rerank_model: Option, + #[serde(default)] + rerank_window: Option, + #[serde(default)] + rerank_threshold: Option, } -/// Update mutable KB metadata only (name/slug/description/status). -/// kind / embedding_model / embedding_dim are immutable after creation. +/// Update mutable KB metadata (name/slug/description/status) plus the +/// rerank trio. kind / embedding_model / embedding_dim are immutable after +/// creation. async fn admin_update_kb( auth: AuthUser, State(state): State, @@ -579,6 +603,7 @@ async fn admin_update_kb( "status must be 'active' or 'archived'".into(), )); } + validate_rerank_params(req.rerank_window, req.rerank_threshold)?; let id = parse_snowflake(&id)?; let tenant = tenant_of(&auth); if knowledge_base::find_kb_by_id(&state.pool, id, &tenant) @@ -600,6 +625,9 @@ async fn admin_update_kb( .map(str::to_owned), slug: req.slug.trim().to_owned(), status: req.status, + rerank_model: req.rerank_model.filter(|m| !m.is_empty()), + rerank_window: req.rerank_window.map(i64::from), + rerank_threshold: req.rerank_threshold, }, &tenant, ) @@ -607,6 +635,21 @@ async fn admin_update_kb( Ok(ApiResponse::success(json!({ "id": id }))) } +/// Shared validation for the rerank trio on create/update. +fn validate_rerank_params(window: Option, threshold: Option) -> AppResult<()> { + if window.is_some_and(|w| w == 0) { + return Err(AppError::BadRequest( + "rerank_window must be a positive integer".into(), + )); + } + if threshold.is_some_and(|t| !(0.0..=1.0).contains(&t)) { + return Err(AppError::BadRequest( + "rerank_threshold must be within 0..=1".into(), + )); + } + Ok(()) +} + /// Delete a KB and everything in it (documents, chunks, FAQs, wiki pages, /// provenance links, vector/FTS indexes, original files). async fn admin_delete_kb( diff --git a/crates/core/src/kb/mod.rs b/crates/core/src/kb/mod.rs index 9a95080a..3a9db17e 100644 --- a/crates/core/src/kb/mod.rs +++ b/crates/core/src/kb/mod.rs @@ -16,6 +16,7 @@ pub mod chunker; pub mod diagnostics; pub mod distill; pub mod eval; +pub mod rerank; /// Runtime singletons shared by handlers and workers (built once at boot, /// `None` when the KB is disabled). @@ -23,6 +24,8 @@ pub struct KbRuntime { pub vector: std::sync::Arc, pub kbsearch: std::sync::Arc, pub embedder: std::sync::Arc, + /// S5 reranker; `None` when no rerank model is configured (§6.1). + pub reranker: Option>, } /// Build the KB runtime from config; `Ok(None)` when the KB is disabled, @@ -39,12 +42,22 @@ pub fn build_kb_runtime( let kbsearch_dir = std::path::Path::new(&config.storage_root_dir).join("kb_search_index"); let kbsearch = kbsearch::KbSearchEngine::open(&kbsearch_dir)?; let embedder: std::sync::Arc = std::sync::Arc::new( - service::ProviderEmbedder::new(router, config.kb.embed_batch_size), + service::ProviderEmbedder::new(router.clone(), config.kb.embed_batch_size), ); + // Rerank support is always wired; enablement is resolved per query from + // the KB rows (per-KB override → global RAISFAST_KB_RERANK_MODEL + // default → passthrough, §6.1.2). + let reranker: std::sync::Arc = std::sync::Arc::new( + rerank::ProviderReranker::new(router.clone(), config.kb.rerank_batch_size), + ); + if let Some(model) = config.kb.rerank_model.as_deref() { + tracing::info!("kb rerank default model: '{model}' (per-KB rows may override)"); + } Ok(Some(std::sync::Arc::new(KbRuntime { vector, kbsearch: std::sync::Arc::new(kbsearch), embedder, + reranker: Some(reranker), }))) } pub mod handler; diff --git a/crates/core/src/kb/models/knowledge_base.rs b/crates/core/src/kb/models/knowledge_base.rs index d6129202..89bccfe2 100644 --- a/crates/core/src/kb/models/knowledge_base.rs +++ b/crates/core/src/kb/models/knowledge_base.rs @@ -24,6 +24,14 @@ pub struct KbKnowledgeBase { pub chunking_config: Option, pub embedding_model: Option, pub embedding_dim: Option, + /// S5 rerank model for this KB (query-time behavior — mutable, unlike + /// the pinned embedding config). Empty = fall back to the global + /// `RAISFAST_KB_RERANK_MODEL` default; neither set = S5 passthrough. + pub rerank_model: Option, + /// Per-KB rerank window override (§6.1.4); `None` = global default. + pub rerank_window: Option, + /// Per-KB rerank score floor override; `None` = global default. + pub rerank_threshold: Option, pub status: String, pub created_at: Timestamp, pub updated_at: Timestamp, @@ -40,6 +48,9 @@ pub struct CreateKbCmd { pub indexing_strategy: Option, pub embedding_model: Option, pub embedding_dim: Option, + pub rerank_model: Option, + pub rerank_window: Option, + pub rerank_threshold: Option, } pub async fn create_kb( @@ -63,6 +74,9 @@ pub async fn create_kb( "indexing_strategy" => cmd.indexing_strategy.clone(), "embedding_model" => cmd.embedding_model.as_deref(), "embedding_dim" => cmd.embedding_dim, + "rerank_model" => cmd.rerank_model.as_deref(), + "rerank_window" => cmd.rerank_window, + "rerank_threshold" => cmd.rerank_threshold, "created_at" => now ], tenant: Some(tenant_id) @@ -72,8 +86,9 @@ pub async fn create_kb( .ok_or_else(|| AppError::Internal(anyhow::anyhow!("kb insert returned no row"))) } -/// Update-KB command — only mutable metadata (name/slug/description/status); -/// kind and embedding config are immutable after creation (WeKnora +/// Update-KB command — mutable metadata (name/slug/description/status) and +/// the rerank trio (query-time behavior, freely editable). kind and +/// embedding config stay immutable after creation (WeKnora /// `vector_store_id` precedent: vectors are pinned to model+dim). #[derive(Debug, Clone)] pub struct UpdateKbCmd { @@ -81,6 +96,9 @@ pub struct UpdateKbCmd { pub description: Option, pub slug: String, pub status: String, + pub rerank_model: Option, + pub rerank_window: Option, + pub rerank_threshold: Option, } pub async fn update_kb( @@ -98,6 +116,9 @@ pub async fn update_kb( "description" => cmd.description.as_deref(), "slug" => cmd.slug.as_str(), "status" => cmd.status.as_str(), + "rerank_model" => cmd.rerank_model.as_deref(), + "rerank_window" => cmd.rerank_window, + "rerank_threshold" => cmd.rerank_threshold, "updated_at" => now ], where: ("id", id), diff --git a/crates/core/src/kb/pipeline/mod.rs b/crates/core/src/kb/pipeline/mod.rs index cbdd79fa..fb3a9f8c 100644 --- a/crates/core/src/kb/pipeline/mod.rs +++ b/crates/core/src/kb/pipeline/mod.rs @@ -10,6 +10,7 @@ mod assemble; mod fusion; mod generate; mod merge; +mod rerank; mod search; mod understand; @@ -101,6 +102,8 @@ pub async fn prepare_answer(deps: &KbDeps, req: &AskRequest) -> AppResult AppResult Ok(v), Err(e) => { trace.fail_stage(&e.to_string()); @@ -228,20 +243,28 @@ fn chunk_title(chunk: &KbChunk) -> String { } /// S2–S7 over a validated KB scope: recall → fuse → top-k → hydrate → -/// wiki boost → merge. Returns `(top_score, context_units)`. Stage -/// summaries are side-recorded into `trace` (DR6 orchestration-level). +/// (rerank) → wiki boost → merge. Returns `(top_score, context_units)`. +/// Stage summaries are side-recorded into `trace` (DR6 orchestration-level). async fn recall_and_merge( deps: &KbDeps, tenant: &str, kbs: &[i64], doc_ids: &[i64], understood: &understand::UnderstoodQuery, + rerank: Option<&rerank::ResolvedRerank>, trace: &mut crate::kb::trace::RunRecorder, ) -> AppResult<(f32, Vec)> { use crate::kb::trace::STAGE_LIST_TOP; // S2+S3+S4 recall → fuse → top-k, per KB scope. let mut candidates = Vec::new(); + // S4 cut size: with rerank, S4 keeps the wider rerank window so S5 sees + // more candidates than the final keep set (§6.1.4 窗口重排; window ≥ + // top_k so a small window never shrinks recall). + let s4_cut = match rerank { + Some(cfg) if deps.reranker.is_some() => cfg.window.max(deps.config.kb.top_k), + _ => deps.config.kb.top_k, + }; // Per-KB raw recall lists for the s2 summary (top-10 each, DR3). type RawRecall = (i64, Vec<(i64, f32)>, Vec<(i64, f32)>); let mut raw_per_kb: Vec = Vec::new(); @@ -285,7 +308,7 @@ async fn recall_and_merge( .dense .retain(|(unit_id, _)| allowed.contains(unit_id)); } - candidates.extend(fusion::fuse_and_cut(recalled, deps.config.kb.top_k)); + candidates.extend(fusion::fuse_and_cut(recalled, s4_cut)); } // Bounded title backfill for the s2 top-10 lists — ids cut by fusion // are exactly the diagnostically interesting ones, so resolve them @@ -372,13 +395,18 @@ async fn recall_and_merge( trace.stage("s4_topk"); trace.end_stage( crate::kb::trace::STAGE_OK, - serde_json::json!({ "kept": hydrated.len() }), + serde_json::json!({ "kept": hydrated.len(), "cut_per_kb": s4_cut }), None, ); - // S5 rerank: v1 passthrough (external rerank provider unimplemented; - // the stage boundary is kept — see kb-technical-design §6 note). - // S6 wiki boost. + // S5 rerank (kb-technical-design §6.1, revised 2026-09-17): per-KB + // resolved config; the rerank score REPLACES the fused RRF score; + // skipped/degraded inside — never fails the query. + trace.stage("s5_rerank"); + rerank::rerank_stage(deps, tenant, rerank, &understood.text, &mut hydrated, trace).await; + + // S6 wiki boost — multiplies `Candidate.score`, i.e. the rerank score + // once S5 ran (乘法提权语义不变, §6.1.4). let boosted: Vec = hydrated .iter() .filter(|c| c.chunk.kind == "wiki_page") @@ -693,10 +721,23 @@ mod tests { } async fn deps() -> KbDeps { + deps_with(|_| {}, None).await + } + + async fn deps_with( + tune: impl FnOnce(&mut crate::config::app::AppConfig), + reranker: Option>, + ) -> KbDeps { let pool = crate::test_pool!(); let mut config = crate::config::app::AppConfig::test_defaults(); config.kb.enabled = true; config.kb.fallback_threshold = 0.05; // deterministic for mocks + if reranker.is_some() { + // Global env default so S5 resolves without a per-KB override; + // tests may override via `tune` or the KB row (§6.1.2). + config.kb.rerank_model = Some("env-default-model".into()); + } + tune(&mut config); let bus = crate::eventbus::EventBus::new(16); let router = crate::llm::service::LlmRouter::with_provider_for_test( Some(pool.clone()), @@ -715,11 +756,118 @@ mod tests { vector: Arc::new(BruteForceIndex::new()), kbsearch: Arc::new(crate::kb::kbsearch::KbSearchEngine::open_in_memory().unwrap()), embedder: Arc::new(MockEmbedder { dim: 4 }), + reranker, router, emitter: crate::event::EventEmitter::eventbus_only(bus), } } + /// Deterministic reranker stub (§6.1.5): scores by keyword membership. + struct KeywordReranker { + keyword: &'static str, + hit: f32, + miss: f32, + } + + #[async_trait::async_trait] + impl crate::kb::rerank::KbReranker for KeywordReranker { + async fn rerank( + &self, + _tenant: &str, + _model: &str, + _query: &str, + docs: &[crate::kb::rerank::RerankDoc], + ) -> AppResult> { + Ok(docs + .iter() + .map(|d| { + if d.text.contains(self.keyword) { + self.hit + } else { + self.miss + } + }) + .collect()) + } + } + + /// Always-failing reranker: exercises the S5 degrade path. + struct FailingReranker; + + #[async_trait::async_trait] + impl crate::kb::rerank::KbReranker for FailingReranker { + async fn rerank( + &self, + _tenant: &str, + _model: &str, + _query: &str, + _docs: &[crate::kb::rerank::RerankDoc], + ) -> AppResult> { + Err(AppError::ServiceUnavailable("rerank down".into())) + } + } + + /// Constant-score reranker that records the resolved model per call + /// (per-KB override assertions, §6.1.2). + struct RecordingReranker { + constant: f32, + models: std::sync::Arc>>, + } + + #[async_trait::async_trait] + impl crate::kb::rerank::KbReranker for RecordingReranker { + async fn rerank( + &self, + _tenant: &str, + model: &str, + _query: &str, + docs: &[crate::kb::rerank::RerankDoc], + ) -> AppResult> { + self.models.lock().unwrap().push(model.to_string()); + Ok(vec![self.constant; docs.len()]) + } + } + + /// KB + one processed document under a distinct slug (rerank tests). + #[allow(clippy::too_many_arguments)] + async fn seeded_rr_kb( + deps: &KbDeps, + slug: &str, + rerank_model: Option<&str>, + rerank_window: Option, + rerank_threshold: Option, + ) -> i64 { + let kb = crate::kb::models::knowledge_base::create_kb( + &deps.pool, + &crate::kb::models::knowledge_base::CreateKbCmd { + name: slug.into(), + description: None, + slug: slug.into(), + kind: "document".into(), + indexing_strategy: None, + embedding_model: Some("m".into()), + embedding_dim: Some(4), + rerank_model: rerank_model.map(str::to_owned), + rerank_window, + rerank_threshold, + }, + "default", + ) + .await + .unwrap(); + let mut markdown = "# 重排\n\n".to_string(); + markdown.push_str(&"raisfast 知识库重排序阶段说明文本。".repeat(60)); + let doc = crate::kb::service::create_online_document( + deps, kb.id, "重排", &markdown, None, "default", + ) + .await + .unwrap(); + crate::kb::service::process_document(deps, doc.id, "default") + .await + .unwrap(); + i64::from(kb.id) + } + #[tokio::test] async fn ask_answered_with_references() { let deps = deps().await; @@ -733,6 +881,9 @@ mod tests { indexing_strategy: None, embedding_model: Some("m".into()), embedding_dim: Some(4), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "default", ) @@ -787,6 +938,9 @@ mod tests { indexing_strategy: None, embedding_model: Some("m".into()), embedding_dim: Some(4), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "default", ) @@ -819,6 +973,9 @@ mod tests { indexing_strategy: None, embedding_model: Some("m".into()), embedding_dim: Some(4), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "default", ) @@ -868,6 +1025,9 @@ mod tests { indexing_strategy: None, embedding_model: Some("m".into()), embedding_dim: Some(4), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, tenant, ) @@ -939,6 +1099,9 @@ mod tests { description: None, slug: "iso-default".into(), status: "disabled".into(), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "default", ) @@ -956,4 +1119,144 @@ mod tests { "inactive kb must fail closed" ); } + + #[tokio::test] + async fn rerank_replaces_scores_and_cuts_to_top_k() { + let deps = deps_with( + |c| { + c.kb.top_k = 2; + c.kb.rerank_model = Some("mock-reranker".into()); + }, + Some(Arc::new(KeywordReranker { + keyword: "PostgreSQL", + hit: 0.9, + miss: 0.1, + })), + ) + .await; + let kb_id = seeded_rr_kb(&deps, "rr-mix", None, None, None).await; + let mut markdown = "# PostgreSQL 部署\n\n".to_string(); + markdown.push_str(&"PostgreSQL 生产部署要点与连接池调优参数。".repeat(60)); + markdown.push_str("\n\n# MySQL 备份\n\n"); + markdown.push_str(&"MySQL 备份策略与权限分配说明。".repeat(60)); + let doc = crate::kb::service::create_online_document( + &deps, + crate::types::snowflake_id::SnowflakeId(kb_id), + "混合", + &markdown, + None, + "default", + ) + .await + .unwrap(); + crate::kb::service::process_document(&deps, doc.id, "default") + .await + .unwrap(); + + let ask = AskRequest { + tenant_id: "default".into(), + kb_ids: vec![kb_id], + doc_ids: Vec::new(), + question: "PostgreSQL 部署要点是什么".into(), + }; + let outcome = prepare_answer(&deps, &ask).await.unwrap(); + assert!( + (outcome.top_score - 0.9).abs() < 1e-6, + "rerank score must replace the fused score, got {}", + outcome.top_score + ); + assert!( + outcome.context_units[0].content.contains("PostgreSQL"), + "top unit after rerank must be the keyword hit" + ); + assert!( + outcome.context_units.len() <= 2, + "S5 must cut back to top_k after rerank" + ); + } + + #[tokio::test] + async fn rerank_threshold_drops_all_and_reports_uncovered() { + let deps = deps_with( + |c| { + c.kb.rerank_threshold = 0.5; + }, + Some(Arc::new(KeywordReranker { + keyword: "\u{0}never-matches", + hit: 0.2, + miss: 0.2, + })), + ) + .await; + let kb_id = seeded_rr_kb(&deps, "rr-floor", None, None, None).await; + let ask = AskRequest { + tenant_id: "default".into(), + kb_ids: vec![kb_id], + doc_ids: Vec::new(), + question: "重排序阶段做什么".into(), + }; + let mut outcome = prepare_answer(&deps, &ask).await.unwrap(); + assert!( + outcome.context_units.is_empty(), + "below-floor candidates must be dropped (no refill)" + ); + assert_eq!(outcome.top_score, 0.0); + finish_answer(&deps, &mut outcome).await.unwrap(); + assert_eq!(outcome.status, "uncovered"); + } + + #[tokio::test] + async fn rerank_failure_degrades_to_rrf_order() { + let deps = deps_with(|_| {}, Some(Arc::new(FailingReranker))).await; + let kb_id = seeded_rr_kb(&deps, "rr-fail", None, None, None).await; + let ask = AskRequest { + tenant_id: "default".into(), + kb_ids: vec![kb_id], + doc_ids: Vec::new(), + question: "重排序阶段做什么".into(), + }; + let mut outcome = prepare_answer(&deps, &ask).await.unwrap(); + assert!( + !outcome.context_units.is_empty(), + "degraded rerank must keep the RRF results" + ); + assert!( + outcome.top_score > 0.0, + "RRF score must survive the fallback" + ); + finish_answer(&deps, &mut outcome).await.unwrap(); + assert_eq!(outcome.status, "answered"); + } + + #[tokio::test] + async fn per_kb_rerank_config_overrides_env_default() { + let models = Arc::new(std::sync::Mutex::new(Vec::new())); + let deps = deps_with( + |_| {}, + Some(Arc::new(RecordingReranker { + constant: 0.2, + models: models.clone(), + })), + ) + .await; + // deps_with seeds the env default model "env-default-model" and + // threshold 0; the KB row overrides both (model + threshold). + let kb_id = seeded_rr_kb(&deps, "rr-override", Some("kb-model"), None, Some(0.5)).await; + let ask = AskRequest { + tenant_id: "default".into(), + kb_ids: vec![kb_id], + doc_ids: Vec::new(), + question: "重排序阶段做什么".into(), + }; + let outcome = prepare_answer(&deps, &ask).await.unwrap(); + assert_eq!( + *models.lock().unwrap(), + vec!["kb-model".to_string()], + "per-KB rerank_model must win over the env default" + ); + assert!( + outcome.context_units.is_empty(), + "per-KB rerank_threshold (0.5) must drop constant-0.2 scores" + ); + } } diff --git a/crates/core/src/kb/pipeline/rerank.rs b/crates/core/src/kb/pipeline/rerank.rs new file mode 100644 index 00000000..e4684724 --- /dev/null +++ b/crates/core/src/kb/pipeline/rerank.rs @@ -0,0 +1,227 @@ +//! S5 rerank — optional relevance rerank over the llm 底座 +//! [抄WK:chat_pipeline/rerank.go 管道位;修订 2026-09-17:模型经 llm 底座 +//! 路由(§10.2),KB 侧只选模型名,见 kb-technical-design §6.1]. +//! +//! Semantics (§6.1.4): the rerank score REPLACES the fused RRF score for +//! ordering, threshold filtering and the S11 fallback verdict; wiki boost +//! (S6) then multiplies the rerank score. Failures degrade to the RRF +//! order — rerank is a gain, never a dependency [抄WK:rerank.go 容错: +//! 错误记日志不中断管道]. + +use serde_json::json; + +use crate::db::DbDriver; +use crate::errors::app_error::{AppError, AppResult}; +use crate::kb::rerank::RerankDoc; +use crate::kb::service::KbDeps; +use crate::kb::trace::RunRecorder; + +use super::Candidate; + +/// Per-query S5 rerank config, resolved from the KB scope (§6.1.2): +/// the first KB with a `rerank_model` provides the whole trio (single +/// model per merged pool, WK RetrievalConfig precedent); null fields on +/// that row fall back to the global env defaults. +#[derive(Debug, Clone)] +pub(super) struct ResolvedRerank { + pub model: String, + pub window: u32, + pub threshold: f32, +} + +/// Raw rerank columns off a KB row (resolution query shape). +#[derive(sqlx::FromRow)] +struct KbRerankRow { + id: i64, + #[sqlx(rename = "rerank_model")] + model: Option, + #[sqlx(rename = "rerank_window")] + window: Option, + #[sqlx(rename = "rerank_threshold")] + threshold: Option, +} + +/// Resolve S5 rerank for a KB scope: first KB row (in scope order) with a +/// non-empty `rerank_model` wins; without any per-KB override the global +/// `RAISFAST_KB_RERANK_MODEL` default applies; neither → `None` +/// (passthrough). Rerank trio is freely editable per KB — query-time +/// behavior, nothing baked (unlike the pinned embedding config). +pub(super) async fn resolve(deps: &KbDeps, kbs: &[i64]) -> AppResult> { + let global = || { + deps.config + .kb + .rerank_model + .as_deref() + .filter(|m| !m.is_empty()) + }; + if deps.reranker.is_none() { + return Ok(None); + } + let global_default = || ResolvedRerank { + model: global().unwrap_or_default().to_string(), + window: deps.config.kb.rerank_window, + threshold: deps.config.kb.rerank_threshold, + }; + if kbs.is_empty() { + return Ok(global().map(|_| global_default())); + } + let placeholders: Vec = (1..=kbs.len()).map(crate::db::Driver::ph).collect(); + let sql = format!( + "SELECT id, rerank_model, rerank_window, rerank_threshold \ + FROM kb_knowledge_bases WHERE id IN ({})", + placeholders.join(", ") + ); + let mut query = sqlx::query_as::<_, KbRerankRow>(crate::db::safe_sql(&sql)); + for id in kbs { + query = query.bind(id); + } + let rows = query + .fetch_all(&deps.pool) + .await + .map_err(|e| AppError::Internal(anyhow::anyhow!(e.to_string())))?; + let by_id: std::collections::HashMap = + rows.into_iter().map(|r| (r.id, r)).collect(); + for kb_id in kbs { + let Some(row) = by_id.get(kb_id) else { + continue; + }; + if row.model.as_deref().is_some_and(|m| !m.is_empty()) { + return Ok(Some(ResolvedRerank { + model: row.model.clone().unwrap_or_default(), + window: row + .window + .map(|w| w.max(1) as u32) + .unwrap_or(deps.config.kb.rerank_window), + threshold: row + .threshold + .map(|t| t as f32) + .unwrap_or(deps.config.kb.rerank_threshold), + })); + } + } + Ok(global().map(|_| global_default())) +} + +/// 送排文本 = `breadcrumb + "\n" + content`,与 S2 双路召回所索引的 FTS +/// 文本同构(重排器看到的与检索器看到的是同一份文本)[§6.1.1]. +fn rerank_text(chunk: &crate::kb::models::chunk::KbChunk) -> String { + match &chunk.breadcrumb { + Some(b) if !b.is_empty() => format!("{b}\n{}", chunk.content), + _ => chunk.content.clone(), + } +} + +/// Bounded top-3 list for the s5 summary (§6.1.4 最小埋点:S5 前后 top-3 +/// 变化即 rerank 开关效果). +fn top3(candidates: &[Candidate]) -> Vec { + candidates + .iter() + .take(3) + .map(|c| { + json!({ + "id": c.unit_id, + "score": (c.score * 10000.0).round() / 10000.0, + "title": super::chunk_title(&c.chunk), + }) + }) + .collect() +} + +/// Degrade without failing the query: keep the incoming (RRF) order and +/// mark the stage degraded (run-level degraded flag via the recorder). +fn degrade(trace: &mut RunRecorder, reason: String, before_top3: Vec) { + tracing::warn!("[kb] rerank degraded to RRF order: {reason}"); + trace.end_stage( + crate::kb::trace::STAGE_DEGRADED, + json!({ "fallback": "rrf", "error": reason, "before_top3": before_top3 }), + None, + ); +} + +/// Stage body: rerank `hydrated` in place when a rerank config resolves, +/// then threshold-filter and cut back to `top_k`. Records +/// skipped/ok/degraded into `trace`; never fails the query. §6.1.4 阈值 +/// 过滤后不足 top_k 不补召回(宁可少上下文,不引入重排器已判低分的噪声). +pub(super) async fn rerank_stage( + deps: &KbDeps, + tenant: &str, + rerank: Option<&ResolvedRerank>, + query: &str, + hydrated: &mut Vec, + trace: &mut RunRecorder, +) { + let Some(reranker) = deps.reranker.clone() else { + trace.end_stage( + crate::kb::trace::STAGE_SKIPPED, + json!({ "reason": "no_reranker_wired" }), + None, + ); + return; + }; + let Some(cfg) = rerank else { + trace.end_stage( + crate::kb::trace::STAGE_SKIPPED, + json!({ "reason": "no_rerank_model" }), + None, + ); + return; + }; + if hydrated.is_empty() { + // 抄WK:rerank.go empty_search_result skip. + trace.end_stage( + crate::kb::trace::STAGE_SKIPPED, + json!({ "reason": "empty_candidates", "model": cfg.model }), + None, + ); + return; + } + let before_top3 = top3(hydrated); + let docs: Vec = hydrated + .iter() + .map(|c| RerankDoc { + unit_id: c.unit_id, + text: rerank_text(&c.chunk), + }) + .collect(); + match reranker.rerank(tenant, &cfg.model, query, &docs).await { + Ok(scores) if scores.len() == hydrated.len() => { + // 分数替换:重排分替换融合分参与排序与 S11 阈值判断. + for (c, s) in hydrated.iter_mut().zip(scores) { + c.score = s; + } + let before_filter = hydrated.len(); + hydrated.retain(|c| c.score >= cfg.threshold); + let dropped_by_threshold = before_filter - hydrated.len(); + hydrated.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + hydrated.truncate(deps.config.kb.top_k as usize); + trace.end_stage( + crate::kb::trace::STAGE_OK, + json!({ + "model": cfg.model, + "threshold": cfg.threshold, + "kept": hydrated.len(), + "dropped_by_threshold": dropped_by_threshold, + "before_top3": before_top3, + "after_top3": top3(hydrated), + }), + None, + ); + } + Ok(scores) => { + degrade( + trace, + format!( + "rerank score count mismatch: {} docs, {} scores", + hydrated.len(), + scores.len() + ), + before_top3, + ); + } + Err(e) => degrade(trace, e.to_string(), before_top3), + } +} diff --git a/crates/core/src/kb/rerank.rs b/crates/core/src/kb/rerank.rs new file mode 100644 index 00000000..91546c92 --- /dev/null +++ b/crates/core/src/kb/rerank.rs @@ -0,0 +1,186 @@ +//! S5 rerank seam — optional relevance reranking over the llm 底座 +//! (kb-technical-design §6.1, revised 2026-09-17). +//! +//! Routing/auth/failover/billing all live in the LLM gateway (`/v1/rerank` +//! relay + `LlmModelType::Rerank` channels, §10.2); the KB side only picks +//! the model name via `RAISFAST_KB_RERANK_MODEL`. Empty model = no reranker +//! wired → S5 passthrough (v1 behavior). + +use crate::errors::app_error::{AppError, AppResult}; + +/// One rerank input unit (§6.1.1). `unit_id` is the `kb_chunks.id` snowflake; +/// `text` is the same breadcrumb+content payload the S2 FTS index sees. +#[derive(Debug, Clone)] +pub struct RerankDoc { + pub unit_id: i64, + pub text: String, +} + +/// Rerank seam: production wraps the llm 底座 facade (§10.2 唯一入口, +/// mirrors [`crate::kb::service::KbEmbedder::embed_for`]'s per-KB model +/// param); tests inject stubs. +#[async_trait::async_trait] +pub trait KbReranker: Send + Sync { + /// Relevance scores (0..1, higher = more relevant), equal length and + /// same order as `docs` — call sites never remap indexes. `model` is + /// the per-KB resolved rerank model (KB row → global default). + async fn rerank( + &self, + tenant: &str, + model: &str, + query: &str, + docs: &[RerankDoc], + ) -> AppResult>; +} + +// Docs per `/rerank` request — §6.1.4 分批:候选超批大小时切片串行. +// Global (transport concern), mirrors `embed_batch_size`. + +/// Retry shape [抄RF:ProviderEmbedder::embed_batched(同源抄WK +/// batchEmbedWithBackoff)]: 5 attempts, backoff 200ms → 3.2s. +const RERANK_RETRY_ATTEMPTS: usize = 5; +const RERANK_RETRY_BASE_DELAY_MS: u64 = 200; + +/// Production reranker over the llm 底座 facade — the ONLY entry +/// (model resolution/routing/failover/billing in the gateway, §10.2). +pub struct ProviderReranker { + router: std::sync::Arc, + /// Docs per `/rerank` request (global, like `embed_batch_size`). + batch_size: usize, +} + +impl ProviderReranker { + pub fn new(router: std::sync::Arc, batch_size: usize) -> Self { + Self { router, batch_size } + } + + /// Slice docs into fixed-size batches and rerank them sequentially, + /// retrying each batch with exponential backoff (§6.1.4 分批). + async fn rerank_batched( + &self, + tenant: &str, + model: &str, + query: &str, + docs: &[RerankDoc], + ) -> AppResult> { + let mut out = Vec::with_capacity(docs.len()); + for batch in docs.chunks(self.batch_size.max(1)) { + let mut delay = RERANK_RETRY_BASE_DELAY_MS; + let mut last_err: Option = None; + let mut scores: Option> = None; + for attempt in 0..RERANK_RETRY_ATTEMPTS { + match self.rerank_once(tenant, model, query, batch).await { + Ok(v) if v.len() == batch.len() => { + scores = Some(v); + break; + } + Ok(v) => { + last_err = Some(format!( + "rerank score count mismatch: batch {} got {}", + batch.len(), + v.len() + )); + } + Err(e) => last_err = Some(e), + } + tracing::warn!( + "rerank batch ({}/{} docs) attempt {}/{} failed: {}", + batch.len(), + docs.len(), + attempt + 1, + RERANK_RETRY_ATTEMPTS, + last_err.as_deref().unwrap_or_default() + ); + if attempt + 1 < RERANK_RETRY_ATTEMPTS { + tokio::time::sleep(std::time::Duration::from_millis(delay)).await; + delay *= 2; + } + } + match scores { + Some(v) => out.extend(v), + None => { + return Err(AppError::ServiceUnavailable(format!( + "rerank: {}", + last_err.unwrap_or_else(|| "unknown error".into()) + ))); + } + } + } + Ok(out) + } + + /// 单批重排:唯一入口 llm 底座(路由/failover/计费/日志,§10.2)。 + async fn rerank_once( + &self, + tenant: &str, + model: &str, + query: &str, + batch: &[RerankDoc], + ) -> Result, String> { + let texts: Vec<&str> = batch.iter().map(|d| d.text.as_str()).collect(); + let results = self + .router + .call(tenant, crate::llm::models::log::LogSource::Kb) + .rerank(model, query, &texts) + .await + .map_err(|e| e.to_string())?; + Ok(backfill_scores(batch.len(), &results)) + } +} + +#[async_trait::async_trait] +impl KbReranker for ProviderReranker { + async fn rerank( + &self, + tenant: &str, + model: &str, + query: &str, + docs: &[RerankDoc], + ) -> AppResult> { + self.rerank_batched(tenant, model, query, docs).await + } +} + +/// Backfill best-first `RerankResult{index, score}` rows into an +/// equal-length, same-order score vector — docs the model did not score +/// default to 0.0 (unranked = irrelevant; §6.1.1 等长同序契约). +fn backfill_scores(len: usize, results: &[raisfast_agent::provider::RerankResult]) -> Vec { + let mut scores = vec![0.0_f32; len]; + for r in results { + if r.index < len { + scores[r.index] = r.relevance_score; + } + } + scores +} + +#[cfg(test)] +mod tests { + use super::*; + + fn result(index: usize, score: f32) -> raisfast_agent::provider::RerankResult { + raisfast_agent::provider::RerankResult { + index, + relevance_score: score, + } + } + + #[test] + fn backfill_maps_by_index_best_first() { + // Best-first rows land on their input positions, not their rank. + let scores = backfill_scores(3, &[result(2, 0.9), result(0, 0.5), result(1, 0.7)]); + assert_eq!(scores, vec![0.5, 0.7, 0.9]); + } + + #[test] + fn backfill_defaults_unscored_to_zero() { + let scores = backfill_scores(3, &[result(0, 0.8)]); + assert_eq!(scores, vec![0.8, 0.0, 0.0]); + } + + #[test] + fn backfill_ignores_out_of_range_index() { + let scores = backfill_scores(2, &[result(0, 0.4), result(9, 1.0)]); + assert_eq!(scores, vec![0.4, 0.0]); + } +} diff --git a/crates/core/src/kb/service.rs b/crates/core/src/kb/service.rs index cef0e717..29bf9463 100644 --- a/crates/core/src/kb/service.rs +++ b/crates/core/src/kb/service.rs @@ -210,6 +210,8 @@ pub struct KbDeps { pub vector: Arc, pub kbsearch: Arc, pub embedder: Arc, + /// S5 reranker; `None` = rerank off (§6.1). + pub reranker: Option>, /// LLM 底座(chat + embedding 唯一入口,§10.2)。 pub router: Arc, pub emitter: EventEmitter, @@ -833,6 +835,7 @@ mod tests { vector: Arc::new(BruteForceIndex::new()), kbsearch: Arc::new(KbSearchEngine::open_in_memory().unwrap()), embedder: Arc::new(MockEmbedder { dim: 4 }), + reranker: None, router, emitter: EventEmitter::eventbus_only(bus), } @@ -849,6 +852,9 @@ mod tests { indexing_strategy: None, embedding_model: Some("test-model".into()), embedding_dim: Some(4), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "default", ) @@ -1304,6 +1310,7 @@ mod faq_tests { vector: Arc::new(BruteForceIndex::new()), kbsearch: Arc::new(crate::kb::kbsearch::KbSearchEngine::open_in_memory().unwrap()), embedder: Arc::new(OneHotEmbedder), + reranker: None, router, emitter: crate::event::EventEmitter::eventbus_only(crate::eventbus::EventBus::new(16)), } @@ -1320,6 +1327,9 @@ mod faq_tests { indexing_strategy: None, embedding_model: Some("m".into()), embedding_dim: Some(4), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "default", ) diff --git a/crates/core/src/kb/trace.rs b/crates/core/src/kb/trace.rs index 1a64fdcc..339efc03 100644 --- a/crates/core/src/kb/trace.rs +++ b/crates/core/src/kb/trace.rs @@ -320,6 +320,9 @@ pub fn config_snapshot(config: &AppConfig) -> Value { "context_budget_tokens": config.kb.context_budget_tokens, "vector_backend": config.kb.vector_backend, "trace_mode": config.kb.trace_mode, + "rerank_model": config.kb.rerank_model, + "rerank_window": config.kb.rerank_window, + "rerank_threshold": config.kb.rerank_threshold, }) } diff --git a/crates/core/src/worker/handlers/kb.rs b/crates/core/src/worker/handlers/kb.rs index 1eecee94..40a85cbd 100644 --- a/crates/core/src/worker/handlers/kb.rs +++ b/crates/core/src/worker/handlers/kb.rs @@ -93,6 +93,7 @@ impl KbProcessDocumentHandler { vector: self.runtime.vector.clone(), kbsearch: self.runtime.kbsearch.clone(), embedder: self.runtime.embedder.clone(), + reranker: self.runtime.reranker.clone(), router: self.llm_router.clone(), emitter: self.emitter.clone(), }; @@ -226,6 +227,7 @@ impl JobHandler for KbDistillWikiHandler { vector: self.runtime.vector.clone(), kbsearch: self.runtime.kbsearch.clone(), embedder: self.runtime.embedder.clone(), + reranker: self.runtime.reranker.clone(), router: self.llm_router.clone(), emitter: self.emitter.clone(), }; diff --git a/crates/core/tests/kb.rs b/crates/core/tests/kb.rs index 4faaab78..deb9d0d0 100644 --- a/crates/core/tests/kb.rs +++ b/crates/core/tests/kb.rs @@ -109,6 +109,7 @@ async fn deps_with_provider(provider: Arc) -> KbDeps { vector: Arc::new(BruteForceIndex::new()), kbsearch: Arc::new(raisfast::kb::kbsearch::KbSearchEngine::open_in_memory().unwrap()), embedder: Arc::new(SumEmbedder(4)), + reranker: None, router, emitter: raisfast::event::EventEmitter::eventbus_only(raisfast::eventbus::EventBus::new( 16, @@ -127,6 +128,9 @@ async fn seed_kb(deps: &KbDeps, name: &str) -> SnowflakeId { indexing_strategy: None, embedding_model: Some("test-model".into()), embedding_dim: Some(4), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "default", ) @@ -918,6 +922,9 @@ async fn s18_update_kb_metadata_and_tenant_guard() { description: Some("新描述".into()), slug: format!("slug-{}", raisfast::utils::id::new_id()), status: "archived".into(), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "default", ) @@ -942,6 +949,9 @@ async fn s18_update_kb_metadata_and_tenant_guard() { description: None, slug: "hijack".into(), status: "active".into(), + rerank_model: None, + rerank_window: None, + rerank_threshold: None, }, "other-tenant", ) @@ -1131,6 +1141,7 @@ async fn s21_chunk_edit_failure_records_failed_run() { vector: good.vector.clone(), kbsearch: good.kbsearch.clone(), embedder: Arc::new(FailingEmbedder), + reranker: None, router: good.router.clone(), emitter: good.emitter.clone(), }; @@ -1184,6 +1195,7 @@ async fn s22_cold_bruteforce_rebuilds_on_first_search() { vector: Arc::new(BruteForceIndex::new()), kbsearch: deps.kbsearch.clone(), embedder: deps.embedder.clone(), + reranker: None, router: deps.router.clone(), emitter: deps.emitter.clone(), }; diff --git a/migrations/mysql/schema.mysql.sql b/migrations/mysql/schema.mysql.sql index 83205703..f1ac43f5 100644 --- a/migrations/mysql/schema.mysql.sql +++ b/migrations/mysql/schema.mysql.sql @@ -1564,6 +1564,9 @@ CREATE TABLE IF NOT EXISTS kb_knowledge_bases ( chunking_config JSON, embedding_model VARCHAR(100), embedding_dim BIGINT, + rerank_model VARCHAR(100), + rerank_window BIGINT, + rerank_threshold DOUBLE, status VARCHAR(20) NOT NULL DEFAULT 'active', created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, diff --git a/migrations/postgres/schema.postgres.sql b/migrations/postgres/schema.postgres.sql index 0f3dadfa..b39d9de8 100644 --- a/migrations/postgres/schema.postgres.sql +++ b/migrations/postgres/schema.postgres.sql @@ -1632,6 +1632,9 @@ CREATE TABLE IF NOT EXISTS kb_knowledge_bases ( chunking_config JSONB, embedding_model TEXT, embedding_dim BIGINT, + rerank_model TEXT, + rerank_window BIGINT, + rerank_threshold DOUBLE PRECISION, status TEXT NOT NULL DEFAULT 'active', created_at TIMESTAMPTZ(0) NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ(0) NOT NULL DEFAULT NOW() diff --git a/migrations/sqlite/schema.sqlite.sql b/migrations/sqlite/schema.sqlite.sql index 0b89163a..904598f0 100644 --- a/migrations/sqlite/schema.sqlite.sql +++ b/migrations/sqlite/schema.sqlite.sql @@ -1625,6 +1625,9 @@ CREATE TABLE IF NOT EXISTS kb_knowledge_bases ( chunking_config TEXT, embedding_model TEXT, embedding_dim BIGINT, + rerank_model TEXT, + rerank_window BIGINT, + rerank_threshold REAL, status TEXT NOT NULL DEFAULT 'active', created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%SZ','now')), updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%SZ','now'))