Merge pull request #69 from philippgille/webassembly

Add experimental WebAssembly binding and example
philippgille · Apr 25, 2024 · 551af6d · 551af6d
2 parents ee7328f + 545e626
commit 551af6d
Show file tree

Hide file tree

Showing 4 changed files with 216 additions and 0 deletions.
diff --git a/examples/README.md b/examples/README.md
@@ -11,3 +11,5 @@
    - This example shows a semantic search application, using `chromem-go` as vector database for finding semantically relevant search results.
    - Loads and searches across ~5,000 arXiv papers in the "Computer Science - Computation and Language" category, which is the relevant one for Natural Language Processing (NLP) related papers.
    - Uses OpenAI for creating the embeddings
+4. [WebAssembly](webassembly)
+   - This example shows how `chromem-go` can be compiled to WebAssembly and then used from JavaScript in a browser
diff --git a/examples/webassembly/README.md b/examples/webassembly/README.md
@@ -0,0 +1,13 @@
+# WebAssembly (WASM)
+
+Go can compile to WebAssembly, which you can then use from JavaScript in a Browser or similar environments (Node, Deno, Bun etc.). You could also target WASI (WebAssembly System Interface) and run it in a standalone runtime (wazero, wasmtime, Wasmer), but in this example we focus on the Browser use case.
+
+1. Compile the `chromem-go` WASM binding to WebAssembly:
+   1. `cd /path/to/chromem-go/wasm`
+   2. `GOOS=js GOARCH=wasm go build -o ../examples/webassembly/chromem-go.wasm`
+2. Copy Go's wrapper JavaScript:
+   1. `cp $(go env GOROOT)/misc/wasm/wasm_exec.js ../examples/webassembly/wasm_exec.js`
+3. Serve the files
+   1. `cd ../examples/webassembly`
+   2. `go run github.com/philippgille/serve@latest -b localhost -p 8080` or similar
+4. Open <http://localhost:8080> in your browser
diff --git a/examples/webassembly/index.html b/examples/webassembly/index.html
@@ -0,0 +1,69 @@
+<!DOCTYPE html>
+<html lang="en">
+
+<head>
+    <meta charset="UTF-8" />
+    <script src="wasm_exec.js"></script>
+    <script>
+        const go = new Go();
+
+        console.log("Initializing WASM...");
+        WebAssembly.instantiateStreaming(fetch("chromem-go.wasm"), go.importObject).then((result) => {
+            console.log("WASM initialized.");
+
+            go.run(result.instance);
+        });
+
+        function initDBWithKey() {
+            console.log("Initializing DB...")
+            const openaiApiKey = document.getElementById("openai-api-key").value;
+
+            const err = initDB(openaiApiKey)
+            if (err) {
+                console.error('Returned error:', err)
+            } else {
+                console.log("DB initialized.")
+            }
+        }
+
+        async function addDocuments() {
+            console.log("Adding documents...")
+            try {
+                await addDocument("1", "The sky is blue because of Rayleigh scattering.");
+                console.log("Document 1 added.")
+                await addDocument("2", "Leaves are green because chlorophyll absorbs red and blue light.");
+                console.log("Document 2 added.")
+                console.log("Documents added.")
+            } catch (err) {
+                console.error('Caught exception', err)
+            }
+        }
+
+        async function queryAndPrint() {
+            console.log("Querying DB...")
+            try {
+                const res = await query("Why is the sky blue?");
+                console.log("DB queried.");
+
+                const outputElement = document.getElementById("output");
+                outputElement.textContent = `ID: ${res.ID}\nSimilarity: ${res.Similarity}\nContent: ${res.Content}\n`;
+            } catch (err) {
+                console.error('Caught exception', err)
+            }
+        }
+
+        async function runWorkflow() {
+            initDBWithKey();
+            await addDocuments();
+            await queryAndPrint();
+        }
+    </script>
+</head>
+
+<body>
+    <input type="text" id="openai-api-key" placeholder="Enter your OpenAI API key">
+    <button onclick="runWorkflow()">Run</button>
+    <p id="output"></p>
+</body>
+
+</html>
diff --git a/wasm/main.go b/wasm/main.go
@@ -0,0 +1,132 @@
+//go:build js
+
+package main
+
+import (
+	"context"
+	"errors"
+	"syscall/js"
+
+	"github.com/philippgille/chromem-go"
+)
+
+var c *chromem.Collection
+
+func main() {
+	js.Global().Set("initDB", js.FuncOf(initDB))
+	js.Global().Set("addDocument", js.FuncOf(addDocument))
+	js.Global().Set("query", js.FuncOf(query))
+
+	select {} // prevent main from exiting
+}
+
+// Exported function to initialize the database and collection.
+// Takes an OpenAI API key as argument.
+func initDB(this js.Value, args []js.Value) interface{} {
+	if len(args) != 1 {
+		return "expected 1 argument with the OpenAI API key"
+	}
+
+	openAIAPIKey := args[0].String()
+	embeddingFunc := chromem.NewEmbeddingFuncOpenAI(openAIAPIKey, chromem.EmbeddingModelOpenAI3Small)
+
+	db := chromem.NewDB()
+	var err error
+	c, err = db.CreateCollection("chromem", nil, embeddingFunc)
+	if err != nil {
+		return err.Error()
+	}
+
+	return nil
+}
+
+// Exported function to add documents to the collection.
+// Takes the document ID and content as arguments.
+func addDocument(this js.Value, args []js.Value) interface{} {
+	ctx := context.Background()
+
+	var id string
+	var content string
+	var err error
+	if len(args) != 2 {
+		err = errors.New("expected 2 arguments with the document ID and content")
+	} else {
+		id = args[0].String()
+		content = args[1].String()
+	}
+
+	handler := js.FuncOf(func(this js.Value, args []js.Value) interface{} {
+		resolve := args[0]
+		reject := args[1]
+		go func() {
+			if err != nil {
+				handleErr(err, reject)
+				return
+			}
+
+			err = c.AddDocument(ctx, chromem.Document{
+				ID:      id,
+				Content: content,
+			})
+			if err != nil {
+				handleErr(err, reject)
+				return
+			}
+			resolve.Invoke()
+		}()
+		return nil
+	})
+
+	promiseConstructor := js.Global().Get("Promise")
+	return promiseConstructor.New(handler)
+}
+
+// Exported function to query the collection
+// Takes the query string and the number of documents to return as argument.
+func query(this js.Value, args []js.Value) interface{} {
+	ctx := context.Background()
+
+	var q string
+	var err error
+	if len(args) != 1 {
+		err = errors.New("expected 1 argument with the query string")
+	} else {
+		q = args[0].String()
+	}
+
+	handler := js.FuncOf(func(this js.Value, args []js.Value) interface{} {
+		resolve := args[0]
+		reject := args[1]
+		go func() {
+			if err != nil {
+				handleErr(err, reject)
+				return
+			}
+
+			res, err := c.Query(ctx, q, 1, nil, nil)
+			if err != nil {
+				handleErr(err, reject)
+				return
+			}
+
+			// Convert response to JS values
+			// TODO: Return more than one result
+			o := js.Global().Get("Object").New()
+			o.Set("ID", res[0].ID)
+			o.Set("Similarity", res[0].Similarity)
+			o.Set("Content", res[0].Content)
+
+			resolve.Invoke(o)
+		}()
+		return nil
+	})
+
+	promiseConstructor := js.Global().Get("Promise")
+	return promiseConstructor.New(handler)
+}
+
+func handleErr(err error, reject js.Value) {
+	errorConstructor := js.Global().Get("Error")
+	errorObject := errorConstructor.New(err.Error())
+	reject.Invoke(errorObject)
+}