From 3f461fc6057f17a4b5ea67ea1ff94f700dcf18cf Mon Sep 17 00:00:00 2001 From: Bogdan Szabo Date: Sun, 25 Oct 2026 09:15:26 +0100 Subject: [PATCH] A search may name a source and an external_id prefix, and both channels recall only those documents before fusion Signed-off-by: Bogdan Szabo --- crates/utopia-core/src/models.rs | 9 +- crates/utopia-search/src/lib.rs | 21 ++- .../utopia-server/src/api/documents_routes.rs | 4 +- crates/utopia-server/src/api/search_routes.rs | 41 +++- .../utopia-server/src/api/sources_routes.rs | 14 +- .../utopia-server/src/api/token_read_tests.rs | 175 +++++++++++++++++- crates/utopia-server/src/main.rs | 1 + crates/utopia-server/src/pipeline_tests.rs | 5 +- crates/utopia-server/src/pushed_keys.rs | 23 +++ crates/utopia-server/src/retrieval.rs | 84 ++++++--- crates/utopia-server/src/retrieval_tests.rs | 83 ++++++++- crates/utopia-server/src/search_recovery.rs | 23 ++- crates/utopia-store/src/documents.rs | 39 +++- .../a_search_reads_the_base_as_it_was.rs | 6 +- ...st_chunk_is_found_however_it_is_reached.rs | 18 +- ...ty-from-the-person-scope-from-the-token.md | 8 +- docs/design/chat-and-mcp.md | 5 +- web/src/docs/ingest.md | 2 + 18 files changed, 497 insertions(+), 64 deletions(-) create mode 100644 crates/utopia-server/src/pushed_keys.rs diff --git a/crates/utopia-core/src/models.rs b/crates/utopia-core/src/models.rs index 178a9fffc..14e8164eb 100644 --- a/crates/utopia-core/src/models.rs +++ b/crates/utopia-core/src/models.rs @@ -479,15 +479,18 @@ pub struct ChunkView { pub filename: String, } -/// 分块带上它来自哪个来源。`external_id` 只有推送来源（api / statements）才有： -/// 推送方给的那个 id，去掉库里加的 `{kind}:` 前缀，好让它拿回去对自己的记录（#1271） +/// 分块带上它来自哪个来源，连同来源的 kind 和文档的 `external_key`：推送来源的 +/// `external_id` 由这两样反推（`sources_routes::pushed_identity`，#1271） #[derive(Debug, Clone, Serialize, sqlx::FromRow)] pub struct SourcedChunk { #[serde(flatten)] #[sqlx(flatten)] pub chunk: ChunkView, pub source_id: Option, - pub external_id: Option, + #[serde(skip)] + pub source_kind: Option, + #[serde(skip)] + pub external_key: Option, } #[derive(Debug, Clone, Serialize, sqlx::FromRow)] diff --git a/crates/utopia-search/src/lib.rs b/crates/utopia-search/src/lib.rs index f518825bd..54ca0b6f2 100644 --- a/crates/utopia-search/src/lib.rs +++ b/crates/utopia-search/src/lib.rs @@ -10,7 +10,7 @@ use std::sync::Mutex; use anyhow::Context; use tantivy::collector::TopDocs; -use tantivy::query::{BooleanQuery, Occur, Query, TermQuery}; +use tantivy::query::{BooleanQuery, Occur, Query, TermQuery, TermSetQuery}; use tantivy::schema::{ Field, IndexRecordOption, Schema, TextFieldIndexing, TextOptions, Value, STORED, STRING, }; @@ -134,7 +134,15 @@ impl SearchIndex { /// BM25 检索（限定 kb）。 /// 查询用与索引完全一致的 jieba analyzer 切词后按 OR 组合—— /// 不能走 QueryParser：CJK 整句会被当成短语查询（要求词连续出现），召回归零。 - pub fn search(&self, kb_id: &str, query: &str, limit: usize) -> anyhow::Result> { + /// + /// `documents`：只在这些文档里找。它和 kb 一样是 `Must`，`limit` 数的就只是允许的块 + pub fn search( + &self, + kb_id: &str, + query: &str, + limit: usize, + documents: Option<&[String]>, + ) -> anyhow::Result> { let mut analyzer = self.analyzer.clone(); let mut stream = analyzer.token_stream(query); let mut term_queries: Vec<(Occur, Box)> = Vec::new(); @@ -161,7 +169,14 @@ impl SearchIndex { IndexRecordOption::Basic, )); let text_query: Box = Box::new(BooleanQuery::new(term_queries)); - let combined = BooleanQuery::new(vec![(Occur::Must, kb_filter), (Occur::Must, text_query)]); + let mut clauses = vec![(Occur::Must, kb_filter), (Occur::Must, text_query)]; + if let Some(documents) = documents { + let terms = documents + .iter() + .map(|id| Term::from_field_text(self.f_document_id, id)); + clauses.push((Occur::Must, Box::new(TermSetQuery::new(terms)))); + } + let combined = BooleanQuery::new(clauses); let searcher = self.reader.searcher(); let collector = TopDocs::with_limit(limit).order_by_score(); diff --git a/crates/utopia-server/src/api/documents_routes.rs b/crates/utopia-server/src/api/documents_routes.rs index 62746b758..f740cde67 100644 --- a/crates/utopia-server/src/api/documents_routes.rs +++ b/crates/utopia-server/src/api/documents_routes.rs @@ -227,11 +227,11 @@ pub async fn by_external_id( ) -> ApiResult { require_reader_kb(&state, &reader, kb_id).await?; let source = super::sources_routes::source_in_kb(&state, kb_id, source_id).await?; - if !super::sources_routes::is_push_kind(&source.kind) { + if !crate::pushed_keys::is_push_kind(&source.kind) { return Err(AppError::NotFound.into()); } let external_id = query.external_id.trim(); - let key = super::sources_routes::pushed_key(&source, external_id); + let key = crate::pushed_keys::pushed_key(&source, external_id); let found = utopia_store::documents::find_by_external_key(&state.pool, source.id, &key) .await? .ok_or(AppError::NotFound)?; diff --git a/crates/utopia-server/src/api/search_routes.rs b/crates/utopia-server/src/api/search_routes.rs index 10ae4066d..daba73234 100644 --- a/crates/utopia-server/src/api/search_routes.rs +++ b/crates/utopia-server/src/api/search_routes.rs @@ -10,7 +10,7 @@ use crate::error::ApiResult; use crate::retrieval; use crate::state::AppState; -/// 不认识的字段答 422：过滤条件还没做，悄悄忽略一个 `filters` 会让调用方以为它生效了 +/// 不认识的字段答 422：悄悄忽略一个拼错的过滤条件，调用方会以为它生效了 #[derive(Deserialize)] #[serde(deny_unknown_fields)] pub struct SearchReq { @@ -23,6 +23,44 @@ pub struct SearchReq { /// 命中的是当时活着的块、当时还没被删的文档 #[serde(default)] pub as_of: Option, + /// 只检索本库这个来源的文档；别的库的来源答 404 + #[serde(default)] + pub source_id: Option, + /// 只检索 `external_id` 以它开头的文档。identity 是推送来源里的，所以要 `source_id` + /// 且它是推送来源，否则 422 + #[serde(default)] + pub external_id_prefix: Option, +} + +async fn source_filter( + state: &AppState, + kb_id: Uuid, + req: &SearchReq, +) -> ApiResult> { + let source = match req.source_id { + Some(source_id) => { + Some(super::sources_routes::source_in_kb(state, kb_id, source_id).await?) + } + None => None, + }; + let Some(prefix) = &req.external_id_prefix else { + return Ok(source.map(|source| retrieval::SourceFilter { + source_id: source.id, + key_prefix: None, + })); + }; + let Some(source) = source.filter(|source| crate::pushed_keys::is_push_kind(&source.kind)) + else { + return Err(AppError::invalid( + "prefix_needs_push_source", + "external_id_prefix needs the source_id of an API or Statements source", + ) + .into()); + }; + Ok(Some(retrieval::SourceFilter { + source_id: source.id, + key_prefix: Some(crate::pushed_keys::pushed_key(&source, prefix)), + })) } pub async fn search( @@ -50,6 +88,7 @@ pub async fn search( top_k, offset: req.offset, as_of, + filter: source_filter(&state, kb_id, &req).await?, }; let page = retrieval::hybrid_scored(&state, kb_id, kb.workspace_id, q, opts, None).await?; Ok(Json(json!({ diff --git a/crates/utopia-server/src/api/sources_routes.rs b/crates/utopia-server/src/api/sources_routes.rs index ddd859dd8..bcb62b728 100644 --- a/crates/utopia-server/src/api/sources_routes.rs +++ b/crates/utopia-server/src/api/sources_routes.rs @@ -12,6 +12,7 @@ use uuid::Uuid; use super::graph_routes::require_kb; use crate::auth::{AuthUser, SessionOrToken}; use crate::error::ApiResult; +use crate::pushed_keys::{is_push_kind, pushed_key}; use crate::state::AppState; /// 生成 api 来源的推送密钥。 @@ -166,19 +167,6 @@ pub async fn create( )) } -/// 有推送密钥的来源：`api` 推文档，`statements` 推陈述（0054）。两种的密钥同一套 -/// 生成、查看、轮换 -pub(crate) fn is_push_kind(kind: &str) -> bool { - matches!(kind, "api" | "statements") -} - -/// The `external_key` a push source files a pushed identity under (`api:…`, `statements:…`). -/// The push handlers write it and the read by `external_id` (#1276) looks it up, so the two -/// cannot drift apart. -pub(crate) fn pushed_key(source: &utopia_core::models::Source, identity: &str) -> String { - format!("{}:{identity}", source.kind) -} - /// 查看推送密钥（Editor；列表响应从不携带，查看走这里）。 /// /// 令牌要 write：密钥能往库里推东西，读令牌拿到它就等于绕过了自己的 scope diff --git a/crates/utopia-server/src/api/token_read_tests.rs b/crates/utopia-server/src/api/token_read_tests.rs index 339e7654f..632405ae9 100644 --- a/crates/utopia-server/src/api/token_read_tests.rs +++ b/crates/utopia-server/src/api/token_read_tests.rs @@ -20,6 +20,7 @@ struct Fixture { entity: Uuid, fact: Uuid, doc: Uuid, + search: Arc, _dir: tempfile::TempDir, } @@ -63,7 +64,7 @@ impl Fixture { let search = Arc::new(utopia_search::SearchIndex::open( &dir.path().join("search"), )?); - let state = AppState::new(pool.clone(), &cfg, search, "test-only".into()); + let state = AppState::new(pool.clone(), &cfg, search.clone(), "test-only".into()); let app = super::router(state, &cfg); Ok(Some(Self { pool, @@ -76,6 +77,7 @@ impl Fixture { entity, fact, doc, + search, _dir: dir, })) } @@ -133,6 +135,22 @@ impl Fixture { Ok(self.app.clone().oneshot(request).await?.status()) } + async fn search( + &self, + token: &str, + body: serde_json::Value, + ) -> anyhow::Result { + let request = Request::post(format!("/api/v1/kbs/{}/search", self.kb)) + .header("authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(Body::from(body.to_string()))?; + let response = self.app.clone().oneshot(request).await?; + let status = response.status().as_u16(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX).await?; + let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap_or_default(); + Ok(json!({ "status": status, "body": body })) + } + async fn cleanup(self) -> anyhow::Result<()> { sqlx::query("DELETE FROM organizations WHERE id = $1") .bind(self.org) @@ -229,3 +247,158 @@ async fn a_search_past_depth_100_or_with_an_unknown_filters_field_gets_422() -> assert_eq!(statuses, [StatusCode::UNPROCESSABLE_ENTITY; 2]); Ok(()) } + +/// 一个 api 来源推了 `identities` 里的每一篇，都讲 harvest；另一个库有个来源 +async fn seed_pushed_and_foreign_sources( + f: &Fixture, + identities: &[&str], +) -> anyhow::Result<(Uuid, Uuid, Uuid)> { + let (pushed, custom, foreign) = (Uuid::now_v7(), Uuid::now_v7(), Uuid::now_v7()); + let (kb, other_kb) = (f.kb, f.other_kb); + sqlx::raw_sql(&format!( + "INSERT INTO sources(id,kb_id,kind,name) VALUES + ('{pushed}','{kb}','api','ci'), ('{custom}','{kb}','custom','poll'), + ('{foreign}','{other_kb}','api','ci');" + )) + .execute(&f.pool) + .await?; + for identity in identities { + let (doc, chunk) = (Uuid::now_v7(), Uuid::now_v7()); + sqlx::query( + "INSERT INTO documents(id,kb_id,source_id,external_key,filename,sha256) + VALUES($1,$2,$3,$4,$5,$5)", + ) + .bind(doc) + .bind(kb) + .bind(pushed) + .bind(format!("api:{identity}")) + .bind(identity) + .execute(&f.pool) + .await?; + sqlx::query( + "INSERT INTO chunks(id,kb_id,document_id,seq,text) VALUES($1,$2,$3,0,'harvest notes')", + ) + .bind(chunk) + .bind(kb) + .bind(doc) + .execute(&f.pool) + .await?; + f.search.reindex_document( + &kb.to_string(), + &doc.to_string(), + &[(chunk.to_string(), "harvest notes".into())], + )?; + } + Ok((pushed, custom, foreign)) +} + +fn found_external_ids(found: &serde_json::Value) -> serde_json::Value { + let mut ids: Vec<_> = found["body"]["results"] + .as_array() + .into_iter() + .flatten() + .map(|hit| hit["external_id"].clone()) + .collect(); + ids.sort_by_key(|id| id.to_string()); + json!({ "status": found["status"], "external_ids": ids }) +} + +#[tokio::test] +async fn a_search_by_source_and_external_id_prefix_docs_finds_only_docs_a_md() -> anyhow::Result<()> +{ + let Some(f) = Fixture::new().await? else { + return Ok(()); + }; + let result = async { + let (pushed, _, _) = + seed_pushed_and_foreign_sources(&f, &["docs/a.md", "code/b.rs"]).await?; + let token = f.read_token(f.kb).await; + let all = f + .search(&token, json!({ "q": "harvest", "source_id": pushed })) + .await?; + let docs = f + .search( + &token, + json!({ "q": "harvest", "source_id": pushed, "external_id_prefix": "docs/" }), + ) + .await?; + anyhow::Ok((all, docs)) + } + .await; + f.cleanup().await?; + let (all, docs) = result?; + assert_eq!( + found_external_ids(&all), + json!({ "status": 200, "external_ids": ["code/b.rs", "docs/a.md"] }) + ); + assert_eq!( + found_external_ids(&docs), + json!({ "status": 200, "external_ids": ["docs/a.md"] }) + ); + Ok(()) +} + +#[tokio::test] +async fn a_search_by_another_bases_source_gets_404_and_a_prefix_without_a_push_source_gets_422( +) -> anyhow::Result<()> { + let Some(f) = Fixture::new().await? else { + return Ok(()); + }; + let result = async { + let (_, custom, foreign) = seed_pushed_and_foreign_sources(&f, &[]).await?; + let token = f.read_token(f.kb).await; + let mut codes = Vec::new(); + for body in [ + json!({ "q": "harvest", "source_id": foreign }), + json!({ "q": "harvest", "external_id_prefix": "docs/" }), + json!({ "q": "harvest", "source_id": custom, "external_id_prefix": "docs/" }), + ] { + let found = f.search(&token, body).await?; + codes.push(json!([found["status"], found["body"]["code"]])); + } + anyhow::Ok(codes) + } + .await; + f.cleanup().await?; + assert_eq!( + result?, + vec![ + json!([404, null]), + json!([422, "prefix_needs_push_source"]), + json!([422, "prefix_needs_push_source"]), + ] + ); + Ok(()) +} + +#[tokio::test] +async fn an_external_id_prefix_docs_underscore_finds_docs_a_md_but_not_docsxa_md( +) -> anyhow::Result<()> { + let Some(f) = Fixture::new().await? else { + return Ok(()); + }; + let result = async { + let (pushed, _, _) = seed_pushed_and_foreign_sources( + &f, + &["docs_a.md", "docsXa.md", "50%off.md", "50 percent off.md"], + ) + .await?; + let token = f.read_token(f.kb).await; + let mut found = Vec::new(); + for prefix in ["docs_", "50%"] { + let body = json!({ "q": "harvest", "source_id": pushed, "external_id_prefix": prefix }); + found.push(found_external_ids(&f.search(&token, body).await?)); + } + anyhow::Ok(found) + } + .await; + f.cleanup().await?; + assert_eq!( + result?, + [ + json!({ "status": 200, "external_ids": ["docs_a.md"] }), + json!({ "status": 200, "external_ids": ["50%off.md"] }), + ] + ); + Ok(()) +} diff --git a/crates/utopia-server/src/main.rs b/crates/utopia-server/src/main.rs index df3141134..3c01813d4 100644 --- a/crates/utopia-server/src/main.rs +++ b/crates/utopia-server/src/main.rs @@ -33,6 +33,7 @@ mod owl_import; mod pack_alignment; mod phrase_alignment; mod pipeline; +mod pushed_keys; mod query_engine; mod rdf; mod readers; diff --git a/crates/utopia-server/src/pipeline_tests.rs b/crates/utopia-server/src/pipeline_tests.rs index d7252dc00..656a6a2d1 100644 --- a/crates/utopia-server/src/pipeline_tests.rs +++ b/crates/utopia-server/src/pipeline_tests.rs @@ -398,7 +398,10 @@ async fn a_large_statements_payload_is_searchable_without_embedding() -> anyhow: payload ); let chunks = utopia_store::documents::chunks_full(&f.pool, doc).await?; - let hits = f.state.search.search(&f.kb.to_string(), "sensor", 1)?; + let hits = f + .state + .search + .search(&f.kb.to_string(), "sensor", 1, None)?; assert_eq!(hits.len(), 1, "the payload remains searchable by keyword"); assert_eq!(hits[0].chunk_id, chunks[0].id.to_string()); f.cleanup().await diff --git a/crates/utopia-server/src/pushed_keys.rs b/crates/utopia-server/src/pushed_keys.rs new file mode 100644 index 000000000..a4c449d4c --- /dev/null +++ b/crates/utopia-server/src/pushed_keys.rs @@ -0,0 +1,23 @@ +//! 推送来源（`api`、`statements`）的文档怎么按推送方给的 identity 存 `external_key`、 +//! 怎么从它取回来。推送的处理器、按 `external_id` 读回和检索都用这一份（#1271） + +/// 有推送密钥的来源：`api` 推文档，`statements` 推陈述（0054）。两种的密钥同一套 +/// 生成、查看、轮换 +pub fn is_push_kind(kind: &str) -> bool { + matches!(kind, "api" | "statements") +} + +/// The `external_key` a push source files a pushed identity under (`api:…`, `statements:…`). +/// The push handlers write it and the read by `external_id` (#1276) looks it up, so the two +/// cannot drift apart. +pub fn pushed_key(source: &utopia_core::models::Source, identity: &str) -> String { + format!("{}:{identity}", source.kind) +} + +/// `pushed_key` 倒过来：推送来源的文档交回推送方给的 identity，别的来源没有（#1271） +pub fn pushed_identity<'k>(kind: &str, key: &'k str) -> Option<&'k str> { + if !is_push_kind(kind) { + return None; + } + key.strip_prefix(kind)?.strip_prefix(':') +} diff --git a/crates/utopia-server/src/retrieval.rs b/crates/utopia-server/src/retrieval.rs index 4a2938107..eb47147fd 100644 --- a/crates/utopia-server/src/retrieval.rs +++ b/crates/utopia-server/src/retrieval.rs @@ -31,6 +31,14 @@ pub struct SearchOpts { pub top_k: usize, pub offset: usize, pub as_of: Option>, + pub filter: Option, +} + +/// 只检索一个来源的文档，给了 `key_prefix` 就只要 `external_key` 以它开头的（#1271）。 +/// 两路各在召回**之内**过滤：融合完再筛，库里别的命中早占满了每路的名额，页就短了 +pub struct SourceFilter { + pub source_id: Uuid, + pub key_prefix: Option, } /// 一条命中连同它的分数（#1271）。`score` 是 RRF 总分，只在同一次查询里能排先后， @@ -39,6 +47,8 @@ pub struct SearchOpts { pub struct ScoredChunk { #[serde(flatten)] pub chunk: SourcedChunk, + /// 推送来源的文档才有：推送方给的那个 id，好让它拿回去对自己的记录 + pub external_id: Option, pub score: f64, pub bm25_rank: Option, pub vector_rank: Option, @@ -66,8 +76,11 @@ pub async fn hybrid( kb_id, workspace_id, query, - RECALL_PER_CHANNEL, - as_of, + Recall { + per_channel: RECALL_PER_CHANNEL, + as_of, + documents: None, + }, embeddings, ) .await?; @@ -89,13 +102,29 @@ pub async fn hybrid_scored( opts: SearchOpts, embeddings: Option<&mut llm_util::EmbedCache>, ) -> AppResult { + let documents = match &opts.filter { + Some(filter) => Some( + utopia_store::documents::ids_in_source( + &state.pool, + kb_id, + filter.source_id, + filter.key_prefix.as_deref(), + opts.as_of, + ) + .await?, + ), + None => None, + }; let (hits, vector_used) = fused_hits( state, kb_id, workspace_id, query, - MAX_DEPTH, - opts.as_of, + Recall { + per_channel: MAX_DEPTH, + as_of: opts.as_of, + documents: documents.as_deref(), + }, embeddings, ) .await?; @@ -113,35 +142,38 @@ pub async fn hybrid_scored( }) } -/// 两路各取 `per_channel` 条、RRF 融合、按记录轴过滤后的整张表，连同向量那一路跑了没有 +/// 每一路怎么召回：取几条、按哪一刻、只在哪些文档里（None 即全库） +#[derive(Clone, Copy)] +struct Recall<'a> { + per_channel: usize, + as_of: Option>, + documents: Option<&'a [Uuid]>, +} + +/// 两路按 `recall` 各取一份、RRF 融合、按记录轴过滤后的整张表，连同向量那一路跑了没有 async fn fused_hits( state: &AppState, kb_id: Uuid, workspace_id: Uuid, query: &str, - per_channel: usize, - as_of: Option>, + recall: Recall<'_>, embeddings: Option<&mut llm_util::EmbedCache>, ) -> AppResult<(Vec, bool)> { let bm25 = { let search = state.search.clone(); let kb = kb_id.to_string(); let query = query.to_string(); - tokio::task::spawn_blocking(move || search.search(&kb, &query, per_channel)) + let documents: Option> = recall + .documents + .map(|ids| ids.iter().map(Uuid::to_string).collect()); + let limit = recall.per_channel; + tokio::task::spawn_blocking(move || search.search(&kb, &query, limit, documents.as_deref())) }; // `join!` 不是 `try_join!`：向量那一路失败是降级不是错误（见 vector_channel）， // `try_join!` 会把模型故障变成检索报错 let (bm25, vector) = tokio::join!( bm25, - vector_channel( - state, - kb_id, - workspace_id, - query, - per_channel, - as_of, - embeddings - ) + vector_channel(state, kb_id, workspace_id, query, recall, embeddings) ); let bm25 = bm25 .map_err(|e| utopia_core::AppError::Other(anyhow::anyhow!("BM25 检索线程退出：{e}")))? @@ -159,14 +191,22 @@ async fn fused_hits( .filter_map(|f| Some((f.id.parse().ok()?, f))) .collect(); let ids: Vec = fused.iter().map(|(id, _)| *id).collect(); - let live = utopia_store::documents::chunks_by_ids(&state.pool, kb_id, &ids, as_of).await?; + let live = + utopia_store::documents::chunks_by_ids(&state.pool, kb_id, &ids, recall.as_of).await?; let mut fused: HashMap = fused.into_iter().collect(); let hits = live .into_iter() .filter_map(|chunk| { let f = fused.remove(&chunk.chunk.id)?; + let external_id = chunk + .source_kind + .as_deref() + .zip(chunk.external_key.as_deref()) + .and_then(|(kind, key)| crate::pushed_keys::pushed_identity(kind, key)) + .map(str::to_string); Some(ScoredChunk { chunk, + external_id, score: f.score, bm25_rank: position(&f.ranks, 0), vector_rank: position(&f.ranks, 1), @@ -196,8 +236,7 @@ async fn vector_channel( kb_id: Uuid, workspace_id: Uuid, query: &str, - per_channel: usize, - as_of: Option>, + recall: Recall<'_>, cache: Option<&mut llm_util::EmbedCache>, ) -> AppResult>> { let vector = match cache.as_deref().and_then(|c| c.get(query)).cloned() { @@ -228,8 +267,9 @@ async fn vector_channel( &state.pool, kb_id, &vector, - per_channel as i64, - as_of, + recall.per_channel as i64, + recall.as_of, + recall.documents, ) .await?; Ok(Some(ids.into_iter().map(|id| id.to_string()).collect())) diff --git a/crates/utopia-server/src/retrieval_tests.rs b/crates/utopia-server/src/retrieval_tests.rs index 5f7d9da64..24287b867 100644 --- a/crates/utopia-server/src/retrieval_tests.rs +++ b/crates/utopia-server/src/retrieval_tests.rs @@ -194,6 +194,7 @@ async fn both_channels_fuse_with_bm25_first() -> anyhow::Result<()> { top_k: 8, offset: 0, as_of: None, + filter: None, }; let page = super::hybrid_scored(&f.state, f.kb, f.ws, "apple harvest", opts, None).await?; let hits: Vec<_> = page @@ -205,7 +206,7 @@ async fn both_channels_fuse_with_bm25_first() -> anyhow::Result<()> { h.bm25_rank, h.vector_rank, h.chunk.source_id, - h.chunk.external_id.as_deref(), + h.external_id.as_deref(), ) }) .collect(); @@ -270,7 +271,7 @@ async fn historical_search_keeps_live_candidates_before_limiting_results() -> an let raw = f .state .search - .search(&f.kb.to_string(), "apple harvest", 24)?; + .search(&f.kb.to_string(), "apple harvest", 24, None)?; assert_eq!(raw.len(), 7); assert_eq!(raw.last().unwrap().chunk_id, f.b.to_string()); let current = super::hybrid(&f.state, f.kb, f.ws, "apple harvest", 6, None, None).await?; @@ -348,6 +349,7 @@ fn page(offset: usize, top_k: usize) -> super::SearchOpts { top_k, offset, as_of: None, + filter: None, } } @@ -394,3 +396,80 @@ async fn a_page_ending_at_depth_100_of_101_hits_has_no_more() -> anyhow::Result< assert_eq!((last.results.len(), last.has_more), (50, false)); Ok(()) } + +/// A source filter works inside both channels (#1271): a hundred better hits from +/// elsewhere fill each channel's candidates, so filtering after fusion would leave +/// nothing, and the six weak hits of the source must still page like any other list. +#[tokio::test] +async fn a_source_filter_pages_through_its_six_hits_under_a_hundred_better_ones( +) -> anyhow::Result<()> { + let Some(f) = fixture().await? else { + return Ok(()); + }; + let _server = seed_ranked_chunks(&f, 100).await?; + let (source, doc) = (Uuid::now_v7(), Uuid::now_v7()); + sqlx::raw_sql(&format!( + "INSERT INTO sources(id,kb_id,kind,name) VALUES ('{source}','{kb}','api','ci'); + INSERT INTO documents(id,kb_id,source_id,external_key,filename,sha256) + VALUES ('{doc}','{kb}','{source}','api:weak','weak.md','w')", + kb = f.kb + )) + .execute(&f.pool) + .await?; + let mut weak = Vec::new(); + for seq in 0..6 { + let id = Uuid::now_v7(); + let text = format!("apple {}", "pear ".repeat(200 + seq)); + sqlx::query( + "INSERT INTO chunks(id,kb_id,document_id,seq,text,embedding) + VALUES($1,$2,$3,$4,$5,$6::text::vector)", + ) + .bind(id) + .bind(f.kb) + .bind(doc) + .bind(seq as i32) + .bind(&text) + .bind(format!("[{},100,0,0]", seq + 1)) + .execute(&f.pool) + .await?; + weak.push((id.to_string(), text)); + } + f.state + .search + .reindex_document(&f.kb.to_string(), &doc.to_string(), &weak)?; + let walk = |offset, top_k| { + let opts = super::SearchOpts { + filter: Some(super::SourceFilter { + source_id: source, + key_prefix: None, + }), + ..page(offset, top_k) + }; + super::hybrid_scored(&f.state, f.kb, f.ws, "apple", opts, None) + }; + let (first, second, whole) = (walk(0, 3).await, walk(3, 3).await, walk(0, 10).await); + f.cleanup().await?; + let (first, second, whole) = (first?, second?, whole?); + let ids = + |hits: &[super::ScoredChunk]| hits.iter().map(|h| h.chunk.chunk.id).collect::>(); + let mut found: Vec<_> = ids(&whole.results).iter().map(Uuid::to_string).collect(); + found.sort(); + let mut expected: Vec<_> = weak.into_iter().map(|(id, _)| id).collect(); + expected.sort(); + assert_eq!(found, expected); + let channels: Vec<_> = whole + .results + .iter() + .map(|h| (h.bm25_rank.is_some(), h.vector_rank.is_some())) + .collect(); + assert_eq!(channels, [(true, true); 6]); + assert_eq!( + [ids(&first.results), ids(&second.results)].concat(), + ids(&whole.results) + ); + assert_eq!( + (first.has_more, second.has_more, whole.has_more), + (true, false, false) + ); + Ok(()) +} diff --git a/crates/utopia-server/src/search_recovery.rs b/crates/utopia-server/src/search_recovery.rs index 0316457ea..3c87e91cc 100644 --- a/crates/utopia-server/src/search_recovery.rs +++ b/crates/utopia-server/src/search_recovery.rs @@ -109,7 +109,12 @@ mod tests { assert!(marker(&index_dir).exists()); assert!(open(&index_dir).is_err(), "a second writer must be refused"); assert_eq!(search.len(), 2, "the live writer keeps its committed data"); - assert_eq!(search.search(&kb.to_string(), "firstneedle", 10)?.len(), 1); + assert_eq!( + search + .search(&kb.to_string(), "firstneedle", 10, None)? + .len(), + 1 + ); drop(search); let search = open(&index_dir)?; @@ -123,9 +128,14 @@ mod tests { rebuild(&index_dir, &search, live)?; assert_eq!(search.len(), 2); assert!(search - .search(&kb.to_string(), "firstneedle", 10)? + .search(&kb.to_string(), "firstneedle", 10, None)? .is_empty()); - assert_eq!(search.search(&kb.to_string(), "thirdneedle", 10)?.len(), 1); + assert_eq!( + search + .search(&kb.to_string(), "thirdneedle", 10, None)? + .len(), + 1 + ); assert!(!marker(&index_dir).exists()); } Ok(()) @@ -154,7 +164,12 @@ mod tests { drop(search); let search = open(&index_dir)?; assert_eq!(search.len(), 2); - assert_eq!(search.search(&kb.to_string(), "legacyneedle", 10)?.len(), 1); + assert_eq!( + search + .search(&kb.to_string(), "legacyneedle", 10, None)? + .len(), + 1 + ); // The database may become empty between attempts. Completing that // recovery must clear its marker instead of rebuilding on every boot. diff --git a/crates/utopia-store/src/documents.rs b/crates/utopia-store/src/documents.rs index 6f1e30fe4..f974577d0 100644 --- a/crates/utopia-store/src/documents.rs +++ b/crates/utopia-store/src/documents.rs @@ -1890,12 +1890,16 @@ async fn set_embeddings_for_input( /// 次序在外层再排一遍（`vector_index::RESORT`）：`relaxed_order` 下索引给的次序 /// 只是大致按距离，而距离并列时精确路径和 HNSW 各排各的——同一问两种计划回的 /// 居首不同（#652）。并列由 id 定 +/// +/// `documents`：只在这些文档的块里找。条件在近邻 CTE 里面，`LIMIT` 数的就只是 +/// 允许的块；放到外层，库里别的块会先占满名额 pub async fn vector_search( pool: &PgPool, kb_id: Uuid, embedding: &[f32], limit: i64, as_of: Option>, + documents: Option<&[Uuid]>, ) -> AppResult> { let dims = embedding.len(); if dims == 0 { @@ -1908,12 +1912,17 @@ pub async fn vector_search( "WITH nearest AS MATERIALIZED ( SELECT id, {distance} AS distance FROM chunks c WHERE c.kb_id = $1 AND c.embedding IS NOT NULL AND {live} - AND {same_dims} + AND {same_dims}{only} ORDER BY {distance} LIMIT $3 ) SELECT id FROM nearest ORDER BY {resort}", live = crate::record_axis::chunk_live_at("c", as_of.map(|_| 4)), + only = if documents.is_some() { + " AND c.document_id = ANY($5)" + } else { + "" + }, same_dims = crate::vector_index::same_dims("c.embedding", dims), distance = crate::vector_index::distance("c.embedding", 2, dims), resort = crate::vector_index::RESORT, @@ -1922,12 +1931,37 @@ pub async fn vector_search( .bind(&query_vec) .bind(limit) .bind(as_of) + .bind(documents) .fetch_all(&mut *tx) .await?; tx.commit().await?; Ok(rows.into_iter().map(|(id,)| id).collect()) } +/// 检索按来源过滤时允许的文档：库里这个来源当时活着的文档，给了 `key_prefix` 就只要 +/// `external_key` 以它开头的。比较用 `starts_with`：路径里的 `_` 和 `%` 进了 LIKE 就是通配符 +pub async fn ids_in_source( + pool: &PgPool, + kb_id: Uuid, + source_id: Uuid, + key_prefix: Option<&str>, + as_of: Option>, +) -> AppResult> { + let rows: Vec<(Uuid,)> = sqlx::query_as(&format!( + "SELECT d.id FROM documents d + WHERE d.kb_id = $1 AND d.source_id = $2 AND {doc_live} + AND ($3::text IS NULL OR starts_with(d.external_key, $3))", + doc_live = crate::record_axis::document_live_at("d", as_of.map(|_| 4)), + )) + .bind(kb_id) + .bind(source_id) + .bind(key_prefix) + .bind(as_of) + .fetch_all(pool) + .await?; + Ok(rows.into_iter().map(|(id,)| id).collect()) +} + /// 按 id 集合取分块（带文档名与来源），保持传入顺序。 /// /// `as_of`：记录轴（0019）——那一刻还活着的块、还没被删的文档。文档的删除 @@ -1940,8 +1974,7 @@ pub async fn chunks_by_ids( ) -> AppResult> { let rows: Vec = sqlx::query_as(&format!( "SELECT c.id, c.document_id, c.seq, c.text, d.filename, d.source_id, - CASE WHEN s.kind IN ('api', 'statements') - THEN substr(d.external_key, length(s.kind) + 2) END AS external_id + s.kind AS source_kind, d.external_key FROM chunks c JOIN documents d ON d.id = c.document_id LEFT JOIN sources s ON s.id = d.source_id WHERE c.kb_id = $1 AND c.id = ANY($2) diff --git a/crates/utopia-store/tests/store/a_search_reads_the_base_as_it_was.rs b/crates/utopia-store/tests/store/a_search_reads_the_base_as_it_was.rs index 6baccb786..cf831ac42 100644 --- a/crates/utopia-store/tests/store/a_search_reads_the_base_as_it_was.rs +++ b/crates/utopia-store/tests/store/a_search_reads_the_base_as_it_was.rs @@ -158,13 +158,15 @@ async fn a_superseded_chunk_keeps_its_vector_and_answers_as_of_then() -> anyhow: // 2. 现在：向量召回给出新块，给不出被顶掉的旧块 let query = vec![1.0_f32, 0.0, 0.0]; - let now_hits = utopia_store::documents::vector_search(&pool, f.kb, &query, 10, None).await?; + let now_hits = + utopia_store::documents::vector_search(&pool, f.kb, &query, 10, None, None).await?; assert!(!now_hits.contains(&f.old_chunk), "现行检索不该翻出旧版本"); assert!(now_hits.contains(&f.new_chunk)); // 3. 三月：那时旧块是现行的，新块还不存在 let then = Some(t("2026-03-01T00:00:00Z")); - let then_hits = utopia_store::documents::vector_search(&pool, f.kb, &query, 10, then).await?; + let then_hits = + utopia_store::documents::vector_search(&pool, f.kb, &query, 10, then, None).await?; assert!( then_hits.contains(&f.old_chunk), "三月该搜得到当时那一版——这一条在丢向量的实现上不可能成立" diff --git a/crates/utopia-store/tests/store/the_nearest_chunk_is_found_however_it_is_reached.rs b/crates/utopia-store/tests/store/the_nearest_chunk_is_found_however_it_is_reached.rs index 8cf1de506..9027893c0 100644 --- a/crates/utopia-store/tests/store/the_nearest_chunk_is_found_however_it_is_reached.rs +++ b/crates/utopia-store/tests/store/the_nearest_chunk_is_found_however_it_is_reached.rs @@ -192,11 +192,18 @@ struct Answers { async fn ask(pool: &PgPool, f: &Fixture) -> anyhow::Result { use utopia_store::documents::vector_search; - let big_now = vector_search(pool, f.big, &query(), 10, None).await?; - let small_now = vector_search(pool, f.small, &query(), 10, None).await?; - let five = vector_search(pool, f.big, &[1.0, 0.0, 0.0, 0.0, 0.0], 10, None).await?; - let big_then = - vector_search(pool, f.big, &query(), 10, Some(t("2026-03-01T00:00:00Z"))).await?; + let big_now = vector_search(pool, f.big, &query(), 10, None, None).await?; + let small_now = vector_search(pool, f.small, &query(), 10, None, None).await?; + let five = vector_search(pool, f.big, &[1.0, 0.0, 0.0, 0.0, 0.0], 10, None, None).await?; + let big_then = vector_search( + pool, + f.big, + &query(), + 10, + Some(t("2026-03-01T00:00:00Z")), + None, + ) + .await?; Ok(Answers { nearest_in_big: big_now.first().copied(), small_limit_10: small_now, @@ -249,6 +256,7 @@ async fn the_answers_agree_with_and_without_the_index() -> anyhow::Result<()> { &query(), 2, Some(t("2026-03-01T00:00:00Z")), + None, ) .await?; assert!( diff --git a/docs/decisions/0014-identity-from-the-person-scope-from-the-token.md b/docs/decisions/0014-identity-from-the-person-scope-from-the-token.md index 9825b075f..76aeb4e63 100644 --- a/docs/decisions/0014-identity-from-the-person-scope-from-the-token.md +++ b/docs/decisions/0014-identity-from-the-person-scope-from-the-token.md @@ -123,6 +123,11 @@ push documents in. Chunk search (`POST /kbs/{id}/search`, added 2026-10-24, #1271) takes a token on the same terms, so the service can ask the base about what it mirrored. Each hit carries its score, its BM25 and vector ranks and its source; `offset` pages through one fused list, at most 100 deep. + `source_id` narrows the search to one source of the base, and a source outside the base answers + 404 like a missing one. `external_id_prefix` narrows it further to a push source's documents + whose `external_id` starts with the prefix, taken literally; without the `source_id` of an + `api` or `statements` source it answers 422 `prefix_needs_push_source`. Both filters apply + inside each channel, before fusion, so other documents cannot fill the candidates first. ## Dead ends @@ -152,7 +157,8 @@ push documents in. and document content" now also counts them. 2026-10-19: the graph overview and neighbourhood join them, and so does reading a pushed document back by its `external_id` (#1276). The ontology and a class's instances follow (#1265), then a document's extractions and the - base's export (#1266). 2026-10-24: chunk search joins them (#1271). + base's export (#1266). 2026-10-24: chunk search joins them (#1271). 2026-10-25: it can be + narrowed to one source and, for a push source, to an `external_id` prefix (#1271). ## Open questions diff --git a/docs/design/chat-and-mcp.md b/docs/design/chat-and-mcp.md index eaa3fbe57..12acc42f7 100644 --- a/docs/design/chat-and-mcp.md +++ b/docs/design/chat-and-mcp.md @@ -79,7 +79,10 @@ RRF score ranks hits within one query only; it is not comparable across queries page fuses 100 candidates per channel where chat fuses 24, so the search page can rank the same query differently from chat and from before #1271. Pages join up only while the base is unchanged between requests, and each channel counts its ranks before the final `as_of` filter, so a -historical page can skip rank numbers. +historical page can skip rank numbers. A `source_id` of the base (another base's source is a 404) +and, for a push source, an `external_id_prefix` limit both channels to that source's documents +before fusion, so the base's other hits cannot take the candidates a filter after fusion would +lose; the prefix compares against the key `pushed_key` builds (#1271). **`remember`** writes a memory document at once; its statements go through open extraction and wait in `pending_facts`; the assistant says the sentence is recorded and its statements will be shown diff --git a/web/src/docs/ingest.md b/web/src/docs/ingest.md index b577f0251..df3a8fa2b 100644 --- a/web/src/docs/ingest.md +++ b/web/src/docs/ingest.md @@ -211,6 +211,8 @@ Content-Type: application/json `q` is required. `top_k` defaults to 10 and is capped at 50; `offset` defaults to 0. `as_of` (an RFC 3339 instant) reads the base as it was then; ranks are counted before `as_of` filters the hits. Any other field is a `422`. +To search one source, add its `source_id`; a source of another base is a `404`. To search part of what you pushed, add `external_id_prefix` as well, for example `"docs/"`: only documents whose `external_id` starts with it are searched, and `_` or `%` in it are plain characters. The prefix needs the `source_id` of an API or Statements source, otherwise it is a `422` (`prefix_needs_push_source`). For one exact `external_id`, read the document back as above. Both filters apply inside each channel before the hits are fused, so paging works the same with them as without. + ```json { "results": [