feat(semantic-index): batch API embedding warmup

packages/node-runtime/src/semantic-index/embedding/api.ts

@@ -10,6 +10,9 @@
 
 import type { EmbeddingProvider } from './types'
 
+/** API embedding 默认一次提交 32 个 chunk，保持 concurrency=1 以降低 429 风险。 */
+export const API_DOCUMENT_BATCH_SIZE = 32
+
 export interface OpenAICompatibleConfig {
   baseUrl: string
   apiKey: string
@@ -37,6 +40,7 @@ function buildEmbeddingsUrl(baseUrl: string): string {
 export class OpenAICompatibleEmbeddingProvider implements EmbeddingProvider {
   readonly modelId: string
   readonly maxTokens: number
+  readonly documentBatchSize = API_DOCUMENT_BATCH_SIZE
 
   private config: OpenAICompatibleConfig
   private fetchFn: FetchFn

packages/node-runtime/src/semantic-index/embedding/local.ts

@@ -47,6 +47,7 @@ export class LocalEmbeddingProvider implements EmbeddingProvider {
   readonly modelId: string
   readonly dim: number
   readonly maxTokens: number
+  readonly documentBatchSize: number
 
   private profile: LocalEmbeddingProfile
   private options: LocalEmbeddingProviderOptions
@@ -58,6 +59,7 @@ export class LocalEmbeddingProvider implements EmbeddingProvider {
     this.modelId = profile.modelId
     this.dim = profile.dim
     this.maxTokens = profile.maxTokens
+    this.documentBatchSize = profile.maxBatchSize ?? 1
   }
 
   private getExtractor(): Promise<FeatureExtractFn> {

packages/node-runtime/src/semantic-index/embedding/providers.test.ts

@@ -82,6 +82,7 @@ test('API provider posts OpenAI-compatible request and parses ordered embeddings
     { fetchFn }
   )
 
+  assert.equal(provider.documentBatchSize, 32)
   const vectors = await provider.embedDocuments(['first', 'second'])
   assert.equal(captured!.url, 'https://api.example.com/v1/embeddings')
   assert.equal(captured!.headers.Authorization, 'Bearer sk-test')

packages/node-runtime/src/semantic-index/embedding/types.ts

@@ -15,6 +15,8 @@ export interface EmbeddingProvider {
   readonly dim: number
   /** 单文本 token 上限（profile cap） */
   readonly maxTokens: number
+  /** warmup 建索引时推荐的 document 批量大小；未声明时按 1 条处理 */
+  readonly documentBatchSize?: number
   /** document 向量化（可批量；Qwen3 内部固定逐条） */
   embedDocuments(texts: string[]): Promise<Float32Array[]>
   /** query 向量化（本地模型会加 queryInstruction） */

packages/node-runtime/src/semantic-index/warmup/runner.test.ts

@@ -52,10 +52,12 @@ class FakeEmbedder implements EmbeddingProvider {
   readonly dim = 4
   readonly maxTokens = 1000
   calls = 0
+  batchSizes: number[] = []
   failAtCall?: number
   afterEmbed?: () => void
 
   async embedDocuments(texts: string[]): Promise<Float32Array[]> {
+    this.batchSizes.push(texts.length)
     return texts.map(() => {
       this.calls++
       if (this.failAtCall && this.calls >= this.failAtCall) throw new Error('boom')
@@ -69,6 +71,10 @@ class FakeEmbedder implements EmbeddingProvider {
   }
 }
 
+class BatchingFakeEmbedder extends FakeEmbedder {
+  readonly documentBatchSize = 32
+}
+
 function setup(messages: ChunkMessageInput[]) {
   const dir = makeTempDir()
   const dbPath = path.join(dir, 'embedding_index.db')
@@ -113,6 +119,20 @@ test('full warmup writes all chunks and marks completed', async () => {
   stateStore.close()
 })
 
+test('warmup batches document embeddings when the provider supports batches', async () => {
+  const { store, stateStore, source } = setup(makeMessages(70))
+  const embedder = new BatchingFakeEmbedder()
+
+  const result = await runWarmup({ dbPathHash: DB_HASH, modelId: MODEL, embedder, store, stateStore, source, config })
+
+  assert.equal(result.status, 'completed')
+  assert.equal(result.chunksWritten, 35)
+  assert.deepEqual(embedder.batchSizes, [32, 3])
+  assert.equal(countStored(store), 35)
+  store.close()
+  stateStore.close()
+})
+
 test('pause then resume completes without duplicate inserts', async () => {
   const { store, stateStore, source } = setup(makeMessages(8))
   const embedder = new FakeEmbedder()

packages/node-runtime/src/semantic-index/warmup/runner.ts

@@ -20,7 +20,7 @@ import { chunkMessages, type ChunkMessageInput, type ChunkSource } from '../chun
 import { CHUNKER_VERSION, STRATEGY_ID, composeChunkId, type ChunkerConfig } from '../chunker-config'
 import type { EmbeddingIndexStore } from '../store'
 import type { SemanticIndexStateStore } from '../session-state-store'
-import type { ChunkRecord } from '../types'
+import type { ChunkInsert, ChunkRecord } from '../types'
 
 /** 对话消息来源（warmup 输入抽象，真实实现读取聊天库） */
 export interface SemanticMessageSource {
@@ -52,6 +52,11 @@ export interface WarmupResult {
   error?: string
 }
 
+function resolveDocumentBatchSize(embedder: EmbeddingProvider): number {
+  const size = embedder.documentBatchSize
+  return Number.isFinite(size) && size && size > 0 ? Math.floor(size) : 1
+}
+
 export async function runWarmup(options: WarmupRunnerOptions): Promise<WarmupResult> {
   const { dbPathHash, modelId, embedder, store, stateStore, source, config, checkStop } = options
 
@@ -111,8 +116,10 @@ export async function runWarmup(options: WarmupRunnerOptions): Promise<WarmupRes
     let storedChunkCount = store.countChunks(dbPathHash, modelId)
     if (resumeIndex < 0) storedChunkCount = 0
 
-    for (const { chunk, endIndex: chunkEndIndex } of chunkRanges) {
-      if (chunkEndIndex <= resumeIndex) continue
+    const pendingRanges = chunkRanges.filter(({ endIndex }) => endIndex > resumeIndex)
+    const batchSize = resolveDocumentBatchSize(embedder)
+    for (let i = 0; i < pendingRanges.length; i += batchSize) {
+      const batchRanges = pendingRanges.slice(i, i + batchSize)
 
       const stop = checkStop?.()
       if (stop) {
@@ -120,40 +127,50 @@ export async function runWarmup(options: WarmupRunnerOptions): Promise<WarmupRes
         return { status: stop, chunksWritten }
       }
 
-      const [vector] = await embedder.embedDocuments([chunk.embeddingInput])
+      const vectors = await embedder.embedDocuments(batchRanges.map(({ chunk }) => chunk.embeddingInput))
+      if (vectors.length !== batchRanges.length) {
+        throw new Error(`embedding provider returned ${vectors.length} vectors for ${batchRanges.length} inputs`)
+      }
       const stopAfterEmbedding = checkStop?.()
       if (stopAfterEmbedding === 'cancelled') {
         stateStore.setIndexStatus(dbPathHash, 'cancelled')
         return { status: 'cancelled', chunksWritten }
       }
 
-      const record: ChunkRecord = {
-        chunkId: composeChunkId(dbPathHash, modelId, chunk.localChunkId),
-        dbPathHash,
-        strategyId: STRATEGY_ID,
-        modelId,
-        dim: vector.length,
-        parentId: chunk.parentId,
-        startMessageId: chunk.startMessageId,
-        endMessageId: chunk.endMessageId,
-        startTs: chunk.startTs,
-        endTs: chunk.endTs,
-        messageCount: chunk.messageCount,
-        rawContentHash: chunk.rawContentHash,
-        embeddingInputHash: chunk.embeddingInputHash,
-        chunkerVersion: CHUNKER_VERSION,
-        chunkerConfigHash,
-        indexedAt: Date.now(),
-        status: 'indexed',
-      }
-      store.insertChunk(record, vector)
-      chunksWritten++
-      storedChunkCount++
+      // API provider 支持一次提交多个文本；本地 Qwen3 仍声明 batch=1，避免 last_token 污染。
+      const indexedAt = Date.now()
+      const inserts: ChunkInsert[] = batchRanges.map(({ chunk }, index) => {
+        const vector = vectors[index]
+        const record: ChunkRecord = {
+          chunkId: composeChunkId(dbPathHash, modelId, chunk.localChunkId),
+          dbPathHash,
+          strategyId: STRATEGY_ID,
+          modelId,
+          dim: vector.length,
+          parentId: chunk.parentId,
+          startMessageId: chunk.startMessageId,
+          endMessageId: chunk.endMessageId,
+          startTs: chunk.startTs,
+          endTs: chunk.endTs,
+          messageCount: chunk.messageCount,
+          rawContentHash: chunk.rawContentHash,
+          embeddingInputHash: chunk.embeddingInputHash,
+          chunkerVersion: CHUNKER_VERSION,
+          chunkerConfigHash,
+          indexedAt,
+          status: 'indexed',
+        }
+        return { record, embedding: vector }
+      })
+      store.insertChunks(inserts)
+      chunksWritten += inserts.length
+      storedChunkCount += inserts.length
 
+      const lastRange = batchRanges[batchRanges.length - 1]
       stateStore.updateProgress(dbPathHash, {
         indexStatus: 'running',
-        indexedMessages: chunkEndIndex + 1,
-        lastIndexedMessageId: chunk.endMessageId,
+        indexedMessages: lastRange.endIndex + 1,
+        lastIndexedMessageId: lastRange.chunk.endMessageId,
         chunkCount: storedChunkCount,
       })
     }