RAGWASM: In-browser Retrieval-Augmented Generation

A RAG pipeline that runs entirely in the browser tab: parse, chunk, embed, retrieve, and generate, built on LiteRT.js, Google's web AI runtime. No server, no data leaves the device.

Stack:

• React 19 + Vite (Bun toolchain)
• TanStack Router for file-based routing (/ ask, /about runtime info)
• Zustand for state management
• LiteParse-WASM for PDF/document text extraction in the browser
• Chonkie for recursive text chunking (WASM)
• EmbeddingGemma on LiteRT.js for .tflite embeddings on WebGPU/WASM
• DuckDB-WASM + VSS for the HNSW vector index in the browser
• LiteRT-LM.js for Gemma text generation (WebGPU)
• Astryx design system (neutral theme)

bun install
bun run dev      # copies the LiteRT / LiteRT-LM WASM into public/ then starts Vite

Two models are involved, see Models: EmbeddingGemma is gated and you supply the .tflite yourself; the Gemma generation model is ungated but large and fetched on first ask. A Chromium-based browser is required for the WebGPU paths.

bun run build      # type-check + production build
bun run typecheck  # tsc only

INGEST:  file ─▶ LiteParse-WASM (text) ─▶ Chonkie (chunks) ─▶ EmbeddingGemma ─▶ DuckDB VSS/HNSW
ASK:     question ─▶ EmbeddingGemma ─▶ DuckDB top-k ─▶ prompt+context ─▶ LiteRT-LM (Gemma) ─▶ streamed, cited answer

All search state, pipeline status, the query, and ranked results, lives in a single Zustand store (src/store/searchStore.ts). Its init() action boots the embedder and DuckDB index and embeds the corpus once; it's idempotent across StrictMode remounts. setQuery() debounces and runs the search. The embedder and index instances are non-reactive module singletons kept out of the store so they don't trigger renders. Components read state through useVectorSearch(), a thin hook that triggers init() and subscribes to the store.

File-based routes live in src/routes/. The @tanstack/router-plugin Vite plugin regenerates src/routeTree.gen.ts on dev/build (also bun run routes); that file is git-ignored. __root.tsx holds the shared AppShell and nav, / is the search page, and /about shows a live LiteRT.js runtime readout. Router devtools load in dev only.

RAGWASM uses two models, both from the Gemma family:

The Gemma tokenizer for embeddings loads for real from the ungated onnx-community/embeddinggemma-300m-ONNX repo, so no license is needed there.

LiteRtEmbedder runs EmbeddingGemma on LiteRT.js and is wired against the real @litertjs/core API: named-signature introspection, masked mean-pooling, manual tensor cleanup. It also handles the EmbeddingGemma specifics for you:

• Task prompts are asymmetric. Queries are prefixed task: search result | query: …, documents get title: none | text: … (src/lib/embeddingGemma.ts).
• The real Gemma tokenizer loads via @huggingface/transformers from the ungated onnx-community/embeddinggemma-300m-ONNX repo (src/lib/ tokenizer.ts). No stub, no license needed for the tokenizer, and it's lazy-loaded so it stays out of the default bundle.
• Matryoshka (MRL) output is truncated to VITE_EMBEDDING_DIM (768 | 512 | 256 | 128) and re-normalized.

The one asset you provide is the model itself. Get the ready-made LiteRT .tflite from litert-community/embeddinggemma-300m (gated, but approval is automatic): accept the terms, download embeddinggemma-300M_seq256_mixed-precision.tflite, rename it to embeddinggemma.tflite, and drop it at public/models/embeddinggemma.tflite. No env variable is needed for the default path. Full steps are in public/models/README.md. To point elsewhere, or to tune the model, set:

# .env.local (all optional)
VITE_EMBEDDING_MODEL_URL=/models/embeddinggemma.tflite   # default shown
VITE_EMBEDDING_DIM=768          # MRL dim: 768 | 512 | 256 | 128
VITE_MAX_SEQ_LEN=256            # EmbeddingGemma supports up to 2048
VITE_TOKENIZER_URL=…            # HF repo id / served dir override
VITE_LLM_MODEL_URL=…            # override the .litertlm generation model
VITE_CHUNK_SIZE=512             # Chonkie target chunk size (chars)

There is no fallback embedder. Until the model loads, the app shows a clear error state naming the expected path. See src/lib/config.ts for all options.

Note: transformers.js is used only for tokenization. The embedding model itself runs on LiteRT.js, by design.

• WASM assets: scripts/copy-wasm.mjs copies the @litertjs/core and @litert-lm/core runtimes to public/litertjs/wasm/ and public/litertlm/wasm/ before dev/build (auto-run via npm lifecycle scripts). Both folders are git-ignored. LiteParse, Chonkie, and DuckDB bundle their own WASM via Vite.
• Cross-origin isolation: vite.config.ts sets COOP/COEP headers so LiteRT can load its multi-threaded WASM build (SharedArrayBuffer). Any production host must send the same headers for threading; without them it still runs single-threaded.
• Memory: LiteRT tensors are manually managed. LiteRtEmbedder deletes every input/output tensor after each run.