feat: Add experimental Supabase.ai API for running inferences

Author: laktek

Cargo.lock

@@ -13,6 +13,7 @@ members = [
   "./crates/sb_graph",
   "./crates/sb_module_loader",
   "./crates/sb_fs",
+  "./crates/sb_ai"
 ]
 resolver = "2"
 

Cargo.toml

@@ -2,6 +2,7 @@
 FROM rust:1.74.1-bookworm as builder
 ARG TARGETPLATFORM
 ARG GIT_V_VERSION
+ARG ONNXRUNTIME_VERSION=1.17.0
 RUN apt-get update && apt-get install -y llvm-dev libclang-dev clang cmake
 WORKDIR /usr/src/edge-runtime
 RUN --mount=type=cache,target=/usr/local/cargo/registry,id=${TARGETPLATFORM} \
@@ -11,10 +12,16 @@ RUN --mount=type=cache,target=/usr/local/cargo/registry,id=${TARGETPLATFORM} --m
     GIT_V_TAG=${GIT_V_VERSION} cargo build --release && \
     cargo strip && \
     mv /usr/src/edge-runtime/target/release/edge-runtime /root
+RUN curl -O https://registry.npmjs.org/onnxruntime-node/-/onnxruntime-node-${ONNXRUNTIME_VERSION}.tgz && tar zxvf onnxruntime-node-${ONNXRUNTIME_VERSION}.tgz && \
+  mv ./package/bin/napi-v3/$TARGETPLATFORM/libonnxruntime.so.${ONNXRUNTIME_VERSION} /root/libonnxruntime.so
 
 
 FROM debian:bookworm-slim
 RUN apt-get update && apt-get install -y libssl-dev && rm -rf /var/lib/apt/lists/*
 RUN apt-get remove -y perl && apt-get autoremove -y
 COPY --from=builder /root/edge-runtime /usr/local/bin/edge-runtime
+COPY --from=builder /root/libonnxruntime.so /usr/local/bin/libonnxruntime.so
+COPY ./models /etc/sb_ai/models
+ENV ORT_DYLIB_PATH=/usr/local/bin/libonnxruntime.so
+ENV SB_AI_MODELS_DIR=/etc/sb_ai/models
 ENTRYPOINT ["edge-runtime"]

Dockerfile

@@ -62,6 +62,7 @@ pin-project = { version = "1.1.3" }
 ctor = { workspace = true }
 deno_canvas.workspace = true
 deno_webgpu.workspace = true
+sb_ai = { version = "0.1.0", path = "../sb_ai" }
 
 [dev-dependencies]
 flaky_test = { version = "0.1.0", path = "../flaky_test" }
@@ -102,4 +103,5 @@ event_worker ={ version = "0.1.0", path = "../event_worker" }
 deno_broadcast_channel.workspace = true
 deno_core.workspace = true
 deno_canvas.workspace = true
-deno_webgpu.workspace = true
+deno_webgpu.workspace = true
+sb_ai = { version = "0.1.0", path = "../sb_ai" }

crates/base/Cargo.toml

@@ -11,6 +11,7 @@ mod supabase_startup_snapshot {
     use deno_http::DefaultHttpPropertyExtractor;
     use event_worker::js_interceptors::sb_events_js_interceptors;
     use event_worker::sb_user_event_worker;
+    use sb_ai::sb_ai;
     use sb_core::http_start::sb_core_http;
     use sb_core::net::sb_core_net;
     use sb_core::permissions::sb_core_permissions;
@@ -196,6 +197,7 @@ mod supabase_startup_snapshot {
             deno_http::deno_http::init_ops_and_esm::<DefaultHttpPropertyExtractor>(),
             deno_io::deno_io::init_ops_and_esm(Some(Default::default())),
             deno_fs::deno_fs::init_ops_and_esm::<Permissions>(fs.clone()),
+            sb_ai::init_ops_and_esm(),
             sb_env::init_ops_and_esm(),
             sb_os::sb_os::init_ops_and_esm(),
             sb_user_workers::init_ops_and_esm(),

crates/base/build.rs

@@ -33,6 +33,7 @@ use crate::snapshot;
 use event_worker::events::{EventMetadata, WorkerEventWithMetadata};
 use event_worker::js_interceptors::sb_events_js_interceptors;
 use event_worker::sb_user_event_worker;
+use sb_ai::sb_ai;
 use sb_core::cache::CacheSetting;
 use sb_core::cert::ValueRootCertStoreProvider;
 use sb_core::external_memory::custom_allocator;
@@ -277,6 +278,7 @@ impl DenoRuntime {
             deno_io::deno_io::init_ops(stdio),
             deno_fs::deno_fs::init_ops::<Permissions>(fs.clone()),
             sb_env_op::init_ops(),
+            sb_ai::init_ops(),
             sb_os::sb_os::init_ops(),
             sb_user_workers::init_ops(),
             sb_user_event_worker::init_ops(),

crates/base/src/deno_runtime.rs

@@ -0,0 +1,21 @@
+[package]
+name = "sb_ai"
+version = "0.1.0"
+authors = ["Supabase <[email protected]>"]
+edition = "2021"
+resolver = "2"
+license = "MIT"
+
+[lib]
+path = "lib.rs"
+
+[dependencies]
+anyhow.workspace = true
+deno_core.workspace = true
+log = { workspace = true }
+serde.workspace = true
+ort = { version = "2.0.0-alpha.4", default-features = false, features = [ "ndarray", "half", "load-dynamic" ] }
+ndarray = "0.15"
+ndarray-linalg = "0.15"
+tokenizers = { version = ">=0.13.4", default-features = false, features = [ "onig" ] }
+rand = "0.8"

crates/sb_ai/Cargo.toml

@@ -0,0 +1,10 @@
+const core = globalThis.Deno.core;
+
+class SupabaseAI {
+	runModel(name, prompt) {
+		const result = core.ops.op_sb_ai_run_model(name, prompt);
+		return result;
+	}
+}
+
+export { SupabaseAI };

crates/sb_ai/ai.js

@@ -0,0 +1,111 @@
+use anyhow::{bail, Error};
+use deno_core::error::AnyError;
+use deno_core::op2;
+use deno_core::OpState;
+use ndarray::{Array1, Array2, Axis, Ix2};
+use ndarray_linalg::norm::{normalize, NormalizeAxis};
+use ort::{inputs, GraphOptimizationLevel, Session, Tensor};
+use std::path::Path;
+use tokenizers::normalizers::bert::BertNormalizer;
+use tokenizers::Tokenizer;
+
+deno_core::extension!(
+    sb_ai,
+    ops = [op_sb_ai_run_model],
+    esm_entry_point = "ext:sb_ai/ai.js",
+    esm = ["ai.js",]
+);
+
+fn run_gte(state: &mut OpState, prompt: String) -> Result<Vec<f32>, Error> {
+    // Create the ONNX Runtime environment, for all sessions created in this process.
+    ort::init().with_name("GTE").commit()?;
+
+    let models_dir = std::env::var("SB_AI_MODELS_DIR").unwrap_or("/etc/sb_ai/models".to_string());
+
+    let mut session = state.try_take::<Session>();
+    if session.is_none() {
+        session = Some(
+            Session::builder()?
+                .with_optimization_level(GraphOptimizationLevel::Disable)?
+                .with_intra_threads(1)?
+                .with_model_from_file(
+                    Path::new(&models_dir)
+                        .join("gte")
+                        .join("gte_small_quantized.onnx"),
+                )?,
+        );
+    }
+    let session = session.unwrap();
+
+    // Load the tokenizer and encode the prompt into a sequence of tokens.
+    let mut tokenizer = state.try_take::<Tokenizer>();
+    if tokenizer.is_none() {
+        tokenizer = Some(
+            Tokenizer::from_file(
+                Path::new(&models_dir)
+                    .join("gte")
+                    .join("gte_small_tokenizer.json"),
+            )
+            .map_err(anyhow::Error::msg)?,
+        )
+    }
+    let mut tokenizer = tokenizer.unwrap();
+
+    let tokenizer_impl = tokenizer
+        .with_normalizer(BertNormalizer::default())
+        .with_padding(None)
+        .with_truncation(None)
+        .map_err(anyhow::Error::msg)?;
+
+    let tokens = tokenizer_impl
+        .encode(prompt, true)
+        .map_err(anyhow::Error::msg)?
+        .get_ids()
+        .iter()
+        .map(|i| *i as i64)
+        .collect::<Vec<_>>();
+
+    let tokens = Array1::from_iter(tokens.iter().cloned());
+
+    let array = tokens.view().insert_axis(Axis(0));
+    let dims = array.raw_dim();
+    let token_type_ids = Array2::<i64>::zeros(dims);
+    let attention_mask = Array2::<i64>::ones(dims);
+    let outputs = session.run(inputs! {
+        "input_ids" => array,
+        "token_type_ids" => token_type_ids,
+        "attention_mask" => attention_mask,
+    }?)?;
+
+    let embeddings: Tensor<f32> = outputs["last_hidden_state"].extract_tensor()?;
+
+    let embeddings_view = embeddings.view();
+    let mean_pool = embeddings_view.mean_axis(Axis(1)).unwrap();
+    let (normalized, _) = normalize(
+        mean_pool.into_dimensionality::<Ix2>().unwrap(),
+        NormalizeAxis::Row,
+    );
+
+    let slice = normalized.view().to_slice().unwrap().to_vec();
+
+    drop(outputs);
+
+    state.put::<Session>(session);
+    state.put::<Tokenizer>(tokenizer);
+
+    Ok(slice)
+}
+
+#[op2]
+#[serde]
+pub fn op_sb_ai_run_model(
+    state: &mut OpState,
+    #[string] name: String,
+    #[string] prompt: String,
+) -> Result<Vec<f32>, AnyError> {
+    if name == "gte" {
+        run_gte(state, prompt)
+    } else {
+        bail!("model not supported")
+    }
+}

crates/sb_ai/lib.rs

@@ -2,7 +2,7 @@ import * as abortSignal from 'ext:deno_web/03_abort_signal.js';
 import * as base64 from 'ext:deno_web/05_base64.js';
 import * as console from 'ext:deno_console/01_console.js';
 import * as crypto from 'ext:deno_crypto/00_crypto.js';
-import { DOMException } from "ext:deno_web/01_dom_exception.js";
+import { DOMException } from 'ext:deno_web/01_dom_exception.js';
 import * as encoding from 'ext:deno_web/08_text_encoding.js';
 import * as event from 'ext:deno_web/02_event.js';
 import * as fetch from 'ext:deno_fetch/26_fetch.js';
@@ -20,6 +20,7 @@ import * as response from 'ext:deno_fetch/23_response.js';
 import * as request from 'ext:deno_fetch/23_request.js';
 import * as globalInterfaces from 'ext:deno_web/04_global_interfaces.js';
 import { SUPABASE_ENV } from 'ext:sb_env/env.js';
+import { SupabaseAI } from 'ext:sb_ai/ai.js';
 import { registerErrors } from 'ext:sb_core_main_js/js/errors.js';
 import {
 	formatException,
@@ -43,12 +44,12 @@ import { SupabaseEventListener } from 'ext:sb_user_event_worker/event_worker.js'
 import * as MainWorker from 'ext:sb_core_main_js/js/main_worker.js';
 import * as DenoWebCompression from 'ext:deno_web/14_compression.js';
 import * as DenoWSStream from 'ext:deno_websocket/02_websocketstream.js';
-import * as eventSource from "ext:deno_fetch/27_eventsource.js";
-import * as WebGPU from "ext:deno_webgpu/00_init.js";
-import * as WebGPUSurface from "ext:deno_webgpu/02_surface.js";
+import * as eventSource from 'ext:deno_fetch/27_eventsource.js';
+import * as WebGPU from 'ext:deno_webgpu/00_init.js';
+import * as WebGPUSurface from 'ext:deno_webgpu/02_surface.js';
 
-import { primordials, core } from "ext:core/mod.js";
-import { op_lazy_load_esm } from "ext:core/ops";
+import { core, primordials } from 'ext:core/mod.js';
+import { op_lazy_load_esm } from 'ext:core/ops';
 const ops = core.ops;
 
 const {
@@ -111,7 +112,7 @@ function ImageWritable(getter) {
 }
 function loadImage() {
 	if (!image) {
-		image = op_lazy_load_esm("ext:deno_canvas/01_image.js");
+		image = op_lazy_load_esm('ext:deno_canvas/01_image.js');
 	}
 }
 
@@ -295,7 +296,13 @@ const deleteDenoApis = (apis) => {
 	});
 };
 
-globalThis.bootstrapSBEdge = (opts, isUserWorker, isEventsWorker, edgeRuntimeVersion, denoVersion) => {
+globalThis.bootstrapSBEdge = (
+	opts,
+	isUserWorker,
+	isEventsWorker,
+	edgeRuntimeVersion,
+	denoVersion,
+) => {
 	// We should delete this after initialization,
 	// Deleting it during bootstrapping can backfire
 	delete globalThis.__bootstrap;
@@ -318,7 +325,7 @@ globalThis.bootstrapSBEdge = (opts, isUserWorker, isEventsWorker, edgeRuntimeVer
 	});
 
 	ObjectDefineProperty(globalThis, 'SUPABASE_VERSION', readOnly(String(edgeRuntimeVersion)));
-	ObjectDefineProperty(globalThis, 'DENO_VERSION', readOnly(denoVersion))
+	ObjectDefineProperty(globalThis, 'DENO_VERSION', readOnly(denoVersion));
 
 	// set these overrides after runtimeStart
 	ObjectDefineProperties(denoOverrides, {
@@ -328,17 +335,29 @@ globalThis.bootstrapSBEdge = (opts, isUserWorker, isEventsWorker, edgeRuntimeVer
 		args: readOnly([]), // args are set to be empty
 		mainModule: getterOnly(() => ops.op_main_module()),
 		version: getterOnly(() => ({
-			deno: `supabase-edge-runtime-${globalThis.SUPABASE_VERSION} (compatible with Deno v${globalThis.DENO_VERSION})`,
+			deno:
+				`supabase-edge-runtime-${globalThis.SUPABASE_VERSION} (compatible with Deno v${globalThis.DENO_VERSION})`,
 			v8: '11.6.189.12',
 			typescript: '5.1.6',
 		})),
 	});
 	ObjectDefineProperty(globalThis, 'Deno', readOnly(denoOverrides));
 
 	setNumCpus(1); // explicitly setting no of CPUs to 1 (since we don't allow workers)
-	setUserAgent(`Deno/${globalThis.DENO_VERSION} (variant; SupabaseEdgeRuntime/${globalThis.SUPABASE_VERSION})`);
+	setUserAgent(
+		`Deno/${globalThis.DENO_VERSION} (variant; SupabaseEdgeRuntime/${globalThis.SUPABASE_VERSION})`,
+	);
 	setLanguage('en');
 
+	const ai = new SupabaseAI();
+	Object.defineProperty(globalThis, 'Supabase_UNSTABLE', {
+		get() {
+			return {
+				ai,
+			};
+		},
+	});
+
 	if (isUserWorker) {
 		delete globalThis.EdgeRuntime;
 
@@ -364,7 +383,7 @@ globalThis.bootstrapSBEdge = (opts, isUserWorker, isEventsWorker, edgeRuntimeVer
 	}
 
 	const nodeBootstrap = globalThis.nodeBootstrap;
-	if(nodeBootstrap) {
+	if (nodeBootstrap) {
 		nodeBootstrap(false, undefined);
 		delete globalThis.nodeBootstrap;
 	}