feat: make embedding dimensions configurable via EMBEDDINGS_DIMENSIONS

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)
 		}

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)
 		}

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{

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

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")
 		}

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 $$

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,131 @@ 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(&currentDim)
+	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"} {
+		if err := reconcileTableEmbedding(ctx, tx, table, columnExists, currentDim, dimensions); err != nil {
+			return 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
+}
+
+// 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,

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)
 	}
@@ -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)
+	db, err := NewPostgreSQL(ctx, testURI, testAuthz, embeddingDimensions)
 	require.NoError(t, err, "Failed to connect to test database")
 
 	// Register cleanup to close connection

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)
 		}