Update CloudFlare worker example to test using streaming LLM

Author: aron

integration/cloudflare-worker/index.js

@@ -1,25 +1,31 @@
 import Replicate from "replicate";
 
-const replicate = new Replicate();
-
 export default {
-  async fetch(_request, _env, _ctx) {
-    const stream = new ReadableStream({
-      async start(controller) {
-        for await (const event of replicate.stream(
-          "replicate/hello-world:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa",
-          {
-            input: {
-              text: "Colin CloudFlare",
-            },
-          }
-        )) {
-          controller.enqueue(`${event}`);
+  async fetch(_request, env, _ctx) {
+    const replicate = new Replicate({ auth: env.REPLICATE_API_TOKEN });
+
+    try {
+      const output = replicate.stream(
+        "replicate/canary:30e22229542eb3f79d4f945dacb58d32001b02cc313ae6f54eef27904edf3272",
+        {
+          input: {
+            text: "Colin CloudFlare",
+          },
         }
-        controller.close();
-      },
-    });
+      );
+      const stream = new ReadableStream({
+        async start(controller) {
+          for await (const event of output) {
+            controller.enqueue(new TextEncoder().encode(`${event}`));
+          }
+          controller.enqueue(new TextEncoder().encode("\n"));
+          controller.close();
+        },
+      });
 
-    return new Response(stream);
+      return new Response(stream);
+    } catch (err) {
+      return Response("", { status: 500 });
+    }
   },
 };

integration/cloudflare-worker/index.test.js

@@ -1,28 +1,36 @@
 // https://developers.cloudflare.com/workers/wrangler/api/#unstable_dev
 import { unstable_dev as dev } from "wrangler";
-import { test, after, before } from "node:test";
+import { test, after, before, describe } from "node:test";
 import assert from "node:assert";
 
-let worker;
+/** @type {import("wrangler").UnstableDevWorker} */
+describe("CloudFlare Worker", () => {
+  let worker;
 
-before(async () => {
-  worker = await dev("index.js", {
-    experimental: { disableExperimentalWarning: true },
+  before(async () => {
+    worker = await dev("./index.js", {
+      port: 3000,
+      experimental: { disableExperimentalWarning: true },
+    });
   });
-});
 
-after(async () => {
-  if (!worker) {
-    // If no worker the before hook failed to run and the process will hang.
-    process.exit(1);
-  }
-  await worker.stop();
-});
+  after(async () => {
+    if (!worker) {
+      // If no worker the before hook failed to run and the process will hang.
+      process.exit(1);
+    }
+    await worker.stop();
+  });
 
-test("worker streams back a response", { timeout: 1000 }, async (t) => {
-  const resp = await worker.fetch("/", { signal: t.signal });
-  const text = await resp.text();
+  test("worker streams back a response", { timeout: 1000 }, async () => {
+    const resp = await worker.fetch();
+    const text = await resp.text();
 
-  assert.ok(resp.ok, "status is 2xx");
-  assert(text.length > 0, "body.length is greater than 0");
+    assert.ok(resp.ok, "status is 2xx");
+    assert(text.length > 0, "body.length is greater than 0");
+    assert(
+      text.includes("Colin CloudFlare"),
+      "body includes stream characters"
+    );
+  });
 });