Add ai_binding_experimental compat flag for AutoRAG/Gateway/ToMarkdown RPC

src/cloudflare/internal/ai-api.ts

@@ -13,10 +13,14 @@ import {
 
 type AiSearchService = object;
 
+const aiBindingExperimental = !!Cloudflare.compatibilityFlags['experimental'];
+
 interface Fetcher {
   fetch: typeof fetch;
   aiSearch: () => AiSearchService;
   gateway: (gatewayId: string) => AiGateway;
+  autorag: (autoragId?: string) => AutoRAG;
+  toMarkdown: () => ToMarkdownService;
 }
 
 interface AiError {
@@ -407,7 +411,9 @@ export class Ai {
     files?: MarkdownDocument | MarkdownDocument[],
     options?: ConversionRequestOptions
   ): ToMarkdownService | Promise<ConversionResponse | ConversionResponse[]> {
-    const service = new ToMarkdownService(this.#fetcher);
+    const service = aiBindingExperimental
+      ? this.#fetcher.toMarkdown()
+      : new ToMarkdownService(this.#fetcher);
 
     if (arguments.length < 1 || !files) return service;
 
@@ -420,13 +426,16 @@ export class Ai {
   }
 
   gateway(gatewayId: string, options?: { beta?: boolean }): AiGateway {
-    if (options?.beta === true) {
+    if (aiBindingExperimental || options?.beta === true) {
       return this.#fetcher.gateway(gatewayId);
     }
     return new AiGateway(this.#fetcher, gatewayId);
   }
 
   autorag(autoragId?: string): AutoRAG {
+    if (aiBindingExperimental) {
+      return this.#fetcher.autorag(autoragId);
+    }
     return new AutoRAG(this.#fetcher, autoragId);
   }
 

src/cloudflare/internal/test/autorag/BUILD.bazel

@@ -7,6 +7,12 @@ wd_test(
     data = glob(["*.js"]),
 )
 
+wd_test(
+    src = "autorag-api-rpc-test.wd-test",
+    args = ["--experimental"],
+    data = glob(["*.js"]),
+)
+
 py_wd_test(
     size = "large",
     src = "python-autorag-api-test.wd-test",

src/cloudflare/internal/test/autorag/autorag-api-rpc-test.wd-test

@@ -0,0 +1,38 @@
+using Workerd = import "/workerd/workerd.capnp";
+
+const unitTests :Workerd.Config = (
+  services = [
+    ( name = "autorag-api-test",
+      worker = (
+        modules = [
+          (name = "worker", esModule = embed "autorag-api-test.js")
+        ],
+        compatibilityFlags = ["experimental", "nodejs_compat",
+                              "service_binding_extra_handlers", "rpc"],
+        bindings = [
+        (
+          name = "ai",
+          wrapped = (
+            moduleName = "cloudflare-internal:ai-api",
+            innerBindings = [(
+              name = "fetcher",
+              service = (
+                name = "autorag-mock",
+                entrypoint = "ServiceEntrypoint"
+              )
+            )],
+          )
+        )
+        ],
+      )
+    ),
+    ( name = "autorag-mock",
+      worker = (
+        compatibilityFlags = ["experimental", "nodejs_compat"],
+        modules = [
+          (name = "worker", esModule = embed "autorag-mock.js")
+        ],
+      )
+    )
+  ]
+);

src/cloudflare/internal/test/autorag/autorag-mock.js

@@ -2,52 +2,88 @@
 // Licensed under the Apache 2.0 license found in the LICENSE file or at:
 //     https://opensource.org/licenses/Apache-2.0
 
-export default {
-  async fetch(request, env, ctx) {
+import { WorkerEntrypoint, RpcTarget } from 'cloudflare:workers';
+
+const SEARCH_DATA = [
+  {
+    file_id: 'file-12345',
+    filename: 'woodchuck_policy.txt',
+    score: 0.85,
+    attributes: {
+      region: 'North America',
+      author: 'Wildlife Department',
+    },
+    content: [
+      {
+        type: 'text',
+        text: 'According to the latest regulations, each passenger is allowed to carry up to two woodchucks.',
+      },
+      {
+        type: 'text',
+        text: 'Ensure that the woodchucks are properly contained during transport.',
+      },
+    ],
+  },
+  {
+    file_id: 'file-67890',
+    filename: 'transport_guidelines.txt',
+    score: 0.75,
+    attributes: {
+      region: 'North America',
+      author: 'Transport Authority',
+    },
+    content: [
+      {
+        type: 'text',
+        text: 'Passengers must adhere to the guidelines set forth by the Transport Authority regarding the transport of woodchucks.',
+      },
+    ],
+  },
+];
+
+const SEARCH_RESPONSE = {
+  object: 'vector_store.search_results.page',
+  search_query: 'How many woodchucks are allowed per passenger?',
+  data: SEARCH_DATA,
+  has_more: false,
+  next_page: null,
+};
+
+/** RPC target returned by autorag(autoragId) */
+class AutoRAGInstance extends RpcTarget {
+  #autoragId;
+
+  constructor(autoragId) {
+    super();
+    this.#autoragId = autoragId;
+  }
+
+  async search(params) {
+    return {
+      ...SEARCH_RESPONSE,
+    };
+  }
+
+  async aiSearch(params) {
+    return {
+      ...SEARCH_RESPONSE,
+      response: 'this is an example result',
+    };
+  }
+}
+
+/** WorkerEntrypoint that supports both fetch (legacy) and RPC (beta) paths */
+export class ServiceEntrypoint extends WorkerEntrypoint {
+  autorag(autoragId) {
+    return new AutoRAGInstance(autoragId);
+  }
+
+  async fetch(request) {
     if (request.method === 'POST') {
       if (request.url.includes('/ai-search')) {
         return Response.json({
           result: {
-            object: 'vector_store.search_results.page',
-            search_query: 'How many woodchucks are allowed per passenger?',
-            data: [
-              {
-                file_id: 'file-12345',
-                filename: 'woodchuck_policy.txt',
-                score: 0.85,
-                attributes: {
-                  region: 'North America',
-                  author: 'Wildlife Department',
-                },
-                content: [
-                  {
-                    type: 'text',
-                    text: 'According to the latest regulations, each passenger is allowed to carry up to two woodchucks.',
-                  },
-                  {
-                    type: 'text',
-                    text: 'Ensure that the woodchucks are properly contained during transport.',
-                  },
-                ],
-              },
-              {
-                file_id: 'file-67890',
-                filename: 'transport_guidelines.txt',
-                score: 0.75,
-                attributes: {
-                  region: 'North America',
-                  author: 'Transport Authority',
-                },
-                content: [
-                  {
-                    type: 'text',
-                    text: 'Passengers must adhere to the guidelines set forth by the Transport Authority regarding the transport of woodchucks.',
-                  },
-                ],
-              },
-            ],
-            has_more: false,
-            next_page: null,
+            ...SEARCH_RESPONSE,
             response: 'this is an example result',
           },
           success: true,
@@ -56,53 +92,19 @@ export default {
 
       if (request.url.includes('/search')) {
         return Response.json({
-          result: {
-            object: 'vector_store.search_results.page',
-            search_query: 'How many woodchucks are allowed per passenger?',
-            data: [
-              {
-                file_id: 'file-12345',
-                filename: 'woodchuck_policy.txt',
-                score: 0.85,
-                attributes: {
-                  region: 'North America',
-                  author: 'Wildlife Department',
-                },
-                content: [
-                  {
-                    type: 'text',
-                    text: 'According to the latest regulations, each passenger is allowed to carry up to two woodchucks.',
-                  },
-                  {
-                    type: 'text',
-                    text: 'Ensure that the woodchucks are properly contained during transport.',
-                  },
-                ],
-              },
-              {
-                file_id: 'file-67890',
-                filename: 'transport_guidelines.txt',
-                score: 0.75,
-                attributes: {
-                  region: 'North America',
-                  author: 'Transport Authority',
-                },
-                content: [
-                  {
-                    type: 'text',
-                    text: 'Passengers must adhere to the guidelines set forth by the Transport Authority regarding the transport of woodchucks.',
-                  },
-                ],
-              },
-            ],
-            has_more: false,
-            next_page: null,
-          },
+          result: SEARCH_RESPONSE,
           success: true,
         });
       }
     }
 
     return Response.json({ success: false }, { status: 500 });
+  }
+}
+
+export default {
+  async fetch(request, env, ctx) {
+    const entrypoint = new ServiceEntrypoint(ctx, env);
+    return entrypoint.fetch(request);
   },
 };