langchaingo/vectorstores/cloudsql/vectorstore_options.go

package cloudsql

import (
	"errors"

	"github.com/tmc/langchaingo/embeddings"
	"github.com/tmc/langchaingo/util/cloudsqlutil"
	"github.com/tmc/langchaingo/vectorstores"
)

const (
	defaultSchemaName         = "public"
	defaultIDColumn           = "langchain_id"
	defaultContentColumn      = "content"
	defaultEmbeddingColumn    = "embedding"
	defaultMetadataJSONColumn = "langchain_metadata"
	defaultK                  = 4
)

// VectorStoreOption is a function for creating new vector store
// with other than the default values.
type VectorStoreOption func(vs *VectorStore)

// WithSchemaName sets the VectorStore's schemaName field.
func WithSchemaName(schemaName string) VectorStoreOption {
	return func(v *VectorStore) {
		v.schemaName = schemaName
	}
}

// WithContentColumn sets VectorStore's the idColumn field.
func WithIDColumn(idColumn string) VectorStoreOption {
	return func(v *VectorStore) {
		v.idColumn = idColumn
	}
}

// WithMetadataJSONColumn sets VectorStore's the metadataJSONColumn field.
func WithMetadataJSONColumn(metadataJSONColumn string) VectorStoreOption {
	return func(v *VectorStore) {
		v.metadataJSONColumn = metadataJSONColumn
	}
}

// WithContentColumn sets the VectorStore's ContentColumn field.
func WithContentColumn(contentColumn string) VectorStoreOption {
	return func(v *VectorStore) {
		v.contentColumn = contentColumn
	}
}

// WithEmbeddingColumn sets the EmbeddingColumn field.
func WithEmbeddingColumn(embeddingColumn string) VectorStoreOption {
	return func(v *VectorStore) {
		v.embeddingColumn = embeddingColumn
	}
}

// WithMetadataColumns sets the VectorStore's MetadataColumns field.
func WithMetadataColumns(metadataColumns []string) VectorStoreOption {
	return func(v *VectorStore) {
		v.metadataColumns = metadataColumns
	}
}

// WithK sets the number of Documents to return from the VectorStore.
func WithK(k int) VectorStoreOption {
	return func(v *VectorStore) {
		v.k = k
	}
}

// WithDistanceStrategy sets the distance strategy used by the VectorStore.
func WithDistanceStrategy(distanceStrategy distanceStrategy) VectorStoreOption {
	return func(v *VectorStore) {
		v.distanceStrategy = distanceStrategy
	}
}

// VectorStoreOption applies the given VectorStore options to the
// VectorStore with a cloudsql Engine.
func applyCloudSQLVectorStoreOptions(engine cloudsqlutil.PostgresEngine,
	embedder embeddings.Embedder,
	tableName string,
	opts ...VectorStoreOption,
) (VectorStore, error) {
	// Check for required values.
	if engine.Pool == nil {
		return VectorStore{}, errors.New("missing vector store engine")
	}
	if embedder == nil {
		return VectorStore{}, errors.New("missing vector store embeder")
	}
	if tableName == "" {
		return VectorStore{}, errors.New("missing vector store table name")
	}
	defaultDistanceStrategy := CosineDistance{}

	vs := &VectorStore{
		engine:             engine,
		embedder:           embedder,
		tableName:          tableName,
		schemaName:         defaultSchemaName,
		idColumn:           defaultIDColumn,
		contentColumn:      defaultContentColumn,
		embeddingColumn:    defaultEmbeddingColumn,
		metadataJSONColumn: defaultMetadataJSONColumn,
		k:                  defaultK,
		distanceStrategy:   defaultDistanceStrategy,
		metadataColumns:    []string{},
	}
	for _, opt := range opts {
		opt(vs)
	}

	return *vs, nil
}

func applyOpts(options ...vectorstores.Option) vectorstores.Options {
	opts := vectorstores.Options{}
	for _, opt := range options {
		opt(&opts)
	}
	return opts
}