alchemy 2.0.0-beta.46 → 2.0.0-beta.47

package/src/Cloudflare/AiGateway/LanguageModel.ts

@@ -0,0 +1,907 @@
+import * as Effect from "effect/Effect";
+import * as Layer from "effect/Layer";
+import * as Stream from "effect/Stream";
+import {
+  AiError,
+  LanguageModel as AiLanguageModel,
+  IdGenerator,
+  Prompt,
+  Response,
+  Tool,
+} from "effect/unstable/ai";
+import * as Sse from "effect/unstable/encoding/Sse";
+import type { RuntimeContext } from "../../RuntimeContext.ts";
+import type { AiGatewayClient } from "./AiGatewayBinding.ts";
+
+/**
+ * Options for constructing an AI Gateway-backed Workers AI LanguageModel.
+ */
+export interface LanguageModelOptions {
+  /** Already-bound AI Gateway client from `AiGatewayBinding.bind(gateway)`. */
+  readonly client: AiGatewayClient;
+  /** Workers AI model id, e.g. `@cf/meta/llama-3.3-70b-instruct-fp8-fast`. */
+  readonly model: string;
+  /** Optional per-call defaults; overridable per request via `providerOptions`. */
+  readonly parameters?: {
+    readonly temperature?: number;
+    readonly maxTokens?: number;
+    readonly topP?: number;
+    readonly topK?: number;
+    readonly seed?: number;
+    readonly frequencyPenalty?: number;
+    readonly presencePenalty?: number;
+  };
+}
+
+/**
+ * Provide a {@link AiLanguageModel.LanguageModel} layer backed by the supplied
+ * AI Gateway client and Workers AI model.
+ */
+export const makeLanguageModelLayer = (
+  options: LanguageModelOptions,
+): Layer.Layer<AiLanguageModel.LanguageModel, never, RuntimeContext> =>
+  Layer.effect(AiLanguageModel.LanguageModel, makeLanguageModel(options));
+
+/**
+ * Build a {@link AiLanguageModel.Service} that proxies generateText/streamText
+ * through the supplied AI Gateway client to a Workers AI model.
+ */
+export const makeLanguageModel = ({
+  client,
+  model,
+  parameters,
+}: LanguageModelOptions): Effect.Effect<
+  AiLanguageModel.Service,
+  never,
+  RuntimeContext
+> =>
+  Effect.gen(function* () {
+    const ai = yield* client.raw;
+    const gatewayId = yield* client.id;
+
+    const callRaw = (
+      body: WorkersAiInputs,
+      method: "generateText" | "streamText",
+    ): Effect.Effect<Response, AiError.AiError> =>
+      Effect.tryPromise({
+        try: () =>
+          ai.run(
+            model as keyof AiModels,
+            body as unknown as AiModels[keyof AiModels]["inputs"],
+            {
+              gateway: { id: gatewayId },
+              returnRawResponse: true,
+            },
+          ),
+        catch: (cause) => toAiError(cause, method),
+      });
+
+    return yield* AiLanguageModel.make({
+      generateText: (options) =>
+        Effect.gen(function* () {
+          const body = toRequestBody({ options, parameters, stream: false });
+          const resp = yield* callRaw(body, "generateText");
+          const json = yield* Effect.tryPromise({
+            try: () => resp.json() as Promise<Record<string, unknown>>,
+            catch: (cause) => toAiError(cause, "generateText"),
+          });
+          return yield* parseGenerateText(json);
+        }),
+      streamText: (options) =>
+        Stream.unwrap(
+          Effect.gen(function* () {
+            const idGen = yield* IdGenerator.IdGenerator;
+            const body = toRequestBody({ options, parameters, stream: true });
+            const resp = yield* callRaw(body, "streamText");
+            return parseStreamText(resp, idGen);
+          }),
+        ),
+    });
+  });
+
+// ---------------------------------------------------------------------------
+// Wire format types (Workers AI request)
+//
+// Workers AI returns two response shapes depending on the model:
+//   - Native:  { response: "...", tool_calls: [...] , usage }
+//   - OpenAI:  { choices: [{ message: { content, tool_calls, reasoning_content } }], usage }
+//
+// We accept both defensively — schemas would over-constrain.
+// ---------------------------------------------------------------------------
+
+interface WorkersAiMessage {
+  readonly role: "system" | "user" | "assistant" | "tool";
+  readonly content?: unknown;
+  readonly name?: string;
+  readonly tool_call_id?: string;
+  readonly tool_calls?: ReadonlyArray<{
+    readonly id: string;
+    readonly type: "function";
+    readonly function: { readonly name: string; readonly arguments: string };
+  }>;
+  readonly reasoning?: string;
+}
+
+interface WorkersAiToolDef {
+  readonly type: "function";
+  readonly function: {
+    readonly name: string;
+    readonly description?: string;
+    readonly parameters: unknown;
+  };
+}
+
+type WorkersAiToolChoice =
+  | "auto"
+  | "required"
+  | "none"
+  | { readonly type: "function"; readonly function: { readonly name: string } };
+
+interface WorkersAiInputs {
+  readonly messages: ReadonlyArray<WorkersAiMessage>;
+  readonly tools?: ReadonlyArray<WorkersAiToolDef>;
+  readonly tool_choice?: WorkersAiToolChoice;
+  readonly stream?: boolean;
+  readonly stream_options?: { readonly include_usage: boolean };
+  readonly max_tokens?: number;
+  readonly temperature?: number;
+  readonly top_p?: number;
+  readonly top_k?: number;
+  readonly random_seed?: number;
+  readonly frequency_penalty?: number;
+  readonly presence_penalty?: number;
+}
+
+// ---------------------------------------------------------------------------
+// Prompt → Workers AI messages (pure, no .push mutation)
+// ---------------------------------------------------------------------------
+
+const uint8ArrayToBase64 = (bytes: Uint8Array): string => {
+  let binary = "";
+  const chunkSize = 8192;
+  for (let i = 0; i < bytes.length; i += chunkSize) {
+    const chunk = bytes.subarray(i, Math.min(i + chunkSize, bytes.length));
+    binary += String.fromCharCode(...chunk);
+  }
+  return btoa(binary);
+};
+
+const fileToImageUrl = (
+  data: string | Uint8Array | URL,
+  mediaType: string,
+): string => {
+  if (data instanceof URL) return data.toString();
+  if (data instanceof Uint8Array) {
+    return `data:${mediaType};base64,${uint8ArrayToBase64(data)}`;
+  }
+  if (data.startsWith("data:") || data.startsWith("http")) return data;
+  return `data:${mediaType};base64,${data}`;
+};
+
+const convertPromptToMessages = (
+  prompt: Prompt.Prompt,
+): ReadonlyArray<WorkersAiMessage> =>
+  prompt.content.flatMap((m): ReadonlyArray<WorkersAiMessage> => {
+    switch (m.role) {
+      case "system":
+        return [{ role: "system", content: m.content }];
+      case "user":
+        return [toUserMessage(m.content)];
+      case "assistant":
+        return [toAssistantMessage(m.content)];
+      case "tool":
+        return m.content.flatMap(toToolMessage);
+    }
+  });
+
+const toUserMessage = (
+  parts: Prompt.UserMessage["content"],
+): WorkersAiMessage => {
+  const text = parts
+    .flatMap((p) => (p.type === "text" ? [p.text] : []))
+    .join("\n");
+  const images = parts.flatMap((p) =>
+    p.type === "file"
+      ? [
+          {
+            type: "image_url" as const,
+            image_url: { url: fileToImageUrl(p.data, p.mediaType) },
+          },
+        ]
+      : [],
+  );
+  if (images.length === 0) return { role: "user", content: text };
+  return {
+    role: "user",
+    content: [
+      ...(text.length > 0 ? [{ type: "text" as const, text }] : []),
+      ...images,
+    ],
+  };
+};
+
+const toAssistantMessage = (
+  parts: Prompt.AssistantMessage["content"],
+): WorkersAiMessage => {
+  const text = parts
+    .flatMap((p) => (p.type === "text" ? [p.text] : []))
+    .join("");
+  const reasoning = parts
+    .flatMap((p) => (p.type === "reasoning" ? [p.text] : []))
+    .join("");
+  const toolCalls = parts.flatMap((p) =>
+    p.type === "tool-call"
+      ? [
+          {
+            id: p.id,
+            type: "function" as const,
+            function: {
+              name: p.name,
+              arguments:
+                typeof p.params === "string"
+                  ? p.params
+                  : JSON.stringify(p.params ?? {}),
+            },
+          },
+        ]
+      : [],
+  );
+  return {
+    role: "assistant",
+    content: text,
+    ...(reasoning ? { reasoning } : {}),
+    ...(toolCalls.length > 0 ? { tool_calls: toolCalls } : {}),
+  };
+};
+
+const toToolMessage = (
+  part: Prompt.ToolMessage["content"][number],
+): ReadonlyArray<WorkersAiMessage> =>
+  part.type === "tool-result"
+    ? [
+        {
+          role: "tool",
+          name: part.name,
+          tool_call_id: part.id,
+          content:
+            typeof part.result === "string"
+              ? part.result
+              : JSON.stringify(part.result),
+        },
+      ]
+    : [];
+
+// ---------------------------------------------------------------------------
+// Tools / tool_choice
+// ---------------------------------------------------------------------------
+
+const prepareTools = (
+  tools: ReadonlyArray<Tool.Any>,
+  toolChoice: AiLanguageModel.ProviderOptions["toolChoice"],
+): {
+  tools?: ReadonlyArray<WorkersAiToolDef>;
+  tool_choice?: WorkersAiToolChoice;
+} => {
+  if (tools.length === 0) return {};
+  const mapped: ReadonlyArray<WorkersAiToolDef> = tools.map((tool) => ({
+    type: "function",
+    function: {
+      name: tool.name,
+      description: Tool.getDescription(tool),
+      parameters: Tool.getJsonSchema(tool),
+    },
+  }));
+
+  if (toolChoice === "auto" || toolChoice == null) {
+    return { tools: mapped, tool_choice: "auto" };
+  }
+  if (toolChoice === "none") return { tools: mapped, tool_choice: "none" };
+  if (toolChoice === "required") {
+    return { tools: mapped, tool_choice: "required" };
+  }
+  if (typeof toolChoice === "object" && "tool" in toolChoice) {
+    return {
+      tools: mapped.filter((t) => t.function.name === toolChoice.tool),
+      tool_choice: "required",
+    };
+  }
+  if (typeof toolChoice === "object" && "oneOf" in toolChoice) {
+    const allowed = new Set(toolChoice.oneOf);
+    return {
+      tools: mapped.filter((t) => allowed.has(t.function.name)),
+      tool_choice: toolChoice.mode === "required" ? "required" : "auto",
+    };
+  }
+  return { tools: mapped, tool_choice: "auto" };
+};
+
+// ---------------------------------------------------------------------------
+// Request body
+// ---------------------------------------------------------------------------
+
+const toRequestBody = ({
+  options,
+  parameters,
+  stream,
+}: {
+  readonly options: AiLanguageModel.ProviderOptions;
+  readonly parameters: LanguageModelOptions["parameters"];
+  readonly stream: boolean;
+}): WorkersAiInputs => {
+  const messages = convertPromptToMessages(options.prompt);
+  const { tools, tool_choice } = prepareTools(
+    options.tools,
+    options.toolChoice,
+  );
+  return {
+    messages,
+    ...(tools !== undefined ? { tools } : {}),
+    ...(tool_choice !== undefined ? { tool_choice } : {}),
+    // `stream_options.include_usage` is the OpenAI-compatible opt-in for
+    // usage tokens to appear in the final streamed chunk. Without it most
+    // Workers AI models omit `usage` from the stream entirely, leaving the
+    // `finish` part with zeroed counts.
+    ...(stream
+      ? { stream: true, stream_options: { include_usage: true } }
+      : {}),
+    ...(parameters?.maxTokens !== undefined
+      ? { max_tokens: parameters.maxTokens }
+      : {}),
+    ...(parameters?.temperature !== undefined
+      ? { temperature: parameters.temperature }
+      : {}),
+    ...(parameters?.topP !== undefined ? { top_p: parameters.topP } : {}),
+    ...(parameters?.topK !== undefined ? { top_k: parameters.topK } : {}),
+    ...(parameters?.seed !== undefined ? { random_seed: parameters.seed } : {}),
+    ...(parameters?.frequencyPenalty !== undefined
+      ? { frequency_penalty: parameters.frequencyPenalty }
+      : {}),
+    ...(parameters?.presencePenalty !== undefined
+      ? { presence_penalty: parameters.presencePenalty }
+      : {}),
+  };
+};
+
+// ---------------------------------------------------------------------------
+// Finish reason / usage mapping
+// ---------------------------------------------------------------------------
+
+const mapFinishReason = (raw: unknown): Response.FinishReason => {
+  switch (raw) {
+    case "stop":
+      return "stop";
+    case "length":
+    case "model_length":
+      return "length";
+    case "tool_calls":
+      return "tool-calls";
+    case "content_filter":
+    case "content-filter":
+      return "content-filter";
+    case "error":
+      return "error";
+    case undefined:
+    case null:
+      return "unknown";
+    default:
+      return "other";
+  }
+};
+
+const mapUsage = (raw: Record<string, unknown> | undefined): Response.Usage => {
+  const usage = (raw?.usage as Record<string, unknown> | undefined) ?? {};
+  const promptTokens = (usage.prompt_tokens as number | undefined) ?? 0;
+  const completionTokens = (usage.completion_tokens as number | undefined) ?? 0;
+  const cached = (
+    usage.prompt_tokens_details as { cached_tokens?: number } | undefined
+  )?.cached_tokens;
+  // Construct an actual `Response.Usage` instance — `Schema.Class<Usage>`
+  // encodes by going through the class constructor / `isInstance` check, so a
+  // plain struct that "matches" the encoded shape isn't enough.
+  return new Response.Usage({
+    inputTokens: {
+      uncached:
+        cached !== undefined
+          ? Math.max(0, promptTokens - cached)
+          : promptTokens,
+      total: promptTokens,
+      cacheRead: cached ?? 0,
+      cacheWrite: 0,
+    },
+    outputTokens: {
+      total: completionTokens,
+      text: 0,
+      reasoning: 0,
+    },
+  });
+};
+
+// ---------------------------------------------------------------------------
+// generateText: JSON → Response.PartEncoded[]
+//
+// Normalize the dual-shape (native + OpenAI) response into a single
+// `DecodedResponse` once, then build the part list with pure spreads.
+// ---------------------------------------------------------------------------
+
+interface DecodedToolCall {
+  readonly rawId: string;
+  readonly name: string;
+  readonly arguments: unknown;
+}
+
+interface DecodedResponse {
+  readonly text: string | undefined;
+  readonly reasoning: string | undefined;
+  readonly toolCalls: ReadonlyArray<DecodedToolCall>;
+  readonly finishReason: string | undefined;
+}
+
+const decodeResponse = (raw: Record<string, unknown>): DecodedResponse => {
+  const choice = (
+    raw.choices as
+      | Array<{
+          message?: {
+            content?: string | null;
+            reasoning_content?: string;
+            reasoning?: string;
+            tool_calls?: ReadonlyArray<Record<string, unknown>>;
+          };
+          finish_reason?: string;
+        }>
+      | undefined
+  )?.[0];
+  const message = choice?.message;
+
+  const openAiText = message?.content;
+  const text =
+    typeof openAiText === "string" && openAiText.length > 0
+      ? openAiText
+      : nativeTextOf(raw.response);
+  const reasoning = message?.reasoning_content ?? message?.reasoning;
+  const rawToolCalls =
+    message?.tool_calls ??
+    (Array.isArray(raw.tool_calls)
+      ? (raw.tool_calls as ReadonlyArray<Record<string, unknown>>)
+      : []);
+
+  return {
+    text,
+    reasoning: reasoning && reasoning.length > 0 ? reasoning : undefined,
+    toolCalls: rawToolCalls.flatMap(decodeToolCall),
+    finishReason:
+      choice?.finish_reason ?? (raw.finish_reason as string | undefined),
+  };
+};
+
+const nativeTextOf = (raw: unknown): string | undefined => {
+  if (raw == null) return undefined;
+  if (typeof raw === "object") return JSON.stringify(raw);
+  const text = String(raw);
+  return text.length > 0 ? text : undefined;
+};
+
+const decodeToolCall = (
+  tc: Record<string, unknown>,
+): ReadonlyArray<DecodedToolCall> => {
+  const fn = tc.function as { name?: string; arguments?: unknown } | undefined;
+  const rawId = (tc.id as string | undefined) ?? "";
+  if (fn?.name) {
+    return [{ rawId, name: fn.name, arguments: fn.arguments ?? "" }];
+  }
+  const flatName = tc.name as string | undefined;
+  if (flatName) {
+    return [{ rawId, name: flatName, arguments: tc.arguments ?? "" }];
+  }
+  return [];
+};
+
+const tryParseJsonArgs = (raw: unknown): unknown => {
+  if (typeof raw !== "string") return raw;
+  try {
+    return JSON.parse(raw);
+  } catch {
+    // Leave as raw string; the framework's tool-result decoder will fail loudly.
+    return raw;
+  }
+};
+
+const parseGenerateText = Effect.fnUntraced(function* (
+  raw: Record<string, unknown>,
+) {
+  const idGen = yield* IdGenerator.IdGenerator;
+  const decoded = decodeResponse(raw);
+
+  const toolCallParts = yield* Effect.forEach(decoded.toolCalls, (tc) =>
+    Effect.gen(function* () {
+      const id = tc.rawId || (yield* idGen.generateId());
+      return {
+        type: "tool-call" as const,
+        id,
+        name: tc.name,
+        params: tryParseJsonArgs(tc.arguments),
+      };
+    }),
+  );
+
+  const finish = mapFinishReason(
+    decoded.finishReason ??
+      (decoded.toolCalls.length > 0 ? "tool_calls" : "stop"),
+  );
+
+  return [
+    ...(decoded.reasoning !== undefined
+      ? [{ type: "reasoning" as const, text: decoded.reasoning }]
+      : []),
+    ...(decoded.text !== undefined && decoded.text.length > 0
+      ? [{ type: "text" as const, text: decoded.text }]
+      : []),
+    ...toolCallParts,
+    {
+      type: "finish" as const,
+      reason: finish,
+      usage: mapUsage(raw),
+      response: undefined,
+    },
+  ] satisfies ReadonlyArray<Response.PartEncoded>;
+});
+
+// ---------------------------------------------------------------------------
+// streamText: SSE byte stream → Stream<Response.StreamPartEncoded>
+//
+// Immutable `StreamState` is threaded through `Stream.mapAccumEffect`. The
+// per-chunk output buffer (`parts: Array<StreamPartEncoded>`) is mutable for
+// performance — it's scoped to one chunk, never escapes the handler, and lets
+// us avoid the O(n²) array-spread that pure threading would force in the hot
+// path. This matches the pattern Effect's own `@effect/ai-*` adapters use.
+// ---------------------------------------------------------------------------
+
+interface StreamState {
+  readonly textId: string | undefined;
+  readonly reasoningId: string | undefined;
+  readonly toolCalls: ReadonlyMap<
+    number,
+    { readonly id: string; readonly name: string }
+  >;
+  readonly lastToolIndex: number | undefined;
+  readonly closedToolIndices: ReadonlySet<number>;
+  readonly usage: Record<string, unknown> | undefined;
+  readonly finishReason: string | undefined;
+  readonly receivedAnyData: boolean;
+  readonly receivedDone: boolean;
+}
+
+const initialStreamState = (): StreamState => ({
+  textId: undefined,
+  reasoningId: undefined,
+  toolCalls: new Map(),
+  lastToolIndex: undefined,
+  closedToolIndices: new Set(),
+  usage: undefined,
+  finishReason: undefined,
+  receivedAnyData: false,
+  receivedDone: false,
+});
+
+type StreamParts = Array<Response.StreamPartEncoded>;
+
+const tryParseJson = (data: string): Record<string, unknown> | undefined => {
+  try {
+    const v = JSON.parse(data);
+    return v && typeof v === "object"
+      ? (v as Record<string, unknown>)
+      : undefined;
+  } catch {
+    return undefined;
+  }
+};
+
+const isNullFinalizationToolCall = (tc: Record<string, unknown>): boolean => {
+  const fn = tc.function as Record<string, unknown> | undefined;
+  const name = fn?.name ?? tc.name ?? null;
+  const args = fn?.arguments ?? tc.arguments ?? null;
+  const id = tc.id ?? null;
+  return !id && !name && (!args || args === "");
+};
+
+const closeReasoning = (
+  state: StreamState,
+  parts: StreamParts,
+): StreamState => {
+  if (state.reasoningId === undefined) return state;
+  parts.push({ type: "reasoning-end", id: state.reasoningId });
+  return { ...state, reasoningId: undefined };
+};
+
+const closeToolCall = (
+  state: StreamState,
+  index: number,
+  parts: StreamParts,
+): StreamState => {
+  if (state.closedToolIndices.has(index)) return state;
+  const tc = state.toolCalls.get(index);
+  if (!tc) return state;
+  parts.push({ type: "tool-params-end", id: tc.id });
+  const closed = new Set(state.closedToolIndices);
+  closed.add(index);
+  return { ...state, closedToolIndices: closed };
+};
+
+const emitTextDelta = (
+  state: StreamState,
+  delta: string,
+  parts: StreamParts,
+  idGen: IdGenerator.Service,
+): Effect.Effect<StreamState> =>
+  Effect.gen(function* () {
+    let s = closeReasoning(state, parts);
+    if (s.textId === undefined) {
+      const id = yield* idGen.generateId();
+      parts.push({ type: "text-start", id });
+      s = { ...s, textId: id };
+    }
+    parts.push({ type: "text-delta", id: s.textId!, delta });
+    return s;
+  });
+
+const emitReasoningDelta = (
+  state: StreamState,
+  delta: string,
+  parts: StreamParts,
+  idGen: IdGenerator.Service,
+): Effect.Effect<StreamState> =>
+  Effect.gen(function* () {
+    let s = state;
+    if (s.reasoningId === undefined) {
+      const id = yield* idGen.generateId();
+      parts.push({ type: "reasoning-start", id });
+      s = { ...s, reasoningId: id };
+    }
+    parts.push({ type: "reasoning-delta", id: s.reasoningId!, delta });
+    return s;
+  });
+
+const handleToolDeltas = (
+  state: StreamState,
+  deltas: ReadonlyArray<Record<string, unknown>>,
+  parts: StreamParts,
+  idGen: IdGenerator.Service,
+): Effect.Effect<StreamState> =>
+  Effect.gen(function* () {
+    let s = state;
+    for (const d of deltas) {
+      if (isNullFinalizationToolCall(d)) {
+        if (s.lastToolIndex !== undefined) {
+          s = closeToolCall(s, s.lastToolIndex, parts);
+        }
+        continue;
+      }
+      const idx = (d.index as number | undefined) ?? 0;
+      const fn = d.function as
+        | { name?: string; arguments?: string }
+        | undefined;
+      const name = fn?.name ?? (d.name as string | undefined) ?? "";
+      const args = fn?.arguments ?? (d.arguments as string | undefined) ?? "";
+      const rawId = (d.id as string | undefined) ?? "";
+
+      const existing = s.toolCalls.get(idx);
+      if (existing === undefined) {
+        if (s.lastToolIndex !== undefined && s.lastToolIndex !== idx) {
+          s = closeToolCall(s, s.lastToolIndex, parts);
+        }
+        const id = rawId || (yield* idGen.generateId());
+        const entry = { id, name };
+        const next = new Map(s.toolCalls);
+        next.set(idx, entry);
+        s = { ...s, toolCalls: next, lastToolIndex: idx };
+        parts.push({ type: "tool-params-start", id, name });
+        if (args.length > 0) {
+          parts.push({ type: "tool-params-delta", id, delta: args });
+        }
+      } else {
+        s = { ...s, lastToolIndex: idx };
+        if (args.length > 0) {
+          parts.push({
+            type: "tool-params-delta",
+            id: existing.id,
+            delta: args,
+          });
+        }
+      }
+    }
+    return s;
+  });
+
+const hasNonZeroUsage = (raw: unknown): boolean => {
+  if (raw == null || typeof raw !== "object") return false;
+  const u = raw as Record<string, unknown>;
+  const prompt = (u.prompt_tokens as number | undefined) ?? 0;
+  const completion = (u.completion_tokens as number | undefined) ?? 0;
+  const total = (u.total_tokens as number | undefined) ?? 0;
+  return prompt > 0 || completion > 0 || total > 0;
+};
+
+const updateChunkMeta = (
+  state: StreamState,
+  chunk: Record<string, unknown>,
+): StreamState => {
+  let s = state;
+  // Workers AI's native stream emits the real usage chunk, then a
+  // "zero-valued terminator" chunk where every count is 0 (it also re-emits
+  // `usage` with all zeros). Treat the zero chunk as a no-op so we keep the
+  // meaningful counts.
+  if (chunk.usage !== undefined && hasNonZeroUsage(chunk.usage)) {
+    s = { ...s, usage: chunk };
+  }
+  const choices = chunk.choices as
+    | Array<{ finish_reason?: string }>
+    | undefined;
+  const finish =
+    choices?.[0]?.finish_reason ?? (chunk.finish_reason as string | undefined);
+  if (finish != null) s = { ...s, finishReason: finish };
+  return s;
+};
+
+const handleNativeText = (
+  state: StreamState,
+  chunk: Record<string, unknown>,
+  parts: StreamParts,
+  idGen: IdGenerator.Service,
+): Effect.Effect<StreamState> => {
+  const native = chunk.response;
+  if (native == null || native === "") return Effect.succeed(state);
+  const text =
+    typeof native === "object" ? JSON.stringify(native) : String(native);
+  if (text.length === 0) return Effect.succeed(state);
+  return emitTextDelta(state, text, parts, idGen);
+};
+
+const handleNativeToolCalls = (
+  state: StreamState,
+  chunk: Record<string, unknown>,
+  parts: StreamParts,
+  idGen: IdGenerator.Service,
+): Effect.Effect<StreamState> => {
+  if (!Array.isArray(chunk.tool_calls)) return Effect.succeed(state);
+  return Effect.gen(function* () {
+    const s = closeReasoning(state, parts);
+    return yield* handleToolDeltas(
+      s,
+      chunk.tool_calls as ReadonlyArray<Record<string, unknown>>,
+      parts,
+      idGen,
+    );
+  });
+};
+
+const handleOpenAiDelta = (
+  state: StreamState,
+  chunk: Record<string, unknown>,
+  parts: StreamParts,
+  idGen: IdGenerator.Service,
+): Effect.Effect<StreamState> => {
+  const delta = (
+    chunk.choices as Array<{ delta?: Record<string, unknown> }> | undefined
+  )?.[0]?.delta;
+  if (!delta) return Effect.succeed(state);
+  return Effect.gen(function* () {
+    let s = state;
+    const reasoning = (delta.reasoning_content ?? delta.reasoning) as
+      | string
+      | undefined;
+    if (reasoning && reasoning.length > 0) {
+      s = yield* emitReasoningDelta(s, reasoning, parts, idGen);
+    }
+    const text = delta.content as string | undefined;
+    if (text && text.length > 0) {
+      s = yield* emitTextDelta(s, text, parts, idGen);
+    }
+    const toolDeltas = delta.tool_calls as
+      | ReadonlyArray<Record<string, unknown>>
+      | undefined;
+    if (Array.isArray(toolDeltas)) {
+      s = closeReasoning(s, parts);
+      s = yield* handleToolDeltas(s, toolDeltas, parts, idGen);
+    }
+    return s;
+  });
+};
+
+const handleStreamChunk = (
+  state: StreamState,
+  data: string,
+  idGen: IdGenerator.Service,
+): Effect.Effect<
+  readonly [StreamState, ReadonlyArray<Response.StreamPartEncoded>]
+> =>
+  Effect.gen(function* () {
+    if (data === "") return [state, []] as const;
+    if (data === "[DONE]") {
+      return [{ ...state, receivedDone: true }, []] as const;
+    }
+    const chunk = tryParseJson(data);
+    if (chunk === undefined) return [state, []] as const;
+
+    const parts: StreamParts = [];
+    let s: StreamState = { ...state, receivedAnyData: true };
+    s = updateChunkMeta(s, chunk);
+    s = yield* handleNativeText(s, chunk, parts, idGen);
+    s = yield* handleNativeToolCalls(s, chunk, parts, idGen);
+    s = yield* handleOpenAiDelta(s, chunk, parts, idGen);
+    return [s, parts] as const;
+  });
+
+const finalizeStream = (
+  state: StreamState,
+): ReadonlyArray<Response.StreamPartEncoded> => {
+  const parts: StreamParts = [];
+  let s = state;
+  for (const [idx] of s.toolCalls) s = closeToolCall(s, idx, parts);
+  s = closeReasoning(s, parts);
+  if (s.textId !== undefined) parts.push({ type: "text-end", id: s.textId });
+
+  // Three cases for the final reason:
+  //  1. The model emitted an explicit `finish_reason`           → map it.
+  //  2. The stream ended cleanly (`[DONE]` seen) but no reason  → "stop".
+  //     Workers AI's native shape never includes `finish_reason`,
+  //     so without this rule every native-mode stream would report
+  //     `unknown` despite completing successfully.
+  //  3. The stream ended abnormally (no `[DONE]`, no reason)    → "error".
+  const reason: Response.FinishReason =
+    s.finishReason !== undefined
+      ? mapFinishReason(s.finishReason)
+      : s.receivedDone
+        ? "stop"
+        : s.receivedAnyData
+          ? "error"
+          : "unknown";
+
+  parts.push({
+    type: "finish",
+    reason,
+    usage: mapUsage(s.usage),
+    response: undefined,
+  });
+  return parts;
+};
+
+const parseStreamText = (
+  resp: Response,
+  idGen: IdGenerator.Service,
+): Stream.Stream<Response.StreamPartEncoded, AiError.AiError> => {
+  const body = resp.body;
+  if (body === null) {
+    return Stream.fromIterable<Response.StreamPartEncoded>(
+      finalizeStream(initialStreamState()),
+    );
+  }
+  return Stream.fromReadableStream<Uint8Array, AiError.AiError>({
+    evaluate: () => body,
+    onError: (cause) => toAiError(cause, "streamText"),
+  }).pipe(
+    Stream.decodeText(),
+    Stream.pipeThroughChannel(Sse.decode<AiError.AiError, unknown>()),
+    Stream.catchTag("Retry", (retry) => Stream.die(retry)),
+    Stream.mapAccumEffect(
+      initialStreamState,
+      (state, event) => handleStreamChunk(state, event.data, idGen),
+      { onHalt: (state) => finalizeStream(state) },
+    ),
+  );
+};
+
+// ---------------------------------------------------------------------------
+// Error mapping
+// ---------------------------------------------------------------------------
+
+const toAiError = (
+  cause: unknown,
+  method: "generateText" | "streamText",
+): AiError.AiError =>
+  AiError.AiError.make({
+    module: "Cloudflare.AiGateway.LanguageModel",
+    method,
+    reason: new AiError.UnknownError({
+      description:
+        cause instanceof Error ? cause.message : "AI Gateway request failed",
+    }),
+  });