Repository: tmc/langchaingo
Pull request #203 (2023-07-18 13:07:07)

embeddings: Separating embeddings and allowing options to be passed
### PR Checklist

- [x] Read the [Contributing documentation](https://github.com/tmc/langchaingo/blob/main/CONTRIBUTING.md).
- [x] Read the [Code of conduct documentation](https://github.com/tmc/langchaingo/blob/main/CODE_OF_CONDUCT.md).
- [x] Name your Pull Request title clearly, concisely, and prefixed with the name of the primarily affected package you changed according to [Good commit messages](https://go.dev/doc/contribute#commit_messages) (such as `memory: add interfaces for X, Y` or `util: add whizzbang helpers`).
- [x] Check that there isn't already a PR that solves the problem the same way to avoid creating a duplicate.
- [x] Provide a description in this PR that addresses **what** the PR is solving, or reference the issue that it solves (e.g. `Fixes #123`).
- [x] Describes the source of new concepts.
- [x] References existing implementations as appropriate.
- [x] Contains test coverage for new functions.
- [x] Passes all [`golangci-lint`](https://golangci-lint.run/) checks.

## Description
- separating embeddings and adding options to be passed
- adding vertexai options to be passed in



Patch:
diff --git a/embeddings/embedding.go b/embeddings/embedding.go
index 7c3af443..029efbfa 100644
--- a/embeddings/embedding.go
+++ b/embeddings/embedding.go
@@ -13,7 +13,7 @@ type Embedder interface {
 	EmbedQuery(ctx context.Context, text string) ([]float64, error)
 }
 
-func maybeRemoveNewLines(texts []string, removeNewLines bool) []string {
+func MaybeRemoveNewLines(texts []string, removeNewLines bool) []string {
 	if !removeNewLines {
 		return texts
 	}
@@ -25,8 +25,8 @@ func maybeRemoveNewLines(texts []string, removeNewLines bool) []string {
 	return texts
 }
 
-// batchTexts splits strings by the length batchSize.
-func batchTexts(texts []string, batchSize int) [][]string {
+// BatchTexts splits strings by the length batchSize.
+func BatchTexts(texts []string, batchSize int) [][]string {
 	batchedTexts := make([][]string, len(texts))
 	for i, text := range texts {
 		runeText := []rune(text)
diff --git a/embeddings/openai.go b/embeddings/openai/openai.go
similarity index 76%
rename from embeddings/openai.go
rename to embeddings/openai/openai.go
index 756023eb..d7de2aee 100644
--- a/embeddings/openai.go
+++ b/embeddings/openai/openai.go
@@ -1,15 +1,13 @@
-package embeddings
+package openai
 
 import (
 	"context"
 	"strings"
 
+	"github.com/tmc/langchaingo/embeddings"
 	"github.com/tmc/langchaingo/llms/openai"
 )
 
-// defaultBatchSize is the default length of batches.
-const defaultBatchSize = 512
-
 // OpenAI is the embedder using the OpenAI api.
 type OpenAI struct {
 	client *openai.LLM
@@ -18,7 +16,7 @@ type OpenAI struct {
 	BatchSize     int
 }
 
-var _ Embedder = OpenAI{}
+var _ embeddings.Embedder = OpenAI{}
 
 // NewOpenAI creates a new OpenAI with options. Options for client, strip new lines and batch.
 func NewOpenAI(opts ...Option) (OpenAI, error) {
@@ -32,12 +30,12 @@ func NewOpenAI(opts ...Option) (OpenAI, error) {
 
 // EmbedDocuments creates one vector embedding for each of the texts.
 func (e OpenAI) EmbedDocuments(ctx context.Context, texts []string) ([][]float64, error) {
-	batchedTexts := batchTexts(
-		maybeRemoveNewLines(texts, e.StripNewLines),
+	batchedTexts := embeddings.BatchTexts(
+		embeddings.MaybeRemoveNewLines(texts, e.StripNewLines),
 		e.BatchSize,
 	)
 
-	embeddings := make([][]float64, 0, len(texts))
+	emb := make([][]float64, 0, len(texts))
 	for _, texts := range batchedTexts {
 		curTextEmbeddings, err := e.client.CreateEmbedding(ctx, texts)
 		if err != nil {
@@ -49,15 +47,15 @@ func (e OpenAI) EmbedDocuments(ctx context.Context, texts []string) ([][]float64
 			textLengths = append(textLengths, len(text))
 		}
 
-		combined, err := combineVectors(curTextEmbeddings, textLengths)
+		combined, err := embeddings.CombineVectors(curTextEmbeddings, textLengths)
 		if err != nil {
 			return nil, err
 		}
 
-		embeddings = append(embeddings, combined)
+		emb = append(emb, combined)
 	}
 
-	return embeddings, nil
+	return emb, nil
 }
 
 // EmbedQuery embeds a single text.
diff --git a/embeddings/options.go b/embeddings/openai/options.go
similarity index 98%
rename from embeddings/options.go
rename to embeddings/openai/options.go
index 00039882..1fd1300f 100644
--- a/embeddings/options.go
+++ b/embeddings/openai/options.go
@@ -1,4 +1,4 @@
-package embeddings
+package openai
 
 import (
 	"github.com/tmc/langchaingo/llms/openai"
diff --git a/embeddings/vector_math.go b/embeddings/vector_math.go
index ce049e02..a0e37e0c 100644
--- a/embeddings/vector_math.go
+++ b/embeddings/vector_math.go
@@ -14,7 +14,7 @@ var (
 	ErrAllTextsLenZero = errors.New("all texts have length 0")
 )
 
-func combineVectors(vectors [][]float64, weights []int) ([]float64, error) {
+func CombineVectors(vectors [][]float64, weights []int) ([]float64, error) {
 	average, err := getAverage(vectors, weights)
 	if err != nil {
 		return nil, err
diff --git a/embeddings/vertexai/options.go b/embeddings/vertexai/options.go
new file mode 100644
index 00000000..eb76eb43
--- /dev/null
+++ b/embeddings/vertexai/options.go
@@ -0,0 +1,55 @@
+package vertexai
+
+import (
+	"github.com/tmc/langchaingo/llms/vertexai"
+)
+
+const (
+	_defaultBatchSize     = 512
+	_defaultStripNewLines = true
+)
+
+// Option is a function type that can be used to modify the client.
+type Option func(p *VertexAIPaLM)
+
+// WithClient is an option for providing the LLM client.
+func WithClient(client vertexai.LLM) Option {
+	return func(p *VertexAIPaLM) {
+		p.client = &client
+	}
+}
+
+// WithStripNewLines is an option for specifying the should it strip new lines.
+func WithStripNewLines(stripNewLines bool) Option {
+	return func(p *VertexAIPaLM) {
+		p.StripNewLines = stripNewLines
+	}
+}
+
+// WithBatchSize is an option for specifying the batch size.
+func WithBatchSize(batchSize int) Option {
+	return func(p *VertexAIPaLM) {
+		p.BatchSize = batchSize
+	}
+}
+
+func applyClientOptions(opts ...Option) (*VertexAIPaLM, error) {
+	v := &VertexAIPaLM{
+		StripNewLines: _defaultStripNewLines,
+		BatchSize:     _defaultBatchSize,
+	}
+
+	for _, opt := range opts {
+		opt(v)
+	}
+
+	if v.client == nil {
+		client, err := vertexai.New()
+		if err != nil {
+			return nil, err
+		}
+		v.client = client
+	}
+
+	return v, nil
+}
diff --git a/embeddings/vertexai_palm.go b/embeddings/vertexai/vertexai_palm.go
similarity index 57%
rename from embeddings/vertexai_palm.go
rename to embeddings/vertexai/vertexai_palm.go
index 910833dc..c1fbe207 100644
--- a/embeddings/vertexai_palm.go
+++ b/embeddings/vertexai/vertexai_palm.go
@@ -1,44 +1,41 @@
-package embeddings
+package vertexai
 
 import (
 	"context"
 	"strings"
 
+	"github.com/tmc/langchaingo/embeddings"
 	"github.com/tmc/langchaingo/llms/vertexai"
 )
 
 // VertexAIPaLM is the embedder using the Google PaLM api.
-type VertexAIPaLM struct {
+type VertexAIPaLM struct { //nolint:revive
 	client *vertexai.LLM
 
 	StripNewLines bool
 	BatchSize     int
 }
 
-var _ Embedder = VertexAIPaLM{}
+var _ embeddings.Embedder = VertexAIPaLM{}
 
-// NewVertexAIPaLM creates a new VertexAI with StripNewLines set to true and batch
-// size set to 512.
-func NewVertexAIPaLM() (*VertexAIPaLM, error) {
-	client, err := vertexai.New()
+// NewVertexAIPaLM creates a new VertexAI with options. Options for client, strip new lines and batch size.
+func NewVertexAIPaLM(opts ...Option) (*VertexAIPaLM, error) {
+	v, err := applyClientOptions(opts...)
 	if err != nil {
 		return nil, err
 	}
-	return &VertexAIPaLM{
-		client:        client,
-		StripNewLines: true,
-		BatchSize:     defaultBatchSize,
-	}, nil
+
+	return v, nil
 }
 
 // EmbedDocuments creates one vector embedding for each of the texts.
 func (e VertexAIPaLM) EmbedDocuments(ctx context.Context, texts []string) ([][]float64, error) {
-	batchedTexts := batchTexts(
-		maybeRemoveNewLines(texts, e.StripNewLines),
+	batchedTexts := embeddings.BatchTexts(
+		embeddings.MaybeRemoveNewLines(texts, e.StripNewLines),
 		e.BatchSize,
 	)
 
-	embeddings := make([][]float64, 0, len(texts))
+	emb := make([][]float64, 0, len(texts))
 	for _, texts := range batchedTexts {
 		curTextEmbeddings, err := e.client.CreateEmbedding(ctx, texts)
 		if err != nil {
@@ -50,15 +47,15 @@ func (e VertexAIPaLM) EmbedDocuments(ctx context.Context, texts []string) ([][]f
 			textLengths = append(textLengths, len(text))
 		}
 
-		combined, err := combineVectors(curTextEmbeddings, textLengths)
+		combined, err := embeddings.CombineVectors(curTextEmbeddings, textLengths)
 		if err != nil {
 			return nil, err
 		}
 
-		embeddings = append(embeddings, combined)
+		emb = append(emb, combined)
 	}
 
-	return embeddings, nil
+	return emb, nil
 }
 
 // EmbedQuery embeds a single text.
@@ -67,10 +64,10 @@ func (e VertexAIPaLM) EmbedQuery(ctx context.Context, text string) ([]float64, e
 		text = strings.ReplaceAll(text, "\n", " ")
 	}
 
-	embeddings, err := e.client.CreateEmbedding(ctx, []string{text})
+	emb, err := e.client.CreateEmbedding(ctx, []string{text})
 	if err != nil {
 		return nil, err
 	}
 
-	return embeddings[0], nil
+	return emb[0], nil
 }


Test patch:
diff --git a/embeddings/embedding_test.go b/embeddings/embedding_test.go
index d06a3823..cce72ba7 100644
--- a/embeddings/embedding_test.go
+++ b/embeddings/embedding_test.go
@@ -27,6 +27,6 @@ func TestBatchTexts(t *testing.T) {
 	}
 
 	for _, tc := range cases {
-		assert.Equal(t, tc.expected, batchTexts(tc.texts, tc.batchSize))
+		assert.Equal(t, tc.expected, BatchTexts(tc.texts, tc.batchSize))
 	}
 }
diff --git a/embeddings/openai_test.go b/embeddings/openai/openai_test.go
similarity index 98%
rename from embeddings/openai_test.go
rename to embeddings/openai/openai_test.go
index 13e91780..6b1abb69 100644
--- a/embeddings/openai_test.go
+++ b/embeddings/openai/openai_test.go
@@ -1,4 +1,4 @@
-package embeddings
+package openai
 
 import (
 	"context"
diff --git a/embeddings/vector_math_test.go b/embeddings/vector_math_test.go
index d96d3d18..4bb1432e 100644
--- a/embeddings/vector_math_test.go
+++ b/embeddings/vector_math_test.go
@@ -22,7 +22,7 @@ func TestCombineVectors(t *testing.T) {
 	}
 
 	for _, tc := range cases {
-		combined, err := combineVectors(tc.vectors, tc.weights)
+		combined, err := CombineVectors(tc.vectors, tc.weights)
 		assert.NoError(t, err)
 		assert.Equal(t, tc.expected, combined)
 	}
diff --git a/embeddings/vertexai/vertexai_palm_test.go b/embeddings/vertexai/vertexai_palm_test.go
new file mode 100644
index 00000000..882522c2
--- /dev/null
+++ b/embeddings/vertexai/vertexai_palm_test.go
@@ -0,0 +1,49 @@
+package vertexai
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/tmc/langchaingo/llms/vertexai"
+)
+
+func TestVertexAIPaLMEmbeddings(t *testing.T) {
+	t.Parallel()
+
+	if gcpProjectID := os.Getenv("GOOGLE_CLOUD_PROJECT"); gcpProjectID == "" {
+		t.Skip("GOOGLE_CLOUD_PROJECT not set")
+	}
+	e, err := NewVertexAIPaLM()
+	require.NoError(t, err)
+
+	_, err = e.EmbedQuery(context.Background(), "Hello world!")
+	require.NoError(t, err)
+
+	embeddings, err := e.EmbedDocuments(context.Background(), []string{"Hello world", "The world is ending", "good bye"})
+	require.NoError(t, err)
+	assert.Len(t, embeddings, 3)
+}
+
+func TestVertexAIPaLMEmbeddingsWithOptions(t *testing.T) {
+	t.Parallel()
+
+	if gcpProjectID := os.Getenv("GOOGLE_CLOUD_PROJECT"); gcpProjectID == "" {
+		t.Skip("GOOGLE_CLOUD_PROJECT not set")
+	}
+
+	client, err := vertexai.New()
+	require.NoError(t, err)
+
+	e, err := NewVertexAIPaLM(WithClient(*client), WithBatchSize(1), WithStripNewLines(false))
+	require.NoError(t, err)
+
+	_, err = e.EmbedQuery(context.Background(), "Hello world!")
+	require.NoError(t, err)
+
+	embeddings, err := e.EmbedDocuments(context.Background(), []string{"Hello world"})
+	require.NoError(t, err)
+	assert.Len(t, embeddings, 1)
+}
diff --git a/embeddings/vertexai_palm_test.go b/embeddings/vertexai_palm_test.go
deleted file mode 100644
index 74b1213d..00000000
--- a/embeddings/vertexai_palm_test.go
+++ /dev/null
@@ -1,27 +0,0 @@
-package embeddings
-
-import (
-	"context"
-	"os"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-)
-
-func TestVertexAIPaLMEmbeddings(t *testing.T) {
-	t.Parallel()
-
-	if gcpProjectID := os.Getenv("GOOGLE_CLOUD_PROJECT"); gcpProjectID == "" {
-		t.Skip("GOOGLE_CLOUD_PROJECT not set")
-	}
-	e, err := NewVertexAIPaLM()
-	require.NoError(t, err)
-
-	_, err = e.EmbedQuery(context.Background(), "Hello world!")
-	require.NoError(t, err)
-
-	embeddings, err := e.EmbedDocuments(context.Background(), []string{"Hello world", "The world is ending", "good bye"})
-	require.NoError(t, err)
-	assert.Len(t, embeddings, 3)
-}
diff --git a/vectorstores/pinecone/pinecone_test.go b/vectorstores/pinecone/pinecone_test.go
index 6abda286..2dbae39e 100644
--- a/vectorstores/pinecone/pinecone_test.go
+++ b/vectorstores/pinecone/pinecone_test.go
@@ -9,7 +9,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 	"github.com/tmc/langchaingo/chains"
-	"github.com/tmc/langchaingo/embeddings"
+	openaiEmbeddings "github.com/tmc/langchaingo/embeddings/openai"
 	"github.com/tmc/langchaingo/llms/openai"
 	"github.com/tmc/langchaingo/schema"
 	"github.com/tmc/langchaingo/vectorstores"
@@ -81,7 +81,7 @@ func TestPineconeStoreRest(t *testing.T) {
 	t.Parallel()
 
 	environment, apiKey, indexName, projectName := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	storer, err := pinecone.New(
@@ -111,7 +111,7 @@ func TestPineconeStoreRestWithScoreThreshold(t *testing.T) {
 	t.Parallel()
 
 	environment, apiKey, indexName, projectName := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	storer, err := pinecone.New(
@@ -158,7 +158,7 @@ func TestSimilaritySearchWithInvalidScoreThreshold(t *testing.T) {
 	t.Parallel()
 
 	environment, apiKey, indexName, projectName := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	storer, err := pinecone.New(
@@ -201,7 +201,7 @@ func TestPineconeAsRetriever(t *testing.T) {
 	t.Parallel()
 
 	environment, apiKey, indexName, projectName := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := pinecone.New(
@@ -246,7 +246,7 @@ func TestPineconeAsRetrieverWithScoreThreshold(t *testing.T) {
 	t.Parallel()
 
 	environment, apiKey, indexName, projectName := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := pinecone.New(
@@ -297,7 +297,7 @@ func TestPineconeAsRetrieverWithMetadataFilterEqualsClause(t *testing.T) {
 	t.Parallel()
 
 	environment, apiKey, indexName, projectName := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := pinecone.New(
@@ -376,7 +376,7 @@ func TestPineconeAsRetrieverWithMetadataFilterInClause(t *testing.T) {
 	t.Parallel()
 
 	environment, apiKey, indexName, projectName := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := pinecone.New(
@@ -456,7 +456,7 @@ func TestPineconeAsRetrieverWithMetadataFilterNotSelected(t *testing.T) {
 	t.Parallel()
 
 	environment, apiKey, indexName, projectName := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := pinecone.New(
@@ -534,7 +534,7 @@ func TestPineconeAsRetrieverWithMetadataFilters(t *testing.T) {
 	t.Parallel()
 
 	environment, apiKey, indexName, projectName := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := pinecone.New(
diff --git a/vectorstores/weaviate/weaviate_test.go b/vectorstores/weaviate/weaviate_test.go
index c4825ae7..9103251c 100644
--- a/vectorstores/weaviate/weaviate_test.go
+++ b/vectorstores/weaviate/weaviate_test.go
@@ -10,7 +10,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 	"github.com/tmc/langchaingo/chains"
-	"github.com/tmc/langchaingo/embeddings"
+	openaiEmbeddings "github.com/tmc/langchaingo/embeddings/openai"
 	"github.com/tmc/langchaingo/llms/openai"
 	"github.com/tmc/langchaingo/schema"
 	"github.com/tmc/langchaingo/vectorstores"
@@ -66,7 +66,7 @@ func TestWeaviateStoreRest(t *testing.T) {
 	t.Parallel()
 
 	scheme, host := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := New(
@@ -101,7 +101,7 @@ func TestWeaviateStoreRestWithScoreThreshold(t *testing.T) {
 	t.Parallel()
 
 	scheme, host := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := New(
@@ -149,7 +149,7 @@ func TestSimilaritySearchWithInvalidScoreThreshold(t *testing.T) {
 	t.Parallel()
 
 	scheme, host := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := New(
@@ -193,7 +193,7 @@ func TestWeaviateAsRetriever(t *testing.T) {
 	t.Parallel()
 
 	scheme, host := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := New(
@@ -240,7 +240,7 @@ func TestWeaviateAsRetrieverWithScoreThreshold(t *testing.T) {
 	t.Parallel()
 
 	scheme, host := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := New(
@@ -292,7 +292,7 @@ func TestWeaviateAsRetrieverWithMetadataFilterEqualsClause(t *testing.T) {
 	t.Parallel()
 
 	scheme, host := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := New(
@@ -380,7 +380,7 @@ func TestWeaviateAsRetrieverWithMetadataFilterNotSelected(t *testing.T) {
 	t.Parallel()
 
 	scheme, host := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := New(
@@ -460,7 +460,7 @@ func TestWeaviateAsRetrieverWithMetadataFilters(t *testing.T) {
 	t.Parallel()
 
 	scheme, host := getValues(t)
-	e, err := embeddings.NewOpenAI()
+	e, err := openaiEmbeddings.NewOpenAI()
 	require.NoError(t, err)
 
 	store, err := New(