Ingest embeds four batches at a time, and both ingests share the loop

crates/utopia-server/src/pipeline.rs

@@ -3,11 +3,26 @@
 
 use crate::llm_util;
 use crate::state::AppState;
-use utopia_core::models::Proposer;
+use futures_util::{stream, StreamExt};
+use utopia_core::models::{LlmSettings, Proposer};
+use utopia_llm::LlmClient;
 use uuid::Uuid;
 
+/// 每次送多少条去嵌入。
+///
+/// 与 ontology_index 的 64 不同，这里没量过，先不动：那边的注释说批大小要拿真实文本量，
+/// 块文本比类标签长得多，16 未必错。并发与批大小两个旋钮别在一刀里同时拧，否则结果读不出来。
 const EMBED_BATCH: usize = 16;
 
+/// 同时在飞的嵌入批数上限（#513）。
+///
+/// 嵌入模型的并发闸门（`model_concurrency`，缺省 10）是部署级共享的：本体补齐
+/// （ontology_index，上限 4）、类型消解、每次提问的查询嵌入都走它。摄入是用户等着看的
+/// 前台活，比后台补齐更有资格用闸门；但它也是能一次带来上千个分块的那一路，不设上限
+/// 就把补齐、消解、检索全饿死——ontology_index 记着一次 `join_all` 把闸门占死、五篇文档
+/// 卡在 embedding 十五分钟的事故。4 与补齐相同：两路同时跑合计 8，闸门还剩 2 给别人。
+const EMBED_JOBS: usize = 4;
+
 pub async fn process_document(state: &AppState, document_id: Uuid) -> anyhow::Result<()> {
     match run(state, document_id).await {
         Ok(()) => Ok(()),
@@ -57,25 +72,10 @@ async fn run(state: &AppState, document_id: Uuid) -> anyhow::Result<()> {
     // 4. embedding（工作区配置了 embedding 模型才做；没配也算 ready，先享受 BM25 搜索）
     let kb_row = utopia_store::kbs::get(&state.pool, doc.kb_id).await?;
     let settings = utopia_store::settings::get(&state.pool, kb_row.workspace_id).await?;
-    if let Some(client) = settings.as_ref().and_then(llm_util::embed_client) {
+    if let Some((settings, client)) = embedder(settings.as_ref()) {
         utopia_store::documents::set_status(&state.pool, document_id, "embedding").await?;
         state.emit_document(doc.kb_id, document_id);
-        let pending =
-            utopia_store::documents::chunks_pending_embedding(&state.pool, document_id).await?;
-        for batch in pending.chunks(EMBED_BATCH) {
-            let texts: Vec<String> = batch.iter().map(|(_, t)| t.clone()).collect();
-            let _permit = match settings.as_ref() {
-                Some(s) => llm_util::acquire_embed(state, s).await,
-                None => None,
-            };
-            let embeddings = client.embed(&texts).await?;
-            if embeddings.len() != batch.len() {
-                anyhow::bail!("Embedding 返回数量不匹配");
-            }
-            let items: Vec<(Uuid, Vec<f32>)> =
-                batch.iter().map(|(id, _)| *id).zip(embeddings).collect();
-            utopia_store::documents::set_embeddings(&state.pool, &items).await?;
-        }
+        embed_pending(state, settings, &client, document_id).await?;
     }
 
     utopia_store::documents::set_ready(&state.pool, document_id, text_len, chunk_count).await?;
@@ -115,23 +115,8 @@ pub async fn memory_ingest(
     let kb_row = utopia_store::kbs::get(&state.pool, doc.kb_id).await?;
     let settings = utopia_store::settings::get(&state.pool, kb_row.workspace_id).await?;
 
-    if let Some(client) = settings.as_ref().and_then(llm_util::embed_client) {
-        let pending =
-            utopia_store::documents::chunks_pending_embedding(&state.pool, document_id).await?;
-        for batch in pending.chunks(EMBED_BATCH) {
-            let texts: Vec<String> = batch.iter().map(|(_, t)| t.clone()).collect();
-            let _permit = match settings.as_ref() {
-                Some(s) => llm_util::acquire_embed(state, s).await,
-                None => None,
-            };
-            let embeddings = client.embed(&texts).await?;
-            if embeddings.len() != batch.len() {
-                anyhow::bail!("Embedding 返回数量不匹配");
-            }
-            let items: Vec<(Uuid, Vec<f32>)> =
-                batch.iter().map(|(id, _)| *id).zip(embeddings).collect();
-            utopia_store::documents::set_embeddings(&state.pool, &items).await?;
-        }
+    if let Some((settings, client)) = embedder(settings.as_ref()) {
+        embed_pending(state, settings, &client, document_id).await?;
     }
 
     let chunks = utopia_store::documents::chunks_full(&state.pool, document_id).await?;
@@ -160,3 +145,62 @@ pub async fn memory_ingest(
     state.emit_document(doc.kb_id, document_id);
     Ok(())
 }
+
+/// 工作区配了嵌入模型才有客户端；设置与客户端一起交出去，闸门许可证要按设置取。
+fn embedder(settings: Option<&LlmSettings>) -> Option<(&LlmSettings, LlmClient)> {
+    let s = settings?;
+    llm_util::embed_client(s).map(|c| (s, c))
+}
+
+/// 把这篇文档还没有向量的分块全部嵌完，返回嵌了多少条。文档摄入与记忆摄入共用——
+/// 从前两处各抄一份同样的循环，且都是一批等完再发下一批。
+///
+/// 批与批之间互不依赖（向量按位置配对，只在一批之内），`EMBED_JOBS` 个批次同时在飞，
+/// 每个批次握一张闸门许可证只到自己写完。
+///
+/// **数量对不上就整批放弃。** 配对是按位置的，少一条就全体错位，把一条的向量写到另一条
+/// 身上；这种错落库之后再也看不出来——正文还在，向量是别人的。
+///
+/// **任一批失败整篇失败。** 串行时错误自然一路 unwind；并发下要明确地把它传出去，让调用方
+/// 把文档标成 failed，而不是留它停在 embedding。已经写进去的向量留着，重跑只补缺的
+/// （`chunks_pending_embedding` 只挑 embedding 为空的）。
+async fn embed_pending(
+    state: &AppState,
+    settings: &LlmSettings,
+    client: &LlmClient,
+    document_id: Uuid,
+) -> anyhow::Result<usize> {
+    let pending =
+        utopia_store::documents::chunks_pending_embedding(&state.pool, document_id).await?;
+    // 批次先切成拥有的数据：借来的切片捕进并发驱动的 future 过不了 Send 边界
+    //（ontology_index 里同一句话）
+    let batches: Vec<Vec<(Uuid, String)>> =
+        pending.chunks(EMBED_BATCH).map(<[_]>::to_vec).collect();
+    let mut batches = stream::iter(batches)
+        .map(|batch| async move {
+            let texts: Vec<String> = batch.iter().map(|(_, t)| t.clone()).collect();
+            let _permit = llm_util::acquire_embed(state, settings).await;
+            let embeddings = client.embed(&texts).await?;
+            if embeddings.len() != batch.len() {
+                anyhow::bail!(
+                    "嵌入返回 {} 条，送去的是 {} 条",
+                    embeddings.len(),
+                    batch.len()
+                );
+            }
+            let items: Vec<(Uuid, Vec<f32>)> =
+                batch.iter().map(|(id, _)| *id).zip(embeddings).collect();
+            utopia_store::documents::set_embeddings(&state.pool, &items).await?;
+            Ok::<usize, anyhow::Error>(items.len())
+        })
+        .buffer_unordered(EMBED_JOBS);
+    let mut done = 0;
+    while let Some(written) = batches.next().await {
+        done += written?;
+    }
+    Ok(done)
+}
+
+#[cfg(test)]
+#[path = "pipeline_tests.rs"]
+mod pipeline_tests;