fix some types and buffering

Author: threepointone

.changeset/fuzzy-pots-smoke.md

@@ -0,0 +1,5 @@
+---
+"workers-ai-provider": patch
+---
+
+fix some types and buffering

examples/ai-chat/src/server/index.ts

@@ -21,7 +21,13 @@ export default {
         model: workersai("@cf/meta/llama-3-8b-instruct"),
         messages: convertToCoreMessages(messages),
       });
-      return result.toDataStreamResponse();
+      return result.toDataStreamResponse({
+        headers: {
+          "Content-Type": "text/x-workers-ai",
+          "content-encoding": "identity",
+          "transfer-encoding": "chunked",
+        },
+      });
     }
 
     return new Response("Not Found!!", { status: 404 });

packages/ai-provider/src/convert-to-workersai-chat-messages.ts

@@ -64,7 +64,7 @@ export function convertToWorkersAIChatMessages(
               break;
             }
             default: {
-              const exhaustiveCheck: never = part;
+              const exhaustiveCheck = part satisfies never;
               throw new Error(`Unsupported part: ${exhaustiveCheck}`);
             }
           }
@@ -96,8 +96,8 @@ export function convertToWorkersAIChatMessages(
         break;
       }
       default: {
-        const _exhaustiveCheck: never = role;
-        throw new Error(`Unsupported role: ${_exhaustiveCheck}`);
+        const exhaustiveCheck = role satisfies never;
+        throw new Error(`Unsupported role: ${exhaustiveCheck}`);
       }
     }
   }

packages/ai-provider/src/index.ts

@@ -58,5 +58,5 @@ export function createWorkersAI(options: WorkersAISettings): WorkersAI {
 
   provider.chat = createChatModel;
 
-  return provider as WorkersAI;
+  return provider;
 }

packages/ai-provider/src/workersai-chat-language-model.ts

@@ -130,8 +130,8 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
       }
 
       default: {
-        const _exhaustiveCheck: never = type;
-        throw new Error(`Unsupported type: ${_exhaustiveCheck}`);
+        const exhaustiveCheck = type satisfies never;
+        throw new Error(`Unsupported type: ${exhaustiveCheck}`);
       }
     }
   }
@@ -141,18 +141,31 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
   ): Promise<Awaited<ReturnType<LanguageModelV1["doGenerate"]>>> {
     const { args, warnings } = this.getArgs(options);
 
-    const response = (await this.config.binding.run(args.model, {
+    const response = await this.config.binding.run(args.model, {
       messages: args.messages,
-    })) as { response: string };
+    });
+
+    if (response instanceof ReadableStream) {
+      throw new Error("This shouldn't happen");
+    }
 
     return {
       text: response.response,
-      finishReason: "stop",
+      // TODO: tool calls
+      // toolCalls: response.tool_calls?.map((toolCall) => ({
+      //   toolCallType: "function",
+      //   toolCallId: toolCall.name, // TODO: what can the id be?
+      //   toolName: toolCall.name,
+      //   args: JSON.stringify(toolCall.arguments || {}),
+      // })),
+      finishReason: "stop", // TODO: mapWorkersAIFinishReason(response.finish_reason),
       rawCall: { rawPrompt: args.messages, rawSettings: args },
       usage: {
+        // TODO: mapWorkersAIUsage(response.usage),
         promptTokens: 0,
         completionTokens: 0,
       },
+      warnings,
     };
   }
 
@@ -163,10 +176,14 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
 
     const decoder = new TextDecoder();
 
-    const response = (await this.config.binding.run(args.model, {
+    const response = await this.config.binding.run(args.model, {
       messages: args.messages,
       stream: true,
-    })) as ReadableStream;
+    });
+
+    if (!(response instanceof ReadableStream)) {
+      throw new Error("This shouldn't happen");
+    }
 
     return {
       stream: response.pipeThrough(
@@ -272,8 +289,8 @@ function prepareToolsAndToolChoice(
         tool_choice: "any",
       };
     default: {
-      const _exhaustiveCheck: never = type;
-      throw new Error(`Unsupported tool choice type: ${_exhaustiveCheck}`);
+      const exhaustiveCheck = type satisfies never;
+      throw new Error(`Unsupported tool choice type: ${exhaustiveCheck}`);
     }
   }
 }