From 497720b859f9c5c6f205b84308f0d75872c60326 Mon Sep 17 00:00:00 2001
From: Marcin Antas <antas.marcin@gmail.com>
Date: Wed, 4 Sep 2024 13:08:26 +0200
Subject: [PATCH] Add support for generative-friendliai module

---
 .cicd/test.sh                                 |  3 +++
 weaviate/templates/_helpers.tpl               |  3 +++
 weaviate/templates/tokenFriendliAISecret.yaml | 12 ++++++++++++
 weaviate/templates/weaviateStatefulset.yaml   |  7 +++++++
 weaviate/values.yaml                          | 12 ++++++++++++
 5 files changed, 37 insertions(+)
 create mode 100644 weaviate/templates/tokenFriendliAISecret.yaml

diff --git a/.cicd/test.sh b/.cicd/test.sh
index d6e7757..97a7551 100755
--- a/.cicd/test.sh
+++ b/.cicd/test.sh
@@ -165,6 +165,8 @@ function check_creates_template() {
   check_modules "--set offload.s3.enabled=true" "value: offload-s3"
   check_modules "--set modules.generative-anthropic.enabled=true" "value: generative-anthropic"
   check_modules "--set modules.generative-anthropic.enabled=true --set modules.generative-anthropic.apiKey=apiKey" "value: generative-anthropic"
+  check_modules "--set modules.generative-friendliai.enabled=true" "value: generative-friendliai"
+  check_modules "--set modules.generative-friendliai.enabled=true --set modules.generative-friendliai.token=token" "value: generative-friendliai"
 
   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"
@@ -208,6 +210,7 @@ function check_creates_template() {
   check_string_existence "--set modules.text2vec-openai.enabled=true --set modules.text2vec-jinaai.enabled=true --set modules.text2vec-jinaai.apiKey=JinaAIApiKey" "name: JINAAI_APIKEY"
   check_string_existence "--set modules.text2vec-openai.enabled=true --set modules.reranker-jinaai.enabled=true --set modules.reranker-jinaai.apiKey=JinaAIApiKey" "name: JINAAI_APIKEY"
   check_string_existence "--set modules.text2vec-openai.enabled=true --set modules.reranker-jinaai.enabled=true --set modules.reranker-jinaai.apiKey=JinaAIApiKey --set modules.text2vec-jinaai.enabled=true --set modules.text2vec-jinaai.apiKey=JinaAIApiKey" "name: JINAAI_APIKEY"
+  check_string_existence "--set modules.text2vec-openai.enabled=true --set modules.generative-friendliai.enabled=true --set modules.generative-friendliai.token=FriendliToken" "name: FRIENDLI_TOKEN"
 
   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 d3f88f9..a8d78d9 100644
--- a/weaviate/templates/_helpers.tpl
+++ b/weaviate/templates/_helpers.tpl
@@ -47,6 +47,9 @@
   {{- if (index .Values "modules" "generative-anthropic" "enabled") -}}
     {{ $modules = append $modules "generative-anthropic" }}
   {{- end -}}
+  {{- if (index .Values "modules" "generative-friendliai" "enabled") -}}
+    {{ $modules = append $modules "generative-friendliai" }}
+  {{- end -}}
   {{- if or (index .Values "modules" "img2vec-neural" "enabled") (index .Values "modules" "img2vec-neural" "inferenceUrl") -}}
     {{ $modules = append $modules "img2vec-neural" }}
   {{- end -}}
diff --git a/weaviate/templates/tokenFriendliAISecret.yaml b/weaviate/templates/tokenFriendliAISecret.yaml
new file mode 100644
index 0000000..ab06881
--- /dev/null
+++ b/weaviate/templates/tokenFriendliAISecret.yaml
@@ -0,0 +1,12 @@
+{{ if and (index .Values "modules" "generative-friendliai" "enabled") (index .Values "modules" "generative-friendliai" "token") }}
+apiVersion: v1
+kind: Secret
+metadata:
+  name: weaviate-friendliai
+  labels:
+    app.kubernetes.io/name: weaviate
+    app.kubernetes.io/managed-by: helm
+type: Opaque
+data:
+  token: {{ index .Values "modules" "generative-friendliai" "token" | b64enc }}
+{{ end }}
diff --git a/weaviate/templates/weaviateStatefulset.yaml b/weaviate/templates/weaviateStatefulset.yaml
index 891bf57..c52e705 100644
--- a/weaviate/templates/weaviateStatefulset.yaml
+++ b/weaviate/templates/weaviateStatefulset.yaml
@@ -274,6 +274,13 @@ spec:
                 name: weaviate-octoai
                 key: apiKey
           {{- end  }}
+          {{- if and (index .Values "modules" "generative-friendliai" "enabled") (index .Values "modules" "generative-friendliai" "token") }}
+          - name: FRIENDLI_TOKEN
+            valueFrom:
+              secretKeyRef:
+                name: weaviate-friendliai
+                key: token
+          {{- end  }}
           {{- if and (index .Values "offload" "s3" "enabled") (index .Values "offload" "s3" "envconfig") }}
             {{- range $key, $value := index .Values "offload" "s3" "envconfig" }}
           - name: {{ $key }}
diff --git a/weaviate/values.yaml b/weaviate/values.yaml
index 7ed7e86..cb00104 100644
--- a/weaviate/values.yaml
+++ b/weaviate/values.yaml
@@ -1374,6 +1374,18 @@ modules:
     # an environment variable
     apiKey: ''
 
+  # The generative-friendliai module uses FriendliAI API.
+  # More information about FriendliAI API can be found here:
+  # https://docs.friendli.ai/openapi/create-chat-completions
+  generative-friendliai:
+
+    # enable if you want to use FriendliAI module
+    enabled: false
+
+    # Set your FriendliAI Token to be passed to Weaviate pod as
+    # an environment variable
+    token: ''
+
   # The img2vec-neural module uses neural networks, to generate
   # a vector representation of the image
   img2vec-neural: