Merge pull request #248 from weaviate/add-support-for-jinaai-reranker

Add support for JinaAI reranker module
weaviate · Jul 24, 2024 · 8b7eb23 · 8b7eb23
2 parents 834d636 + 835d1e2
commit 8b7eb23
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 3 deletions.
diff --git a/.cicd/test.sh b/.cicd/test.sh
@@ -141,6 +141,9 @@ function check_creates_template() {
   check_modules "--set modules.multi2vec-bind.enabled=true" "value: multi2vec-bind"
   check_modules "--set modules.multi2vec-bind.enabled=true --set modules.text-spellcheck.enabled=true" "value: text-spellcheck,multi2vec-bind"
   check_modules "--set modules.text2vec-jinaai.enabled=true" "value: text2vec-jinaai"
+  check_modules "--set modules.text2vec-jinaai.enabled=true --set modules.text2vec-jinaai.apiKey=apiKey" "value: text2vec-jinaai"
+  check_modules "--set modules.reranker-jinaai.enabled=true" "value: reranker-jinaai"
+  check_modules "--set modules.reranker-jinaai.enabled=true --set modules.reranker-jinaai.apiKey=apiKey" "value: reranker-jinaai"
   check_modules "--set modules.text2vec-aws.enabled=true" "value: text2vec-aws"
   check_modules "--set modules.generative-aws.enabled=true" "value: generative-aws"
   check_modules "--set modules.text2vec-aws.enabled=true --set modules.generative-aws.enabled=true" "value: generative-aws,text2vec-aws"
@@ -202,6 +205,9 @@ function check_creates_template() {
   check_string_existence "--set offload.s3.enabled=true --set modules.generative-aws.enabled=true --set offload.s3.envconfig.OFFLOAD_S3_PATH=/custom-path --set modules.text2vec-aws.enabled=true" "name: OFFLOAD_S3_PATH"
   check_string_existence "--set initContainers.ensureFileOwnershipContainer.enabled=true --set containerSecurityContext.runAsUser=1000 --set containerSecurityContext.fsGroup=2000" "name: ensure-file-ownership"
   check_string_existence "--set modules.text2vec-openai.enabled=true --set modules.generative-anthropic.enabled=true --set modules.generative-anthropic.apiKey=azureApiKey" "name: ANTHROPIC_APIKEY"
+  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_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
@@ -107,6 +107,9 @@
   {{- if (index .Values "modules" "reranker-voyageai" "enabled") -}}
     {{ $modules = append $modules "reranker-voyageai" }}
   {{- end -}}
+  {{- if (index .Values "modules" "reranker-jinaai" "enabled") -}}
+    {{ $modules = append $modules "reranker-jinaai" }}
+  {{- end -}}
   {{- if (index .Values "offload" "s3" "enabled") -}}
     {{ $modules = append $modules "offload-s3" }}
   {{- end -}}

diff --git a/weaviate/templates/apiKeyJinaAISecret.yaml b/weaviate/templates/apiKeyJinaAISecret.yaml
@@ -1,4 +1,20 @@
-{{ if and (index .Values "modules" "text2vec-jinaai" "enabled") (index .Values "modules" "text2vec-jinaai" "apiKey") }}
+{{- $t2vJinaAI :=  and (index .Values "modules" "text2vec-jinaai" "enabled") (index .Values "modules" "text2vec-jinaai" "apiKey") }}
+{{- $rerankerJinaAI := and (index .Values "modules" "reranker-jinaai" "enabled") (index .Values "modules" "reranker-jinaai" "apiKey") }}
+{{- $jinaai := or ($t2vJinaAI) ($rerankerJinaAI) }}
+{{- if $jinaai }}
+
+{{- $apiKeys := list }}
+{{- if $t2vJinaAI }}
+  {{- $apiKeys = append $apiKeys (index .Values "modules" "text2vec-jinaai" "apiKey") -}}
+{{- end }}
+{{- if $rerankerJinaAI }}
+  {{- $apiKeys = append $apiKeys (index .Values "modules" "reranker-jinaai" "apiKey") -}}
+{{- end }}
+
+{{- if gt (len ($apiKeys | uniq)) 1 -}}
+  {{- fail "JinaAI modules activated, but their keys differ. Use the same key on all activated JinaAI modules." }}
+{{- end }}
+
 apiVersion: v1
 kind: Secret
 metadata:
@@ -8,5 +24,5 @@ metadata:
     app.kubernetes.io/managed-by: {{ .Release.Service }}
 type: Opaque
 data:
-  apiKey: {{ index .Values "modules" "text2vec-jinaai" "apiKey" | b64enc }}
+  apiKey: {{ index $apiKeys 0 | b64enc }}
 {{ end }}
diff --git a/weaviate/templates/weaviateStatefulset.yaml b/weaviate/templates/weaviateStatefulset.yaml
@@ -232,7 +232,7 @@ spec:
                 name: weaviate-palm
                 key: apiKey
           {{- end  }}
-          {{- if and (index .Values "modules" "text2vec-jinaai" "enabled") (index .Values "modules" "text2vec-jinaai" "apiKey") }}
+          {{- if or (and (index .Values "modules" "text2vec-jinaai" "enabled") (index .Values "modules" "text2vec-jinaai" "apiKey")) (and (index .Values "modules" "reranker-jinaai" "enabled") (index .Values "modules" "reranker-jinaai" "apiKey")) }}
           - name: JINAAI_APIKEY
             valueFrom:
               secretKeyRef:

diff --git a/weaviate/values.yaml b/weaviate/values.yaml
@@ -1462,6 +1462,19 @@ modules:
     # an environment variable
     apiKey: ''
 
+  # The reranker-jinaai module uses JinaAI API
+  # to dynamically compute a score for the relevance 
+  # of the query with each of the initial search results.
+  # More information about Cohere API can be found here: https://jina.ai/reranker
+  reranker-jinaai:
+
+    # enable if you want to use JinaAI module
+    enabled: false
+
+    # Set your JinaAI API Key to be passed to Weaviate pod as
+    # an environment variable
+    apiKey: ''
+
   # The reranker-transformers module uses Cross-Encoders for 
   # sentence pair scoring and sentence pair classification tasks.
   # More information about Cross-Encoders can be found here: