refactor(ai): add protocol body extensions

2026-09-15 05:16:26 +00:00 · 2026-09-14 17:53:23 -05:00
5 changed files with 110 additions and 44 deletions
@@ -16,33 +16,30 @@ const Options = Schema.Struct({
  ),
 })
 export type OptionsInput = typeof Options.Type & Pick<AnthropicMessages.OptionsInput, "outputConfig">
-export const protocol = Protocol.make({
+export const protocol = Protocol.withBody(AnthropicMessages.protocol, {
  id: "alibaba-messages",
-  body: {
-    schema: Schema.Struct({
-      ...AnthropicMessages.AnthropicMessagesBody.fields,
-      thinking: Schema.optional(Schema.Struct({ type: Schema.String, budget_tokens: Schema.optional(Schema.Int) })),
-    }),
-    from: Effect.fn("AlibabaMessages.fromRequest")(function* (req) {
-      const opts = yield* ProviderShared.validateWith(Schema.decodeUnknownEffect(Options))(req.providerOptions ?? {})
-      // Model Studio accepts enabled thinking without Anthropic's mandatory token budget.
-      return {
-        ...(yield* AnthropicMessages.protocol.body.from(
-          LLMRequest.update(req, {
-            providerOptions: { ...req.providerOptions, thinking: undefined },
-          }),
-        )),
-        thinking:
-          opts.thinking === undefined
-            ? undefined
-            : {
-                type: opts.thinking.type,
-                budget_tokens: opts.thinking.budgetTokens ?? opts.thinking.budget_tokens,
-              },
-      }
-    }),
-  },
-  stream: AnthropicMessages.protocol.stream,
+  schema: Schema.Struct({
+    ...AnthropicMessages.AnthropicMessagesBody.fields,
+    thinking: Schema.optional(Schema.Struct({ type: Schema.String, budget_tokens: Schema.optional(Schema.Int) })),
+  }),
+  from: Effect.fn("AlibabaMessages.fromRequest")(function* (req, fromBase) {
+    const opts = yield* ProviderShared.validateWith(Schema.decodeUnknownEffect(Options))(req.providerOptions ?? {})
+    // Model Studio accepts enabled thinking without Anthropic's mandatory token budget.
+    return {
+      ...(yield* fromBase(
+        LLMRequest.update(req, {
+          providerOptions: { ...req.providerOptions, thinking: undefined },
+        }),
+      )),
+      thinking:
+        opts.thinking === undefined
+          ? undefined
+          : {
+              type: opts.thinking.type,
+              budget_tokens: opts.thinking.budgetTokens ?? opts.thinking.budget_tokens,
+            },
+    }
+  }),
 })

 export * as AlibabaMessages from "./alibaba-messages.js"
@@ -13,7 +13,7 @@ import { JsonObject, optionalArray, optionalNull, ProviderShared } from "./share
 import { OpenAIImage } from "./utils/openai-image.js"
 import { ResponsesHostedTools } from "./utils/responses-hosted-tools.js"
 import { ToolSchemaProjection } from "./utils/tool-schema.js"
-import { OpenResponsesChannel } from "./open-responses-channel.js"
+import { OpenResponsesChannel, type Options } from "./open-responses-channel.js"
 import { ResponsesCompaction } from "./utils/responses-compaction.js"
 import { ResponsesCheckpoint } from "./utils/responses-checkpoint.js"

@@ -96,7 +96,7 @@ const OpenAIResponsesToolChoice = Schema.Union([
  Schema.Struct({ type: Schema.tag("image_generation") }),
 ])

-const OpenAIResponsesInputItem = Schema.Union([
+export const OpenAIResponsesInputItem = Schema.Union([
  OpenResponses.InputItem,
  OpenAIResponsesHostedToolItem,
  OpenResponses.ConfigurationUpdate,
@@ -117,7 +117,7 @@ const OpenAIResponsesCoreFields = {
  ),
 }

-const OpenAIResponsesBody = Schema.Struct({
+export const OpenAIResponsesBody = Schema.Struct({
  ...OpenAIResponsesCoreFields,
  stream: Schema.Literal(true),
 })
@@ -322,7 +322,9 @@ const endpoint = Endpoint.path<OpenAIResponsesBody>(PATH, { baseURL: DEFAULT_BAS
 const auth = Auth.none

 export const httpTransport = HttpTransport.sseJson.with<OpenAIResponsesBody>()
-export const channelTransport = OpenResponsesChannel.transport<OpenAIResponsesBody>
+export const channelTransport = <Body extends OpenAIResponsesBody = OpenAIResponsesBody>(
+  input: Options,
+) => OpenResponsesChannel.transport<Body>(input)
 export const transport = channelTransport({
  id: ADAPTER,
  name: NAME,
@@ -19,21 +19,21 @@ const Body = Schema.Struct({
  thinking: Options.fields.thinking,
 })

-const fromRequest = Effect.fn("ZAIMessages.fromRequest")(function* (request: LLMRequest) {
-  const options = yield* ProviderShared.validateWith(Schema.decodeUnknownEffect(Options))(request.providerOptions ?? {})
-  // Z.AI accepts enabled thinking without Anthropic's mandatory token budget.
-  const body = yield* AnthropicMessages.protocol.body.from(
-    LLMRequest.update(request, {
-      providerOptions: { ...request.providerOptions, thinking: undefined },
-    }),
-  )
-  return { ...body, thinking: options.thinking }
-})
-
-export const protocol = Protocol.make({
+export const protocol = Protocol.withBody(AnthropicMessages.protocol, {
  id: "zai-messages",
-  body: { schema: Body, from: fromRequest },
-  stream: AnthropicMessages.protocol.stream,
+  schema: Body,
+  from: Effect.fn("ZAIMessages.fromRequest")(function* (request, fromBase) {
+    const options = yield* ProviderShared.validateWith(Schema.decodeUnknownEffect(Options))(
+      request.providerOptions ?? {},
+    )
+    // Z.AI accepts enabled thinking without Anthropic's mandatory token budget.
+    const body = yield* fromBase(
+      LLMRequest.update(request, {
+        providerOptions: { ...request.providerOptions, thinking: undefined },
+      }),
+    )
+    return { ...body, thinking: options.thinking }
+  }),
 })

 export * as ZAIMessages from "./zai-messages.js"
@@ -52,6 +52,15 @@ export interface ProtocolBody<Body> {
  readonly from: (request: LLMRequest) => Effect.Effect<Body, AIError>
 }

+export interface ProtocolBodyExtension<Body, NextBody> {
+  /** Stable id for the derived wire protocol. Defaults to the base protocol id. */
+  readonly id?: ProtocolID
+  /** Schema for the derived provider-native body. */
+  readonly schema: Schema.Codec<NextBody, unknown>
+  /** Build the derived body, optionally invoking the base protocol with this or another request. */
+  readonly from: (request: LLMRequest, base: ProtocolBody<Body>["from"]) => Effect.Effect<NoInfer<NextBody>, AIError>
+}
+
 export interface ProtocolStream<Frame, Event, State> {
  /** Schema for one decoded streaming event, decoded from a transport frame. */
  readonly event: Schema.Codec<Event, Frame>
@@ -81,6 +90,20 @@ export const make = <Body, Frame, Event, State>(
  input: Protocol<Body, Frame, Event, State>,
 ): Protocol<Body, Frame, Event, State> => input

+/** Derive a protocol with a different request body while retaining its response stream. */
+export const withBody = <Body, NextBody, Frame, Event, State>(
+  protocol: Protocol<Body, Frame, Event, State>,
+  extension: ProtocolBodyExtension<Body, NextBody>,
+): Protocol<NextBody, Frame, Event, State> =>
+  make({
+    ...protocol,
+    id: extension.id ?? protocol.id,
+    body: {
+      schema: extension.schema,
+      from: (request) => extension.from(request, protocol.body.from),
+    },
+  })
+
 export const jsonEvent = <const S extends Schema.Top>(schema: S) => Schema.fromJsonString(schema)

 export * as Protocol from "./protocol.js"
@@ -0,0 +1,44 @@
+import { describe, expect, test } from "bun:test"
+import { Effect, Schema } from "effect"
+import { LLM } from "../src/index.js"
+import { OpenAIChat } from "../src/protocols/openai-chat.js"
+import { Protocol } from "../src/route.js"
+import { LanguageModel } from "../src/schema/index.js"
+
+describe("Protocol.withBody", () => {
+  const BaseBody = Schema.Struct({ prompt: Schema.String })
+  const ExtendedBody = Schema.Struct({ ...BaseBody.fields, priority: Schema.String })
+  const stream = {
+    event: Schema.String,
+    initial: () => undefined,
+    step: () => Effect.succeed([undefined, []] as const),
+  }
+  const base = Protocol.make({
+    id: "base",
+    body: {
+      schema: BaseBody,
+      from: (request) => Effect.succeed({ prompt: String(request.model.id) }),
+    },
+    stream,
+    supportsEffortUpdates: () => true,
+  })
+
+  test("derives a typed body while retaining the response protocol", async () => {
+    const derived = Protocol.withBody(base, {
+      schema: ExtendedBody,
+      from: (request, fromBase) => fromBase(request).pipe(Effect.map((body) => ({ ...body, priority: "high" }))),
+    })
+    const model = LanguageModel.make({
+      id: "model",
+      provider: "test",
+      route: OpenAIChat.route,
+    })
+
+    expect(await Effect.runPromise(derived.body.from(LLM.request({ model, prompt: "Hello" })))).toEqual({
+      prompt: "model",
+      priority: "high",
+    })
+    expect(derived.stream).toBe(stream)
+    expect(derived.supportsEffortUpdates).toBe(base.supportsEffortUpdates)
+  })
+})