Add the API key via a parameter instead of environment variable

FluxCapacitor2 · Nov 18, 2024 · e7270e1 · e7270e1
1 parent 2aaf3b7
commit e7270e1
Show file tree

Hide file tree

Showing 5 changed files with 6 additions and 4 deletions.
diff --git a/app/config/config.go b/app/config/config.go
@@ -17,6 +17,7 @@ type Config struct {
 	} `yaml:"db"`
 	Embeddings struct {
 		OpenAIBaseURL string `yaml:"openaiBaseUrl"`
+		APIKey        string `yaml:"apiKey"`
 		Model         string
 		Dimensions    int
 		ChunkSize     int `yaml:"chunkSize"`

diff --git a/app/embedding/embedding.go b/app/embedding/embedding.go
@@ -9,9 +9,9 @@ import (
 	"github.com/tmc/langchaingo/llms/openai"
 )
 
-func GetEmbeddings(openAIBaseURL string, model string, chunk string) ([]float32, error) {
+func GetEmbeddings(openAIBaseURL string, model string, apiKey string, chunk string) ([]float32, error) {
 
-	llm, err := openai.New(openai.WithBaseURL(openAIBaseURL), openai.WithEmbeddingModel(model))
+	llm, err := openai.New(openai.WithBaseURL(openAIBaseURL), openai.WithEmbeddingModel(model), openai.WithToken(apiKey))
 	if err != nil {
 		return nil, fmt.Errorf("error setting up LLM for embedding: %v", err)
 	}

diff --git a/app/processqueue.go b/app/processqueue.go
@@ -82,7 +82,7 @@ func processEmbedQueue(db database.Database, config *config.Config, src config.S
 		}
 	}
 
-	vector, err := embedding.GetEmbeddings(config.Embeddings.OpenAIBaseURL, config.Embeddings.Model, item.Content)
+	vector, err := embedding.GetEmbeddings(config.Embeddings.OpenAIBaseURL, config.Embeddings.Model, config.Embeddings.APIKey, item.Content)
 	if err != nil {
 		fmt.Printf("error getting embeddings: %v\n", err)
 		markFailure()

diff --git a/app/server/server.go b/app/server/server.go
@@ -134,7 +134,7 @@ func Start(db database.Database, config *config.Config) {
 
 		if q != "" && src != nil && len(src) > 0 {
 
-			vector, err := embedding.GetEmbeddings(config.Embeddings.OpenAIBaseURL, config.Embeddings.Model, q)
+			vector, err := embedding.GetEmbeddings(config.Embeddings.OpenAIBaseURL, config.Embeddings.Model, config.Embeddings.APIKey, q)
 			if err != nil {
 				response = &httpResponse{
 					status:  500,

diff --git a/config-sample.yml b/config-sample.yml
@@ -51,6 +51,7 @@ embeddings:
   # openaiBaseUrl: https://api.openai.com/v1/
   # model: text-embedding-3-small
   # dimensions: 1536
+  # apiKey: sk-*************************************
 
   # You can also use any OpenAI-compatible API, like a local Ollama server:
   openaiBaseUrl: http://localhost:11434/v1/