Repository: tmc/langchaingo
Pull request #173 (2023-07-06 17:46:40)

feat: Embeddings - extending the create method to accept options
### 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
- adding ability to pass in the embedding options when creting the OpenAI embedder
- this resolves issue: https://github.com/tmc/langchaingo/issues/159


Patch:
diff --git a/embeddings/openai.go b/embeddings/openai.go
index 023d2e74..756023eb 100644
--- a/embeddings/openai.go
+++ b/embeddings/openai.go
@@ -20,19 +20,14 @@ type OpenAI struct {
 
 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()
+// NewOpenAI creates a new OpenAI with options. Options for client, strip new lines and batch.
+func NewOpenAI(opts ...Option) (OpenAI, error) {
+	o, err := applyClientOptions(opts...)
 	if err != nil {
 		return OpenAI{}, err
 	}
 
-	return OpenAI{
-		client:        client,
-		StripNewLines: true,
-		BatchSize:     defaultBatchSize,
-	}, nil
+	return o, nil
 }
 
 // EmbedDocuments creates one vector embedding for each of the texts.
diff --git a/embeddings/options.go b/embeddings/options.go
new file mode 100644
index 00000000..00039882
--- /dev/null
+++ b/embeddings/options.go
@@ -0,0 +1,55 @@
+package embeddings
+
+import (
+	"github.com/tmc/langchaingo/llms/openai"
+)
+
+const (
+	_defaultBatchSize     = 512
+	_defaultStripNewLines = true
+)
+
+// Option is a function type that can be used to modify the client.
+type Option func(p *OpenAI)
+
+// WithClient is an option for providing the LLM client.
+func WithClient(client openai.LLM) Option {
+	return func(p *OpenAI) {
+		p.client = &client
+	}
+}
+
+// WithStripNewLines is an option for specifying the should it strip new lines.
+func WithStripNewLines(stripNewLines bool) Option {
+	return func(p *OpenAI) {
+		p.StripNewLines = stripNewLines
+	}
+}
+
+// WithBatchSize is an option for specifying the batch size.
+func WithBatchSize(batchSize int) Option {
+	return func(p *OpenAI) {
+		p.BatchSize = batchSize
+	}
+}
+
+func applyClientOptions(opts ...Option) (OpenAI, error) {
+	o := &OpenAI{
+		StripNewLines: _defaultStripNewLines,
+		BatchSize:     _defaultBatchSize,
+	}
+
+	for _, opt := range opts {
+		opt(o)
+	}
+
+	if o.client == nil {
+		client, err := openai.New()
+		if err != nil {
+			return OpenAI{}, err
+		}
+		o.client = client
+	}
+
+	return *o, nil
+}


Test patch:
diff --git a/embeddings/openai_test.go b/embeddings/openai_test.go
index 88aabdac..13e91780 100644
--- a/embeddings/openai_test.go
+++ b/embeddings/openai_test.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/tmc/langchaingo/llms/openai"
 )
 
 func TestOpenaiEmbeddings(t *testing.T) {
@@ -25,3 +26,24 @@ func TestOpenaiEmbeddings(t *testing.T) {
 	require.NoError(t, err)
 	assert.Len(t, embeddings, 3)
 }
+
+func TestOpenaiEmbeddingsWithOptions(t *testing.T) {
+	t.Parallel()
+
+	if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" {
+		t.Skip("OPENAI_API_KEY not set")
+	}
+
+	client, err := openai.New()
+	require.NoError(t, err)
+
+	e, err := NewOpenAI(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)
+}