Repository: tmc/langchaingo
Pull request #75 (2023-04-29 17:39:03)

embeddings: refactor and move out of exp
 - Add embedder interface
https://github.com/hwchase17/langchainjs/blob/main/langchain/src/embeddings/base.ts
https://github.com/hwchase17/langchain/blob/master/langchain/embeddings/base.py
 - Add embedder implemetaion using openAI
https://github.com/hwchase17/langchain/blob/master/langchain/embeddings/openai.py
https://github.com/hwchase17/langchainjs/blob/main/langchain/src/embeddings/openai.ts

### 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 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 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.
- [x]  I am an existing contributor.



Patch:
diff --git a/exp/embeddings/doc.go b/embeddings/doc.go
similarity index 56%
rename from exp/embeddings/doc.go
rename to embeddings/doc.go
index a5d11239..c7ed6205 100644
--- a/exp/embeddings/doc.go
+++ b/embeddings/doc.go
@@ -1,2 +1,3 @@
-// Package embeddings contains a standard interface and implementations for creating numerical representation of textual data
+// Package embeddings contains a standard interface and implementations
+// for creating numerical representation of textual data
 package embeddings
diff --git a/embeddings/embedding.go b/embeddings/embedding.go
new file mode 100644
index 00000000..b63ef382
--- /dev/null
+++ b/embeddings/embedding.go
@@ -0,0 +1,11 @@
+package embeddings
+
+import "context"
+
+// Embedder is the interface for creating vector embeddings from texts.
+type Embedder interface {
+	// EmbedDocuments returns a vector for each text.
+	EmbedDocuments(ctx context.Context, texts []string) ([][]float64, error)
+	// EmbedQuery embeds a single text.
+	EmbedQuery(ctx context.Context, text string) ([]float64, error)
+}
diff --git a/embeddings/openai.go b/embeddings/openai.go
new file mode 100644
index 00000000..eec695d9
--- /dev/null
+++ b/embeddings/openai.go
@@ -0,0 +1,111 @@
+package embeddings
+
+import (
+	"context"
+	"strings"
+
+	"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
+
+	StripNewLines bool
+	BatchSize     int
+}
+
+var _ Embedder = OpenAI{}
+
+// NewOpenAI creates a new OpenAI with StripNewLines set to true and batch
+// size set to 512.
+func NewOpenAI() (OpenAI, error) {
+	client, err := openai.New()
+	if err != nil {
+		return OpenAI{}, err
+	}
+
+	return OpenAI{
+		client:        client,
+		StripNewLines: true,
+		BatchSize:     defaultBatchSize,
+	}, nil
+}
+
+// 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),
+		e.BatchSize,
+	)
+
+	embeddings := make([][]float64, 0, len(texts))
+	for _, texts := range batchedTexts {
+		curTextEmbeddings, err := e.client.CreateEmbedding(ctx, texts)
+		if err != nil {
+			return nil, err
+		}
+
+		textLengths := make([]int, 0, len(texts))
+		for _, text := range texts {
+			textLengths = append(textLengths, len(text))
+		}
+
+		combined, err := combineVectors(curTextEmbeddings, textLengths)
+		if err != nil {
+			return nil, err
+		}
+
+		embeddings = append(embeddings, combined)
+	}
+
+	return embeddings, nil
+}
+
+// EmbedQuery embeds a single text.
+func (e OpenAI) EmbedQuery(ctx context.Context, text string) ([]float64, error) {
+	if e.StripNewLines {
+		text = strings.ReplaceAll(text, "\n", " ")
+	}
+
+	embeddings, err := e.client.CreateEmbedding(ctx, []string{text})
+	if err != nil {
+		return nil, err
+	}
+
+	return embeddings[0], nil
+}
+
+func maybeRemoveNewLines(texts []string, removeNewLines bool) []string {
+	if !removeNewLines {
+		return texts
+	}
+
+	for i := 0; i < len(texts); i++ {
+		texts[i] = strings.ReplaceAll(texts[i], "\n", " ")
+	}
+
+	return texts
+}
+
+// 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)
+
+		for j := 0; j < len(runeText); j += batchSize {
+			if j+batchSize >= len(runeText) {
+				batchedTexts[i] = append(batchedTexts[i], string(runeText[j:]))
+				break
+			}
+
+			batchedTexts[i] = append(batchedTexts[i], string(runeText[j:j+batchSize]))
+		}
+	}
+
+	return batchedTexts
+}
diff --git a/embeddings/vector_math.go b/embeddings/vector_math.go
new file mode 100644
index 00000000..ce049e02
--- /dev/null
+++ b/embeddings/vector_math.go
@@ -0,0 +1,83 @@
+package embeddings
+
+import (
+	"errors"
+	"math"
+)
+
+var (
+	// ErrVectorsNotSameSize is returned if the vectors returned from the
+	// embeddings api have different sizes.
+	ErrVectorsNotSameSize = errors.New("vectors gotten not the same size")
+	// ErrAllTextsLenZero is returned if all texts to be embedded has the combined
+	// length of zero.
+	ErrAllTextsLenZero = errors.New("all texts have length 0")
+)
+
+func combineVectors(vectors [][]float64, weights []int) ([]float64, error) {
+	average, err := getAverage(vectors, weights)
+	if err != nil {
+		return nil, err
+	}
+
+	averageNorm := getNorm(average)
+	for i := 0; i < len(average); i++ {
+		average[i] /= averageNorm
+	}
+
+	return average, nil
+}
+
+// getAverage does the following calculation:
+//
+//	avg = sum(vectors * weights) / sum(weights).
+func getAverage(vectors [][]float64, weights []int) ([]float64, error) {
+	// Check that all vectors are the same size and get that size.
+	vectorLen := -1
+	for _, vector := range vectors {
+		if vectorLen == -1 {
+			vectorLen = len(vector)
+			continue
+		}
+
+		if len(vector) != vectorLen {
+			return nil, ErrVectorsNotSameSize
+		}
+	}
+
+	if vectorLen == -1 {
+		return []float64{}, nil
+	}
+
+	// Get the sum of the weights.
+	weightSum := 0
+	for _, weight := range weights {
+		weightSum += weight
+	}
+
+	if weightSum == 0 {
+		return nil, ErrAllTextsLenZero
+	}
+
+	average := make([]float64, vectorLen)
+	for i := 0; i < vectorLen; i++ {
+		for j := 0; j < len(vectors); j++ {
+			average[i] += vectors[j][i] * float64(weights[j])
+		}
+	}
+
+	for i := 0; i < len(average); i++ {
+		average[i] /= float64(weightSum)
+	}
+
+	return average, nil
+}
+
+func getNorm(v []float64) float64 {
+	var sum float64
+	for i := 0; i < len(v); i++ {
+		sum += v[i] * v[i]
+	}
+
+	return math.Sqrt(sum)
+}
diff --git a/exp/embeddings/embeddings.go b/exp/embeddings/embeddings.go
deleted file mode 100644
index de54e285..00000000
--- a/exp/embeddings/embeddings.go
+++ /dev/null
@@ -1,9 +0,0 @@
-package embeddings
-
-// Embeddings is the interface for creating vector embeddings from texts.
-type Embeddings interface {
-	// Returns vector for each text
-	EmbedDocuments(texts []string) ([][]float64, error)
-	// Embeds a single text
-	EmbedQuery(text string) ([]float64, error)
-}
diff --git a/exp/embeddings/openai.go b/exp/embeddings/openai.go
deleted file mode 100644
index eb17b0fc..00000000
--- a/exp/embeddings/openai.go
+++ /dev/null
@@ -1,80 +0,0 @@
-package embeddings
-
-import (
-	"math"
-	"strings"
-
-	"github.com/tmc/langchaingo/llms/openai"
-)
-
-type OpenAI struct {
-	client        *openai.LLM
-	StripNewLines bool
-	BatchSize     int
-}
-
-func NewOpenAI() (OpenAI, error) {
-	client, err := openai.New()
-	if err != nil {
-		return OpenAI{}, err
-	}
-
-	return OpenAI{
-		client:        client,
-		StripNewLines: true,
-		BatchSize:     512,
-	}, nil
-}
-
-func (e OpenAI) EmbedDocuments(texts []string) ([][]float64, error) {
-	removedNewLines := make([]string, 0)
-
-	for i := 0; i < len(texts); i++ {
-		curText := texts[i]
-		if e.StripNewLines {
-			curText = strings.ReplaceAll(curText, "\n", " ")
-		}
-
-		removedNewLines = append(removedNewLines, curText)
-	}
-
-	subPrompts := chunkArray(removedNewLines, e.BatchSize)
-
-	embeddings := make([][]float64, 0)
-	for i := 0; i < len(subPrompts); i++ {
-		curEmbeddings, err := e.client.CreateEmbedding(subPrompts[i])
-		if err != nil {
-			return [][]float64{}, err
-		}
-
-		embeddings = append(embeddings, curEmbeddings...)
-	}
-
-	return embeddings, nil
-}
-
-func (e OpenAI) EmbedQuery(text string) ([]float64, error) {
-	if e.StripNewLines {
-		text = strings.ReplaceAll(text, "\n", " ")
-	}
-
-	embeddings, err := e.client.CreateEmbedding([]string{text})
-	if err != nil {
-		return []float64{}, err
-	}
-
-	return embeddings[0], nil
-}
-
-func chunkArray[T any](arr []T, chunkSize int) [][]T {
-	var chunks [][]T
-	var chunkIndex int
-	for i, elem := range arr {
-		chunkIndex = int(math.Floor(float64(i) / float64(chunkSize)))
-		if chunkIndex >= len(chunks) {
-			chunks = append(chunks, []T{})
-		}
-		chunks[chunkIndex] = append(chunks[chunkIndex], elem)
-	}
-	return chunks
-}
diff --git a/exp/vector_stores/pinecone/pinecone.go b/exp/vector_stores/pinecone/pinecone.go
index 6cd6baeb..cdc6c829 100644
--- a/exp/vector_stores/pinecone/pinecone.go
+++ b/exp/vector_stores/pinecone/pinecone.go
@@ -7,7 +7,7 @@ import (
 	"os"
 
 	"github.com/google/uuid"
-	"github.com/tmc/langchaingo/exp/embeddings"
+	"github.com/tmc/langchaingo/embeddings"
 	"github.com/tmc/langchaingo/exp/vector_stores/pinecone/internal/pineconeClient"
 	"github.com/tmc/langchaingo/schema"
 )
