Merge pull request #9 from jeasonstudio/feat-embedding-model

feat: add embedding model

.changeset/sweet-donuts-prove.md

@@ -0,0 +1,5 @@
+---
+"chrome-ai": minor
+---
+
+feat: add embedding model for client side embed with ai sdk

package.json

@@ -35,7 +35,8 @@
     "access": "public"
   },
   "dependencies": {
-    "@ai-sdk/provider": "^0.0.10",
+    "@ai-sdk/provider": "^0.0.11",
+    "@mediapipe/tasks-text": "^0.10.14",
     "debug": "^4.3.5"
   },
   "devDependencies": {
@@ -46,16 +47,16 @@
     "@radix-ui/react-dropdown-menu": "^2.1.0",
     "@radix-ui/react-label": "^2.1.0",
     "@radix-ui/react-select": "^2.1.0",
-    "@radix-ui/react-slot": "^1.1.0",
     "@radix-ui/react-slider": "^1.2.0",
+    "@radix-ui/react-slot": "^1.1.0",
     "@radix-ui/react-tooltip": "^1.1.2",
     "@tailwindcss/typography": "^0.5.13",
     "@types/debug": "^4.1.12",
     "@types/node": "^20",
     "@types/react": "^18",
     "@types/react-dom": "^18",
     "@vitest/coverage-v8": "^1.6.0",
-    "ai": "^3.1.31",
+    "ai": "^3.2.16",
     "autoprefixer": "^10.4.19",
     "class-variance-authority": "^0.7.0",
     "clsx": "^2.1.1",

src/embedding-model.test.ts

@@ -0,0 +1,53 @@
+import { describe, it, expect, vi } from 'vitest';
+import { ChromeAIEmbeddingModel, chromeEmbedding } from './embedding-model';
+import { embed } from 'ai';
+
+vi.mock('@mediapipe/tasks-text', async () => ({
+  FilesetResolver: {
+    forTextTasks: vi.fn(async () => ({
+      wasmLoaderPath: 'wasmLoaderPath',
+      wasmBinaryPath: 'wasmBinaryPath',
+    })),
+  },
+  TextEmbedder: {
+    createFromOptions: vi.fn(async () => ({
+      embed: vi.fn((text: string) => ({
+        embeddings: [
+          { floatEmbedding: text === 'undefined' ? undefined : [1, 2] },
+        ],
+      })),
+    })),
+  },
+}));
+
+describe('embedding-model', () => {
+  it('should instantiation anyways', async () => {
+    expect(new ChromeAIEmbeddingModel()).toBeInstanceOf(ChromeAIEmbeddingModel);
+    expect(chromeEmbedding()).toBeInstanceOf(ChromeAIEmbeddingModel);
+  });
+  it('should embed', async () => {
+    const model = chromeEmbedding();
+    expect(
+      await embed({
+        model,
+        value: 'test',
+      })
+    ).toMatchObject({ embedding: [1, 2] });
+
+    expect(
+      await embed({
+        model,
+        value: 'test2',
+      })
+    ).toMatchObject({ embedding: [1, 2] });
+  });
+
+  it('should embed result empty', async () => {
+    expect(
+      await embed({
+        model: chromeEmbedding({ l2Normalize: true }),
+        value: 'undefined',
+      })
+    ).toMatchObject({ embedding: [] });
+  });
+});