feat: migrate to upstream clickhouse image

Signed-off-by: Philip Miglinci <pmig@glasskube.com>

Author: pmig

glasskube/trieve/clickhouse/clickhouse.yaml

@@ -49,8 +49,10 @@ spec:
                   mountPath: /opt/user_functions/
                   readOnly: true
               env:
+                - name: CLICKHOUSE_ALWAYS_RUN_INITDB_SCRIPTS
+                  value: "true"
                 - name: EMBEDDING_SERVER_URL
-                  value: "" # todo
+                  value: "http://trieve-embedding-bgem3" # todo
                 - name: CLICKHOUSE_ADMIN_PASSWORD
                   value: "password" # todo
 

glasskube/trieve/clickhouse/embedding_function.xml

@@ -0,0 +1,14 @@
+<functions>
+    <function>
+        <name>embed_p</name>
+        <type>executable_pool</type>
+        <pool_size>10</pool_size>
+        <send_chunk_header>true</send_chunk_header>
+        <format>TabSeparated</format>
+        <return_type>Array(Float32)</return_type>
+        <argument>
+            <type>String</type>
+        </argument>
+        <command>embed.py</command>
+    </function>
+</functions>

glasskube/trieve/clickhouse/functions.yaml

@@ -1,4 +1,15 @@
 resources:
-  - scripts.yaml
-  - functions.yaml
-  - clickhouse.yaml
+  - clickhouse.yaml
+
+configMapGenerator:
+  - name: trieve-clickhouse-functions
+    files:
+      - embedding_function.xml
+    options:
+      disableNameSuffixHash: true
+
+  - name: trieve-clickhouse-scripts
+    files:
+      - load_embed.sh
+    options:
+      disableNameSuffixHash: true

glasskube/trieve/clickhouse/kustomization.yaml

@@ -0,0 +1,67 @@
+#!/usr/bin/env bash
+echo "STARTING TRIEVE CLICKHOUSE INIT SCRIPT"
+
+apt -qq update
+apt -qq install -y python3 python3-pip
+
+cat > requirements.txt << EOL
+tenacity==8.4.1
+requests==2.32.3
+EOL
+
+pip install -r requirements.txt
+
+cat > /opt/user_scripts/embed.py << EOL
+#!/usr/bin/python3
+import sys
+import requests
+import os
+
+request_timeout = 3
+
+
+def completion_with_backoff(model_input):
+  url = os.getenv("EMBEDDING_SERVER_URL")
+  parameters = {"model": "dense-embeddings", "input": f"Search for {model_input}"}
+  headers = {
+    "Content-Type": "application/json",
+  }
+  try:
+    response = requests.post(
+      f"{url}/embeddings?api-version=2023-05-15",
+      headers=headers,
+      json=parameters,
+    )
+    response.raise_for_status()
+    return [embedding["embedding"] for embedding in response.json()["data"]][0]
+  except requests.exceptions.RequestException as e:
+    raise Exception("Failed to send message to embedding server") from e
+  except Exception as e:
+    raise Exception("Failed to get text from embeddings") from e
+
+
+def embed(text):
+  if text == "":
+    return "NULL"
+  try:
+    response = completion_with_backoff(text)
+    return response
+  except:
+    return "ERROR"
+
+
+for size in sys.stdin:
+  try:
+    # collect a batch for performance
+    for row in range(0, int(size)):
+      print(embed(sys.stdin.readline().strip()))
+    sys.stdout.flush()
+  except Exception as e:
+    print(f"ERROR: {e}")
+    sys.stdout.flush()
+    continue
+EOL
+
+chmod -v +x /opt/user_scripts/embed.py
+
+echo "FINISHED TRIEVE CLICKHOUSE INIT SCRIPT"