Open AI module: provide API key at query time

adapters/handlers/rest/middlewares.go

@@ -12,8 +12,10 @@
 package rest
 
 import (
+	"context"
 	"fmt"
 	"net/http"
+	"strings"
 
 	"github.com/rs/cors"
 	"github.com/semi-technologies/weaviate/adapters/handlers/rest/state"
@@ -91,6 +93,7 @@ func makeSetupGlobalMiddleware(appState *state.State) func(http.Handler) http.Ha
 		handler = addLiveAndReadyness(handler)
 		handler = addHandleRoot(handler)
 		handler = makeAddModuleHandlers(appState.Modules)(handler)
+		handler = addInjectHeadersIntoContext(handler)
 
 		return handler
 	}
@@ -122,6 +125,25 @@ func addPreflight(next http.Handler) http.Handler {
 	})
 }
 
+func addInjectHeadersIntoContext(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ctx := r.Context()
+		changed := false
+		for k, v := range r.Header {
+			if strings.HasPrefix(k, "X-") {
+				ctx = context.WithValue(ctx, k, v)
+				changed = true
+			}
+		}
+
+		if changed {
+			next.ServeHTTP(w, r.Clone(ctx))
+		} else {
+			next.ServeHTTP(w, r)
+		}
+	})
+}
+
 func addLiveAndReadyness(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.String() == "/v1/.well-known/live" {

modules/text2vec-openai/clients/vectorizer.go

@@ -87,7 +87,11 @@ func (v *vectorizer) vectorize(ctx context.Context, input string,
 	if err != nil {
 		return nil, errors.Wrap(err, "create POST request")
 	}
-	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", v.apiKey))
+	apiKey, err := v.getApiKey(ctx)
+	if err != nil {
+		return nil, errors.Wrapf(err, "OpenAI API Key")
+	}
+	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", apiKey))
 	req.Header.Add("Content-Type", "application/json")
 
 	res, err := v.httpClient.Do(req)
@@ -124,6 +128,20 @@ func (v *vectorizer) vectorize(ctx context.Context, input string,
 	}, nil
 }
 
+func (v *vectorizer) getApiKey(ctx context.Context) (string, error) {
+	if len(v.apiKey) > 0 {
+		return v.apiKey, nil
+	}
+	apiKey := ctx.Value("X-Openai-Api-Key")
+	if apiKeyHeader, ok := apiKey.([]string); ok &&
+		len(apiKeyHeader) > 0 && len(apiKeyHeader[0]) > 0 {
+		return apiKeyHeader[0], nil
+	}
+	return "", errors.New("no api key found " +
+		"neither in request header: X-OpenAI-Api-Key " +
+		"nor in environment variable under OPENAI_APIKEY")
+}
+
 func (v *vectorizer) docUrl(docType, model string) string {
 	return v.urlBuilder.documentUrl(docType, model)
 }

modules/text2vec-openai/clients/vectorizer_test.go

@@ -98,6 +98,86 @@ func TestClient(t *testing.T) {
 		require.NotNil(t, err)
 		assert.Equal(t, err.Error(), "failed with status: 500 error: nope, not gonna happen")
 	})
+
+	t.Run("when OpenAI key is passed using X-Openai-Api-Key header", func(t *testing.T) {
+		server := httptest.NewServer(&fakeHandler{t: t})
+		defer server.Close()
+		c := &vectorizer{
+			apiKey:     "",
+			httpClient: &http.Client{},
+			urlBuilder: &openAIUrlBuilder{
+				origin:   server.URL,
+				pathMask: "/v1/engines/%s-search-%s-%s-001/embeddings",
+			},
+			logger: nullLogger(),
+		}
+		ctxWithValue := context.WithValue(context.Background(),
+			"X-Openai-Api-Key", []string{"some-key"})
+
+		expected := &ent.VectorizationResult{
+			Text:       "This is my text",
+			Vector:     []float32{0.1, 0.2, 0.3},
+			Dimensions: 3,
+		}
+		res, err := c.Vectorize(ctxWithValue, "This is my text",
+			ent.VectorizationConfig{
+				Type:  "text",
+				Model: "ada",
+			})
+
+		require.Nil(t, err)
+		assert.Equal(t, expected, res)
+	})
+
+	t.Run("when OpenAI key is empty", func(t *testing.T) {
+		server := httptest.NewServer(&fakeHandler{t: t})
+		defer server.Close()
+		c := &vectorizer{
+			apiKey:     "",
+			httpClient: &http.Client{},
+			urlBuilder: &openAIUrlBuilder{
+				origin:   server.URL,
+				pathMask: "/v1/engines/%s-search-%s-%s-001/embeddings",
+			},
+			logger: nullLogger(),
+		}
+		ctx, cancel := context.WithDeadline(context.Background(), time.Now())
+		defer cancel()
+
+		_, err := c.Vectorize(ctx, "This is my text", ent.VectorizationConfig{})
+
+		require.NotNil(t, err)
+		assert.Equal(t, err.Error(), "OpenAI API Key: no api key found "+
+			"neither in request header: X-OpenAI-Api-Key "+
+			"nor in environment variable under OPENAI_APIKEY")
+	})
+
+	t.Run("when X-Openai-Api-Key header is passed but empty", func(t *testing.T) {
+		server := httptest.NewServer(&fakeHandler{t: t})
+		defer server.Close()
+		c := &vectorizer{
+			apiKey:     "",
+			httpClient: &http.Client{},
+			urlBuilder: &openAIUrlBuilder{
+				origin:   server.URL,
+				pathMask: "/v1/engines/%s-search-%s-%s-001/embeddings",
+			},
+			logger: nullLogger(),
+		}
+		ctxWithValue := context.WithValue(context.Background(),
+			"X-Openai-Api-Key", []string{""})
+
+		_, err := c.Vectorize(ctxWithValue, "This is my text",
+			ent.VectorizationConfig{
+				Type:  "text",
+				Model: "ada",
+			})
+
+		require.NotNil(t, err)
+		assert.Equal(t, err.Error(), "OpenAI API Key: no api key found "+
+			"neither in request header: X-OpenAI-Api-Key "+
+			"nor in environment variable under OPENAI_APIKEY")
+	})
 }
 
 type fakeHandler struct {

modules/text2vec-openai/module.go

@@ -98,10 +98,6 @@ func (m *OpenAIModule) InitDependency(modules []modulecapabilities.Module) error
 func (m *OpenAIModule) initVectorizer(ctx context.Context,
 	logger logrus.FieldLogger) error {
 	apiKey := os.Getenv("OPENAI_APIKEY")
-	if apiKey == "" {
-		return errors.Errorf("required variable OPENAI_APIKEY is not set")
-	}
-
 	client := clients.New(apiKey, logger)
 
 	m.vectorizer = vectorizer.New(client)