workers-ai-provider 0.0.12 → 0.1.0

package/dist/index.d.ts

@@ -9,6 +9,7 @@ interface WorkersAIChatSettings {
     safePrompt?: boolean;
     /**
      * Optionally set Cloudflare AI Gateway options.
+     * @deprecated
      */
     gateway?: GatewayOptions;
 }
@@ -16,13 +17,15 @@ interface WorkersAIChatSettings {
 /**
  * The names of the BaseAiTextGeneration models.
  */
-type TextGenerationModels = {
-    [K in keyof AiModels]: AiModels[K] extends BaseAiTextGeneration ? K : never;
-}[keyof AiModels];
+type TextGenerationModels = Exclude<value2key<AiModels, BaseAiTextGeneration>, value2key<AiModels, BaseAiTextToImage>>;
+type value2key<T, V> = {
+    [K in keyof T]: T[K] extends V ? K : never;
+}[keyof T];
 
 type WorkersAIChatConfig = {
     provider: string;
     binding: Ai;
+    gateway?: GatewayOptions;
 };
 declare class WorkersAIChatLanguageModel implements LanguageModelV1 {
     readonly specificationVersion = "v1";
@@ -56,6 +59,10 @@ interface WorkersAISettings {
     ```
      **/
     binding: Ai;
+    /**
+     * Optionally set Cloudflare AI Gateway options.
+     */
+    gateway?: GatewayOptions;
 }
 /**
  * Create a Workers AI provider instance.

package/dist/index.js

@@ -96,6 +96,20 @@ function convertToWorkersAIChatMessages(prompt) {
 
 // src/workersai-chat-language-model.ts
 import { events } from "fetch-event-stream";
+
+// 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/workersai-chat-language-model.ts
 var WorkersAIChatLanguageModel = class {
   constructor(modelId, settings, config) {
     __publicField(this, "specificationVersion", "v1");
@@ -193,11 +207,12 @@ var WorkersAIChatLanguageModel = class {
       args.model,
       {
         messages: args.messages,
-        tools: args.tools
+        max_tokens: args.max_tokens,
+        temperature: args.temperature,
+        tools: args.tools,
+        top_p: args.top_p
       },
-      {
-        gateway: this.settings.gateway
-      }
+      { gateway: this.config.gateway ?? this.settings.gateway }
     );
     if (output instanceof ReadableStream) {
       throw new Error("This shouldn't happen");
@@ -213,63 +228,56 @@ var WorkersAIChatLanguageModel = class {
       finishReason: "stop",
       // TODO: mapWorkersAIFinishReason(response.finish_reason),
       rawCall: { rawPrompt: args.messages, rawSettings: args },
-      usage: {
-        // TODO: mapWorkersAIUsage(response.usage),
-        promptTokens: 0,
-        completionTokens: 0
-      },
+      usage: mapWorkersAIUsage(output),
       warnings
     };
   }
   async doStream(options) {
     const { args, warnings } = this.getArgs(options);
-    const decoder = new TextDecoder();
-    const response = await this.config.binding.run(args.model, {
-      messages: args.messages,
-      stream: true,
-      tools: args.tools
-    });
+    const response = await this.config.binding.run(
+      args.model,
+      {
+        messages: args.messages,
+        max_tokens: args.max_tokens,
+        stream: true,
+        temperature: args.temperature,
+        tools: args.tools,
+        top_p: args.top_p
+      },
+      { gateway: this.config.gateway ?? this.settings.gateway }
+    );
     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: response.pipeThrough(
-        new TransformStream({
-          async transform(chunk, controller) {
-            const chunkToText = decoder.decode(chunk);
-            const chunks = events(new Response(chunkToText));
-            for await (const singleChunk of chunks) {
-              if (!singleChunk.data) {
-                continue;
-              }
-              if (singleChunk.data === "[DONE]") {
-                controller.enqueue({
-                  type: "finish",
-                  finishReason: "stop",
-                  usage: {
-                    promptTokens: 0,
-                    completionTokens: 0
-                  }
-                });
-                return;
-              }
-              const data = JSON.parse(singleChunk.data);
-              controller.enqueue({
-                type: "text-delta",
-                textDelta: data.response ?? "DATALOSS"
-              });
+      stream: new ReadableStream({
+        async start(controller) {
+          for await (const event of chunkEvent) {
+            if (!event.data) {
+              continue;
             }
-            controller.enqueue({
-              type: "finish",
-              finishReason: "stop",
-              usage: {
-                promptTokens: 0,
-                completionTokens: 0
-              }
+            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();
+        }
+      }),
       rawCall: { rawPrompt: args.messages, rawSettings: args },
       warnings
     };
@@ -326,7 +334,8 @@ function prepareToolsAndToolChoice(mode) {
 function createWorkersAI(options) {
   const createChatModel = (modelId, settings = {}) => new WorkersAIChatLanguageModel(modelId, settings, {
     provider: "workersai.chat",
-    binding: options.binding
+    binding: options.binding,
+    gateway: options.gateway
   });
   const provider = function(modelId, settings) {
     if (new.target) {

package/dist/index.js.map

@@ -1 +1 @@
-{"version":3,"sources":["../src/workersai-chat-language-model.ts","../src/convert-to-workersai-chat-messages.ts","../src/index.ts"],"sourcesContent":["import {\n  type LanguageModelV1,\n  type LanguageModelV1CallWarning,\n  type LanguageModelV1StreamPart,\n  UnsupportedFunctionalityError,\n} from \"@ai-sdk/provider\";\nimport { z } from \"zod\";\nimport { convertToWorkersAIChatMessages } from \"./convert-to-workersai-chat-messages\";\nimport type { WorkersAIChatSettings } from \"./workersai-chat-settings\";\nimport type { TextGenerationModels } from \"./workersai-models\";\n\nimport { events } from \"fetch-event-stream\";\n\ntype WorkersAIChatConfig = {\n  provider: string;\n  binding: Ai;\n};\n\nexport class WorkersAIChatLanguageModel implements LanguageModelV1 {\n  readonly specificationVersion = \"v1\";\n  readonly defaultObjectGenerationMode = \"json\";\n\n  readonly modelId: TextGenerationModels;\n  readonly settings: WorkersAIChatSettings;\n\n  private readonly config: WorkersAIChatConfig;\n\n  constructor(\n    modelId: TextGenerationModels,\n    settings: WorkersAIChatSettings,\n    config: WorkersAIChatConfig\n  ) {\n    this.modelId = modelId;\n    this.settings = settings;\n    this.config = config;\n  }\n\n  get provider(): string {\n    return this.config.provider;\n  }\n\n  private getArgs({\n    mode,\n    prompt,\n    maxTokens,\n    temperature,\n    topP,\n    frequencyPenalty,\n    presencePenalty,\n    seed,\n  }: Parameters<LanguageModelV1[\"doGenerate\"]>[0]) {\n    const type = mode.type;\n\n    const warnings: LanguageModelV1CallWarning[] = [];\n\n    if (frequencyPenalty != null) {\n      warnings.push({\n        type: \"unsupported-setting\",\n        setting: \"frequencyPenalty\",\n      });\n    }\n\n    if (presencePenalty != null) {\n      warnings.push({\n        type: \"unsupported-setting\",\n        setting: \"presencePenalty\",\n      });\n    }\n\n    const baseArgs = {\n      // model id:\n      model: this.modelId,\n\n      // model specific settings:\n      safe_prompt: this.settings.safePrompt,\n\n      // standardized settings:\n      max_tokens: maxTokens,\n      temperature,\n      top_p: topP,\n      random_seed: seed,\n\n      // messages:\n      messages: convertToWorkersAIChatMessages(prompt),\n    };\n\n    switch (type) {\n      case \"regular\": {\n        return {\n          args: { ...baseArgs, ...prepareToolsAndToolChoice(mode) },\n          warnings,\n        };\n      }\n\n      case \"object-json\": {\n        return {\n          args: {\n            ...baseArgs,\n            response_format: { type: \"json_object\" },\n            tools: undefined,\n          },\n          warnings,\n        };\n      }\n\n      case \"object-tool\": {\n        return {\n          args: {\n            ...baseArgs,\n            tool_choice: \"any\",\n            tools: [{ type: \"function\", function: mode.tool }],\n          },\n          warnings,\n        };\n      }\n\n      // @ts-expect-error - this is unreachable code\n      // TODO: fixme\n      case \"object-grammar\": {\n        throw new UnsupportedFunctionalityError({\n          functionality: \"object-grammar mode\",\n        });\n      }\n\n      default: {\n        const exhaustiveCheck = type satisfies never;\n        throw new Error(`Unsupported type: ${exhaustiveCheck}`);\n      }\n    }\n  }\n\n  async doGenerate(\n    options: Parameters<LanguageModelV1[\"doGenerate\"]>[0]\n  ): Promise<Awaited<ReturnType<LanguageModelV1[\"doGenerate\"]>>> {\n    const { args, warnings } = this.getArgs(options);\n\n    const output = await this.config.binding.run(\n      args.model,\n      {\n        messages: args.messages,\n        tools: args.tools,\n      },\n      {\n        gateway: this.settings.gateway,\n      }\n    );\n\n    if (output instanceof ReadableStream) {\n      throw new Error(\"This shouldn't happen\");\n    }\n\n    return {\n      text: output.response,\n      toolCalls: output.tool_calls?.map((toolCall) => ({\n        toolCallType: \"function\",\n        toolCallId: toolCall.name,\n        toolName: toolCall.name,\n        args: JSON.stringify(toolCall.arguments || {}),\n      })),\n      finishReason: \"stop\", // TODO: mapWorkersAIFinishReason(response.finish_reason),\n      rawCall: { rawPrompt: args.messages, rawSettings: args },\n      usage: {\n\t    // TODO: mapWorkersAIUsage(response.usage),\n  \t\tpromptTokens: 0,\n\t    completionTokens: 0,\n      },\n      warnings,\n    };\n  }\n\n  async doStream(\n    options: Parameters<LanguageModelV1[\"doStream\"]>[0]\n  ): Promise<Awaited<ReturnType<LanguageModelV1[\"doStream\"]>>> {\n    const { args, warnings } = this.getArgs(options);\n\n    const decoder = new TextDecoder();\n\n    const response = await this.config.binding.run(args.model, {\n      messages: args.messages,\n      stream: true,\n      tools: args.tools,\n    });\n\n    if (!(response instanceof ReadableStream)) {\n      throw new Error(\"This shouldn't happen\");\n    }\n\n    return {\n      stream: response.pipeThrough(\n        new TransformStream<Uint8Array, LanguageModelV1StreamPart>({\n          async transform(chunk, controller) {\n            const chunkToText = decoder.decode(chunk);\n            const chunks = events(new Response(chunkToText));\n            for await (const singleChunk of chunks) {\n              if (!singleChunk.data) {\n                continue;\n              }\n              if (singleChunk.data === \"[DONE]\") {\n                controller.enqueue({\n                  type: \"finish\",\n                  finishReason: \"stop\",\n                  usage: {\n                    promptTokens: 0,\n                    completionTokens: 0,\n                  },\n                });\n                return;\n              }\n              const data = JSON.parse(singleChunk.data);\n\n              controller.enqueue({\n                type: \"text-delta\",\n                textDelta: data.response ?? \"DATALOSS\",\n              });\n            }\n            controller.enqueue({\n              type: \"finish\",\n              finishReason: \"stop\",\n              usage: {\n                promptTokens: 0,\n                completionTokens: 0,\n              },\n            });\n          },\n        })\n      ),\n      rawCall: { rawPrompt: args.messages, rawSettings: args },\n      warnings,\n    };\n  }\n}\n// limited version of the schema, focussed on what is needed for the implementation\n// this approach limits breakages when the API changes and increases efficiency\nconst workersAIChatResponseSchema = z.object({\n  response: z.string(),\n});\n\n// limited version of the schema, focussed on what is needed for the implementation\n// this approach limits breakages when the API changes and increases efficiency\nconst workersAIChatChunkSchema = z.instanceof(Uint8Array);\n\nfunction prepareToolsAndToolChoice(\n  mode: Parameters<LanguageModelV1[\"doGenerate\"]>[0][\"mode\"] & {\n    type: \"regular\";\n  }\n) {\n  // when the tools array is empty, change it to undefined to prevent errors:\n  const tools = mode.tools?.length ? mode.tools : undefined;\n\n  if (tools == null) {\n    return { tools: undefined, tool_choice: undefined };\n  }\n\n  const mappedTools = tools.map((tool) => ({\n    type: \"function\",\n    function: {\n      name: tool.name,\n      // @ts-expect-error - description is not a property of tool\n      description: tool.description,\n      // @ts-expect-error - parameters is not a property of tool\n      parameters: tool.parameters,\n    },\n  }));\n\n  const toolChoice = mode.toolChoice;\n\n  if (toolChoice == null) {\n    return { tools: mappedTools, tool_choice: undefined };\n  }\n\n  const type = toolChoice.type;\n\n  switch (type) {\n    case \"auto\":\n      return { tools: mappedTools, tool_choice: type };\n    case \"none\":\n      return { tools: mappedTools, tool_choice: type };\n    case \"required\":\n      return { tools: mappedTools, tool_choice: \"any\" };\n\n    // workersAI does not support tool mode directly,\n    // so we filter the tools and force the tool choice through 'any'\n    case \"tool\":\n      return {\n        tools: mappedTools.filter(\n          (tool) => tool.function.name === toolChoice.toolName\n        ),\n        tool_choice: \"any\",\n      };\n    default: {\n      const exhaustiveCheck = type satisfies never;\n      throw new Error(`Unsupported tool choice type: ${exhaustiveCheck}`);\n    }\n  }\n}\n","import {\n  type LanguageModelV1Prompt,\n  UnsupportedFunctionalityError,\n} from \"@ai-sdk/provider\";\nimport type { WorkersAIChatPrompt } from \"./workersai-chat-prompt\";\n\n// TODO\nexport function convertToWorkersAIChatMessages(\n  prompt: LanguageModelV1Prompt\n): WorkersAIChatPrompt {\n  const messages: WorkersAIChatPrompt = [];\n\n  for (const { role, content } of prompt) {\n    switch (role) {\n      case \"system\": {\n        messages.push({ role: \"system\", content });\n        break;\n      }\n\n      case \"user\": {\n        messages.push({\n          role: \"user\",\n          content: content\n            .map((part) => {\n              switch (part.type) {\n                case \"text\": {\n                  return part.text;\n                }\n                case \"image\": {\n                  throw new UnsupportedFunctionalityError({\n                    functionality: \"image-part\",\n                  });\n                }\n              }\n            })\n            .join(\"\"),\n        });\n        break;\n      }\n\n      case \"assistant\": {\n        let text = \"\";\n        const toolCalls: Array<{\n          id: string;\n          type: \"function\";\n          function: { name: string; arguments: string };\n        }> = [];\n\n        for (const part of content) {\n          switch (part.type) {\n            case \"text\": {\n              text += part.text;\n              break;\n            }\n            case \"tool-call\": {\n              toolCalls.push({\n                id: part.toolCallId,\n                type: \"function\",\n                function: {\n                  name: part.toolName,\n                  arguments: JSON.stringify(part.args),\n                },\n              });\n              break;\n            }\n            default: {\n              const exhaustiveCheck = part satisfies never;\n              throw new Error(`Unsupported part: ${exhaustiveCheck}`);\n            }\n          }\n        }\n\n        messages.push({\n          role: \"assistant\",\n          content: text,\n          tool_calls:\n            toolCalls.length > 0\n              ? toolCalls.map(({ function: { name, arguments: args } }) => ({\n                  id: \"null\",\n                  type: \"function\",\n                  function: { name, arguments: args },\n                }))\n              : undefined,\n        });\n\n        break;\n      }\n      case \"tool\": {\n        for (const toolResponse of content) {\n          messages.push({\n            role: \"tool\",\n            name: toolResponse.toolName,\n            content: JSON.stringify(toolResponse.result),\n          });\n        }\n        break;\n      }\n      default: {\n        const exhaustiveCheck = role satisfies never;\n        throw new Error(`Unsupported role: ${exhaustiveCheck}`);\n      }\n    }\n  }\n\n  return messages;\n}\n","import { WorkersAIChatLanguageModel } from \"./workersai-chat-language-model\";\nimport type { WorkersAIChatSettings } from \"./workersai-chat-settings\";\nimport type { TextGenerationModels } from \"./workersai-models\";\n\nexport interface WorkersAI {\n  (\n    modelId: TextGenerationModels,\n    settings?: WorkersAIChatSettings\n  ): WorkersAIChatLanguageModel;\n\n  /**\n   * Creates a model for text generation.\n   **/\n  chat(\n    modelId: TextGenerationModels,\n    settings?: WorkersAIChatSettings\n  ): WorkersAIChatLanguageModel;\n}\n\nexport interface WorkersAISettings {\n  /**\n   * Provide an `env.AI` binding to use for the AI inference.\n   * You can set up an AI bindings in your Workers project\n   * by adding the following this to `wrangler.toml`:\n\n  ```toml\n[ai]\nbinding = \"AI\"\n  ```\n   **/\n  binding: Ai;\n}\n\n/**\n * Create a Workers AI provider instance.\n **/\nexport function createWorkersAI(options: WorkersAISettings): WorkersAI {\n  const createChatModel = (\n    modelId: TextGenerationModels,\n    settings: WorkersAIChatSettings = {}\n  ) =>\n    new WorkersAIChatLanguageModel(modelId, settings, {\n      provider: \"workersai.chat\",\n      binding: options.binding,\n    });\n\n  const provider = function (\n    modelId: TextGenerationModels,\n    settings?: WorkersAIChatSettings\n  ) {\n    if (new.target) {\n      throw new Error(\n        \"The WorkersAI model function cannot be called with the new keyword.\"\n      );\n    }\n\n    return createChatModel(modelId, settings);\n  };\n\n  provider.chat = createChatModel;\n\n  return provider;\n}\n"],"mappings":";;;;;AAAA;AAAA,EAIE,iCAAAA;AAAA,OACK;AACP,SAAS,SAAS;;;ACNlB;AAAA,EAEE;AAAA,OACK;AAIA,SAAS,+BACd,QACqB;AACrB,QAAM,WAAgC,CAAC;AAEvC,aAAW,EAAE,MAAM,QAAQ,KAAK,QAAQ;AACtC,YAAQ,MAAM;AAAA,MACZ,KAAK,UAAU;AACb,iBAAS,KAAK,EAAE,MAAM,UAAU,QAAQ,CAAC;AACzC;AAAA,MACF;AAAA,MAEA,KAAK,QAAQ;AACX,iBAAS,KAAK;AAAA,UACZ,MAAM;AAAA,UACN,SAAS,QACN,IAAI,CAAC,SAAS;AACb,oBAAQ,KAAK,MAAM;AAAA,cACjB,KAAK,QAAQ;AACX,uBAAO,KAAK;AAAA,cACd;AAAA,cACA,KAAK,SAAS;AACZ,sBAAM,IAAI,8BAA8B;AAAA,kBACtC,eAAe;AAAA,gBACjB,CAAC;AAAA,cACH;AAAA,YACF;AAAA,UACF,CAAC,EACA,KAAK,EAAE;AAAA,QACZ,CAAC;AACD;AAAA,MACF;AAAA,MAEA,KAAK,aAAa;AAChB,YAAI,OAAO;AACX,cAAM,YAID,CAAC;AAEN,mBAAW,QAAQ,SAAS;AAC1B,kBAAQ,KAAK,MAAM;AAAA,YACjB,KAAK,QAAQ;AACX,sBAAQ,KAAK;AACb;AAAA,YACF;AAAA,YACA,KAAK,aAAa;AAChB,wBAAU,KAAK;AAAA,gBACb,IAAI,KAAK;AAAA,gBACT,MAAM;AAAA,gBACN,UAAU;AAAA,kBACR,MAAM,KAAK;AAAA,kBACX,WAAW,KAAK,UAAU,KAAK,IAAI;AAAA,gBACrC;AAAA,cACF,CAAC;AACD;AAAA,YACF;AAAA,YACA,SAAS;AACP,oBAAM,kBAAkB;AACxB,oBAAM,IAAI,MAAM,qBAAqB,eAAe,EAAE;AAAA,YACxD;AAAA,UACF;AAAA,QACF;AAEA,iBAAS,KAAK;AAAA,UACZ,MAAM;AAAA,UACN,SAAS;AAAA,UACT,YACE,UAAU,SAAS,IACf,UAAU,IAAI,CAAC,EAAE,UAAU,EAAE,MAAM,WAAW,KAAK,EAAE,OAAO;AAAA,YAC1D,IAAI;AAAA,YACJ,MAAM;AAAA,YACN,UAAU,EAAE,MAAM,WAAW,KAAK;AAAA,UACpC,EAAE,IACF;AAAA,QACR,CAAC;AAED;AAAA,MACF;AAAA,MACA,KAAK,QAAQ;AACX,mBAAW,gBAAgB,SAAS;AAClC,mBAAS,KAAK;AAAA,YACZ,MAAM;AAAA,YACN,MAAM,aAAa;AAAA,YACnB,SAAS,KAAK,UAAU,aAAa,MAAM;AAAA,UAC7C,CAAC;AAAA,QACH;AACA;AAAA,MACF;AAAA,MACA,SAAS;AACP,cAAM,kBAAkB;AACxB,cAAM,IAAI,MAAM,qBAAqB,eAAe,EAAE;AAAA,MACxD;AAAA,IACF;AAAA,EACF;AAEA,SAAO;AACT;;;AD9FA,SAAS,cAAc;AAOhB,IAAM,6BAAN,MAA4D;AAAA,EASjE,YACE,SACA,UACA,QACA;AAZF,wBAAS,wBAAuB;AAChC,wBAAS,+BAA8B;AAEvC,wBAAS;AACT,wBAAS;AAET,wBAAiB;AAOf,SAAK,UAAU;AACf,SAAK,WAAW;AAChB,SAAK,SAAS;AAAA,EAChB;AAAA,EAEA,IAAI,WAAmB;AACrB,WAAO,KAAK,OAAO;AAAA,EACrB;AAAA,EAEQ,QAAQ;AAAA,IACd;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,EACF,GAAiD;AAC/C,UAAM,OAAO,KAAK;AAElB,UAAM,WAAyC,CAAC;AAEhD,QAAI,oBAAoB,MAAM;AAC5B,eAAS,KAAK;AAAA,QACZ,MAAM;AAAA,QACN,SAAS;AAAA,MACX,CAAC;AAAA,IACH;AAEA,QAAI,mBAAmB,MAAM;AAC3B,eAAS,KAAK;AAAA,QACZ,MAAM;AAAA,QACN,SAAS;AAAA,MACX,CAAC;AAAA,IACH;AAEA,UAAM,WAAW;AAAA;AAAA,MAEf,OAAO,KAAK;AAAA;AAAA,MAGZ,aAAa,KAAK,SAAS;AAAA;AAAA,MAG3B,YAAY;AAAA,MACZ;AAAA,MACA,OAAO;AAAA,MACP,aAAa;AAAA;AAAA,MAGb,UAAU,+BAA+B,MAAM;AAAA,IACjD;AAEA,YAAQ,MAAM;AAAA,MACZ,KAAK,WAAW;AACd,eAAO;AAAA,UACL,MAAM,EAAE,GAAG,UAAU,GAAG,0BAA0B,IAAI,EAAE;AAAA,UACxD;AAAA,QACF;AAAA,MACF;AAAA,MAEA,KAAK,eAAe;AAClB,eAAO;AAAA,UACL,MAAM;AAAA,YACJ,GAAG;AAAA,YACH,iBAAiB,EAAE,MAAM,cAAc;AAAA,YACvC,OAAO;AAAA,UACT;AAAA,UACA;AAAA,QACF;AAAA,MACF;AAAA,MAEA,KAAK,eAAe;AAClB,eAAO;AAAA,UACL,MAAM;AAAA,YACJ,GAAG;AAAA,YACH,aAAa;AAAA,YACb,OAAO,CAAC,EAAE,MAAM,YAAY,UAAU,KAAK,KAAK,CAAC;AAAA,UACnD;AAAA,UACA;AAAA,QACF;AAAA,MACF;AAAA;AAAA;AAAA,MAIA,KAAK,kBAAkB;AACrB,cAAM,IAAIC,+BAA8B;AAAA,UACtC,eAAe;AAAA,QACjB,CAAC;AAAA,MACH;AAAA,MAEA,SAAS;AACP,cAAM,kBAAkB;AACxB,cAAM,IAAI,MAAM,qBAAqB,eAAe,EAAE;AAAA,MACxD;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAM,WACJ,SAC6D;AAC7D,UAAM,EAAE,MAAM,SAAS,IAAI,KAAK,QAAQ,OAAO;AAE/C,UAAM,SAAS,MAAM,KAAK,OAAO,QAAQ;AAAA,MACvC,KAAK;AAAA,MACL;AAAA,QACE,UAAU,KAAK;AAAA,QACf,OAAO,KAAK;AAAA,MACd;AAAA,MACA;AAAA,QACE,SAAS,KAAK,SAAS;AAAA,MACzB;AAAA,IACF;AAEA,QAAI,kBAAkB,gBAAgB;AACpC,YAAM,IAAI,MAAM,uBAAuB;AAAA,IACzC;AAEA,WAAO;AAAA,MACL,MAAM,OAAO;AAAA,MACb,WAAW,OAAO,YAAY,IAAI,CAAC,cAAc;AAAA,QAC/C,cAAc;AAAA,QACd,YAAY,SAAS;AAAA,QACrB,UAAU,SAAS;AAAA,QACnB,MAAM,KAAK,UAAU,SAAS,aAAa,CAAC,CAAC;AAAA,MAC/C,EAAE;AAAA,MACF,cAAc;AAAA;AAAA,MACd,SAAS,EAAE,WAAW,KAAK,UAAU,aAAa,KAAK;AAAA,MACvD,OAAO;AAAA;AAAA,QAET,cAAc;AAAA,QACb,kBAAkB;AAAA,MACjB;AAAA,MACA;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAM,SACJ,SAC2D;AAC3D,UAAM,EAAE,MAAM,SAAS,IAAI,KAAK,QAAQ,OAAO;AAE/C,UAAM,UAAU,IAAI,YAAY;AAEhC,UAAM,WAAW,MAAM,KAAK,OAAO,QAAQ,IAAI,KAAK,OAAO;AAAA,MACzD,UAAU,KAAK;AAAA,MACf,QAAQ;AAAA,MACR,OAAO,KAAK;AAAA,IACd,CAAC;AAED,QAAI,EAAE,oBAAoB,iBAAiB;AACzC,YAAM,IAAI,MAAM,uBAAuB;AAAA,IACzC;AAEA,WAAO;AAAA,MACL,QAAQ,SAAS;AAAA,QACf,IAAI,gBAAuD;AAAA,UACzD,MAAM,UAAU,OAAO,YAAY;AACjC,kBAAM,cAAc,QAAQ,OAAO,KAAK;AACxC,kBAAM,SAAS,OAAO,IAAI,SAAS,WAAW,CAAC;AAC/C,6BAAiB,eAAe,QAAQ;AACtC,kBAAI,CAAC,YAAY,MAAM;AACrB;AAAA,cACF;AACA,kBAAI,YAAY,SAAS,UAAU;AACjC,2BAAW,QAAQ;AAAA,kBACjB,MAAM;AAAA,kBACN,cAAc;AAAA,kBACd,OAAO;AAAA,oBACL,cAAc;AAAA,oBACd,kBAAkB;AAAA,kBACpB;AAAA,gBACF,CAAC;AACD;AAAA,cACF;AACA,oBAAM,OAAO,KAAK,MAAM,YAAY,IAAI;AAExC,yBAAW,QAAQ;AAAA,gBACjB,MAAM;AAAA,gBACN,WAAW,KAAK,YAAY;AAAA,cAC9B,CAAC;AAAA,YACH;AACA,uBAAW,QAAQ;AAAA,cACjB,MAAM;AAAA,cACN,cAAc;AAAA,cACd,OAAO;AAAA,gBACL,cAAc;AAAA,gBACd,kBAAkB;AAAA,cACpB;AAAA,YACF,CAAC;AAAA,UACH;AAAA,QACF,CAAC;AAAA,MACH;AAAA,MACA,SAAS,EAAE,WAAW,KAAK,UAAU,aAAa,KAAK;AAAA,MACvD;AAAA,IACF;AAAA,EACF;AACF;AAGA,IAAM,8BAA8B,EAAE,OAAO;AAAA,EAC3C,UAAU,EAAE,OAAO;AACrB,CAAC;AAID,IAAM,2BAA2B,EAAE,WAAW,UAAU;AAExD,SAAS,0BACP,MAGA;AAEA,QAAM,QAAQ,KAAK,OAAO,SAAS,KAAK,QAAQ;AAEhD,MAAI,SAAS,MAAM;AACjB,WAAO,EAAE,OAAO,QAAW,aAAa,OAAU;AAAA,EACpD;AAEA,QAAM,cAAc,MAAM,IAAI,CAAC,UAAU;AAAA,IACvC,MAAM;AAAA,IACN,UAAU;AAAA,MACR,MAAM,KAAK;AAAA;AAAA,MAEX,aAAa,KAAK;AAAA;AAAA,MAElB,YAAY,KAAK;AAAA,IACnB;AAAA,EACF,EAAE;AAEF,QAAM,aAAa,KAAK;AAExB,MAAI,cAAc,MAAM;AACtB,WAAO,EAAE,OAAO,aAAa,aAAa,OAAU;AAAA,EACtD;AAEA,QAAM,OAAO,WAAW;AAExB,UAAQ,MAAM;AAAA,IACZ,KAAK;AACH,aAAO,EAAE,OAAO,aAAa,aAAa,KAAK;AAAA,IACjD,KAAK;AACH,aAAO,EAAE,OAAO,aAAa,aAAa,KAAK;AAAA,IACjD,KAAK;AACH,aAAO,EAAE,OAAO,aAAa,aAAa,MAAM;AAAA;AAAA;AAAA,IAIlD,KAAK;AACH,aAAO;AAAA,QACL,OAAO,YAAY;AAAA,UACjB,CAAC,SAAS,KAAK,SAAS,SAAS,WAAW;AAAA,QAC9C;AAAA,QACA,aAAa;AAAA,MACf;AAAA,IACF,SAAS;AACP,YAAM,kBAAkB;AACxB,YAAM,IAAI,MAAM,iCAAiC,eAAe,EAAE;AAAA,IACpE;AAAA,EACF;AACF;;;AElQO,SAAS,gBAAgB,SAAuC;AACrE,QAAM,kBAAkB,CACtB,SACA,WAAkC,CAAC,MAEnC,IAAI,2BAA2B,SAAS,UAAU;AAAA,IAChD,UAAU;AAAA,IACV,SAAS,QAAQ;AAAA,EACnB,CAAC;AAEH,QAAM,WAAW,SACf,SACA,UACA;AACA,QAAI,YAAY;AACd,YAAM,IAAI;AAAA,QACR;AAAA,MACF;AAAA,IACF;AAEA,WAAO,gBAAgB,SAAS,QAAQ;AAAA,EAC1C;AAEA,WAAS,OAAO;AAEhB,SAAO;AACT;","names":["UnsupportedFunctionalityError","UnsupportedFunctionalityError"]}
+{"version":3,"sources":["../src/workersai-chat-language-model.ts","../src/convert-to-workersai-chat-messages.ts","../src/map-workersai-usage.ts","../src/index.ts"],"sourcesContent":["import {\n  type LanguageModelV1,\n  type LanguageModelV1CallWarning,\n  type LanguageModelV1StreamPart,\n  UnsupportedFunctionalityError,\n} from \"@ai-sdk/provider\";\nimport { z } from \"zod\";\nimport { convertToWorkersAIChatMessages } from \"./convert-to-workersai-chat-messages\";\nimport type { WorkersAIChatSettings } from \"./workersai-chat-settings\";\nimport type { TextGenerationModels } from \"./workersai-models\";\n\nimport { events } from \"fetch-event-stream\";\nimport { mapWorkersAIUsage } from \"./map-workersai-usage\";\n\ntype WorkersAIChatConfig = {\n  provider: string;\n  binding: Ai;\n  gateway?: GatewayOptions;\n};\n\nexport class WorkersAIChatLanguageModel implements LanguageModelV1 {\n  readonly specificationVersion = \"v1\";\n  readonly defaultObjectGenerationMode = \"json\";\n\n  readonly modelId: TextGenerationModels;\n  readonly settings: WorkersAIChatSettings;\n\n  private readonly config: WorkersAIChatConfig;\n\n  constructor(\n    modelId: TextGenerationModels,\n    settings: WorkersAIChatSettings,\n    config: WorkersAIChatConfig\n  ) {\n    this.modelId = modelId;\n    this.settings = settings;\n    this.config = config;\n  }\n\n  get provider(): string {\n    return this.config.provider;\n  }\n\n  private getArgs({\n    mode,\n    prompt,\n    maxTokens,\n    temperature,\n    topP,\n    frequencyPenalty,\n    presencePenalty,\n    seed,\n  }: Parameters<LanguageModelV1[\"doGenerate\"]>[0]) {\n    const type = mode.type;\n\n    const warnings: LanguageModelV1CallWarning[] = [];\n\n    if (frequencyPenalty != null) {\n      warnings.push({\n        type: \"unsupported-setting\",\n        setting: \"frequencyPenalty\",\n      });\n    }\n\n    if (presencePenalty != null) {\n      warnings.push({\n        type: \"unsupported-setting\",\n        setting: \"presencePenalty\",\n      });\n    }\n\n    const baseArgs = {\n      // model id:\n      model: this.modelId,\n\n      // model specific settings:\n      safe_prompt: this.settings.safePrompt,\n\n      // standardized settings:\n      max_tokens: maxTokens,\n      temperature,\n      top_p: topP,\n      random_seed: seed,\n\n      // messages:\n      messages: convertToWorkersAIChatMessages(prompt),\n    };\n\n    switch (type) {\n      case \"regular\": {\n        return {\n          args: { ...baseArgs, ...prepareToolsAndToolChoice(mode) },\n          warnings,\n        };\n      }\n\n      case \"object-json\": {\n        return {\n          args: {\n            ...baseArgs,\n            response_format: { type: \"json_object\" },\n            tools: undefined,\n          },\n          warnings,\n        };\n      }\n\n      case \"object-tool\": {\n        return {\n          args: {\n            ...baseArgs,\n            tool_choice: \"any\",\n            tools: [{ type: \"function\", function: mode.tool }],\n          },\n          warnings,\n        };\n      }\n\n      // @ts-expect-error - this is unreachable code\n      // TODO: fixme\n      case \"object-grammar\": {\n        throw new UnsupportedFunctionalityError({\n          functionality: \"object-grammar mode\",\n        });\n      }\n\n      default: {\n        const exhaustiveCheck = type satisfies never;\n        throw new Error(`Unsupported type: ${exhaustiveCheck}`);\n      }\n    }\n  }\n\n  async doGenerate(\n    options: Parameters<LanguageModelV1[\"doGenerate\"]>[0]\n  ): Promise<Awaited<ReturnType<LanguageModelV1[\"doGenerate\"]>>> {\n    const { args, warnings } = this.getArgs(options);\n\n    const output = await this.config.binding.run(\n      args.model,\n      {\n        messages: args.messages,\n        max_tokens: args.max_tokens,\n        temperature: args.temperature,\n        tools: args.tools,\n        top_p: args.top_p,\n      },\n      { gateway: this.config.gateway ?? this.settings.gateway }\n    );\n\n    if (output instanceof ReadableStream) {\n      throw new Error(\"This shouldn't happen\");\n    }\n\n    return {\n      text: output.response,\n      toolCalls: output.tool_calls?.map((toolCall) => ({\n        toolCallType: \"function\",\n        toolCallId: toolCall.name,\n        toolName: toolCall.name,\n        args: JSON.stringify(toolCall.arguments || {}),\n      })),\n      finishReason: \"stop\", // TODO: mapWorkersAIFinishReason(response.finish_reason),\n      rawCall: { rawPrompt: args.messages, rawSettings: args },\n      usage: mapWorkersAIUsage(output),\n      warnings,\n    };\n  }\n\n  async doStream(\n    options: Parameters<LanguageModelV1[\"doStream\"]>[0]\n  ): Promise<Awaited<ReturnType<LanguageModelV1[\"doStream\"]>>> {\n    const { args, warnings } = this.getArgs(options);\n\n    const response = await this.config.binding.run(\n      args.model,\n      {\n        messages: args.messages,\n        max_tokens: args.max_tokens,\n        stream: true,\n        temperature: args.temperature,\n        tools: args.tools,\n        top_p: args.top_p,\n      },\n      { gateway: this.config.gateway ?? this.settings.gateway }\n    );\n\n    if (!(response instanceof ReadableStream)) {\n      throw new Error(\"This shouldn't happen\");\n    }\n\n    const chunkEvent = events(new Response(response));\n    let usage = { promptTokens: 0, completionTokens: 0 };\n\n    return {\n      stream: new ReadableStream<LanguageModelV1StreamPart>({\n        async start(controller) {\n          for await (const event of chunkEvent) {\n            if (!event.data) {\n              continue;\n            }\n            if (event.data === \"[DONE]\") {\n              break;\n            }\n            const chunk = JSON.parse(event.data);\n            if (chunk.usage) {\n              usage = mapWorkersAIUsage(chunk);\n            }\n            chunk.response.length &&\n              controller.enqueue({\n                type: \"text-delta\",\n                textDelta: chunk.response,\n              });\n          }\n          controller.enqueue({\n            type: \"finish\",\n            finishReason: \"stop\",\n            usage: usage,\n          });\n          controller.close();\n        },\n      }),\n      rawCall: { rawPrompt: args.messages, rawSettings: args },\n      warnings,\n    };\n  }\n}\n// limited version of the schema, focussed on what is needed for the implementation\n// this approach limits breakages when the API changes and increases efficiency\nconst workersAIChatResponseSchema = z.object({\n  response: z.string(),\n});\n\n// limited version of the schema, focussed on what is needed for the implementation\n// this approach limits breakages when the API changes and increases efficiency\nconst workersAIChatChunkSchema = z.instanceof(Uint8Array);\n\nfunction prepareToolsAndToolChoice(\n  mode: Parameters<LanguageModelV1[\"doGenerate\"]>[0][\"mode\"] & {\n    type: \"regular\";\n  }\n) {\n  // when the tools array is empty, change it to undefined to prevent errors:\n  const tools = mode.tools?.length ? mode.tools : undefined;\n\n  if (tools == null) {\n    return { tools: undefined, tool_choice: undefined };\n  }\n\n  const mappedTools = tools.map((tool) => ({\n    type: \"function\",\n    function: {\n      name: tool.name,\n      // @ts-expect-error - description is not a property of tool\n      description: tool.description,\n      // @ts-expect-error - parameters is not a property of tool\n      parameters: tool.parameters,\n    },\n  }));\n\n  const toolChoice = mode.toolChoice;\n\n  if (toolChoice == null) {\n    return { tools: mappedTools, tool_choice: undefined };\n  }\n\n  const type = toolChoice.type;\n\n  switch (type) {\n    case \"auto\":\n      return { tools: mappedTools, tool_choice: type };\n    case \"none\":\n      return { tools: mappedTools, tool_choice: type };\n    case \"required\":\n      return { tools: mappedTools, tool_choice: \"any\" };\n\n    // workersAI does not support tool mode directly,\n    // so we filter the tools and force the tool choice through 'any'\n    case \"tool\":\n      return {\n        tools: mappedTools.filter(\n          (tool) => tool.function.name === toolChoice.toolName\n        ),\n        tool_choice: \"any\",\n      };\n    default: {\n      const exhaustiveCheck = type satisfies never;\n      throw new Error(`Unsupported tool choice type: ${exhaustiveCheck}`);\n    }\n  }\n}\n","import {\n  type LanguageModelV1Prompt,\n  UnsupportedFunctionalityError,\n} from \"@ai-sdk/provider\";\nimport type { WorkersAIChatPrompt } from \"./workersai-chat-prompt\";\n\n// TODO\nexport function convertToWorkersAIChatMessages(\n  prompt: LanguageModelV1Prompt\n): WorkersAIChatPrompt {\n  const messages: WorkersAIChatPrompt = [];\n\n  for (const { role, content } of prompt) {\n    switch (role) {\n      case \"system\": {\n        messages.push({ role: \"system\", content });\n        break;\n      }\n\n      case \"user\": {\n        messages.push({\n          role: \"user\",\n          content: content\n            .map((part) => {\n              switch (part.type) {\n                case \"text\": {\n                  return part.text;\n                }\n                case \"image\": {\n                  throw new UnsupportedFunctionalityError({\n                    functionality: \"image-part\",\n                  });\n                }\n              }\n            })\n            .join(\"\"),\n        });\n        break;\n      }\n\n      case \"assistant\": {\n        let text = \"\";\n        const toolCalls: Array<{\n          id: string;\n          type: \"function\";\n          function: { name: string; arguments: string };\n        }> = [];\n\n        for (const part of content) {\n          switch (part.type) {\n            case \"text\": {\n              text += part.text;\n              break;\n            }\n            case \"tool-call\": {\n              toolCalls.push({\n                id: part.toolCallId,\n                type: \"function\",\n                function: {\n                  name: part.toolName,\n                  arguments: JSON.stringify(part.args),\n                },\n              });\n              break;\n            }\n            default: {\n              const exhaustiveCheck = part satisfies never;\n              throw new Error(`Unsupported part: ${exhaustiveCheck}`);\n            }\n          }\n        }\n\n        messages.push({\n          role: \"assistant\",\n          content: text,\n          tool_calls:\n            toolCalls.length > 0\n              ? toolCalls.map(({ function: { name, arguments: args } }) => ({\n                  id: \"null\",\n                  type: \"function\",\n                  function: { name, arguments: args },\n                }))\n              : undefined,\n        });\n\n        break;\n      }\n      case \"tool\": {\n        for (const toolResponse of content) {\n          messages.push({\n            role: \"tool\",\n            name: toolResponse.toolName,\n            content: JSON.stringify(toolResponse.result),\n          });\n        }\n        break;\n      }\n      default: {\n        const exhaustiveCheck = role satisfies never;\n        throw new Error(`Unsupported role: ${exhaustiveCheck}`);\n      }\n    }\n  }\n\n  return messages;\n}\n","export function mapWorkersAIUsage(output:  AiTextGenerationOutput | AiTextToImageOutput) {\n\tconst usage = (output as { usage: { prompt_tokens: number, completion_tokens: number} }).usage ?? {\n\t\tprompt_tokens: 0,\n\t\tcompletion_tokens: 0,\n\t};\n\n\treturn {\n\t\tpromptTokens: usage.prompt_tokens,\n\t\tcompletionTokens: usage.completion_tokens\n\t}\n}\n","import { WorkersAIChatLanguageModel } from \"./workersai-chat-language-model\";\nimport type { WorkersAIChatSettings } from \"./workersai-chat-settings\";\nimport type { TextGenerationModels } from \"./workersai-models\";\n\nexport interface WorkersAI {\n  (\n    modelId: TextGenerationModels,\n    settings?: WorkersAIChatSettings\n  ): WorkersAIChatLanguageModel;\n\n  /**\n   * Creates a model for text generation.\n   **/\n  chat(\n    modelId: TextGenerationModels,\n    settings?: WorkersAIChatSettings\n  ): WorkersAIChatLanguageModel;\n}\n\nexport interface WorkersAISettings {\n  /**\n   * Provide an `env.AI` binding to use for the AI inference.\n   * You can set up an AI bindings in your Workers project\n   * by adding the following this to `wrangler.toml`:\n\n  ```toml\n[ai]\nbinding = \"AI\"\n  ```\n   **/\n  binding: Ai;\n  /**\n   * Optionally set Cloudflare AI Gateway options.\n   */\n  gateway?: GatewayOptions;\n}\n\n/**\n * Create a Workers AI provider instance.\n **/\nexport function createWorkersAI(options: WorkersAISettings): WorkersAI {\n  const createChatModel = (\n    modelId: TextGenerationModels,\n    settings: WorkersAIChatSettings = {}\n  ) =>\n    new WorkersAIChatLanguageModel(modelId, settings, {\n      provider: \"workersai.chat\",\n      binding: options.binding,\n      gateway: options.gateway,\n    });\n\n  const provider = function (\n    modelId: TextGenerationModels,\n    settings?: WorkersAIChatSettings\n  ) {\n    if (new.target) {\n      throw new Error(\n        \"The WorkersAI model function cannot be called with the new keyword.\"\n      );\n    }\n\n    return createChatModel(modelId, settings);\n  };\n\n  provider.chat = createChatModel;\n\n  return provider;\n}\n"],"mappings":";;;;;AAAA;AAAA,EAIE,iCAAAA;AAAA,OACK;AACP,SAAS,SAAS;;;ACNlB;AAAA,EAEE;AAAA,OACK;AAIA,SAAS,+BACd,QACqB;AACrB,QAAM,WAAgC,CAAC;AAEvC,aAAW,EAAE,MAAM,QAAQ,KAAK,QAAQ;AACtC,YAAQ,MAAM;AAAA,MACZ,KAAK,UAAU;AACb,iBAAS,KAAK,EAAE,MAAM,UAAU,QAAQ,CAAC;AACzC;AAAA,MACF;AAAA,MAEA,KAAK,QAAQ;AACX,iBAAS,KAAK;AAAA,UACZ,MAAM;AAAA,UACN,SAAS,QACN,IAAI,CAAC,SAAS;AACb,oBAAQ,KAAK,MAAM;AAAA,cACjB,KAAK,QAAQ;AACX,uBAAO,KAAK;AAAA,cACd;AAAA,cACA,KAAK,SAAS;AACZ,sBAAM,IAAI,8BAA8B;AAAA,kBACtC,eAAe;AAAA,gBACjB,CAAC;AAAA,cACH;AAAA,YACF;AAAA,UACF,CAAC,EACA,KAAK,EAAE;AAAA,QACZ,CAAC;AACD;AAAA,MACF;AAAA,MAEA,KAAK,aAAa;AAChB,YAAI,OAAO;AACX,cAAM,YAID,CAAC;AAEN,mBAAW,QAAQ,SAAS;AAC1B,kBAAQ,KAAK,MAAM;AAAA,YACjB,KAAK,QAAQ;AACX,sBAAQ,KAAK;AACb;AAAA,YACF;AAAA,YACA,KAAK,aAAa;AAChB,wBAAU,KAAK;AAAA,gBACb,IAAI,KAAK;AAAA,gBACT,MAAM;AAAA,gBACN,UAAU;AAAA,kBACR,MAAM,KAAK;AAAA,kBACX,WAAW,KAAK,UAAU,KAAK,IAAI;AAAA,gBACrC;AAAA,cACF,CAAC;AACD;AAAA,YACF;AAAA,YACA,SAAS;AACP,oBAAM,kBAAkB;AACxB,oBAAM,IAAI,MAAM,qBAAqB,eAAe,EAAE;AAAA,YACxD;AAAA,UACF;AAAA,QACF;AAEA,iBAAS,KAAK;AAAA,UACZ,MAAM;AAAA,UACN,SAAS;AAAA,UACT,YACE,UAAU,SAAS,IACf,UAAU,IAAI,CAAC,EAAE,UAAU,EAAE,MAAM,WAAW,KAAK,EAAE,OAAO;AAAA,YAC1D,IAAI;AAAA,YACJ,MAAM;AAAA,YACN,UAAU,EAAE,MAAM,WAAW,KAAK;AAAA,UACpC,EAAE,IACF;AAAA,QACR,CAAC;AAED;AAAA,MACF;AAAA,MACA,KAAK,QAAQ;AACX,mBAAW,gBAAgB,SAAS;AAClC,mBAAS,KAAK;AAAA,YACZ,MAAM;AAAA,YACN,MAAM,aAAa;AAAA,YACnB,SAAS,KAAK,UAAU,aAAa,MAAM;AAAA,UAC7C,CAAC;AAAA,QACH;AACA;AAAA,MACF;AAAA,MACA,SAAS;AACP,cAAM,kBAAkB;AACxB,cAAM,IAAI,MAAM,qBAAqB,eAAe,EAAE;AAAA,MACxD;AAAA,IACF;AAAA,EACF;AAEA,SAAO;AACT;;;AD9FA,SAAS,cAAc;;;AEXhB,SAAS,kBAAkB,QAAuD;AACxF,QAAM,QAAS,OAA0E,SAAS;AAAA,IACjG,eAAe;AAAA,IACf,mBAAmB;AAAA,EACpB;AAEA,SAAO;AAAA,IACN,cAAc,MAAM;AAAA,IACpB,kBAAkB,MAAM;AAAA,EACzB;AACD;;;AFUO,IAAM,6BAAN,MAA4D;AAAA,EASjE,YACE,SACA,UACA,QACA;AAZF,wBAAS,wBAAuB;AAChC,wBAAS,+BAA8B;AAEvC,wBAAS;AACT,wBAAS;AAET,wBAAiB;AAOf,SAAK,UAAU;AACf,SAAK,WAAW;AAChB,SAAK,SAAS;AAAA,EAChB;AAAA,EAEA,IAAI,WAAmB;AACrB,WAAO,KAAK,OAAO;AAAA,EACrB;AAAA,EAEQ,QAAQ;AAAA,IACd;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,EACF,GAAiD;AAC/C,UAAM,OAAO,KAAK;AAElB,UAAM,WAAyC,CAAC;AAEhD,QAAI,oBAAoB,MAAM;AAC5B,eAAS,KAAK;AAAA,QACZ,MAAM;AAAA,QACN,SAAS;AAAA,MACX,CAAC;AAAA,IACH;AAEA,QAAI,mBAAmB,MAAM;AAC3B,eAAS,KAAK;AAAA,QACZ,MAAM;AAAA,QACN,SAAS;AAAA,MACX,CAAC;AAAA,IACH;AAEA,UAAM,WAAW;AAAA;AAAA,MAEf,OAAO,KAAK;AAAA;AAAA,MAGZ,aAAa,KAAK,SAAS;AAAA;AAAA,MAG3B,YAAY;AAAA,MACZ;AAAA,MACA,OAAO;AAAA,MACP,aAAa;AAAA;AAAA,MAGb,UAAU,+BAA+B,MAAM;AAAA,IACjD;AAEA,YAAQ,MAAM;AAAA,MACZ,KAAK,WAAW;AACd,eAAO;AAAA,UACL,MAAM,EAAE,GAAG,UAAU,GAAG,0BAA0B,IAAI,EAAE;AAAA,UACxD;AAAA,QACF;AAAA,MACF;AAAA,MAEA,KAAK,eAAe;AAClB,eAAO;AAAA,UACL,MAAM;AAAA,YACJ,GAAG;AAAA,YACH,iBAAiB,EAAE,MAAM,cAAc;AAAA,YACvC,OAAO;AAAA,UACT;AAAA,UACA;AAAA,QACF;AAAA,MACF;AAAA,MAEA,KAAK,eAAe;AAClB,eAAO;AAAA,UACL,MAAM;AAAA,YACJ,GAAG;AAAA,YACH,aAAa;AAAA,YACb,OAAO,CAAC,EAAE,MAAM,YAAY,UAAU,KAAK,KAAK,CAAC;AAAA,UACnD;AAAA,UACA;AAAA,QACF;AAAA,MACF;AAAA;AAAA;AAAA,MAIA,KAAK,kBAAkB;AACrB,cAAM,IAAIC,+BAA8B;AAAA,UACtC,eAAe;AAAA,QACjB,CAAC;AAAA,MACH;AAAA,MAEA,SAAS;AACP,cAAM,kBAAkB;AACxB,cAAM,IAAI,MAAM,qBAAqB,eAAe,EAAE;AAAA,MACxD;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAM,WACJ,SAC6D;AAC7D,UAAM,EAAE,MAAM,SAAS,IAAI,KAAK,QAAQ,OAAO;AAE/C,UAAM,SAAS,MAAM,KAAK,OAAO,QAAQ;AAAA,MACvC,KAAK;AAAA,MACL;AAAA,QACE,UAAU,KAAK;AAAA,QACf,YAAY,KAAK;AAAA,QACjB,aAAa,KAAK;AAAA,QAClB,OAAO,KAAK;AAAA,QACZ,OAAO,KAAK;AAAA,MACd;AAAA,MACA,EAAE,SAAS,KAAK,OAAO,WAAW,KAAK,SAAS,QAAQ;AAAA,IAC1D;AAEA,QAAI,kBAAkB,gBAAgB;AACpC,YAAM,IAAI,MAAM,uBAAuB;AAAA,IACzC;AAEA,WAAO;AAAA,MACL,MAAM,OAAO;AAAA,MACb,WAAW,OAAO,YAAY,IAAI,CAAC,cAAc;AAAA,QAC/C,cAAc;AAAA,QACd,YAAY,SAAS;AAAA,QACrB,UAAU,SAAS;AAAA,QACnB,MAAM,KAAK,UAAU,SAAS,aAAa,CAAC,CAAC;AAAA,MAC/C,EAAE;AAAA,MACF,cAAc;AAAA;AAAA,MACd,SAAS,EAAE,WAAW,KAAK,UAAU,aAAa,KAAK;AAAA,MACvD,OAAO,kBAAkB,MAAM;AAAA,MAC/B;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAM,SACJ,SAC2D;AAC3D,UAAM,EAAE,MAAM,SAAS,IAAI,KAAK,QAAQ,OAAO;AAE/C,UAAM,WAAW,MAAM,KAAK,OAAO,QAAQ;AAAA,MACzC,KAAK;AAAA,MACL;AAAA,QACE,UAAU,KAAK;AAAA,QACf,YAAY,KAAK;AAAA,QACjB,QAAQ;AAAA,QACR,aAAa,KAAK;AAAA,QAClB,OAAO,KAAK;AAAA,QACZ,OAAO,KAAK;AAAA,MACd;AAAA,MACA,EAAE,SAAS,KAAK,OAAO,WAAW,KAAK,SAAS,QAAQ;AAAA,IAC1D;AAEA,QAAI,EAAE,oBAAoB,iBAAiB;AACzC,YAAM,IAAI,MAAM,uBAAuB;AAAA,IACzC;AAEA,UAAM,aAAa,OAAO,IAAI,SAAS,QAAQ,CAAC;AAChD,QAAI,QAAQ,EAAE,cAAc,GAAG,kBAAkB,EAAE;AAEnD,WAAO;AAAA,MACL,QAAQ,IAAI,eAA0C;AAAA,QACpD,MAAM,MAAM,YAAY;AACtB,2BAAiB,SAAS,YAAY;AACpC,gBAAI,CAAC,MAAM,MAAM;AACf;AAAA,YACF;AACA,gBAAI,MAAM,SAAS,UAAU;AAC3B;AAAA,YACF;AACA,kBAAM,QAAQ,KAAK,MAAM,MAAM,IAAI;AACnC,gBAAI,MAAM,OAAO;AACf,sBAAQ,kBAAkB,KAAK;AAAA,YACjC;AACA,kBAAM,SAAS,UACb,WAAW,QAAQ;AAAA,cACjB,MAAM;AAAA,cACN,WAAW,MAAM;AAAA,YACnB,CAAC;AAAA,UACL;AACA,qBAAW,QAAQ;AAAA,YACjB,MAAM;AAAA,YACN,cAAc;AAAA,YACd;AAAA,UACF,CAAC;AACD,qBAAW,MAAM;AAAA,QACnB;AAAA,MACF,CAAC;AAAA,MACD,SAAS,EAAE,WAAW,KAAK,UAAU,aAAa,KAAK;AAAA,MACvD;AAAA,IACF;AAAA,EACF;AACF;AAGA,IAAM,8BAA8B,EAAE,OAAO;AAAA,EAC3C,UAAU,EAAE,OAAO;AACrB,CAAC;AAID,IAAM,2BAA2B,EAAE,WAAW,UAAU;AAExD,SAAS,0BACP,MAGA;AAEA,QAAM,QAAQ,KAAK,OAAO,SAAS,KAAK,QAAQ;AAEhD,MAAI,SAAS,MAAM;AACjB,WAAO,EAAE,OAAO,QAAW,aAAa,OAAU;AAAA,EACpD;AAEA,QAAM,cAAc,MAAM,IAAI,CAAC,UAAU;AAAA,IACvC,MAAM;AAAA,IACN,UAAU;AAAA,MACR,MAAM,KAAK;AAAA;AAAA,MAEX,aAAa,KAAK;AAAA;AAAA,MAElB,YAAY,KAAK;AAAA,IACnB;AAAA,EACF,EAAE;AAEF,QAAM,aAAa,KAAK;AAExB,MAAI,cAAc,MAAM;AACtB,WAAO,EAAE,OAAO,aAAa,aAAa,OAAU;AAAA,EACtD;AAEA,QAAM,OAAO,WAAW;AAExB,UAAQ,MAAM;AAAA,IACZ,KAAK;AACH,aAAO,EAAE,OAAO,aAAa,aAAa,KAAK;AAAA,IACjD,KAAK;AACH,aAAO,EAAE,OAAO,aAAa,aAAa,KAAK;AAAA,IACjD,KAAK;AACH,aAAO,EAAE,OAAO,aAAa,aAAa,MAAM;AAAA;AAAA;AAAA,IAIlD,KAAK;AACH,aAAO;AAAA,QACL,OAAO,YAAY;AAAA,UACjB,CAAC,SAAS,KAAK,SAAS,SAAS,WAAW;AAAA,QAC9C;AAAA,QACA,aAAa;AAAA,MACf;AAAA,IACF,SAAS;AACP,YAAM,kBAAkB;AACxB,YAAM,IAAI,MAAM,iCAAiC,eAAe,EAAE;AAAA,IACpE;AAAA,EACF;AACF;;;AG1PO,SAAS,gBAAgB,SAAuC;AACrE,QAAM,kBAAkB,CACtB,SACA,WAAkC,CAAC,MAEnC,IAAI,2BAA2B,SAAS,UAAU;AAAA,IAChD,UAAU;AAAA,IACV,SAAS,QAAQ;AAAA,IACjB,SAAS,QAAQ;AAAA,EACnB,CAAC;AAEH,QAAM,WAAW,SACf,SACA,UACA;AACA,QAAI,YAAY;AACd,YAAM,IAAI;AAAA,QACR;AAAA,MACF;AAAA,IACF;AAEA,WAAO,gBAAgB,SAAS,QAAQ;AAAA,EAC1C;AAEA,WAAS,OAAO;AAEhB,SAAO;AACT;","names":["UnsupportedFunctionalityError","UnsupportedFunctionalityError"]}

package/package.json

@@ -2,7 +2,7 @@
   "name": "workers-ai-provider",
   "description": "Workers AI Provider for the vercel AI SDK",
   "type": "module",
-  "version": "0.0.12",
+  "version": "0.1.0",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "repository": {

package/src/index.ts

@@ -29,6 +29,10 @@ binding = "AI"
   ```
    **/
   binding: Ai;
+  /**
+   * Optionally set Cloudflare AI Gateway options.
+   */
+  gateway?: GatewayOptions;
 }
 
 /**
@@ -42,6 +46,7 @@ export function createWorkersAI(options: WorkersAISettings): WorkersAI {
     new WorkersAIChatLanguageModel(modelId, settings, {
       provider: "workersai.chat",
       binding: options.binding,
+      gateway: options.gateway,
     });
 
   const provider = function (

package/src/map-workersai-usage.ts

@@ -0,0 +1,11 @@
+export function mapWorkersAIUsage(output:  AiTextGenerationOutput | AiTextToImageOutput) {
+	const usage = (output as { usage: { prompt_tokens: number, completion_tokens: number} }).usage ?? {
+		prompt_tokens: 0,
+		completion_tokens: 0,
+	};
+
+	return {
+		promptTokens: usage.prompt_tokens,
+		completionTokens: usage.completion_tokens
+	}
+}

package/src/workersai-chat-language-model.ts

@@ -10,10 +10,12 @@ import type { WorkersAIChatSettings } from "./workersai-chat-settings";
 import type { TextGenerationModels } from "./workersai-models";
 
 import { events } from "fetch-event-stream";
+import { mapWorkersAIUsage } from "./map-workersai-usage";
 
 type WorkersAIChatConfig = {
   provider: string;
   binding: Ai;
+  gateway?: GatewayOptions;
 };
 
 export class WorkersAIChatLanguageModel implements LanguageModelV1 {
@@ -138,11 +140,12 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
       args.model,
       {
         messages: args.messages,
+        max_tokens: args.max_tokens,
+        temperature: args.temperature,
         tools: args.tools,
+        top_p: args.top_p,
       },
-      {
-        gateway: this.settings.gateway,
-      }
+      { gateway: this.config.gateway ?? this.settings.gateway }
     );
 
     if (output instanceof ReadableStream) {
@@ -159,11 +162,7 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
       })),
       finishReason: "stop", // TODO: mapWorkersAIFinishReason(response.finish_reason),
       rawCall: { rawPrompt: args.messages, rawSettings: args },
-      usage: {
-	    // TODO: mapWorkersAIUsage(response.usage),
-  		promptTokens: 0,
-	    completionTokens: 0,
-      },
+      usage: mapWorkersAIUsage(output),
       warnings,
     };
   }
@@ -173,57 +172,54 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
   ): Promise<Awaited<ReturnType<LanguageModelV1["doStream"]>>> {
     const { args, warnings } = this.getArgs(options);
 
-    const decoder = new TextDecoder();
-
-    const response = await this.config.binding.run(args.model, {
-      messages: args.messages,
-      stream: true,
-      tools: args.tools,
-    });
+    const response = await this.config.binding.run(
+      args.model,
+      {
+        messages: args.messages,
+        max_tokens: args.max_tokens,
+        stream: true,
+        temperature: args.temperature,
+        tools: args.tools,
+        top_p: args.top_p,
+      },
+      { gateway: this.config.gateway ?? this.settings.gateway }
+    );
 
     if (!(response instanceof ReadableStream)) {
       throw new Error("This shouldn't happen");
     }
 
-    return {
-      stream: response.pipeThrough(
-        new TransformStream<Uint8Array, LanguageModelV1StreamPart>({
-          async transform(chunk, controller) {
-            const chunkToText = decoder.decode(chunk);
-            const chunks = events(new Response(chunkToText));
-            for await (const singleChunk of chunks) {
-              if (!singleChunk.data) {
-                continue;
-              }
-              if (singleChunk.data === "[DONE]") {
-                controller.enqueue({
-                  type: "finish",
-                  finishReason: "stop",
-                  usage: {
-                    promptTokens: 0,
-                    completionTokens: 0,
-                  },
-                });
-                return;
-              }
-              const data = JSON.parse(singleChunk.data);
+    const chunkEvent = events(new Response(response));
+    let usage = { promptTokens: 0, completionTokens: 0 };
 
+    return {
+      stream: new ReadableStream<LanguageModelV1StreamPart>({
+        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: data.response ?? "DATALOSS",
+                textDelta: chunk.response,
               });
-            }
-            controller.enqueue({
-              type: "finish",
-              finishReason: "stop",
-              usage: {
-                promptTokens: 0,
-                completionTokens: 0,
-              },
-            });
-          },
-        })
-      ),
+          }
+          controller.enqueue({
+            type: "finish",
+            finishReason: "stop",
+            usage: usage,
+          });
+          controller.close();
+        },
+      }),
       rawCall: { rawPrompt: args.messages, rawSettings: args },
       warnings,
     };

package/src/workersai-chat-settings.ts

@@ -7,6 +7,7 @@ export interface WorkersAIChatSettings {
   safePrompt?: boolean;
   /**
    * Optionally set Cloudflare AI Gateway options.
+   * @deprecated
    */
-  gateway?: GatewayOptions
+  gateway?: GatewayOptions;
 }

package/src/workersai-models.ts

@@ -1,6 +1,9 @@
 /**
  * The names of the BaseAiTextGeneration models.
  */
-export type TextGenerationModels = {
-  [K in keyof AiModels]: AiModels[K] extends BaseAiTextGeneration ? K : never;
-}[keyof AiModels];
+export type TextGenerationModels = Exclude<
+  value2key<AiModels, BaseAiTextGeneration>,
+  value2key<AiModels, BaseAiTextToImage>
+>;
+
+type value2key<T, V> = { [K in keyof T]: T[K] extends V ? K : never }[keyof T];