From bd6eb98a46240d2a164b411f7e2c2cada0cbb86a Mon Sep 17 00:00:00 2001 From: Simon Zhu Date: Wed, 25 Feb 2026 16:36:00 -0500 Subject: [PATCH 1/2] feat: make embedding dimensions configurable via EMBEDDINGS_DIMENSIONS The EMBEDDINGS_DIMENSIONS env var exists but was locked to 1536 by a hardcoded validation check, making it impossible to use embedding providers with different dimensions (e.g., Voyage AI at 1024). Changes: - Remove the hardcoded 1536 validation in config/validate.go - Add ensureVectorDimensions() that runs after migrations on startup - It compares the configured dimension with the actual pgvector column type (via pg_attribute.atttypmod) and, if they differ, atomically: 1. Drops HNSW indexes 2. Clears existing embeddings (they must be regenerated) 3. Alters both vector columns to the new dimension 4. Recreates HNSW indexes - Pass embeddingDimensions into NewPostgreSQL (0 = skip reconciliation) This means operators can switch providers by setting a single env var: EMBEDDINGS_DIMENSIONS=1024 EMBEDDINGS_PROVIDER=voyageai EMBEDDINGS_MODEL=voyage-3 On next startup the schema is reconciled automatically. Existing deployments using the default 1536 are unaffected (no-op reconciliation). --- internal/cli/export.go | 2 +- internal/cli/import.go | 2 +- internal/registry/config/validate.go | 5 - .../migrations/001_initial_schema.sql | 18 +-- internal/registry/database/postgres.go | 106 +++++++++++++++++- internal/registry/database/testutil.go | 4 +- internal/registry/registry_app.go | 2 +- 7 files changed, 115 insertions(+), 24 deletions(-) diff --git a/internal/cli/export.go b/internal/cli/export.go index 25fdb6a4..c286c371 100644 --- a/internal/cli/export.go +++ b/internal/cli/export.go @@ -41,7 +41,7 @@ var ExportCmd = &cobra.Command{ // so that the authn middleware extracts the session and stores in the context. (which the db can use to authorize queries) authz := auth.Authorizer{Authz: nil} - db, err := database.NewPostgreSQL(ctx, cfg.DatabaseURL, authz) + db, err := database.NewPostgreSQL(ctx, cfg.DatabaseURL, authz, 0) if err != nil { return fmt.Errorf("failed to connect to database: %w", err) } diff --git a/internal/cli/import.go b/internal/cli/import.go index 08c6fa70..e133b357 100644 --- a/internal/cli/import.go +++ b/internal/cli/import.go @@ -55,7 +55,7 @@ var ImportCmd = &cobra.Command{ // so that the authn middleware extracts the session and stores in the context. (which the db can use to authorize queries) authz := auth.Authorizer{Authz: nil} - db, err := database.NewPostgreSQL(ctx, cfg.DatabaseURL, authz) + db, err := database.NewPostgreSQL(ctx, cfg.DatabaseURL, authz, 0) if err != nil { return fmt.Errorf("failed to connect to database: %w", err) } diff --git a/internal/registry/config/validate.go b/internal/registry/config/validate.go index b9142972..cfbd709a 100644 --- a/internal/registry/config/validate.go +++ b/internal/registry/config/validate.go @@ -3,7 +3,6 @@ package config import "fmt" // Validate performs runtime validations on the loaded configuration. -// It is intentionally strict for embeddings to avoid runtime pgvector errors. func Validate(cfg *Config) error { if cfg == nil { return fmt.Errorf("config is nil") @@ -12,10 +11,6 @@ func Validate(cfg *Config) error { if cfg.Embeddings.Dimensions <= 0 { return fmt.Errorf("embeddings dimensions must be positive (got %d)", cfg.Embeddings.Dimensions) } - // Database schema currently provisions vector(1536). Reject mismatches early. - if cfg.Embeddings.Dimensions != 1536 { - return fmt.Errorf("embeddings dimensions must equal 1536 to match database schema (got %d)", cfg.Embeddings.Dimensions) - } if cfg.Embeddings.Model == "" { return fmt.Errorf("embeddings model must be specified when embeddings are enabled") } diff --git a/internal/registry/database/migrations/001_initial_schema.sql b/internal/registry/database/migrations/001_initial_schema.sql index 52b75dfc..8218ffe4 100644 --- a/internal/registry/database/migrations/001_initial_schema.sql +++ b/internal/registry/database/migrations/001_initial_schema.sql @@ -37,14 +37,14 @@ CREATE TABLE servers ( published_date TIMESTAMP WITH TIME ZONE, unpublished_date TIMESTAMP WITH TIME ZONE, - -- Semantic embedding columns for vector search - semantic_embedding vector(1536), + -- Semantic embedding metadata (the vector column itself is created by + -- ensureVectorDimensions at startup, sized to EMBEDDINGS_DIMENSIONS) semantic_embedding_provider TEXT, semantic_embedding_model TEXT, semantic_embedding_dimensions INTEGER, semantic_embedding_checksum TEXT, semantic_embedding_generated_at TIMESTAMPTZ, - + -- Primary key CONSTRAINT servers_pkey PRIMARY KEY (server_name, version) ); @@ -65,9 +65,6 @@ CREATE UNIQUE INDEX idx_unique_latest_per_server ON servers (server_name) WHERE CREATE INDEX idx_servers_json_remotes ON servers USING GIN((value->'remotes')); CREATE INDEX idx_servers_json_packages ON servers USING GIN((value->'packages')); --- HNSW index for semantic embedding similarity search -CREATE INDEX idx_servers_semantic_embedding_hnsw ON servers USING hnsw (semantic_embedding vector_cosine_ops); - -- Check constraints for servers ALTER TABLE servers ADD CONSTRAINT check_status_valid CHECK (status IN ('active', 'deprecated', 'deleted')); @@ -181,14 +178,14 @@ CREATE TABLE agents ( published_date TIMESTAMP WITH TIME ZONE, unpublished_date TIMESTAMP WITH TIME ZONE, - -- Semantic embedding columns for vector search - semantic_embedding vector(1536), + -- Semantic embedding metadata (the vector column itself is created by + -- ensureVectorDimensions at startup, sized to EMBEDDINGS_DIMENSIONS) semantic_embedding_provider TEXT, semantic_embedding_model TEXT, semantic_embedding_dimensions INTEGER, semantic_embedding_checksum TEXT, semantic_embedding_generated_at TIMESTAMPTZ, - + -- Primary key CONSTRAINT agents_pkey PRIMARY KEY (agent_name, version) ); @@ -205,9 +202,6 @@ CREATE INDEX idx_agents_published ON agents (published); -- Ensure only one version per agent is marked as latest CREATE UNIQUE INDEX idx_unique_latest_per_agent ON agents (agent_name) WHERE is_latest = true; --- HNSW index for semantic embedding similarity search -CREATE INDEX idx_agents_semantic_embedding_hnsw ON agents USING hnsw (semantic_embedding vector_cosine_ops); - -- Trigger function to auto-update updated_at CREATE OR REPLACE FUNCTION update_agents_updated_at() RETURNS TRIGGER AS $$ diff --git a/internal/registry/database/postgres.go b/internal/registry/database/postgres.go index d96374cd..b675ed5d 100644 --- a/internal/registry/database/postgres.go +++ b/internal/registry/database/postgres.go @@ -44,8 +44,10 @@ func (db *PostgreSQL) getExecutor(tx pgx.Tx) Executor { return db.pool } -// NewPostgreSQL creates a new instance of the PostgreSQL database -func NewPostgreSQL(ctx context.Context, connectionURI string, authz auth.Authorizer) (*PostgreSQL, error) { +// NewPostgreSQL creates a new instance of the PostgreSQL database. +// embeddingDimensions configures the pgvector column size for semantic search. +// Pass 0 to skip dimension reconciliation (e.g., when embeddings are disabled). +func NewPostgreSQL(ctx context.Context, connectionURI string, authz auth.Authorizer, embeddingDimensions int) (*PostgreSQL, error) { // Parse connection config for pool settings config, err := pgxpool.ParseConfig(connectionURI) if err != nil { @@ -81,12 +83,112 @@ func NewPostgreSQL(ctx context.Context, connectionURI string, authz auth.Authori return nil, fmt.Errorf("failed to run database migrations: %w", err) } + // After migrations, ensure the semantic_embedding vector columns exist with + // the configured dimension. The initial schema deliberately omits these + // columns so the dimension is never hardcoded in SQL. This function creates + // them on first run and reconciles the dimension if it changes later. + if embeddingDimensions > 0 { + if err := ensureVectorDimensions(ctx, conn.Conn(), embeddingDimensions); err != nil { + return nil, fmt.Errorf("failed to reconcile vector dimensions: %w", err) + } + } + return &PostgreSQL{ pool: pool, authz: authz, }, nil } +// ensureVectorDimensions creates or reconciles the semantic_embedding vector +// columns to match the configured EMBEDDINGS_DIMENSIONS. The initial SQL +// migration deliberately omits these columns so the dimension is never +// hardcoded in SQL. This function: +// - Creates the columns + HNSW indexes on first startup (fresh install) +// - Alters the columns if the configured dimension changes (provider switch) +// - No-ops if the columns already exist with the correct dimension +func ensureVectorDimensions(ctx context.Context, conn *pgx.Conn, dimensions int) error { + // Check whether the column exists and, if so, its current dimension. + // pgvector stores dimension+4 in atttypmod (the 4 accounts for the typmod header). + var currentDim int + columnExists := true + err := conn.QueryRow(ctx, ` + SELECT a.atttypmod - 4 + FROM pg_attribute a + JOIN pg_class c ON a.attrelid = c.oid + JOIN pg_namespace n ON c.relnamespace = n.oid + WHERE n.nspname = 'public' + AND c.relname = 'servers' + AND a.attname = 'semantic_embedding' + AND a.atttypmod > 0 + `).Scan(¤tDim) + if err != nil { + // Column doesn't exist yet (fresh install or pre-embeddings schema). + columnExists = false + } + + if columnExists && currentDim == dimensions { + return nil // Already correct, nothing to do. + } + + tx, err := conn.Begin(ctx) + if err != nil { + return fmt.Errorf("failed to begin transaction: %w", err) + } + defer func() { + if err := tx.Rollback(ctx); err != nil && !errors.Is(err, pgx.ErrTxClosed) { + log.Printf("Failed to rollback dimension migration: %v", err) + } + }() + + for _, table := range []string{"servers", "agents"} { + idx := fmt.Sprintf("idx_%s_semantic_embedding_hnsw", table) + + if columnExists { + // Column exists with wrong dimension: drop index, clear data, alter type. + log.Printf("Reconciling %s.semantic_embedding: %d -> %d dimensions", table, currentDim, dimensions) + + if _, err := tx.Exec(ctx, fmt.Sprintf("DROP INDEX IF EXISTS %s", idx)); err != nil { + return fmt.Errorf("failed to drop index %s: %w", idx, err) + } + if _, err := tx.Exec(ctx, fmt.Sprintf("UPDATE %s SET semantic_embedding = NULL", table)); err != nil { + return fmt.Errorf("failed to clear embeddings in %s: %w", table, err) + } + alter := fmt.Sprintf("ALTER TABLE %s ALTER COLUMN semantic_embedding TYPE vector(%d)", table, dimensions) + if _, err := tx.Exec(ctx, alter); err != nil { + return fmt.Errorf("failed to alter %s.semantic_embedding to vector(%d): %w", table, dimensions, err) + } + } else { + // Column doesn't exist: add it. + log.Printf("Creating %s.semantic_embedding as vector(%d)", table, dimensions) + + addCol := fmt.Sprintf("ALTER TABLE %s ADD COLUMN semantic_embedding vector(%d)", table, dimensions) + if _, err := tx.Exec(ctx, addCol); err != nil { + return fmt.Errorf("failed to add semantic_embedding to %s: %w", table, err) + } + } + + // Create (or recreate) the HNSW index. + createIdx := fmt.Sprintf( + "CREATE INDEX IF NOT EXISTS %s ON %s USING hnsw (semantic_embedding vector_cosine_ops)", + idx, table, + ) + if _, err := tx.Exec(ctx, createIdx); err != nil { + return fmt.Errorf("failed to create index %s: %w", idx, err) + } + } + + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("failed to commit vector dimension setup: %w", err) + } + + if columnExists { + log.Printf("Vector dimensions reconciled to %d (embeddings cleared; regeneration required)", dimensions) + } else { + log.Printf("Vector columns created with %d dimensions", dimensions) + } + return nil +} + func (db *PostgreSQL) ListServers( ctx context.Context, tx pgx.Tx, diff --git a/internal/registry/database/testutil.go b/internal/registry/database/testutil.go index 0a253a8f..5fbe662a 100644 --- a/internal/registry/database/testutil.go +++ b/internal/registry/database/testutil.go @@ -106,7 +106,7 @@ func ensureTemplateDB(ctx context.Context, adminConn *pgx.Conn) error { // Connect to template and run migrations (always) to keep it up-to-date // Create a permissive authz for tests testAuthz := createTestAuthz() - templateDB, err := NewPostgreSQL(ctx, templateURI, testAuthz) + templateDB, err := NewPostgreSQL(ctx, templateURI, testAuthz, 0) if err != nil { return fmt.Errorf("failed to connect to template database: %w", err) } @@ -180,7 +180,7 @@ func NewTestDB(t *testing.T) database.Database { // Create a permissive authz for tests testAuthz := createTestAuthz() - db, err := NewPostgreSQL(ctx, testURI, testAuthz) + db, err := NewPostgreSQL(ctx, testURI, testAuthz, 0) require.NoError(t, err, "Failed to connect to test database") // Register cleanup to close connection diff --git a/internal/registry/registry_app.go b/internal/registry/registry_app.go index 6c01709e..d74c2b0d 100644 --- a/internal/registry/registry_app.go +++ b/internal/registry/registry_app.go @@ -84,7 +84,7 @@ func App(_ context.Context, opts ...types.AppOptions) error { return fmt.Errorf("failed to create database via factory: %w", err) } } else { - baseDB, err := internaldb.NewPostgreSQL(ctx, cfg.DatabaseURL, authz) + baseDB, err := internaldb.NewPostgreSQL(ctx, cfg.DatabaseURL, authz, cfg.Embeddings.Dimensions) if err != nil { return fmt.Errorf("failed to connect to PostgreSQL: %w", err) } From 2c5d9ec0316e48e8a6edd125a7a9ade7339ed3f6 Mon Sep 17 00:00:00 2001 From: Simon Zhu Date: Thu, 5 Mar 2026 04:57:58 -0500 Subject: [PATCH 2/2] fix: resolve CI lint and test failures for configurable embedding dimensions Refactor ensureVectorDimensions to eliminate nestif complexity by extracting per-table logic into focused helper functions. Fix TestListServersSemanticSearch by adding NewTestDBWithEmbeddings that creates the semantic_embedding column with the correct dimension for embedding tests. Co-Authored-By: Claude Opus 4.6 --- internal/registry/api/handlers/v0/servers.go | 3 +- .../registry/api/handlers/v0/servers_test.go | 13 +-- internal/registry/database/postgres.go | 85 ++++++++++++------- internal/registry/database/testutil.go | 9 +- 4 files changed, 62 insertions(+), 48 deletions(-) diff --git a/internal/registry/api/handlers/v0/servers.go b/internal/registry/api/handlers/v0/servers.go index e17dabe8..375e9951 100644 --- a/internal/registry/api/handlers/v0/servers.go +++ b/internal/registry/api/handlers/v0/servers.go @@ -128,8 +128,7 @@ func RegisterServersEndpoints(api huma.API, pathPrefix string, registry service. }, nil }) - var tags []string - tags = []string{"servers"} + var tags = []string{"servers"} // List servers endpoint huma.Register(api, huma.Operation{ diff --git a/internal/registry/api/handlers/v0/servers_test.go b/internal/registry/api/handlers/v0/servers_test.go index 4b35f11f..71252aed 100644 --- a/internal/registry/api/handlers/v0/servers_test.go +++ b/internal/registry/api/handlers/v0/servers_test.go @@ -22,7 +22,6 @@ import ( "github.com/agentregistry-dev/agentregistry/pkg/registry/database" "github.com/danielgtaylor/huma/v2" "github.com/danielgtaylor/huma/v2/adapters/humago" - "github.com/jackc/pgx/v5" apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0" "github.com/modelcontextprotocol/registry/pkg/model" "github.com/stretchr/testify/assert" @@ -136,8 +135,7 @@ func TestListServersEndpoint(t *testing.T) { func TestListServersSemanticSearch(t *testing.T) { ctx := context.Background() - db := internaldb.NewTestDB(t) - ensureVectorExtension(t, db) + db := internaldb.NewTestDBWithEmbeddings(t, 3) cfg := config.NewConfig() cfg.Embeddings.Enabled = true @@ -436,15 +434,6 @@ func TestGetServerVersionEndpoint(t *testing.T) { } } -func ensureVectorExtension(t *testing.T, db database.Database) { - t.Helper() - err := db.InTransaction(context.Background(), func(ctx context.Context, tx pgx.Tx) error { - _, execErr := tx.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector") - return execErr - }) - require.NoError(t, err, "failed to ensure pgvector extension for tests") -} - type stubEmbeddingProvider struct { mu sync.Mutex vectors map[string][]float32 diff --git a/internal/registry/database/postgres.go b/internal/registry/database/postgres.go index b675ed5d..f714a337 100644 --- a/internal/registry/database/postgres.go +++ b/internal/registry/database/postgres.go @@ -141,39 +141,8 @@ func ensureVectorDimensions(ctx context.Context, conn *pgx.Conn, dimensions int) }() for _, table := range []string{"servers", "agents"} { - idx := fmt.Sprintf("idx_%s_semantic_embedding_hnsw", table) - - if columnExists { - // Column exists with wrong dimension: drop index, clear data, alter type. - log.Printf("Reconciling %s.semantic_embedding: %d -> %d dimensions", table, currentDim, dimensions) - - if _, err := tx.Exec(ctx, fmt.Sprintf("DROP INDEX IF EXISTS %s", idx)); err != nil { - return fmt.Errorf("failed to drop index %s: %w", idx, err) - } - if _, err := tx.Exec(ctx, fmt.Sprintf("UPDATE %s SET semantic_embedding = NULL", table)); err != nil { - return fmt.Errorf("failed to clear embeddings in %s: %w", table, err) - } - alter := fmt.Sprintf("ALTER TABLE %s ALTER COLUMN semantic_embedding TYPE vector(%d)", table, dimensions) - if _, err := tx.Exec(ctx, alter); err != nil { - return fmt.Errorf("failed to alter %s.semantic_embedding to vector(%d): %w", table, dimensions, err) - } - } else { - // Column doesn't exist: add it. - log.Printf("Creating %s.semantic_embedding as vector(%d)", table, dimensions) - - addCol := fmt.Sprintf("ALTER TABLE %s ADD COLUMN semantic_embedding vector(%d)", table, dimensions) - if _, err := tx.Exec(ctx, addCol); err != nil { - return fmt.Errorf("failed to add semantic_embedding to %s: %w", table, err) - } - } - - // Create (or recreate) the HNSW index. - createIdx := fmt.Sprintf( - "CREATE INDEX IF NOT EXISTS %s ON %s USING hnsw (semantic_embedding vector_cosine_ops)", - idx, table, - ) - if _, err := tx.Exec(ctx, createIdx); err != nil { - return fmt.Errorf("failed to create index %s: %w", idx, err) + if err := reconcileTableEmbedding(ctx, tx, table, columnExists, currentDim, dimensions); err != nil { + return err } } @@ -189,6 +158,56 @@ func ensureVectorDimensions(ctx context.Context, conn *pgx.Conn, dimensions int) return nil } +// reconcileTableEmbedding creates or alters the semantic_embedding column for a +// single table and (re)creates its HNSW index. +func reconcileTableEmbedding(ctx context.Context, tx pgx.Tx, table string, columnExists bool, currentDim, dimensions int) error { + idx := fmt.Sprintf("idx_%s_semantic_embedding_hnsw", table) + if columnExists { + return alterEmbeddingColumn(ctx, tx, table, idx, currentDim, dimensions) + } + return addEmbeddingColumn(ctx, tx, table, idx, dimensions) +} + +// alterEmbeddingColumn changes an existing semantic_embedding column to a new dimension. +func alterEmbeddingColumn(ctx context.Context, tx pgx.Tx, table, idx string, currentDim, dimensions int) error { + log.Printf("Reconciling %s.semantic_embedding: %d -> %d dimensions", table, currentDim, dimensions) + + if _, err := tx.Exec(ctx, fmt.Sprintf("DROP INDEX IF EXISTS %s", idx)); err != nil { + return fmt.Errorf("failed to drop index %s: %w", idx, err) + } + if _, err := tx.Exec(ctx, fmt.Sprintf("UPDATE %s SET semantic_embedding = NULL", table)); err != nil { + return fmt.Errorf("failed to clear embeddings in %s: %w", table, err) + } + alter := fmt.Sprintf("ALTER TABLE %s ALTER COLUMN semantic_embedding TYPE vector(%d)", table, dimensions) + if _, err := tx.Exec(ctx, alter); err != nil { + return fmt.Errorf("failed to alter %s.semantic_embedding to vector(%d): %w", table, dimensions, err) + } + return createEmbeddingIndex(ctx, tx, table, idx) +} + +// addEmbeddingColumn adds a new semantic_embedding column with the given dimension. +func addEmbeddingColumn(ctx context.Context, tx pgx.Tx, table, idx string, dimensions int) error { + log.Printf("Creating %s.semantic_embedding as vector(%d)", table, dimensions) + + addCol := fmt.Sprintf("ALTER TABLE %s ADD COLUMN semantic_embedding vector(%d)", table, dimensions) + if _, err := tx.Exec(ctx, addCol); err != nil { + return fmt.Errorf("failed to add semantic_embedding to %s: %w", table, err) + } + return createEmbeddingIndex(ctx, tx, table, idx) +} + +// createEmbeddingIndex creates the HNSW index for semantic embedding similarity search. +func createEmbeddingIndex(ctx context.Context, tx pgx.Tx, table, idx string) error { + createIdx := fmt.Sprintf( + "CREATE INDEX IF NOT EXISTS %s ON %s USING hnsw (semantic_embedding vector_cosine_ops)", + idx, table, + ) + if _, err := tx.Exec(ctx, createIdx); err != nil { + return fmt.Errorf("failed to create index %s: %w", idx, err) + } + return nil +} + func (db *PostgreSQL) ListServers( ctx context.Context, tx pgx.Tx, diff --git a/internal/registry/database/testutil.go b/internal/registry/database/testutil.go index 5fbe662a..ad5da48b 100644 --- a/internal/registry/database/testutil.go +++ b/internal/registry/database/testutil.go @@ -132,6 +132,13 @@ func ensureVectorExtension(ctx context.Context, uri string) error { // The template database has migrations pre-applied, so each test is fast. // Requires PostgreSQL to be running on localhost:5432 (e.g., via docker-compose). func NewTestDB(t *testing.T) database.Database { + return NewTestDBWithEmbeddings(t, 0) +} + +// NewTestDBWithEmbeddings is like NewTestDB but also creates the semantic_embedding +// vector columns with the given dimension. Use this for tests that exercise +// embedding/semantic-search functionality. Pass 0 to skip (same as NewTestDB). +func NewTestDBWithEmbeddings(t *testing.T, embeddingDimensions int) database.Database { t.Helper() ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) @@ -180,7 +187,7 @@ func NewTestDB(t *testing.T) database.Database { // Create a permissive authz for tests testAuthz := createTestAuthz() - db, err := NewPostgreSQL(ctx, testURI, testAuthz, 0) + db, err := NewPostgreSQL(ctx, testURI, testAuthz, embeddingDimensions) require.NoError(t, err, "Failed to connect to test database") // Register cleanup to close connection