From 3004a8755021305e8508d124e840b4b7523519ed Mon Sep 17 00:00:00 2001 From: Eva Date: Tue, 5 May 2026 17:29:46 +0700 Subject: [PATCH] fix: support Voyage 2048d schema setup --- src/core/ai/dims.ts | 23 +++++++++++++++++++---- src/core/ai/recipes/voyage.ts | 2 +- src/core/pglite-schema.ts | 17 ++++++++++++----- src/core/postgres-engine.ts | 3 ++- src/core/vector-index.ts | 16 ++++++++++++++++ test/ai/gateway.test.ts | 12 +++++++++++- test/ai/schema-templating.test.ts | 14 ++++++++++++++ 7 files changed, 75 insertions(+), 12 deletions(-) create mode 100644 src/core/vector-index.ts diff --git a/src/core/ai/dims.ts b/src/core/ai/dims.ts index 31e2d5bd5..d08c46053 100644 --- a/src/core/ai/dims.ts +++ b/src/core/ai/dims.ts @@ -11,12 +11,23 @@ import type { Implementation } from './types.ts'; +const VOYAGE_OUTPUT_DIMENSION_MODELS = new Set([ + 'voyage-4-large', + 'voyage-4', + 'voyage-4-lite', + 'voyage-3-large', + 'voyage-3.5', + 'voyage-3.5-lite', + 'voyage-code-3', +]); + /** * Build the providerOptions blob for embedMany() that pins output dimensions. * * Matryoshka providers (OpenAI text-embedding-3, Gemini embedding-001) can be - * asked to return reduced-dim vectors. openai-compatible + anthropic do not - * take a dimension parameter. + * asked to return reduced-dim vectors. Anthropic does not take a dimension + * parameter. Most openai-compatible providers do not either, but Voyage's + * OpenAI-compatible embeddings endpoint accepts `output_dimension`. */ export function dimsProviderOptions( implementation: Implementation, @@ -41,8 +52,12 @@ export function dimsProviderOptions( // Anthropic has no embedding model. return undefined; case 'openai-compatible': - // Ollama, LM Studio, vLLM, Voyage-compat, LiteLLM: no standard dim param. - // Users pick a model that natively outputs the dims they want. + // Most openai-compatible providers (Ollama, LM Studio, vLLM, LiteLLM) + // do not expose a standard dimensions knob. Voyage's compat endpoint is + // the exception: it accepts output_dimension and defaults to 1024 dims. + if (VOYAGE_OUTPUT_DIMENSION_MODELS.has(modelId)) { + return { openaiCompatible: { output_dimension: dims } }; + } return undefined; } } diff --git a/src/core/ai/recipes/voyage.ts b/src/core/ai/recipes/voyage.ts index 003ead0ed..2ce1a2ef2 100644 --- a/src/core/ai/recipes/voyage.ts +++ b/src/core/ai/recipes/voyage.ts @@ -16,7 +16,7 @@ export const voyage: Recipe = { }, touchpoints: { embedding: { - models: ['voyage-3-large', 'voyage-3', 'voyage-3-lite'], + models: ['voyage-4-large', 'voyage-4', 'voyage-3.5', 'voyage-3-large', 'voyage-3', 'voyage-3-lite'], default_dims: 1024, cost_per_1m_tokens_usd: 0.18, price_last_verified: '2026-04-20', diff --git a/src/core/pglite-schema.ts b/src/core/pglite-schema.ts index 60b17079b..abc0859d7 100644 --- a/src/core/pglite-schema.ts +++ b/src/core/pglite-schema.ts @@ -16,6 +16,8 @@ * test/edge-bundle.test.ts has a drift detection test. */ +import { applyChunkEmbeddingIndexPolicy } from './vector-index.ts'; + const PGLITE_SCHEMA_SQL_TEMPLATE = ` -- GBrain PGLite schema (local embedded Postgres) @@ -212,8 +214,8 @@ CREATE TABLE IF NOT EXISTS config ( INSERT INTO config (key, value) VALUES ('version', '1'), ('engine', 'pglite'), - ('embedding_model', 'text-embedding-3-large'), - ('embedding_dimensions', '1536'), + ('embedding_model', '__EMBEDDING_MODEL__'), + ('embedding_dimensions', '__EMBEDDING_DIMS__'), ('chunk_strategy', 'semantic') ON CONFLICT (key) DO NOTHING; @@ -532,9 +534,14 @@ DROP FUNCTION IF EXISTS update_page_search_vector_from_timeline(); * Defaults preserve v0.13 behavior (1536d + text-embedding-3-large). */ export function getPGLiteSchema(dims: number = 1536, model: string = 'text-embedding-3-large'): string { - return PGLITE_SCHEMA_SQL_TEMPLATE - .replace(/__EMBEDDING_DIMS__/g, String(dims)) - .replace(/__EMBEDDING_MODEL__/g, model); + const parsedDims = Number(dims); + if (!Number.isInteger(parsedDims) || parsedDims <= 0) { + throw new Error(`Invalid embedding dimensions: ${dims}`); + } + const sanitizedModel = String(model).replace(/'/g, "''"); + return applyChunkEmbeddingIndexPolicy(PGLITE_SCHEMA_SQL_TEMPLATE, parsedDims) + .replace(/__EMBEDDING_DIMS__/g, String(parsedDims)) + .replace(/__EMBEDDING_MODEL__/g, sanitizedModel); } /** Back-compat: pre-computed default-1536 schema for existing callers. */ diff --git a/src/core/postgres-engine.ts b/src/core/postgres-engine.ts index 4a25a3f84..e74c06bfc 100644 --- a/src/core/postgres-engine.ts +++ b/src/core/postgres-engine.ts @@ -4,6 +4,7 @@ import { MAX_SEARCH_LIMIT, clampSearchLimit } from './engine.ts'; import { runMigrations } from './migrate.ts'; import { SCHEMA_SQL } from './schema-embedded.ts'; import { verifySchema } from './schema-verify.ts'; +import { applyChunkEmbeddingIndexPolicy } from './vector-index.ts'; import type { Page, PageInput, PageFilters, PageType, Chunk, ChunkInput, StaleChunkRow, @@ -110,7 +111,7 @@ export class PostgresEngine implements BrainEngine { model = gw.getEmbeddingModel().split(':').slice(1).join(':') || model; } catch { /* gateway not yet configured — use defaults */ } - const sql = SCHEMA_SQL + const sql = applyChunkEmbeddingIndexPolicy(SCHEMA_SQL, dims) .replace(/vector\(1536\)/g, `vector(${dims})`) .replace(/'text-embedding-3-large'/g, `'${model}'`); diff --git a/src/core/vector-index.ts b/src/core/vector-index.ts new file mode 100644 index 000000000..7ba319507 --- /dev/null +++ b/src/core/vector-index.ts @@ -0,0 +1,16 @@ +export const PGVECTOR_HNSW_VECTOR_MAX_DIMS = 2000; + +const CHUNK_EMBEDDING_HNSW_INDEX = + 'CREATE INDEX IF NOT EXISTS idx_chunks_embedding ON content_chunks USING hnsw (embedding vector_cosine_ops);'; + +export function chunkEmbeddingIndexSql(dims: number): string { + if (dims <= PGVECTOR_HNSW_VECTOR_MAX_DIMS) return CHUNK_EMBEDDING_HNSW_INDEX; + return [ + '-- idx_chunks_embedding skipped: pgvector HNSW vector indexes support', + `-- at most ${PGVECTOR_HNSW_VECTOR_MAX_DIMS} dimensions; exact vector scans remain available.`, + ].join('\n'); +} + +export function applyChunkEmbeddingIndexPolicy(sql: string, dims: number): string { + return sql.replaceAll(CHUNK_EMBEDDING_HNSW_INDEX, chunkEmbeddingIndexSql(dims)); +} diff --git a/test/ai/gateway.test.ts b/test/ai/gateway.test.ts index 83d9b586e..e9f781047 100644 --- a/test/ai/gateway.test.ts +++ b/test/ai/gateway.test.ts @@ -151,8 +151,18 @@ describe('dims.dimsProviderOptions', () => { expect(opts).toBeUndefined(); }); - test('openai-compatible returns undefined (no standard dim param)', () => { + test('openai-compatible returns undefined for providers without a dim param', () => { const opts = dimsProviderOptions('openai-compatible', 'nomic-embed-text', 768); expect(opts).toBeUndefined(); }); + + test('Voyage openai-compatible returns output_dimension', () => { + const opts = dimsProviderOptions('openai-compatible', 'voyage-4-large', 2048); + expect(opts).toEqual({ openaiCompatible: { output_dimension: 2048 } }); + }); + + test('Voyage model without flexible dimensions returns undefined', () => { + const opts = dimsProviderOptions('openai-compatible', 'voyage-3-lite', 1024); + expect(opts).toBeUndefined(); + }); }); diff --git a/test/ai/schema-templating.test.ts b/test/ai/schema-templating.test.ts index bb140f8cb..2b6092875 100644 --- a/test/ai/schema-templating.test.ts +++ b/test/ai/schema-templating.test.ts @@ -14,6 +14,8 @@ describe('getPGLiteSchema', () => { const sql = getPGLiteSchema(768, 'gemini-embedding-001'); expect(sql).toMatch(/vector\(768\)/); expect(sql).toMatch(/'gemini-embedding-001'/); + expect(sql).toMatch(/\('embedding_model', 'gemini-embedding-001'\)/); + expect(sql).toMatch(/\('embedding_dimensions', '768'\)/); expect(sql).not.toMatch(/vector\(1536\)/); }); @@ -21,6 +23,18 @@ describe('getPGLiteSchema', () => { const sql = getPGLiteSchema(1024, 'voyage-3-large'); expect(sql).toMatch(/vector\(1024\)/); expect(sql).toMatch(/'voyage-3-large'/); + expect(sql).toMatch(/\('embedding_model', 'voyage-3-large'\)/); + expect(sql).toMatch(/\('embedding_dimensions', '1024'\)/); + expect(sql).toContain('idx_chunks_embedding ON content_chunks USING hnsw'); + }); + + test('Voyage 2048d skips unsupported HNSW index but keeps vector column', () => { + const sql = getPGLiteSchema(2048, 'voyage-4-large'); + expect(sql).toMatch(/vector\(2048\)/); + expect(sql).toMatch(/'voyage-4-large'/); + expect(sql).toMatch(/\('embedding_dimensions', '2048'\)/); + expect(sql).not.toContain('idx_chunks_embedding ON content_chunks USING hnsw'); + expect(sql).toContain('exact vector scans remain available'); }); test('PGLITE_SCHEMA_SQL back-compat constant is the default-dim schema', () => {