@@ -18,12 +18,12 @@ var ErrMissingToken = errors.New("missing the Pinecone API key, set it in the PI
 
 type Client struct {
 	client     pineconeClient.Client
-	embeddings embeddings.Embeddings
+	embeddings embeddings.Embedder
 	textKey    string
 }
 
 // Environment for project is found in the pinecone console. Index name must not be larger then 45 characters.
-func NewPinecone(embeddings embeddings.Embeddings, environment, indexName string, dimensions int) (Client, error) {
+func NewPinecone(embeddings embeddings.Embedder, environment, indexName string, dimensions int) (Client, error) {
 	token := os.Getenv(pineconeEnvVrName)
 	if token == "" {
 		return Client{}, ErrMissingToken
@@ -60,7 +60,7 @@ func (p Client) AddDocuments(documents []schema.Document, documentIds []string,
 		texts = append(texts, documents[i].PageContent)
 	}
 
-	vectorData, err := p.embeddings.EmbedDocuments(texts)
+	vectorData, err := p.embeddings.EmbedDocuments(context.TODO(), texts)
 	if err != nil {
 		return err
 	}
@@ -85,7 +85,7 @@ func (p Client) AddDocuments(documents []schema.Document, documentIds []string,
 }
 
 func (p Client) SimilaritySearch(query string, numDocuments int, nameSpace string) ([]schema.Document, error) {
-	vector, err := p.embeddings.EmbedQuery(query)
+	vector, err := p.embeddings.EmbedQuery(context.TODO(), query)
 	if err != nil {
 		return []schema.Document{}, err
 	}
diff --git a/llms/openai/openaillm.go b/llms/openai/openaillm.go
index e30df5d2..5f2e8fdd 100644
--- a/llms/openai/openaillm.go
+++ b/llms/openai/openaillm.go
@@ -62,9 +62,9 @@ func (o *LLM) Chat(prompt string) (string, error) {
 	return result.Text, nil
 }
 
-// CreateEmbedding creates an embedding for the given input text.
-func (o *LLM) CreateEmbedding(inputTexts []string) ([][]float64, error) {
-	embeddings, err := o.client.CreateEmbedding(context.TODO(), &openaiclient.EmbeddingRequest{
+// CreateEmbedding creates embeddings for the given input texts.
+func (o *LLM) CreateEmbedding(ctx context.Context, inputTexts []string) ([][]float64, error) {
+	embeddings, err := o.client.CreateEmbedding(ctx, &openaiclient.EmbeddingRequest{
 		Input: inputTexts,
 	})
 


Test patch:
diff --git a/embeddings/openai_test.go b/embeddings/openai_test.go
new file mode 100644
index 00000000..57ce4485
--- /dev/null
+++ b/embeddings/openai_test.go
@@ -0,0 +1,52 @@
+package embeddings
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestOpenaiEmbeddings(t *testing.T) {
+	t.Parallel()
+
+	if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" {
+		t.Skip("OPENAI_API_KEY not set")
+	}
+	e, err := NewOpenAI()
+	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 TestBatchTexts(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		texts     []string
+		batchSize int
+		expected  [][]string
+	}{
+		{
+			texts:     []string{"foo bar zoo"},
+			batchSize: 4,
+			expected:  [][]string{{"foo ", "bar ", "zoo"}},
+		},
+		{
+			texts:     []string{"foo bar zoo", "foo"},
+			batchSize: 7,
+			expected:  [][]string{{"foo bar", " zoo"}, {"foo"}},
+		},
+	}
+
+	for _, tc := range cases {
+		assert.Equal(t, tc.expected, batchTexts(tc.texts, tc.batchSize))
+	}
+}
diff --git a/embeddings/vector_math_test.go b/embeddings/vector_math_test.go
new file mode 100644
index 00000000..d96d3d18
--- /dev/null
+++ b/embeddings/vector_math_test.go
@@ -0,0 +1,83 @@
+package embeddings
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCombineVectors(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		vectors  [][]float64
+		weights  []int
+		expected []float64
+	}{
+		{
+			vectors:  [][]float64{{10, 5, 10, -3}, {14, 12, 2, 3}},
+			weights:  []int{4, 6},
+			expected: []float64{0.7605787665953052, 0.5643003752158716, 0.3189523859915796, 0.0368021983836438},
+		},
+	}
+
+	for _, tc := range cases {
+		combined, err := combineVectors(tc.vectors, tc.weights)
+		assert.NoError(t, err)
+		assert.Equal(t, tc.expected, combined)
+	}
+}
+
+func TestGetAverage(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		vectors  [][]float64
+		weights  []int
+		expected []float64
+	}{
+		{
+			vectors:  [][]float64{{10, 5, 10}, {-10, -10, 20}},
+			weights:  []int{1, 1},
+			expected: []float64{0, -2.5, 15},
+		},
+		{
+			vectors:  [][]float64{{10, 5, 10}, {-10, -10, 20}},
+			weights:  []int{9, 1},
+			expected: []float64{8, 3.5, 11},
+		},
+		{
+			vectors:  [][]float64{{79, 26, -2}, {4, 78, 23}},
+			weights:  []int{3, 7},
+			expected: []float64{26.5, 62.4, 15.5},
+		},
+	}
+
+	for _, tc := range cases {
+		average, err := getAverage(tc.vectors, tc.weights)
+		assert.NoError(t, err)
+		assert.Equal(t, tc.expected, average)
+	}
+}
+
+func TestGetNorm(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		vector   []float64
+		expected float64
+	}{
+		{
+			vector:   []float64{5, 5, 5, 5},
+			expected: 10.0,
+		},
+		{
+			vector:   []float64{3, 4},
+			expected: 5.0,
+		},
+	}
+
+	for _, tc := range cases {
+		assert.Equal(t, tc.expected, getNorm(tc.vector))
+	}
+}
diff --git a/exp/embeddings/openai_test.go b/exp/embeddings/openai_test.go
deleted file mode 100644
index bd5fe109..00000000
--- a/exp/embeddings/openai_test.go
+++ /dev/null
@@ -1,26 +0,0 @@
-package embeddings
-
-import (
-	"os"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-)
-
-func TestOpenaiEmbeddings(t *testing.T) {
-	t.Parallel()
-
-	if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" {
-		t.Skip("OPENAI_API_KEY not set")
-	}
-	e, err := NewOpenAI()
-	require.NoError(t, err)
-
-	_, err = e.EmbedQuery("Hello world!")
-	require.NoError(t, err)
-
-	embeddings, err := e.EmbedDocuments([]string{"Hello world", "The world is ending", "bye bye"})
-	require.NoError(t, err)
-	assert.Len(t, embeddings, 3)
-}