chore: add withClient to packages using openai (#86)

henomis · Jun 13, 2023 · 1cabb13 · 1cabb13
1 parent 4090215
commit 1cabb13
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 8 deletions.
diff --git a/embedder/openai/openai.go b/embedder/openai/openai.go
@@ -71,8 +71,8 @@ func New(model Model) *openAIEmbedder {
 	}
 }
 
-func (o *openAIEmbedder) WithAPIKey(apiKey string) *openAIEmbedder {
-	o.openAIClient = openai.NewClient(apiKey)
+func (o *openAIEmbedder) WithClient(client *openai.Client) *openAIEmbedder {
+	o.openAIClient = client
 	return o
 }
 

diff --git a/llm/openai/openai.go b/llm/openai/openai.go
@@ -95,8 +95,8 @@ func (o *openAI) WithStop(stop []string) *openAI {
 	return o
 }
 
-func (o *openAI) WithAPIKey(apiKey string) *openAI {
-	o.openAIClient = openai.NewClient(apiKey)
+func (o *openAI) WithClient(client *openai.Client) *openAI {
+	o.openAIClient = client
 	return o
 }
 

diff --git a/loader/whisper.go b/loader/whisper.go
@@ -27,8 +27,8 @@ func NewWhisperLoader(filename string) *whisperLoader {
 	}
 }
 
-func (w *whisperLoader) WithOpenAIApiKey(openAIApiKey string) *whisperLoader {
-	w.openAIClient = openai.NewClient(openAIApiKey)
+func (w *whisperLoader) WithClient(client *openai.Client) *whisperLoader {
+	w.openAIClient = client
 	return w
 }
 

diff --git a/prompt/whisper.go b/prompt/whisper.go
@@ -39,6 +39,11 @@ func NewPromptFromAudioFile(ctx context.Context, filePath string, audioResponseF
 	}, nil
 }
 
+func (p *whisperPrompt) WithClient(client *openai.Client) *whisperPrompt {
+	p.openAIClient = client
+	return p
+}
+
 func (p *whisperPrompt) Format(input types.M) error {
 	return nil
 }

diff --git a/transformer/dall-e.go b/transformer/dall-e.go
@@ -46,8 +46,8 @@ func NewDallE() *DallE {
 	}
 }
 
-func (d *DallE) WithAPIKey(apiKey string) *DallE {
-	d.openAIClient = openai.NewClient(apiKey)
+func (d *DallE) WithClient(client *openai.Client) *DallE {
+	d.openAIClient = client
 	return d
 }