workers-ai-provider 0.3.2 → 0.4.1

package/dist/index.d.ts

@@ -4,17 +4,12 @@ type StringLike = string | {
     toString(): string;
 };
 
-type WorkersAIChatSettings = {
+type AutoRAGChatSettings = {
     /**
      * Whether to inject a safety prompt before all conversations.
      * Defaults to `false`.
      */
     safePrompt?: boolean;
-    /**
-     * Optionally set Cloudflare AI Gateway options.
-     * @deprecated
-     */
-    gateway?: GatewayOptions;
 } & {
     /**
      * Passthrough settings that are provided directly to the run function.
@@ -31,6 +26,42 @@ type value2key<T, V> = {
     [K in keyof T]: T[K] extends V ? K : never;
 }[keyof T];
 
+type AutoRAGChatConfig = {
+    provider: string;
+    binding: AutoRAG;
+    gateway?: GatewayOptions;
+};
+declare class AutoRAGChatLanguageModel implements LanguageModelV1 {
+    readonly specificationVersion = "v1";
+    readonly defaultObjectGenerationMode = "json";
+    readonly modelId: TextGenerationModels;
+    readonly settings: AutoRAGChatSettings;
+    private readonly config;
+    constructor(modelId: TextGenerationModels, settings: AutoRAGChatSettings, config: AutoRAGChatConfig);
+    get provider(): string;
+    private getArgs;
+    doGenerate(options: Parameters<LanguageModelV1["doGenerate"]>[0]): Promise<Awaited<ReturnType<LanguageModelV1["doGenerate"]>>>;
+    doStream(options: Parameters<LanguageModelV1["doStream"]>[0]): Promise<Awaited<ReturnType<LanguageModelV1["doStream"]>>>;
+}
+
+type WorkersAIChatSettings = {
+    /**
+     * Whether to inject a safety prompt before all conversations.
+     * Defaults to `false`.
+     */
+    safePrompt?: boolean;
+    /**
+     * Optionally set Cloudflare AI Gateway options.
+     * @deprecated
+     */
+    gateway?: GatewayOptions;
+} & {
+    /**
+     * Passthrough settings that are provided directly to the run function.
+     */
+    [key: string]: StringLike;
+};
+
 type WorkersAIChatConfig = {
     provider: string;
     binding: Ai;
@@ -111,5 +142,19 @@ interface WorkersAI {
  * Create a Workers AI provider instance.
  */
 declare function createWorkersAI(options: WorkersAISettings): WorkersAI;
+type AutoRAGSettings = {
+    binding: AutoRAG;
+};
+interface AutoRAGProvider {
+    (options?: AutoRAGChatSettings): AutoRAGChatLanguageModel;
+    /**
+     * Creates a model for text generation.
+     **/
+    chat(settings?: AutoRAGChatSettings): AutoRAGChatLanguageModel;
+}
+/**
+ * Create a Workers AI provider instance.
+ */
+declare function createAutoRAG(options: AutoRAGSettings): AutoRAGProvider;
 
-export { type WorkersAI, type WorkersAISettings, createWorkersAI };
+export { type AutoRAGProvider, type AutoRAGSettings, type WorkersAI, type WorkersAISettings, createAutoRAG, createWorkersAI };

package/dist/index.js

@@ -9,51 +9,7 @@ var __privateGet = (obj, member, getter) => (__accessCheck(obj, member, "read fr
 var __privateAdd = (obj, member, value) => member.has(obj) ? __typeError("Cannot add the same private member more than once") : member instanceof WeakSet ? member.add(obj) : member.set(obj, value);
 var __privateSet = (obj, member, value, setter) => (__accessCheck(obj, member, "write to private field"), setter ? setter.call(obj, value) : member.set(obj, value), value);
 
-// src/utils.ts
-function createRun(config) {
-  const { accountId, apiKey } = config;
-  return async function run(model, inputs, options) {
-    const { gateway, prefix, extraHeaders, returnRawResponse, ...passthroughOptions } = options || {};
-    const urlParams = new URLSearchParams();
-    for (const [key, value] of Object.entries(passthroughOptions)) {
-      try {
-        const valueStr = value.toString();
-        if (!valueStr) {
-          continue;
-        }
-        urlParams.append(key, valueStr);
-      } catch (error) {
-        throw new Error(
-          `Value for option '${key}' is not able to be coerced into a string.`
-        );
-      }
-    }
-    const url = `https://api.cloudflare.com/client/v4/accounts/${accountId}/ai/run/${model}${urlParams ? `?${urlParams}` : ""}`;
-    const headers = {
-      "Content-Type": "application/json",
-      Authorization: `Bearer ${apiKey}`
-    };
-    const body = JSON.stringify(inputs);
-    const response = await fetch(url, {
-      method: "POST",
-      headers,
-      body
-    });
-    if (returnRawResponse) {
-      return response;
-    }
-    if (inputs.stream === true) {
-      if (response.body) {
-        return response.body;
-      }
-      throw new Error("No readable body available for streaming.");
-    }
-    const data = await response.json();
-    return data.result;
-  };
-}
-
-// src/workersai-chat-language-model.ts
+// src/autorag-chat-language-model.ts
 import {
   UnsupportedFunctionalityError
 } from "@ai-sdk/provider";
@@ -157,6 +113,18 @@ function convertToWorkersAIChatMessages(prompt) {
   return { messages, images };
 }
 
+// src/map-workersai-usage.ts
+function mapWorkersAIUsage(output) {
+  const usage = output.usage ?? {
+    prompt_tokens: 0,
+    completion_tokens: 0
+  };
+  return {
+    promptTokens: usage.prompt_tokens,
+    completionTokens: usage.completion_tokens
+  };
+}
+
 // ../../node_modules/.pnpm/fetch-event-stream@0.1.5/node_modules/fetch-event-stream/esm/deps/jsr.io/@std/streams/0.221.0/text_line_stream.js
 var _currentLine;
 var TextLineStream = class extends TransformStream {
@@ -250,19 +218,274 @@ async function* events(res, signal) {
   }
 }
 
-// src/map-workersai-usage.ts
-function mapWorkersAIUsage(output) {
-  const usage = output.usage ?? {
-    prompt_tokens: 0,
-    completion_tokens: 0
-  };
-  return {
-    promptTokens: usage.prompt_tokens,
-    completionTokens: usage.completion_tokens
+// src/streaming.ts
+function getMappedStream(response) {
+  const chunkEvent = events(response);
+  let usage = { promptTokens: 0, completionTokens: 0 };
+  return new ReadableStream({
+    async start(controller) {
+      for await (const event of chunkEvent) {
+        if (!event.data) {
+          continue;
+        }
+        if (event.data === "[DONE]") {
+          break;
+        }
+        const chunk = JSON.parse(event.data);
+        if (chunk.usage) {
+          usage = mapWorkersAIUsage(chunk);
+        }
+        chunk.response?.length && controller.enqueue({
+          type: "text-delta",
+          textDelta: chunk.response
+        });
+      }
+      controller.enqueue({
+        type: "finish",
+        finishReason: "stop",
+        usage
+      });
+      controller.close();
+    }
+  });
+}
+
+// src/utils.ts
+function createRun(config) {
+  const { accountId, apiKey } = config;
+  return async function run(model, inputs, options) {
+    const { gateway, prefix, extraHeaders, returnRawResponse, ...passthroughOptions } = options || {};
+    const urlParams = new URLSearchParams();
+    for (const [key, value] of Object.entries(passthroughOptions)) {
+      try {
+        const valueStr = value.toString();
+        if (!valueStr) {
+          continue;
+        }
+        urlParams.append(key, valueStr);
+      } catch (error) {
+        throw new Error(
+          `Value for option '${key}' is not able to be coerced into a string.`
+        );
+      }
+    }
+    const url = `https://api.cloudflare.com/client/v4/accounts/${accountId}/ai/run/${model}${urlParams ? `?${urlParams}` : ""}`;
+    const headers = {
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${apiKey}`
+    };
+    const body = JSON.stringify(inputs);
+    const response = await fetch(url, {
+      method: "POST",
+      headers,
+      body
+    });
+    if (returnRawResponse) {
+      return response;
+    }
+    if (inputs.stream === true) {
+      if (response.body) {
+        return response.body;
+      }
+      throw new Error("No readable body available for streaming.");
+    }
+    const data = await response.json();
+    return data.result;
   };
 }
+function prepareToolsAndToolChoice(mode) {
+  const tools = mode.tools?.length ? mode.tools : void 0;
+  if (tools == null) {
+    return { tools: void 0, tool_choice: void 0 };
+  }
+  const mappedTools = tools.map((tool) => ({
+    type: "function",
+    function: {
+      name: tool.name,
+      // @ts-expect-error - description is not a property of tool
+      description: tool.description,
+      // @ts-expect-error - parameters is not a property of tool
+      parameters: tool.parameters
+    }
+  }));
+  const toolChoice = mode.toolChoice;
+  if (toolChoice == null) {
+    return { tools: mappedTools, tool_choice: void 0 };
+  }
+  const type = toolChoice.type;
+  switch (type) {
+    case "auto":
+      return { tools: mappedTools, tool_choice: type };
+    case "none":
+      return { tools: mappedTools, tool_choice: type };
+    case "required":
+      return { tools: mappedTools, tool_choice: "any" };
+    // workersAI does not support tool mode directly,
+    // so we filter the tools and force the tool choice through 'any'
+    case "tool":
+      return {
+        tools: mappedTools.filter((tool) => tool.function.name === toolChoice.toolName),
+        tool_choice: "any"
+      };
+    default: {
+      const exhaustiveCheck = type;
+      throw new Error(`Unsupported tool choice type: ${exhaustiveCheck}`);
+    }
+  }
+}
+function lastMessageWasUser(messages) {
+  return messages.length > 0 && messages[messages.length - 1].role === "user";
+}
+function processToolCalls(output) {
+  if (output.tool_calls && Array.isArray(output.tool_calls)) {
+    return output.tool_calls.map((toolCall) => {
+      if (toolCall.function && toolCall.id) {
+        return {
+          toolCallType: "function",
+          toolCallId: toolCall.id,
+          toolName: toolCall.function.name,
+          args: typeof toolCall.function.arguments === "string" ? toolCall.function.arguments : JSON.stringify(toolCall.function.arguments || {})
+        };
+      }
+      return {
+        toolCallType: "function",
+        toolCallId: toolCall.name,
+        toolName: toolCall.name,
+        args: typeof toolCall.arguments === "string" ? toolCall.arguments : JSON.stringify(toolCall.arguments || {})
+      };
+    });
+  }
+  return [];
+}
+
+// src/autorag-chat-language-model.ts
+var AutoRAGChatLanguageModel = class {
+  constructor(modelId, settings, config) {
+    __publicField(this, "specificationVersion", "v1");
+    __publicField(this, "defaultObjectGenerationMode", "json");
+    __publicField(this, "modelId");
+    __publicField(this, "settings");
+    __publicField(this, "config");
+    this.modelId = modelId;
+    this.settings = settings;
+    this.config = config;
+  }
+  get provider() {
+    return this.config.provider;
+  }
+  getArgs({
+    mode,
+    prompt,
+    frequencyPenalty,
+    presencePenalty
+  }) {
+    const type = mode.type;
+    const warnings = [];
+    if (frequencyPenalty != null) {
+      warnings.push({
+        type: "unsupported-setting",
+        setting: "frequencyPenalty"
+      });
+    }
+    if (presencePenalty != null) {
+      warnings.push({
+        type: "unsupported-setting",
+        setting: "presencePenalty"
+      });
+    }
+    const baseArgs = {
+      // model id:
+      model: this.modelId,
+      // messages:
+      messages: convertToWorkersAIChatMessages(prompt)
+    };
+    switch (type) {
+      case "regular": {
+        return {
+          args: { ...baseArgs, ...prepareToolsAndToolChoice(mode) },
+          warnings
+        };
+      }
+      case "object-json": {
+        return {
+          args: {
+            ...baseArgs,
+            response_format: {
+              type: "json_schema",
+              json_schema: mode.schema
+            },
+            tools: void 0
+          },
+          warnings
+        };
+      }
+      case "object-tool": {
+        return {
+          args: {
+            ...baseArgs,
+            tool_choice: "any",
+            tools: [{ type: "function", function: mode.tool }]
+          },
+          warnings
+        };
+      }
+      // @ts-expect-error - this is unreachable code
+      // TODO: fixme
+      case "object-grammar": {
+        throw new UnsupportedFunctionalityError({
+          functionality: "object-grammar mode"
+        });
+      }
+      default: {
+        const exhaustiveCheck = type;
+        throw new Error(`Unsupported type: ${exhaustiveCheck}`);
+      }
+    }
+  }
+  async doGenerate(options) {
+    const { args, warnings } = this.getArgs(options);
+    const { messages } = convertToWorkersAIChatMessages(options.prompt);
+    const output = await this.config.binding.aiSearch({
+      query: messages.map(({ content, role }) => `${role}: ${content}`).join("\n\n")
+    });
+    return {
+      text: output.response,
+      toolCalls: processToolCalls(output),
+      finishReason: "stop",
+      // TODO: mapWorkersAIFinishReason(response.finish_reason),
+      rawCall: { rawPrompt: args.messages, rawSettings: args },
+      usage: mapWorkersAIUsage(output),
+      warnings,
+      sources: output.data.map(({ file_id, filename, score }) => ({
+        id: file_id,
+        sourceType: "url",
+        url: filename,
+        providerMetadata: {
+          attributes: { score }
+        }
+      }))
+    };
+  }
+  async doStream(options) {
+    const { args, warnings } = this.getArgs(options);
+    const { messages } = convertToWorkersAIChatMessages(options.prompt);
+    const query = messages.map(({ content, role }) => `${role}: ${content}`).join("\n\n");
+    const response = await this.config.binding.aiSearch({
+      query,
+      stream: true
+    });
+    return {
+      stream: getMappedStream(response),
+      rawCall: { rawPrompt: args.messages, rawSettings: args },
+      warnings
+    };
+  }
+};
 
 // src/workersai-chat-language-model.ts
+import {
+  UnsupportedFunctionalityError as UnsupportedFunctionalityError2
+} from "@ai-sdk/provider";
 var WorkersAIChatLanguageModel = class {
   constructor(modelId, settings, config) {
     __publicField(this, "specificationVersion", "v1");
@@ -344,7 +567,7 @@ var WorkersAIChatLanguageModel = class {
       // @ts-expect-error - this is unreachable code
       // TODO: fixme
       case "object-grammar": {
-        throw new UnsupportedFunctionalityError({
+        throw new UnsupportedFunctionalityError2({
           functionality: "object-grammar mode"
         });
       }
@@ -357,9 +580,7 @@ var WorkersAIChatLanguageModel = class {
   async doGenerate(options) {
     const { args, warnings } = this.getArgs(options);
     const { gateway, safePrompt, ...passthroughOptions } = this.settings;
-    const { messages, images } = convertToWorkersAIChatMessages(
-      options.prompt
-    );
+    const { messages, images } = convertToWorkersAIChatMessages(options.prompt);
     if (images.length !== 0 && images.length !== 1) {
       throw new Error("Multiple images are not yet supported as input");
     }
@@ -385,12 +606,7 @@ var WorkersAIChatLanguageModel = class {
     }
     return {
       text: typeof output.response === "object" && output.response !== null ? JSON.stringify(output.response) : output.response,
-      toolCalls: output.tool_calls?.map((toolCall) => ({
-        toolCallType: "function",
-        toolCallId: toolCall.name,
-        toolName: toolCall.name,
-        args: JSON.stringify(toolCall.arguments || {})
-      })),
+      toolCalls: processToolCalls(output),
       finishReason: "stop",
       // TODO: mapWorkersAIFinishReason(response.finish_reason),
       rawCall: { rawPrompt: messages, rawSettings: args },
@@ -400,9 +616,7 @@ var WorkersAIChatLanguageModel = class {
   }
   async doStream(options) {
     const { args, warnings } = this.getArgs(options);
-    const { messages, images } = convertToWorkersAIChatMessages(
-      options.prompt
-    );
+    const { messages, images } = convertToWorkersAIChatMessages(options.prompt);
     if (args.tools?.length && lastMessageWasUser(messages)) {
       const response2 = await this.doGenerate(options);
       if (response2 instanceof ReadableStream) {
@@ -462,85 +676,13 @@ var WorkersAIChatLanguageModel = class {
     if (!(response instanceof ReadableStream)) {
       throw new Error("This shouldn't happen");
     }
-    const chunkEvent = events(new Response(response));
-    let usage = { promptTokens: 0, completionTokens: 0 };
     return {
-      stream: new ReadableStream({
-        async start(controller) {
-          for await (const event of chunkEvent) {
-            if (!event.data) {
-              continue;
-            }
-            if (event.data === "[DONE]") {
-              break;
-            }
-            const chunk = JSON.parse(event.data);
-            if (chunk.usage) {
-              usage = mapWorkersAIUsage(chunk);
-            }
-            chunk.response?.length && controller.enqueue({
-              type: "text-delta",
-              textDelta: chunk.response
-            });
-          }
-          controller.enqueue({
-            type: "finish",
-            finishReason: "stop",
-            usage
-          });
-          controller.close();
-        }
-      }),
+      stream: getMappedStream(new Response(response)),
       rawCall: { rawPrompt: messages, rawSettings: args },
       warnings
     };
   }
 };
-function prepareToolsAndToolChoice(mode) {
-  const tools = mode.tools?.length ? mode.tools : void 0;
-  if (tools == null) {
-    return { tools: void 0, tool_choice: void 0 };
-  }
-  const mappedTools = tools.map((tool) => ({
-    type: "function",
-    function: {
-      name: tool.name,
-      // @ts-expect-error - description is not a property of tool
-      description: tool.description,
-      // @ts-expect-error - parameters is not a property of tool
-      parameters: tool.parameters
-    }
-  }));
-  const toolChoice = mode.toolChoice;
-  if (toolChoice == null) {
-    return { tools: mappedTools, tool_choice: void 0 };
-  }
-  const type = toolChoice.type;
-  switch (type) {
-    case "auto":
-      return { tools: mappedTools, tool_choice: type };
-    case "none":
-      return { tools: mappedTools, tool_choice: type };
-    case "required":
-      return { tools: mappedTools, tool_choice: "any" };
-    // workersAI does not support tool mode directly,
-    // so we filter the tools and force the tool choice through 'any'
-    case "tool":
-      return {
-        tools: mappedTools.filter(
-          (tool) => tool.function.name === toolChoice.toolName
-        ),
-        tool_choice: "any"
-      };
-    default: {
-      const exhaustiveCheck = type;
-      throw new Error(`Unsupported tool choice type: ${exhaustiveCheck}`);
-    }
-  }
-}
-function lastMessageWasUser(messages) {
-  return messages.length > 0 && messages[messages.length - 1].role === "user";
-}
 
 // src/workersai-image-model.ts
 var WorkersAIImageModel = class {
@@ -666,7 +808,23 @@ function createWorkersAI(options) {
   provider.imageModel = createImageModel;
   return provider;
 }
+function createAutoRAG(options) {
+  const binding = options.binding;
+  const createChatModel = (settings = {}) => new AutoRAGChatLanguageModel("@cf/meta/llama-3.3-70b-instruct-fp8-fast", settings, {
+    provider: "autorag.chat",
+    binding
+  });
+  const provider = (settings) => {
+    if (new.target) {
+      throw new Error("The WorkersAI model function cannot be called with the new keyword.");
+    }
+    return createChatModel(settings);
+  };
+  provider.chat = createChatModel;
+  return provider;
+}
 export {
+  createAutoRAG,
   createWorkersAI
 };
 //# sourceMappingURL=index.js.map