diff --git a/src/core/embedding-dim-check.ts b/src/core/embedding-dim-check.ts index e863f748e..f4f1e7ee8 100644 --- a/src/core/embedding-dim-check.ts +++ b/src/core/embedding-dim-check.ts @@ -14,7 +14,7 @@ */ import type { BrainEngine } from './engine.ts'; -import { PGVECTOR_HNSW_VECTOR_MAX_DIMS } from './vector-index.ts'; +import { PGVECTOR_HNSW_VECTOR_MAX_DIMS, hnswMaxDimsForType } from './vector-index.ts'; import { gbrainPath } from './config.ts'; import { resolveRecipe } from './ai/model-resolver.ts'; import type { Recipe } from './ai/types.ts'; @@ -609,6 +609,17 @@ export function buildFactsAlterRecipe( const opclass = columnType === 'halfvec' ? 'halfvec_cosine_ops' : 'vector_cosine_ops'; const targetType = columnType === 'halfvec' ? `halfvec(${configuredDims})` : `vector(${configuredDims})`; const dimsChanged = columnDims !== configuredDims; + const hnswMaxDims = hnswMaxDimsForType(columnType); + const indexLines = configuredDims <= hnswMaxDims + ? [ + `CREATE INDEX idx_facts_embedding_hnsw`, + ` ON facts USING hnsw (embedding ${opclass})`, + ` WHERE embedding IS NOT NULL AND expired_at IS NULL;`, + ] + : [ + `-- Skip reindex. ${columnType}(${configuredDims}) exceeds pgvector's HNSW cap of ${hnswMaxDims};`, + `-- fact similarity falls back to exact scans.`, + ]; return [ `-- ALTER ${columnType}(${columnDims}) → ${columnType}(${configuredDims}) on indexed column.`, `-- HOLD a maintenance window: this rewrites every row's embedding.`, @@ -629,9 +640,7 @@ export function buildFactsAlterRecipe( : []), `ALTER TABLE facts ALTER COLUMN embedding TYPE ${targetType}`, ` USING embedding::${targetType};`, - `CREATE INDEX idx_facts_embedding_hnsw`, - ` ON facts USING hnsw (embedding ${opclass})`, - ` WHERE embedding IS NOT NULL AND expired_at IS NULL;`, + ...indexLines, ].join('\n'); } diff --git a/src/core/migrate.ts b/src/core/migrate.ts index 124b79c8b..9dd889af9 100644 --- a/src/core/migrate.ts +++ b/src/core/migrate.ts @@ -1,6 +1,7 @@ import type { BrainEngine } from './engine.ts'; import { slugifyPath } from './sync.ts'; import { getFtsLanguage } from './fts-language.ts'; +import { hnswMaxDimsForType } from './vector-index.ts'; /** * Schema migrations — run automatically on initSchema(). @@ -2276,11 +2277,19 @@ export const MIGRATIONS: Migration[] = [ useHalfvec = true; } - const vecType = useHalfvec ? 'HALFVEC' : 'VECTOR'; + const columnType = useHalfvec ? 'halfvec' : 'vector'; + const vecType = columnType.toUpperCase(); // HNSW operator class must match the column type: // VECTOR(n) → vector_cosine_ops // HALFVEC(n) → halfvec_cosine_ops const opclass = useHalfvec ? 'halfvec_cosine_ops' : 'vector_cosine_ops'; + const hnswMaxDims = hnswMaxDimsForType(columnType); + const factsEmbeddingIndexSql = embeddingDim <= hnswMaxDims + ? `CREATE INDEX IF NOT EXISTS idx_facts_embedding_hnsw + ON facts USING hnsw (embedding ${opclass}) + WHERE embedding IS NOT NULL AND expired_at IS NULL;` + : `-- idx_facts_embedding_hnsw skipped: pgvector HNSW ${columnType} indexes support + -- at most ${hnswMaxDims} dimensions; exact vector scans remain available.`; // FK to sources is added in a separate ALTER TABLE rather than inline // on the column. Inline `REFERENCES` worked on PGLite but silently // got dropped by postgres.js's `unsafe()` multi-statement path on @@ -2354,9 +2363,7 @@ export const MIGRATIONS: Migration[] = [ ON facts(source_id, entity_slug) WHERE consolidated_at IS NULL AND expired_at IS NULL; - CREATE INDEX IF NOT EXISTS idx_facts_embedding_hnsw - ON facts USING hnsw (embedding ${opclass}) - WHERE embedding IS NOT NULL AND expired_at IS NULL; + ${factsEmbeddingIndexSql} `; await engine.runMigration(40, factsDDL); @@ -2870,8 +2877,16 @@ export const MIGRATIONS: Migration[] = [ useHalfvec = true; } - const vecType = useHalfvec ? 'HALFVEC' : 'VECTOR'; + const columnType = useHalfvec ? 'halfvec' : 'vector'; + const vecType = columnType.toUpperCase(); const opclass = useHalfvec ? 'halfvec_cosine_ops' : 'vector_cosine_ops'; + const hnswMaxDims = hnswMaxDimsForType(columnType); + const queryCacheEmbeddingIndexSql = embeddingDim <= hnswMaxDims + ? `CREATE INDEX IF NOT EXISTS idx_query_cache_embedding_hnsw + ON query_cache USING hnsw (embedding ${opclass}) + WHERE embedding IS NOT NULL;` + : `-- idx_query_cache_embedding_hnsw skipped: pgvector HNSW ${columnType} indexes support + -- at most ${hnswMaxDims} dimensions; exact vector scans remain available.`; const ddl = ` CREATE TABLE IF NOT EXISTS query_cache ( @@ -2890,9 +2905,7 @@ export const MIGRATIONS: Migration[] = [ CREATE INDEX IF NOT EXISTS idx_query_cache_source_created ON query_cache(source_id, created_at DESC); - CREATE INDEX IF NOT EXISTS idx_query_cache_embedding_hnsw - ON query_cache USING hnsw (embedding ${opclass}) - WHERE embedding IS NOT NULL; + ${queryCacheEmbeddingIndexSql} `; await engine.runMigration(55, ddl); diff --git a/src/core/vector-index.ts b/src/core/vector-index.ts index 6110ea9dc..b89e9a63a 100644 --- a/src/core/vector-index.ts +++ b/src/core/vector-index.ts @@ -17,6 +17,7 @@ import type { BrainEngine } from './engine.ts'; export const PGVECTOR_HNSW_VECTOR_MAX_DIMS = 2000; +export const PGVECTOR_HNSW_HALFVEC_MAX_DIMS = 4000; const CHUNK_EMBEDDING_HNSW_INDEX = 'CREATE INDEX IF NOT EXISTS idx_chunks_embedding ON content_chunks USING hnsw (embedding vector_cosine_ops);'; @@ -29,6 +30,10 @@ export function chunkEmbeddingIndexSql(dims: number): string { ].join('\n'); } +export function hnswMaxDimsForType(columnType: 'vector' | 'halfvec'): number { + return columnType === 'halfvec' ? PGVECTOR_HNSW_HALFVEC_MAX_DIMS : PGVECTOR_HNSW_VECTOR_MAX_DIMS; +} + export function applyChunkEmbeddingIndexPolicy(sql: string, dims: number): string { return sql.replaceAll(CHUNK_EMBEDDING_HNSW_INDEX, chunkEmbeddingIndexSql(dims)); } diff --git a/test/embedding-dim-check-facts.test.ts b/test/embedding-dim-check-facts.test.ts index 2b0a9a8fc..8c5652729 100644 --- a/test/embedding-dim-check-facts.test.ts +++ b/test/embedding-dim-check-facts.test.ts @@ -122,9 +122,9 @@ describe('buildFactsAlterRecipe', () => { }); test('vector recipe uses vector_cosine_ops + vector(N) USING cast', () => { - const recipe = buildFactsAlterRecipe(1024, 2048, 'vector'); - expect(recipe).toContain('vector(2048)'); - expect(recipe).toContain('USING embedding::vector(2048)'); + const recipe = buildFactsAlterRecipe(1024, 1536, 'vector'); + expect(recipe).toContain('vector(1536)'); + expect(recipe).toContain('USING embedding::vector(1536)'); expect(recipe).toContain('vector_cosine_ops'); expect(recipe).not.toContain('halfvec_cosine_ops'); }); @@ -163,6 +163,14 @@ describe('buildFactsAlterRecipe', () => { expect(recipe).not.toContain('UPDATE facts SET embedding = NULL'); expect(recipe).toContain('USING embedding::vector(1536)'); }); + + test('halfvec recipe skips HNSW rebuild above pgvector cap', () => { + const recipe = buildFactsAlterRecipe(1536, 4096, 'halfvec'); + expect(recipe).toContain('halfvec(4096)'); + expect(recipe).toContain('Skip reindex'); + expect(recipe).toContain("exceeds pgvector's HNSW cap of 4000"); + expect(recipe).not.toMatch(/CREATE INDEX idx_facts_embedding_hnsw[\s\S]*USING hnsw/); + }); }); describe('FactsEmbeddingDimMismatchError', () => { diff --git a/test/facts-migration-dim.test.ts b/test/facts-migration-dim.test.ts index 6560e2f59..dca59f559 100644 --- a/test/facts-migration-dim.test.ts +++ b/test/facts-migration-dim.test.ts @@ -11,6 +11,7 @@ import { describe, test, expect, beforeAll, afterAll } from 'bun:test'; import { PGLiteEngine } from '../src/core/pglite-engine.ts'; +import { configureGateway, resetGateway } from '../src/core/ai/gateway.ts'; let engine: PGLiteEngine; @@ -93,4 +94,60 @@ describe('migration v45 facts column shape', () => { ); expect(after[0].udt_name).toBe(before[0].udt_name); }); + +}); + +describe('migration v45/v55 large-dim HNSW policy', () => { + let largeDimEngine: PGLiteEngine; + + beforeAll(async () => { + configureGateway({ + embedding_model: 'litellm:custom-4096d', + embedding_dimensions: 4096, + env: { ...process.env }, + }); + + largeDimEngine = new PGLiteEngine(); + await largeDimEngine.connect({}); + await largeDimEngine.initSchema(); + }); + + afterAll(async () => { + await largeDimEngine.disconnect(); + resetGateway(); + }); + + test('4096d init skips unsupported HNSW indexes but keeps vector columns', async () => { + const formatRows = await largeDimEngine.executeRaw<{ format_type: string }>( + `SELECT format_type(atttypid, atttypmod) AS format_type + FROM pg_attribute + WHERE attrelid = 'facts'::regclass AND attname = 'embedding'`, + ); + expect(formatRows[0]?.format_type).toMatch(/(halfvec|vector)\(4096\)/); + + const indexRows = await largeDimEngine.executeRaw<{ exists: boolean }>( + `SELECT EXISTS ( + SELECT 1 FROM pg_indexes + WHERE tablename = 'facts' + AND indexname = 'idx_facts_embedding_hnsw' + ) AS exists`, + ); + expect(indexRows[0]?.exists).toBe(false); + + const queryCacheFormatRows = await largeDimEngine.executeRaw<{ format_type: string }>( + `SELECT format_type(atttypid, atttypmod) AS format_type + FROM pg_attribute + WHERE attrelid = 'query_cache'::regclass AND attname = 'embedding'`, + ); + expect(queryCacheFormatRows[0]?.format_type).toMatch(/(halfvec|vector)\(4096\)/); + + const queryCacheIndexRows = await largeDimEngine.executeRaw<{ exists: boolean }>( + `SELECT EXISTS ( + SELECT 1 FROM pg_indexes + WHERE tablename = 'query_cache' + AND indexname = 'idx_query_cache_embedding_hnsw' + ) AS exists`, + ); + expect(queryCacheIndexRows[0]?.exists).toBe(false); + }, 60000); });