Refactor LLM route-first provider API (#28523)

packages/llm/src/route/transport/http.ts

@@ -1,20 +1,13 @@
 import { Effect, Stream } from "effect"
 import { Headers, HttpClientRequest } from "effect/unstable/http"
-import { Auth, type Auth as AuthDef } from "../auth"
-import { type Endpoint, render as renderEndpoint } from "../endpoint"
-import type { Framing } from "../framing"
-import type { Transport } from "./index"
+import { Auth } from "../auth"
+import { render as renderEndpoint } from "../endpoint"
+import { Framing, type Framing as FramingDef } from "../framing"
+import type { Transport, TransportPrepareInput } from "./index"
 import * as ProviderShared from "../../protocols/shared"
 import { mergeJsonRecords, type LLMRequest } from "../../schema"
 
-export interface JsonRequestInput<Body> {
-  readonly body: Body
-  readonly request: LLMRequest
-  readonly endpoint: Endpoint<Body>
-  readonly auth: AuthDef
-  readonly encodeBody: (body: Body) => string
-  readonly headers?: (input: { readonly request: LLMRequest }) => Record<string, string>
-}
+export type JsonRequestInput<Body> = TransportPrepareInput<Body>
 
 export interface JsonRequestParts<Body = unknown> {
   readonly url: string
@@ -25,7 +18,7 @@ export interface JsonRequestParts<Body = unknown> {
 
 export interface HttpPrepared<Frame> {
   readonly request: HttpClientRequest.HttpClientRequest
-  readonly framing: Framing<Frame>
+  readonly framing: FramingDef<Frame>
 }
 
 const applyQuery = (url: string, query: Record<string, string> | undefined) => {
@@ -52,28 +45,21 @@ export const jsonRequestParts = <Body>(input: JsonRequestInput<Body>) =>
       input.request.http?.query,
     )
     const body = yield* bodyWithOverlay(input.body, input.request, input.encodeBody)
-    const headers = yield* Auth.toEffect(Auth.isAuth(input.request.model.auth) ? input.request.model.auth : input.auth)(
-      {
-        request: input.request,
-        method: "POST",
-        url,
-        body: body.bodyText,
-        headers: Headers.fromInput({
-          ...(input.headers?.({ request: input.request }) ?? {}),
-          ...input.request.model.headers,
-          ...input.request.http?.headers,
-        }),
-      },
-    )
+    const headers = yield* Auth.toEffect(input.auth)({
+      request: input.request,
+      method: "POST",
+      url,
+      body: body.bodyText,
+      headers: Headers.fromInput({
+        ...input.headers?.({ request: input.request }),
+        ...input.request.http?.headers,
+      }),
+    })
     return { url, jsonBody: body.jsonBody, bodyText: body.bodyText, headers }
   })
 
-export interface HttpJsonInput<Body, Frame> {
-  readonly endpoint: Endpoint<Body>
-  readonly auth?: AuthDef
-  readonly framing: Framing<Frame>
-  readonly encodeBody: (body: Body) => string
-  readonly headers?: (input: { readonly request: LLMRequest }) => Record<string, string>
+export interface HttpJsonInput<_Body, Frame> {
+  readonly framing: FramingDef<Frame>
 }
 
 export type HttpJsonPatch<Body, Frame> = Partial<HttpJsonInput<Body, Frame>>
@@ -85,14 +71,9 @@ export interface HttpJsonTransport<Body, Frame> extends Transport<Body, HttpPrep
 export const httpJson = <Body, Frame>(input: HttpJsonInput<Body, Frame>): HttpJsonTransport<Body, Frame> => ({
   id: "http-json",
   with: (patch) => httpJson({ ...input, ...patch }),
-  prepare: (body, request) =>
+  prepare: (prepareInput) =>
     jsonRequestParts({
-      body,
-      request,
-      endpoint: input.endpoint,
-      auth: input.auth ?? Auth.bearer(),
-      encodeBody: input.encodeBody,
-      headers: input.headers,
+      ...prepareInput,
     }).pipe(
       Effect.map((parts) => ({
         request: ProviderShared.jsonPost({ url: parts.url, body: parts.bodyText, headers: parts.headers }),
@@ -109,8 +90,8 @@ export const httpJson = <Body, Frame>(input: HttpJsonInput<Body, Frame>): HttpJs
               response.stream.pipe(
                 Stream.mapError((error) =>
                   ProviderShared.eventError(
-                    `${request.model.provider}/${request.model.route}`,
-                    `Failed to read ${request.model.provider}/${request.model.route} stream`,
+                    `${request.model.provider}/${request.model.route.id}`,
+                    `Failed to read ${request.model.provider}/${request.model.route.id} stream`,
                     ProviderShared.errorText(error),
                   ),
                 ),
@@ -120,3 +101,8 @@ export const httpJson = <Body, Frame>(input: HttpJsonInput<Body, Frame>): HttpJs
         ),
     ),
 })
+
+export const sseJson = {
+  id: "http-json/sse",
+  with: <Body>() => httpJson<Body, string>({ framing: Framing.sse }),
+} as const

packages/llm/src/route/transport/index.ts

@@ -1,4 +1,6 @@
 import type { Effect, Stream } from "effect"
+import type { Endpoint } from "../endpoint"
+import type { Auth } from "../auth"
 import type { Interface as RequestExecutorInterface } from "../executor"
 import type { Interface as WebSocketExecutorInterface } from "./websocket"
 import type { LLMError, LLMRequest } from "../../schema"
@@ -10,7 +12,7 @@ export interface TransportRuntime {
 
 export interface Transport<Body, Prepared, Frame> {
   readonly id: string
-  readonly prepare: (body: Body, request: LLMRequest) => Effect.Effect<Prepared, LLMError>
+  readonly prepare: (input: TransportPrepareInput<Body>) => Effect.Effect<Prepared, LLMError>
   readonly frames: (
     prepared: Prepared,
     request: LLMRequest,
@@ -18,5 +20,14 @@ export interface Transport<Body, Prepared, Frame> {
   ) => Stream.Stream<Frame, LLMError>
 }
 
+export interface TransportPrepareInput<Body> {
+  readonly body: Body
+  readonly request: LLMRequest
+  readonly endpoint: Endpoint<Body>
+  readonly auth: Auth
+  readonly encodeBody: (body: Body) => string
+  readonly headers?: (input: { readonly request: LLMRequest }) => Record<string, string>
+}
+
 export * as HttpTransport from "./http"
 export { WebSocketExecutor, WebSocketTransport } from "./websocket"

packages/llm/src/route/transport/websocket.ts

@@ -1,7 +1,6 @@
-import { Cause, Context, Effect, Queue, Stream } from "effect"
+import { Cause, Context, Effect, Layer, Queue, Stream } from "effect"
 import { Headers } from "effect/unstable/http"
-import { Auth, type Auth as AuthDef } from "../auth"
-import type { Endpoint } from "../endpoint"
+import { Auth } from "../auth"
 import { LLMError, TransportReason, type LLMRequest } from "../../schema"
 import * as HttpTransport from "./http"
 import type { Transport } from "./index"
@@ -135,6 +134,8 @@ export const open = (input: WebSocketRequest) =>
       }),
   }).pipe(Effect.flatMap((ws) => fromWebSocket(ws, input)))
 
+export const layer: Layer.Layer<Service> = Layer.succeed(Service, Service.of({ open }))
+
 export const fromWebSocket = (
   ws: globalThis.WebSocket,
   input: WebSocketRequest,
@@ -213,12 +214,8 @@ export interface JsonPrepared {
 }
 
 export interface JsonInput<Body, Message> {
-  readonly endpoint: Endpoint<Body>
-  readonly auth?: AuthDef
-  readonly encodeBody: (body: Body) => string
   readonly toMessage: (body: Body | Record<string, unknown>) => Effect.Effect<Message, LLMError>
   readonly encodeMessage: (message: Message) => string
-  readonly headers?: (input: { readonly request: LLMRequest }) => Record<string, string>
 }
 
 export type JsonPatch<Body, Message> = Partial<JsonInput<Body, Message>>
@@ -230,15 +227,10 @@ export interface JsonTransport<Body, Message> extends Transport<Body, JsonPrepar
 export const json = <Body, Message>(input: JsonInput<Body, Message>): JsonTransport<Body, Message> => ({
   id: "websocket-json",
   with: (patch) => json({ ...input, ...patch }),
-  prepare: (body, request) =>
+  prepare: (prepareInput) =>
     Effect.gen(function* () {
       const parts = yield* HttpTransport.jsonRequestParts({
-        body,
-        request,
-        endpoint: input.endpoint,
-        auth: input.auth ?? Auth.bearer(),
-        encodeBody: input.encodeBody,
-        headers: input.headers,
+        ...prepareInput,
       })
       return {
         url: yield* webSocketUrl(parts.url),
@@ -270,8 +262,14 @@ export const json = <Body, Message>(input: JsonInput<Body, Message>): JsonTransp
   },
 })
 
+export const jsonTransport = {
+  id: "websocket-json",
+  with: json,
+} as const
+
 export const WebSocketExecutor = {
   Service,
+  layer,
   open,
   fromWebSocket,
   messageText,
@@ -279,4 +277,5 @@ export const WebSocketExecutor = {
 
 export const WebSocketTransport = {
   json,
+  jsonTransport,
 } as const