diff --git a/packages/app/src-tauri/src/vector/mod.rs b/packages/app/src-tauri/src/vector/mod.rs index cf8b107d8..e418ba800 100644 --- a/packages/app/src-tauri/src/vector/mod.rs +++ b/packages/app/src-tauri/src/vector/mod.rs @@ -1,9 +1,10 @@ use crate::storage; use anyhow::Result; -use rusqlite::{Connection, params}; +use rusqlite::{params, Connection}; use serde::{Deserialize, Serialize}; use std::path::PathBuf; use std::sync::Mutex; +use std::time::Duration; use tauri::{AppHandle, Manager}; #[derive(Debug, Clone, Serialize, Deserialize)] @@ -25,6 +26,12 @@ pub struct VectorDB { dimension: usize, } +fn dimension_from_vec_schema(sql: &str) -> Option { + let start = sql.find("float[")? + "float[".len(); + let end = sql[start..].find(']')? + start; + sql[start..end].parse().ok() +} + fn parse_embedding_blob(blob: &[u8]) -> Vec { let chunk_size = std::mem::size_of::(); if blob.len() % chunk_size != 0 { @@ -60,7 +67,7 @@ impl VectorDB { ))); } - conn.execute("PRAGMA busy_timeout=5000", [])?; + conn.busy_timeout(Duration::from_millis(5000))?; let _: String = conn.query_row("PRAGMA journal_mode=WAL", [], |row| row.get(0))?; conn.execute( @@ -72,54 +79,53 @@ impl VectorDB { [], )?; - // Check if vec_embeddings already exists with a different dimension. - // If so, drop and recreate to avoid dimension mismatch. + // A vec0 table's dimension is fixed at creation. On startup, opening a + // persisted table must never discard vectors merely because the fallback + // dimension changed (for example, a 1024d remote model vs. 384d builtin). + // The vectorization path can explicitly reinitialize an *empty* table. let table_exists: bool = conn.query_row( "SELECT COUNT(*) > 0 FROM sqlite_master WHERE type='table' AND name='vec_embeddings'", [], |row| row.get(0), )?; - if table_exists { - // Probe actual dimension by checking the schema via sqlite_master - // sqlite-vec stores the dimension in the table's SQL definition - let existing_sql: Option = conn.query_row( + let actual_dimension = if table_exists { + let existing_sql: String = conn.query_row( "SELECT sql FROM sqlite_master WHERE type='table' AND name='vec_embeddings'", [], |row| row.get(0), - ).ok(); - - let needs_recreate = if let Some(sql) = existing_sql { - // sql looks like: CREATE VIRTUAL TABLE vec_embeddings USING vec0(embedding float[4096]) - let dim_str = format!("float[{}]", dimension); - !sql.contains(&dim_str) - } else { - true - }; - - if needs_recreate { - println!("[VectorDB] Existing vec_embeddings has wrong dimension, recreating for {}", dimension); - conn.execute("DROP TABLE IF EXISTS vec_embeddings", [])?; - conn.execute("DELETE FROM id_mapping", [])?; - } - } + )?; + dimension_from_vec_schema(&existing_sql).ok_or_else(|| { + anyhow::anyhow!("Could not determine existing vec_embeddings dimension") + })? + } else { + conn.execute( + &format!( + "CREATE VIRTUAL TABLE vec_embeddings USING vec0( + embedding float[{}] + )", + dimension + ), + [], + )?; + dimension + }; conn.execute( - &format!( - "CREATE VIRTUAL TABLE IF NOT EXISTS vec_embeddings USING vec0( - embedding float[{}] - )", - dimension - ), + "CREATE INDEX IF NOT EXISTS idx_id_mapping_book ON id_mapping(book_id)", [], )?; - conn.execute("CREATE INDEX IF NOT EXISTS idx_id_mapping_book ON id_mapping(book_id)", [])?; - let version: String = conn.query_row("SELECT vec_version()", [], |row| row.get(0))?; - println!("[VectorDB] sqlite-vec version: {}, dimension: {}", version, dimension); - - Ok(Self { conn, dimension }) + println!( + "[VectorDB] sqlite-vec version: {}, dimension: {}", + version, actual_dimension + ); + + Ok(Self { + conn, + dimension: actual_dimension, + }) } pub fn insert(&self, records: &[VectorRecord]) -> Result<()> { @@ -293,6 +299,29 @@ impl VectorDB { } } +#[cfg(test)] +mod tests { + use super::dimension_from_vec_schema; + + #[test] + fn reads_dimension_from_existing_vec0_schema() { + assert_eq!( + dimension_from_vec_schema( + "CREATE VIRTUAL TABLE vec_embeddings USING vec0(embedding float[1024])" + ), + Some(1024), + ); + } + + #[test] + fn rejects_a_schema_without_an_embedding_dimension() { + assert_eq!( + dimension_from_vec_schema("CREATE TABLE vec_embeddings (id TEXT)"), + None + ); + } +} + pub struct VectorDBState { pub db: Mutex>, } diff --git a/packages/app/src/lib/rag/vectorize-trigger.ts b/packages/app/src/lib/rag/vectorize-trigger.ts index 39025e113..331fa7c13 100644 --- a/packages/app/src/lib/rag/vectorize-trigger.ts +++ b/packages/app/src/lib/rag/vectorize-trigger.ts @@ -42,6 +42,7 @@ export async function triggerVectorizeBook( url: selected.url, apiKey: selected.apiKey, modelId: selected.modelId, + dimension: selected.dimension, }; })(), }; diff --git a/packages/app/src/main.tsx b/packages/app/src/main.tsx index 98f915a80..1d90de191 100644 --- a/packages/app/src/main.tsx +++ b/packages/app/src/main.tsx @@ -9,10 +9,16 @@ import { createRoot } from "react-dom/client"; import App from "./App"; import "./styles/globals.css"; import { setEmbeddingWorkerFactory, setStreamingFetch } from "@readany/core/ai"; -import { BUILTIN_EMBEDDING_MODELS } from "@readany/core/ai/builtin-embedding-models"; import { onLibraryChanged } from "@readany/core/events/library-events"; import { installFeedbackLogCapture, setFeedbackWorkerUrl } from "@readany/core/feedback"; -import { setVectorDB } from "@readany/core/rag"; +import { + createBuiltinEmbeddingService, + EmbeddingService, + normalizeEmbeddingEndpoint, + clearSearchConfiguration, + configureSearch, + setVectorDB, +} from "@readany/core/rag"; import { setPlatformService } from "@readany/core/services"; import { fetch as tauriFetch } from "@tauri-apps/plugin-http"; import { TauriPlatformService } from "./lib/platform/tauri-platform-service"; @@ -46,6 +52,53 @@ setEmbeddingWorkerFactory( new Worker(new URL("@readany/core/ai/embedding-worker", import.meta.url), { type: "module" }), ); +/** + * The vectorization pipeline configures its embedding source independently from + * search. Keep the query-side service aligned with the active vector-model setting + * so Reader Agent ragSearch can generate query embeddings as well. + */ +function configureRagSearchFromVectorModelStore(): void { + const state = useVectorModelStore.getState(); + if (!state.vectorModelEnabled) { + clearSearchConfiguration(); + return; + } + + if (state.vectorModelMode === "builtin" && state.selectedBuiltinModelId) { + configureSearch(createBuiltinEmbeddingService(state.selectedBuiltinModelId)); + return; + } + + const remoteModel = state.getSelectedVectorModel(); + if (state.vectorModelMode === "remote" && remoteModel) { + configureSearch( + new EmbeddingService({ + model: { + id: remoteModel.modelId, + name: remoteModel.name || remoteModel.modelId, + dimensions: remoteModel.dimension ?? 0, + maxTokens: 8192, + provider: "openai", + }, + apiKey: remoteModel.apiKey || "local", + baseUrl: remoteModel.url, + }), + { + kind: "remote", + modelId: remoteModel.modelId, + endpoint: normalizeEmbeddingEndpoint(remoteModel.url), + dimensions: remoteModel.dimension ?? 0, + }, + ); + return; + } + + clearSearchConfiguration(); +} + +configureRagSearchFromVectorModelStore(); +useVectorModelStore.subscribe(configureRagSearchFromVectorModelStore); + // Set vector database reference (initialized in Rust setup) const tauriVectorDB = new TauriVectorDB(); setVectorDB(tauriVectorDB); @@ -53,31 +106,6 @@ console.log("[VectorDB] TauriVectorDB reference set"); const desktopDataRootReady = syncLegacyDesktopLibraryRootConfig().catch(console.error); -// Align vector DB dimension with the currently selected model -(async () => { - try { - await desktopDataRootReady; - const { vectorModelMode, selectedBuiltinModelId, getSelectedVectorModel } = - useVectorModelStore.getState(); - let dimension: number | undefined; - - if (vectorModelMode === "builtin" && selectedBuiltinModelId) { - const model = BUILTIN_EMBEDDING_MODELS.find((m) => m.id === selectedBuiltinModelId); - dimension = model?.dimension; - } else if (vectorModelMode === "remote") { - const remoteModel = getSelectedVectorModel(); - dimension = remoteModel?.dimension; - } - - if (dimension && dimension !== 384) { - await tauriVectorDB.reinit(dimension); - console.log(`[VectorDB] Aligned dimension to ${dimension}`); - } - } catch (err) { - console.warn("[VectorDB] Failed to align dimension on startup:", err); - } -})(); - // Ensure i18n is fully initialized before rendering i18nReady.then(() => { desktopDataRootReady.catch(console.error); diff --git a/packages/cli/src/rag-config.ts b/packages/cli/src/rag-config.ts index 1a872637a..e10c16ea1 100644 --- a/packages/cli/src/rag-config.ts +++ b/packages/cli/src/rag-config.ts @@ -96,13 +96,19 @@ export async function configureRagSearchForCli( const key = `${model.url}\n${model.modelId}\n${model.apiKey}`; if (configuredEmbeddingKey === key) return { embeddingConfigured: true }; - const { EmbeddingService, configureSearch } = await import("@readany/core/rag"); + const { EmbeddingService, configureSearch, normalizeEmbeddingEndpoint } = await import("@readany/core/rag"); configureSearch( new EmbeddingService({ model: toEmbeddingModel(model), apiKey: model.apiKey || "local", baseUrl: model.url, }), + { + kind: "remote", + modelId: model.modelId, + endpoint: normalizeEmbeddingEndpoint(model.url), + dimensions: model.dimension ?? 0, + }, ); configuredEmbeddingKey = key; return { embeddingConfigured: true }; diff --git a/packages/core/src/ai/tools/rag-tools.ts b/packages/core/src/ai/tools/rag-tools.ts index 5937e0e2b..b6245f1bc 100644 --- a/packages/core/src/ai/tools/rag-tools.ts +++ b/packages/core/src/ai/tools/rag-tools.ts @@ -187,6 +187,12 @@ export function createRagSearchTool(bookId: string): ToolDefinition { returnedResults: truncatedResults.length, totalTokens, tokenBudget: MAX_TOTAL_TOKENS, + ...(results[0]?.vectorStatus + ? { + vectorStatus: results[0].vectorStatus, + vectorError: results[0].vectorError, + } + : {}), }; }, }; diff --git a/packages/core/src/db/chunk-queries.ts b/packages/core/src/db/chunk-queries.ts index 4f4431566..0d7214098 100644 --- a/packages/core/src/db/chunk-queries.ts +++ b/packages/core/src/db/chunk-queries.ts @@ -1,4 +1,4 @@ -import type { Chunk } from "../types"; +import type { Chunk, VectorIndexProvenance } from "../types"; import { getDB, getLocalDB, serializeEmbedding, deserializeEmbedding } from "./db-core"; export async function getChunks(bookId: string): Promise { @@ -57,11 +57,57 @@ export async function deleteChunks(bookId: string): Promise { await database.execute("DELETE FROM chunks WHERE book_id = ?", [bookId]); } +export async function getVectorIndexProvenance(bookId: string): Promise { + const database = await getLocalDB(); + const rows = await database.select<{ + book_id: string; + model_kind: "builtin" | "remote"; + model_id: string; + endpoint: string | null; + dimensions: number; + created_at: number; + }>("SELECT * FROM vector_index_provenance WHERE book_id = ?", [bookId]); + const row = rows[0]; + if (!row) return null; + return { + bookId: row.book_id, + kind: row.model_kind, + modelId: row.model_id, + endpoint: row.endpoint || undefined, + dimensions: row.dimensions, + createdAt: row.created_at, + }; +} + +export async function setVectorIndexProvenance(provenance: VectorIndexProvenance): Promise { + const database = await getLocalDB(); + await database.execute( + `INSERT OR REPLACE INTO vector_index_provenance + (book_id, model_kind, model_id, endpoint, dimensions, created_at) + VALUES (?, ?, ?, ?, ?, ?)`, + [ + provenance.bookId, + provenance.kind, + provenance.modelId, + provenance.endpoint || null, + provenance.dimensions, + provenance.createdAt, + ], + ); +} + +export async function deleteVectorIndexProvenance(bookId: string): Promise { + const database = await getLocalDB(); + await database.execute("DELETE FROM vector_index_provenance WHERE book_id = ?", [bookId]); +} + export async function clearVectorizationFlagsWithoutLocalChunks(): Promise { const database = await getDB(); const localDatabase = await getLocalDB(); const rows = await localDatabase.select<{ book_id: string }>( - "SELECT DISTINCT book_id FROM chunks", + `SELECT DISTINCT chunks.book_id + FROM chunks + INNER JOIN vector_index_provenance ON vector_index_provenance.book_id = chunks.book_id`, ); const bookIds = rows.map((row) => row.book_id).filter((bookId) => !!bookId); diff --git a/packages/core/src/db/database.ts b/packages/core/src/db/database.ts index b42428869..f29b1bbab 100644 --- a/packages/core/src/db/database.ts +++ b/packages/core/src/db/database.ts @@ -103,6 +103,9 @@ export { getChunks, insertChunks, deleteChunks, + getVectorIndexProvenance, + setVectorIndexProvenance, + deleteVectorIndexProvenance, clearVectorizationFlagsWithoutLocalChunks, } from "./chunk-queries"; diff --git a/packages/core/src/db/db-core.test.ts b/packages/core/src/db/db-core.test.ts new file mode 100644 index 000000000..059e7574a --- /dev/null +++ b/packages/core/src/db/db-core.test.ts @@ -0,0 +1,15 @@ +import { describe, expect, it } from "vitest"; +import { deserializeEmbedding, serializeEmbedding } from "./db-core"; + +describe("embedding serialization", () => { + it("decodes the JSON byte-array TEXT form persisted by Tauri SQL", () => { + const original = [0.125, -0.5, 1.25]; + const bytes = serializeEmbedding(original)!; + + expect(deserializeEmbedding(JSON.stringify(Array.from(bytes)))).toEqual(original); + }); + + it("rejects malformed byte lengths instead of constructing a partial float", () => { + expect(deserializeEmbedding("[1,2,3]")).toBeUndefined(); + }); +}); diff --git a/packages/core/src/db/db-core.ts b/packages/core/src/db/db-core.ts index b83f1b282..1ce4f4598 100644 --- a/packages/core/src/db/db-core.ts +++ b/packages/core/src/db/db-core.ts @@ -764,6 +764,19 @@ export async function initLocalDatabase(): Promise { ) `); + // Local-only metadata: embeddings are not synced, so the model identity that + // produced them must live beside the local chunk/vector indexes. + await database.execute(` + CREATE TABLE IF NOT EXISTS vector_index_provenance ( + book_id TEXT PRIMARY KEY, + model_kind TEXT NOT NULL, + model_id TEXT NOT NULL, + endpoint TEXT, + dimensions INTEGER NOT NULL, + created_at INTEGER NOT NULL + ) + `); + try { await database.execute( "ALTER TABLE chunks ADD COLUMN updated_at INTEGER NOT NULL DEFAULT 0", @@ -885,9 +898,24 @@ export function serializeEmbedding(embedding?: number[]): Uint8Array | null { /** Deserialize bytes back to float32 embedding array */ export function deserializeEmbedding(data: unknown): number[] | undefined { if (!data) return undefined; - // Data comes as an array of bytes from the SQL plugin - const bytes = data instanceof Uint8Array ? data : new Uint8Array(data as ArrayBuffer); + // Tauri's SQL plugin serializes Uint8Array parameters as a JSON byte array + // when writing to SQLite. That leaves a TEXT value such as "[63,172,...]" + // rather than a BLOB. Decode that persisted form before falling back to the + // native byte-array/BLOB forms used by other platforms. + let bytes: Uint8Array; + if (typeof data === "string") { + try { + const parsed = JSON.parse(data); + if (!Array.isArray(parsed)) return undefined; + bytes = new Uint8Array(parsed); + } catch { + return undefined; + } + } else { + bytes = data instanceof Uint8Array ? data : new Uint8Array(data as ArrayBuffer); + } if (bytes.length === 0) return undefined; + if (bytes.byteLength % Float32Array.BYTES_PER_ELEMENT !== 0) return undefined; const view = new Float32Array(bytes.buffer, bytes.byteOffset, bytes.byteLength / 4); return Array.from(view); } diff --git a/packages/core/src/rag/builtin-embedding-service.test.ts b/packages/core/src/rag/builtin-embedding-service.test.ts new file mode 100644 index 000000000..91abe34c7 --- /dev/null +++ b/packages/core/src/rag/builtin-embedding-service.test.ts @@ -0,0 +1,23 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { setLocalEmbeddingEngine } from "../ai/local-embedding-service"; +import { createBuiltinEmbeddingService } from "./builtin-embedding-service"; + +describe("createBuiltinEmbeddingService", () => { + const load = vi.fn(async () => undefined); + + beforeEach(() => { + setLocalEmbeddingEngine({ + init: () => undefined, + load, + generate: async (modelId, texts) => texts.map(() => (modelId === "bge-small-zh-v1.5" ? [1, 2] : [])), + dispose: async () => undefined, + clearCache: async () => undefined, + }); + }); + + it("generates query embeddings with the selected builtin model", async () => { + const service = createBuiltinEmbeddingService("bge-small-zh-v1.5"); + await expect(service.embed("查询文本")).resolves.toEqual([1, 2]); + expect(load).toHaveBeenCalledWith("bge-small-zh-v1.5", "Xenova/bge-small-zh-v1.5", undefined); + }); +}); diff --git a/packages/core/src/rag/builtin-embedding-service.ts b/packages/core/src/rag/builtin-embedding-service.ts new file mode 100644 index 000000000..a9bfb6c2d --- /dev/null +++ b/packages/core/src/rag/builtin-embedding-service.ts @@ -0,0 +1,28 @@ +import { generateLocalEmbeddings, loadEmbeddingPipeline } from "../ai/local-embedding-service"; +import { BUILTIN_EMBEDDING_MODELS } from "../ai/builtin-embedding-models"; +import type { QueryEmbeddingService } from "./search"; + +/** + * Adapts the platform's configured builtin embedding engine to RAG query search. + * Vectorization and query embedding therefore use the same Transformers.js pipeline. + */ +export function createBuiltinEmbeddingService(builtinModelId: string): QueryEmbeddingService { + const model = BUILTIN_EMBEDDING_MODELS.find((candidate) => candidate.id === builtinModelId); + if (!model) throw new Error(`Unknown built-in embedding model: ${builtinModelId}`); + + return { + provenance: { + kind: "builtin", + modelId: model.id, + dimensions: model.dimension, + }, + async embed(text: string): Promise { + await loadEmbeddingPipeline(builtinModelId); + const [embedding] = await generateLocalEmbeddings(builtinModelId, [text]); + if (!embedding?.length) { + throw new Error(`Built-in embedding model ${builtinModelId} returned no embedding.`); + } + return embedding; + }, + }; +} diff --git a/packages/core/src/rag/embedding-provenance.ts b/packages/core/src/rag/embedding-provenance.ts new file mode 100644 index 000000000..bd8c85e93 --- /dev/null +++ b/packages/core/src/rag/embedding-provenance.ts @@ -0,0 +1,11 @@ +import { normalizeEmbeddingEndpointUrl } from "../utils/api"; + +/** + * Persist the same canonical request URL that the embedding client uses. This + * treats an API base URL, a trailing slash, and an explicit /embeddings URL as + * the same endpoint identity. + */ +export function normalizeEmbeddingEndpoint(endpoint: string): string { + const trimmed = endpoint.trim(); + return trimmed ? normalizeEmbeddingEndpointUrl(trimmed) : ""; +} diff --git a/packages/core/src/rag/embedding-service.test.ts b/packages/core/src/rag/embedding-service.test.ts new file mode 100644 index 000000000..4ceb49a86 --- /dev/null +++ b/packages/core/src/rag/embedding-service.test.ts @@ -0,0 +1,31 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { EmbeddingService } from "./embedding-service"; + +describe("EmbeddingService endpoint handling", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("does not append embeddings twice when given a complete request URL", async () => { + const fetchMock = vi.spyOn(globalThis, "fetch").mockResolvedValue( + Response.json({ data: [{ index: 0, embedding: [0.1, 0.2] }] }), + ); + const service = new EmbeddingService({ + model: { + id: "text-embedding-test", + name: "test", + provider: "openai", + dimensions: 2, + maxTokens: 8192, + }, + apiKey: "test-key", + baseUrl: "https://example.test/api-openai/v1/embeddings", + }); + + await expect(service.embed("query")).resolves.toEqual([0.1, 0.2]); + expect(fetchMock).toHaveBeenCalledWith( + "https://example.test/api-openai/v1/embeddings", + expect.objectContaining({ method: "POST" }), + ); + }); +}); diff --git a/packages/core/src/rag/embedding-service.ts b/packages/core/src/rag/embedding-service.ts index 9c9cad499..014fb2831 100644 --- a/packages/core/src/rag/embedding-service.ts +++ b/packages/core/src/rag/embedding-service.ts @@ -3,7 +3,7 @@ */ import type { EmbeddingModel } from "../types"; -import { buildOpenAICompatibleUrl } from "../utils/api"; +import { normalizeEmbeddingEndpointUrl } from "../utils/api"; export interface EmbeddingConfig { model: EmbeddingModel; @@ -92,7 +92,11 @@ export class EmbeddingService { } private async callOpenAI(texts: string[]): Promise { - const url = buildOpenAICompatibleUrl(this.config.baseUrl, "embeddings"); + // Vector-model settings persist the complete embeddings request URL so the + // vectorization path can call it directly. Accept that form here as well as + // a plain OpenAI-compatible base URL; appending `/embeddings` unconditionally + // would otherwise produce `.../embeddings/embeddings` for Reader queries. + const url = normalizeEmbeddingEndpointUrl(this.config.baseUrl); const response = await this.fetchWithRetry(url, { method: "POST", diff --git a/packages/core/src/rag/index.ts b/packages/core/src/rag/index.ts index eb441266b..bde1b66bc 100644 --- a/packages/core/src/rag/index.ts +++ b/packages/core/src/rag/index.ts @@ -20,6 +20,8 @@ export type { export { EmbeddingService } from "./embedding-service"; export type { EmbeddingConfig } from "./embedding-service"; +export { createBuiltinEmbeddingService } from "./builtin-embedding-service"; +export { normalizeEmbeddingEndpoint } from "./embedding-provenance"; export { getEmbeddingModels, @@ -36,6 +38,8 @@ export { invalidateChunkCache, clearChunkCache, } from "./search"; +export type { QueryEmbeddingService } from "./search"; +export type { EmbeddingProvenance, VectorIndexProvenance } from "../types"; // Tokenizer exports export { tokenize, tokenizeQuery, getTokenFrequencies } from "./tokenizer"; diff --git a/packages/core/src/rag/search.test.ts b/packages/core/src/rag/search.test.ts new file mode 100644 index 000000000..ab751f52e --- /dev/null +++ b/packages/core/src/rag/search.test.ts @@ -0,0 +1,95 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { Chunk } from "../types"; +import { getChunks, getVectorIndexProvenance } from "../db/database"; +import { clearChunkCache, clearSearchConfiguration, configureSearch, search } from "./search"; + +vi.mock("../db/database", () => ({ + getChunks: vi.fn(), + getVectorIndexProvenance: vi.fn(), +})); + +const chunk: Chunk = { + id: "chunk-1", + bookId: "book-1", + chapterIndex: 0, + chapterTitle: "Chapter 1", + content: "A semantic search result about astronomy.", + tokenCount: 7, + startCfi: "", + endCfi: "", + embedding: [0.1, 0.2], +}; + +describe("RAG index provenance", () => { + beforeEach(() => { + vi.clearAllMocks(); + clearChunkCache(); + clearSearchConfiguration(); + vi.mocked(getChunks).mockResolvedValue([chunk]); + }); + + it("rejects vector search when the active query model differs from the book index", async () => { + vi.mocked(getVectorIndexProvenance).mockResolvedValue({ + bookId: "book-1", + kind: "builtin", + modelId: "all-MiniLM-L6-v2", + dimensions: 384, + createdAt: 1, + }); + configureSearch({ + provenance: { kind: "builtin", modelId: "bge-small-zh-v1.5", dimensions: 512 }, + embed: vi.fn(), + }); + + await expect( + search({ query: "astronomy", bookId: "book-1", mode: "vector", topK: 5, threshold: 0 }), + ).rejects.toThrow("Vector index model mismatch"); + }); + + it("marks hybrid results when vector retrieval is unavailable instead of silently returning BM25", async () => { + vi.mocked(getVectorIndexProvenance).mockResolvedValue(null); + configureSearch({ + provenance: { kind: "builtin", modelId: "all-MiniLM-L6-v2", dimensions: 384 }, + embed: vi.fn(), + }); + + const results = await search({ + query: "astronomy", + bookId: "book-1", + mode: "hybrid", + topK: 5, + threshold: 0, + }); + + expect(results).toHaveLength(1); + expect(results[0]).toMatchObject({ + matchType: "bm25", + vectorStatus: "unavailable", + vectorError: expect.stringContaining("no embedding provenance"), + }); + }); + + it("accepts remote endpoints that differ only by a trailing slash", async () => { + vi.mocked(getVectorIndexProvenance).mockResolvedValue({ + bookId: "book-1", + kind: "remote", + modelId: "text-embedding-v3", + endpoint: "https://dashscope.aliyuncs.com/compatible-mode/v1", + dimensions: 2, + createdAt: 1, + }); + configureSearch({ + provenance: { + kind: "remote", + modelId: "text-embedding-v3", + endpoint: "https://dashscope.aliyuncs.com/compatible-mode/v1/", + dimensions: 2, + }, + embed: vi.fn().mockResolvedValue([0.1, 0.2]), + }); + + await expect( + search({ query: "astronomy", bookId: "book-1", mode: "vector", topK: 5, threshold: 0 }), + ).resolves.toHaveLength(1); + }); +}); diff --git a/packages/core/src/rag/search.ts b/packages/core/src/rag/search.ts index 5a2ae1317..6a193eb69 100644 --- a/packages/core/src/rag/search.ts +++ b/packages/core/src/rag/search.ts @@ -1,4 +1,4 @@ -import { getChunks } from "../db/database"; +import { getChunks, getVectorIndexProvenance } from "../db/database"; /** * Hybrid search — vector + BM25 with configurable weighting * @@ -8,20 +8,64 @@ import { getChunks } from "../db/database"; * - In-memory caching for chunks and indexes * - Graceful fallback when vector search fails */ -import type { Chunk, SearchQuery, SearchResult } from "../types"; +import type { Chunk, EmbeddingProvenance, SearchQuery, SearchResult, VectorIndexProvenance } from "../types"; import { cosineSimilarity } from "./embedding"; -import type { EmbeddingService } from "./embedding-service"; +import { normalizeEmbeddingEndpoint } from "./embedding-provenance"; import { type InvertedIndex, buildInvertedIndex, searchInvertedIndex } from "./inverted-index"; import { tokenize, tokenizeQuery } from "./tokenizer"; import { getVectorDB, hasVectorDB } from "./vector-db"; -let embeddingService: EmbeddingService | null = null; +/** Minimal contract needed by vector search. Both remote and builtin engines implement it. */ +export interface QueryEmbeddingService { + embed(text: string): Promise; + provenance?: EmbeddingProvenance; +} + +let embeddingService: QueryEmbeddingService | null = null; /** Configure the embedding service for vector search */ -export function configureSearch(service: EmbeddingService): void { +export function configureSearch(service: QueryEmbeddingService, provenance?: EmbeddingProvenance): void { + if (provenance && !service.provenance) { + service.provenance = provenance; + } embeddingService = service; } +function provenanceMatches( + query: EmbeddingProvenance, + index: VectorIndexProvenance, +): boolean { + return ( + query.kind === index.kind && + query.modelId === index.modelId && + // Some remote providers do not expose dimensions in their saved settings. + // In that case the actual query vector is checked immediately after embedding. + (query.dimensions === 0 || index.dimensions === 0 || query.dimensions === index.dimensions) && + (query.kind === "builtin" || + normalizeEmbeddingEndpoint(query.endpoint || "") === normalizeEmbeddingEndpoint(index.endpoint || "")) + ); +} + +async function assertCompatibleIndex(bookId: string): Promise { + if (!embeddingService?.provenance) { + throw new Error("Embedding service has no model identity; cannot safely query this vector index."); + } + const index = await getVectorIndexProvenance(bookId); + if (!index) { + throw new Error( + "This book's vector index has no embedding provenance. Re-vectorize it before semantic search.", + ); + } + if (!provenanceMatches(embeddingService.provenance, index)) { + throw new Error( + `Vector index model mismatch: book uses ${index.kind}:${index.modelId} (${index.dimensions}d), ` + + `but the active query model is ${embeddingService.provenance.kind}:${embeddingService.provenance.modelId} ` + + `(${embeddingService.provenance.dimensions}d). Re-vectorize this book with the active model.`, + ); + } + return index; +} + /** Clear configured embedding service for vector search */ export function clearSearchConfiguration(): void { embeddingService = null; @@ -106,8 +150,16 @@ async function vectorSearch(query: SearchQuery): Promise { throw new Error("Embedding service not configured. Call configureSearch() first."); } + const indexProvenance = await assertCompatibleIndex(query.bookId); + // Get query embedding const queryEmbedding = await embeddingService.embed(query.query); + if (indexProvenance.dimensions > 0 && queryEmbedding.length !== indexProvenance.dimensions) { + throw new Error( + `Vector index dimension mismatch: book uses ${indexProvenance.dimensions}d, ` + + `but the query model returned ${queryEmbedding.length}d. Re-vectorize this book with the active model.`, + ); + } // Try vector database first (sqlite-vec) if (hasVectorDB()) { @@ -195,6 +247,13 @@ async function hybridSearch(query: SearchQuery): Promise { vectorResults = await vectorSearch(expandedQuery); } catch (err) { console.warn("[Search] Vector search failed, falling back to BM25 only:", err); + const vectorError = err instanceof Error ? err.message : String(err); + bm25Results = await bm25Search(expandedQuery); + return bm25Results.slice(0, query.topK).map((result) => ({ + ...result, + vectorStatus: "unavailable" as const, + vectorError, + })); } bm25Results = await bm25Search(expandedQuery); diff --git a/packages/core/src/rag/vectorize-trigger.test.ts b/packages/core/src/rag/vectorize-trigger.test.ts new file mode 100644 index 000000000..ffa32b7d4 --- /dev/null +++ b/packages/core/src/rag/vectorize-trigger.test.ts @@ -0,0 +1,14 @@ +import { describe, expect, it } from "vitest"; +import { canStoreInSharedVectorDB } from "./vectorize-trigger"; + +describe("shared sqlite-vec dimension guard", () => { + it("preserves a 384d book's acceleration index when a 1024d book is indexed", () => { + expect(canStoreInSharedVectorDB({ totalVectors: 120, dimension: 384 }, 1024)).toBe(false); + }); + + it("allows both books to remain searchable through their persisted chunk embeddings", () => { + expect(canStoreInSharedVectorDB({ totalVectors: 120, dimension: 384 }, 384)).toBe(true); + expect(canStoreInSharedVectorDB({ totalVectors: 120, dimension: 1024 }, 384)).toBe(false); + expect(canStoreInSharedVectorDB({ totalVectors: 120, dimension: 1024 }, 1024)).toBe(true); + }); +}); diff --git a/packages/core/src/rag/vectorize-trigger.ts b/packages/core/src/rag/vectorize-trigger.ts index b5c462a42..c893abcd6 100644 --- a/packages/core/src/rag/vectorize-trigger.ts +++ b/packages/core/src/rag/vectorize-trigger.ts @@ -1,7 +1,12 @@ import { BUILTIN_EMBEDDING_MODELS } from "../ai/builtin-embedding-models"; import { generateLocalEmbeddings, loadEmbeddingPipeline } from "../ai/local-embedding-service"; -import { deleteChunks, insertChunks } from "../db/database"; -import type { VectorizeProgress } from "../types"; +import { + deleteChunks, + deleteVectorIndexProvenance, + insertChunks, + setVectorIndexProvenance, +} from "../db/database"; +import type { EmbeddingProvenance, VectorizeProgress } from "../types"; /** * Vectorize Trigger — high-level service that orchestrates book vectorization. * Connects: chapter data → chunking → embedding → database indexing → state update. @@ -13,6 +18,7 @@ import type { VectorizeProgress } from "../types"; */ import { eventBus } from "../utils/event-bus"; import { chunkContent } from "./chunker"; +import { normalizeEmbeddingEndpoint } from "./embedding-provenance"; import type { ChapterData } from "./rag-types"; import { requestRemoteEmbeddingBatch } from "./remote-embedding"; import { invalidateChunkCache } from "./search"; @@ -30,6 +36,7 @@ export interface VectorizeTriggerConfig { url: string; apiKey: string; modelId: string; + dimension?: number; } | null; } @@ -45,6 +52,14 @@ export interface VectorizeTriggerCallbacks { /** Yield to the event loop so UI can repaint */ const yieldToUI = () => new Promise((r) => setTimeout(r, 0)); +/** sqlite-vec currently owns one global vector dimension. */ +export function canStoreInSharedVectorDB( + stats: { totalVectors: number; dimension: number }, + embeddingDimension: number, +): boolean { + return stats.totalVectors === 0 || stats.dimension === embeddingDimension; +} + /** * Trigger full vectorization for a book. * 1. Chunks chapter text into manageable pieces @@ -154,6 +169,7 @@ export async function triggerVectorizeBook( await yieldToUI(); await deleteChunks(bookId); + await deleteVectorIndexProvenance(bookId); // Insert in batches of 50 to avoid huge single transaction const insertBatchSize = 50; for (let i = 0; i < allChunks.length; i += insertBatchSize) { @@ -179,17 +195,34 @@ export async function triggerVectorizeBook( ]; }); + let storedInVectorDb = false; if (vectorRecords.length > 0) { // Detect actual embedding dimension and reinit vector DB if needed const detectedDimension = vectorRecords[0].embedding.length; - if (detectedDimension > 0 && vectorDB.reinit) { + const stats = await vectorDB.getStats(); + let canInsertIntoVectorDb = canStoreInSharedVectorDB(stats, detectedDimension); + if (detectedDimension > 0 && vectorDB.reinit && stats.totalVectors === 0) { await vectorDB.reinit(detectedDimension); + canInsertIntoVectorDb = true; + } else if (stats.dimension !== detectedDimension) { + // sqlite-vec currently has one global vector dimension. Never drop + // other books' acceleration indexes just to insert this one: search + // will safely use the persisted chunk embeddings for this book. + console.warn( + `[Vectorize] Skipping sqlite-vec insert for ${bookId}: ${detectedDimension}d index is incompatible with existing ${stats.dimension}d database.`, + ); + } + if (canInsertIntoVectorDb) { + await vectorDB.insert(vectorRecords); + storedInVectorDb = true; } - - await vectorDB.insert(vectorRecords); } - console.log(`[Vectorize] Stored ${vectorRecords.length} vectors in sqlite-vec`); + console.log( + storedInVectorDb + ? `[Vectorize] Stored ${vectorRecords.length} vectors in sqlite-vec` + : `[Vectorize] Persisted ${vectorRecords.length} vectors in chunk storage only`, + ); } else { console.warn("[Vectorize] Vector database not ready, skipping vector storage"); } @@ -198,6 +231,9 @@ export async function triggerVectorizeBook( } } + const provenance = getEmbeddingProvenance(config, allChunks[0]?.embedding?.length ?? 0); + await setVectorIndexProvenance({ bookId, ...provenance, createdAt: Date.now() }); + // Invalidate search cache so next query picks up new embeddings invalidateChunkCache(bookId); @@ -228,6 +264,25 @@ export async function triggerVectorizeBook( } } +function getEmbeddingProvenance( + config: VectorizeTriggerConfig, + actualDimensions: number, +): EmbeddingProvenance { + if (config.vectorModelMode === "builtin") { + const model = BUILTIN_EMBEDDING_MODELS.find((candidate) => candidate.id === config.selectedBuiltinModelId); + if (!model) throw new Error("Cannot save provenance for an unknown built-in embedding model."); + return { kind: "builtin", modelId: model.id, dimensions: actualDimensions || model.dimension }; + } + + if (!config.remoteModel) throw new Error("Cannot save provenance without a selected remote embedding model."); + return { + kind: "remote", + modelId: config.remoteModel.modelId, + endpoint: normalizeEmbeddingEndpoint(config.remoteModel.url), + dimensions: actualDimensions || config.remoteModel.dimension || 0, + }; +} + /** Generate embeddings using a built-in Transformers.js model (via Web Worker) */ async function generateBuiltinEmbeddings( chunks: Array<{ content: string; embedding?: number[] }>, diff --git a/packages/core/src/types/rag.ts b/packages/core/src/types/rag.ts index 3ccbbe5fb..bf2df05d1 100644 --- a/packages/core/src/types/rag.ts +++ b/packages/core/src/types/rag.ts @@ -18,6 +18,9 @@ export interface SearchResult { score: number; matchType: "vector" | "bm25" | "hybrid"; highlights?: string[]; // matched text segments + /** Present when hybrid search deliberately returned BM25-only results. */ + vectorStatus?: "unavailable"; + vectorError?: string; } export type SearchMode = "hybrid" | "vector" | "bm25"; @@ -38,6 +41,22 @@ export interface EmbeddingModel { provider: "openai" | "local"; } +/** + * Stable, non-secret identity of the embedding space used to build a book index. + * API keys deliberately never belong here. + */ +export interface EmbeddingProvenance { + kind: "builtin" | "remote"; + modelId: string; + endpoint?: string; + dimensions: number; +} + +export interface VectorIndexProvenance extends EmbeddingProvenance { + bookId: string; + createdAt: number; +} + export interface VectorConfig { model: EmbeddingModel; chunkSize: number; // default 300 tokens