fix: polyfill for cloudflare worker (run-llama#919)

himself65 · Jun 10, 2024 · bf25ff6 · bf25ff6
1 parent 32ad099
commit bf25ff6
Show file tree

Hide file tree

Showing 7 changed files with 114 additions and 59 deletions.
diff --git a/.changeset/curly-singers-help.md b/.changeset/curly-singers-help.md
@@ -0,0 +1,7 @@
+---
+"llamaindex": patch
+"@llamaindex/core-e2e": patch
+"@llamaindex/cloudflare-worker-agent-test": patch
+---
+
+fix: polyfill for cloudflare worker
diff --git a/packages/core/e2e/examples/cloudflare-worker-agent/package.json b/packages/core/e2e/examples/cloudflare-worker-agent/package.json
@@ -12,13 +12,13 @@
     "cf-typegen": "wrangler types"
   },
   "devDependencies": {
-    "@cloudflare/vitest-pool-workers": "^0.2.6",
-    "@cloudflare/workers-types": "^4.20240502.0",
+    "@cloudflare/vitest-pool-workers": "^0.4.3",
+    "@cloudflare/workers-types": "^4.20240605.0",
     "@vitest/runner": "1.3.0",
     "@vitest/snapshot": "1.3.0",
     "typescript": "^5.4.5",
     "vitest": "1.3.0",
-    "wrangler": "^3.53.1"
+    "wrangler": "^3.60.1"
   },
   "dependencies": {
     "llamaindex": "workspace:*"

diff --git a/packages/core/e2e/examples/cloudflare-worker-agent/src/index.ts b/packages/core/e2e/examples/cloudflare-worker-agent/src/index.ts
@@ -10,10 +10,12 @@ export default {
     const agent = new OpenAIAgent({
       tools: [],
     });
+    console.log(1);
     const responseStream = await agent.chat({
       stream: true,
       message: "Hello? What is the weather today?",
     });
+    console.log(2);
     const textEncoder = new TextEncoder();
     const response = responseStream.pipeThrough<Uint8Array>(
       // @ts-expect-error: see https://github.com/cloudflare/workerd/issues/2067

diff --git a/packages/core/src/embeddings/CloudflareWorkerEmbedding.ts b/packages/core/src/embeddings/CloudflareWorkerEmbedding.ts
@@ -0,0 +1,14 @@
+import type { ImageType } from "../Node.js";
+import { MultiModalEmbedding } from "./MultiModalEmbedding.js";
+
+/**
+ * Cloudflare worker doesn't support image embeddings for now
+ */
+export class CloudflareWorkerMultiModalEmbedding extends MultiModalEmbedding {
+  getImageEmbedding(images: ImageType): Promise<number[]> {
+    throw new Error("Method not implemented.");
+  }
+  getTextEmbedding(text: string): Promise<number[]> {
+    throw new Error("Method not implemented.");
+  }
+}
diff --git a/packages/core/src/internal/settings/image-embed-model.ts b/packages/core/src/internal/settings/image-embed-model.ts
@@ -0,0 +1,8 @@
+export async function getImageEmbedModel() {
+  if (globalThis.navigator?.userAgent === "Cloudflare-Workers") {
+    return (await import("../../embeddings/CloudflareWorkerEmbedding.js"))
+      .CloudflareWorkerMultiModalEmbedding;
+  } else {
+    return (await import("../../embeddings/ClipEmbedding.js")).ClipEmbedding;
+  }
+}
diff --git a/packages/core/src/storage/StorageContext.ts b/packages/core/src/storage/StorageContext.ts
@@ -1,6 +1,6 @@
 import { path } from "@llamaindex/env";
 import { ModalityType, ObjectType } from "../Node.js";
-import { ClipEmbedding } from "../embeddings/ClipEmbedding.js";
+import { getImageEmbedModel } from "../internal/settings/image-embed-model.js";
 import {
   DEFAULT_IMAGE_VECTOR_NAMESPACE,
   DEFAULT_NAMESPACE,
@@ -44,7 +44,7 @@ export async function storageContextFromDefaults({
     }
     if (storeImages && !(ModalityType.IMAGE in vectorStores)) {
       vectorStores[ModalityType.IMAGE] = new SimpleVectorStore({
-        embedModel: new ClipEmbedding(),
+        embedModel: new (await getImageEmbedModel())(),
       });
     }
   } else {
@@ -60,7 +60,7 @@ export async function storageContextFromDefaults({
     if (storeImages && !(ObjectType.IMAGE in vectorStores)) {
       vectorStores[ModalityType.IMAGE] = await SimpleVectorStore.fromPersistDir(
         path.join(persistDir, DEFAULT_IMAGE_VECTOR_NAMESPACE),
-        new ClipEmbedding(),
+        new (await getImageEmbedModel())(),
       );
     }
   }