A search may name a source and an external_id prefix, and both channels recall only those documents before fusion

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<Uuid>,
-    pub external_id: Option<String>,
+    #[serde(skip)]
+    pub source_kind: Option<String>,
+    #[serde(skip)]
+    pub external_key: Option<String>,
 }
 
 #[derive(Debug, Clone, Serialize, sqlx::FromRow)]

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<Vec<Hit>> {
+    ///
+    /// `documents`：只在这些文档里找。它和 kb 一样是 `Must`，`limit` 数的就只是允许的块
+    pub fn search(
+        &self,
+        kb_id: &str,
+        query: &str,
+        limit: usize,
+        documents: Option<&[String]>,
+    ) -> anyhow::Result<Vec<Hit>> {
         let mut analyzer = self.analyzer.clone();
         let mut stream = analyzer.token_stream(query);
         let mut term_queries: Vec<(Occur, Box<dyn Query>)> = Vec::new();
@@ -161,7 +169,14 @@ impl SearchIndex {
             IndexRecordOption::Basic,
         ));
         let text_query: Box<dyn Query> = 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();

crates/utopia-server/src/api/documents_routes.rs

@@ -227,11 +227,11 @@ pub async fn by_external_id(
 ) -> ApiResult<Response> {
     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)?;

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<String>,
+    /// 只检索本库这个来源的文档；别的库的来源答 404
+    #[serde(default)]
+    pub source_id: Option<Uuid>,
+    /// 只检索 `external_id` 以它开头的文档。identity 是推送来源里的，所以要 `source_id`
+    /// 且它是推送来源，否则 422
+    #[serde(default)]
+    pub external_id_prefix: Option<String>,
+}
+
+async fn source_filter(
+    state: &AppState,
+    kb_id: Uuid,
+    req: &SearchReq,
+) -> ApiResult<Option<retrieval::SourceFilter>> {
+    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!({

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

crates/utopia-server/src/api/token_read_tests.rs

@@ -20,6 +20,7 @@ struct Fixture {
     entity: Uuid,
     fact: Uuid,
     doc: Uuid,
+    search: Arc<utopia_search::SearchIndex>,
     _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<serde_json::Value> {
+        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(())
+}

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;

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