From 900e0436005cc3d5abad195e814971c32b4d086e Mon Sep 17 00:00:00 2001
From: Marcin Antas <antas.marcin@gmail.com>
Date: Thu, 5 Sep 2024 08:53:31 +0200
Subject: [PATCH] Add support for text2vec-mistral module

---
 .cicd/test.sh                               |  3 +++
 weaviate/templates/_helpers.tpl             |  3 +++
 weaviate/templates/apiKeyMistralSecret.yaml | 22 ++++++++++++++++++---
 weaviate/templates/weaviateStatefulset.yaml |  2 +-
 weaviate/values.yaml                        | 12 +++++++++++
 5 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/.cicd/test.sh b/.cicd/test.sh
index d8afb5b..e8f082b 100755
--- a/.cicd/test.sh
+++ b/.cicd/test.sh
@@ -169,6 +169,7 @@ function check_creates_template() {
   check_modules "--set modules.generative-friendliai.enabled=true --set modules.generative-friendliai.token=token" "value: generative-friendliai"
   check_modules "--set modules.text2vec-databricks.enabled=true" "value: text2vec-databricks"
   check_modules "--set modules.generative-databricks.enabled=true" "value: generative-databricks"
+  check_modules "--set modules.text2vec-mistral.enabled=true" "value: text2vec-mistral"
 
   check_modules "--set modules.text2vec-openai.enabled=true --set modules.text2vec-openai.azureApiKey=azureApiKey" "value: text2vec-openai"
   check_modules "--set modules.qna-openai.enabled=true --set modules.qna-openai.azureApiKey=azureApiKey" "value: qna-openai"
@@ -216,6 +217,8 @@ function check_creates_template() {
   check_string_existence "--set modules.text2vec-databricks.enabled=true --set modules.text2vec-databricks.token=databricksToken" "name: DATABRICKS_TOKEN"
   check_string_existence "--set modules.generative-databricks.enabled=true --set modules.generative-databricks.token=databricksToken" "name: DATABRICKS_TOKEN"
   check_string_existence "--set modules.generative-databricks.enabled=true --set modules.generative-databricks.token=databricksToken --set modules.text2vec-databricks.enabled=true --set modules.text2vec-databricks.token=databricksToken" "name: DATABRICKS_TOKEN"
+  check_string_existence "--set modules.text2vec-mistral.enabled=true --set modules.text2vec-mistral.apiKey=mistralApiKey" "name: MISTRAL_APIKEY"
+  check_string_existence "--set modules.generative-mistral.enabled=true --set modules.generative-mistral.apiKey=mistralApiKey --set modules.text2vec-mistral.enabled=true --set modules.text2vec-mistral.apiKey=mistralApiKey" "name: MISTRAL_APIKEY"
 
   check_setting_has_value "--set replicas=3 --set env.RAFT_BOOTSTRAP_EXPECT=3" "name: RAFT_JOIN" "value: \"weaviate-0,weaviate-1,weaviate-2\""
   check_setting_has_value "--set replicas=3 --set env.RAFT_BOOTSTRAP_EXPECT=3" "name: RAFT_BOOTSTRAP_EXPECT" "value: \"3\""
diff --git a/weaviate/templates/_helpers.tpl b/weaviate/templates/_helpers.tpl
index 0d20ae2..0757a83 100644
--- a/weaviate/templates/_helpers.tpl
+++ b/weaviate/templates/_helpers.tpl
@@ -104,6 +104,9 @@
   {{- if (index .Values "modules" "text2vec-databricks" "enabled") -}}
     {{ $modules = append $modules "text2vec-databricks" }}
   {{- end -}}
+  {{- if (index .Values "modules" "text2vec-mistral" "enabled") -}}
+    {{ $modules = append $modules "text2vec-mistral" }}
+  {{- end -}}
   {{- if (index .Values "modules" "ref2vec-centroid" "enabled") -}}
     {{ $modules = append $modules "ref2vec-centroid" }}
   {{- end -}}
diff --git a/weaviate/templates/apiKeyMistralSecret.yaml b/weaviate/templates/apiKeyMistralSecret.yaml
index 9ab7b11..412283b 100644
--- a/weaviate/templates/apiKeyMistralSecret.yaml
+++ b/weaviate/templates/apiKeyMistralSecret.yaml
@@ -1,12 +1,28 @@
-{{ if and (index .Values "modules" "generative-mistral" "enabled") (index .Values "modules" "generative-mistral" "apiKey") }}
+{{-  $t2vMistral := and (index .Values "modules" "text2vec-mistral" "enabled") (index .Values "modules" "text2vec-mistral" "apiKey") -}}
+{{-  $generativeMistral := and (index .Values "modules" "generative-mistral" "enabled") (index .Values "modules" "generative-mistral" "apiKey") }}
+{{-  $mistral := or ($t2vMistral) ($generativeMistral) }}
+{{- if $mistral }}
+
+{{- $apiKeys := list }}
+{{- if $t2vMistral }}
+  {{- $apiKeys = append $apiKeys (index .Values "modules" "text2vec-mistral" "apiKey") -}}
+{{- end }}
+{{- if $generativeMistral }}
+  {{- $apiKeys = append $apiKeys (index .Values "modules" "generative-mistral" "apiKey") -}}
+{{- end }}
+
+{{- if gt (len ($apiKeys | uniq)) 1 -}}
+  {{- fail "Mistral modules activated, but their keys differ. Use the same key on all activated Mistral modules." }}
+{{- end }}
+
 apiVersion: v1
 kind: Secret
 metadata:
   name: weaviate-mistral
   labels:
     app.kubernetes.io/name: weaviate
-    app.kubernetes.io/managed-by: helm
+    app.kubernetes.io/managed-by: {{ .Release.Service }}
 type: Opaque
 data:
-  apiKey: {{ index .Values "modules" "generative-mistral" "apiKey" | b64enc }}
+  apiKey: {{ index $apiKeys 0 | b64enc }}
 {{ end }}
diff --git a/weaviate/templates/weaviateStatefulset.yaml b/weaviate/templates/weaviateStatefulset.yaml
index a15b0e3..1157845 100644
--- a/weaviate/templates/weaviateStatefulset.yaml
+++ b/weaviate/templates/weaviateStatefulset.yaml
@@ -260,7 +260,7 @@ spec:
                 name: weaviate-anthropic
                 key: apiKey
           {{- end  }}
-          {{- if and (index .Values "modules" "generative-mistral" "enabled") (index .Values "modules" "generative-mistral" "apiKey") }}
+          {{- if or (and (index .Values "modules" "text2vec-mistral" "enabled") (index .Values "modules" "text2vec-mistral" "apiKey")) (and (index .Values "modules" "generative-mistral" "enabled") (index .Values "modules" "generative-mistral" "apiKey")) }}
           - name: MISTRAL_APIKEY
             valueFrom:
               secretKeyRef:
diff --git a/weaviate/values.yaml b/weaviate/values.yaml
index c00d3e6..085911b 100644
--- a/weaviate/values.yaml
+++ b/weaviate/values.yaml
@@ -1008,6 +1008,18 @@ modules:
     # an environment variable
     apiKey: ''
 
+  # The text2vec-mistral module uses MISTRAL AI Embeddings API
+  # More information about MISTRAL AI's Embeddings API can be found here:
+  # https://docs.mistral.ai/api/#operation/createEmbedding
+  text2vec-mistral:
+
+    # enable if you want to use MISTRAL AI module
+    enabled: false
+
+    # Set your MISTRAL API Key to be passed to Weaviate pod as
+    # an environment variable
+    apiKey: ''
+
   # The text2vec-databricks module uses Databricks Embeddings API
   # to dynamically compute vector embeddings based on the
   # sentence's context.