feat(ai): configure chat max tokens field (#39909)

This commit is contained in:
Aiden Cline 2026-07-31 11:14:08 -05:00 committed by GitHub
commit 3468aa0140
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 44 additions and 84 deletions

View file

@ -107,6 +107,7 @@ export const bodyFields = {
stream_options: Schema.optional(Schema.Struct({ include_usage: Schema.Boolean })), stream_options: Schema.optional(Schema.Struct({ include_usage: Schema.Boolean })),
store: Schema.optional(Schema.Boolean), store: Schema.optional(Schema.Boolean),
reasoning_effort: Schema.optional(OpenAIOptions.OpenAIReasoningEffort), reasoning_effort: Schema.optional(OpenAIOptions.OpenAIReasoningEffort),
max_completion_tokens: Schema.optional(Schema.Number),
max_tokens: Schema.optional(Schema.Number), max_tokens: Schema.optional(Schema.Number),
temperature: Schema.optional(Schema.Number), temperature: Schema.optional(Schema.Number),
top_p: Schema.optional(Schema.Number), top_p: Schema.optional(Schema.Number),
@ -415,6 +416,7 @@ const fromRequest = Effect.fn("OpenAIChat.fromRequest")(function* (request: LLMR
) )
const generation = request.generation const generation = request.generation
const toolSchemaCompatibility = request.model.compatibility?.toolSchema const toolSchemaCompatibility = request.model.compatibility?.toolSchema
const maxTokensField = request.model.compatibility?.maxTokensField ?? "max_tokens"
return { return {
model: request.model.id, model: request.model.id,
messages: yield* lowerMessages(request), messages: yield* lowerMessages(request),
@ -427,7 +429,9 @@ const fromRequest = Effect.fn("OpenAIChat.fromRequest")(function* (request: LLMR
tool_choice: request.toolChoice ? yield* lowerToolChoice(request.toolChoice) : undefined, tool_choice: request.toolChoice ? yield* lowerToolChoice(request.toolChoice) : undefined,
stream: true as const, stream: true as const,
stream_options: { include_usage: true }, stream_options: { include_usage: true },
max_tokens: generation?.maxTokens, ...(maxTokensField === "max_completion_tokens"
? { max_completion_tokens: generation?.maxTokens }
: { max_tokens: generation?.maxTokens }),
temperature: generation?.temperature, temperature: generation?.temperature,
top_p: generation?.topP, top_p: generation?.topP,
frequency_penalty: generation?.frequencyPenalty, frequency_penalty: generation?.frequencyPenalty,

View file

@ -28,57 +28,9 @@ const applyQuery = (url: string, query: Record<string, string> | undefined) => {
return next.toString() return next.toString()
} }
const PROTOCOL_BODY_OVERLAY_DENYLIST = new Set([
"anthropic_version",
"content",
"contents",
"frequencyPenalty",
"frequency_penalty",
"generationConfig",
"inferenceConfig",
"input",
"maxTokens",
"max_tokens",
"messages",
"model",
"presencePenalty",
"presence_penalty",
"responseFormat",
"response_format",
"seed",
"stop",
"stopSequences",
"stop_sequences",
"stream",
"streamOptions",
"stream_options",
"system",
"systemInstruction",
"system_instruction",
"temperature",
"thinking",
"toolChoice",
"toolConfig",
"tool_choice",
"tool_config",
"tools",
"topK",
"topP",
"top_k",
"top_p",
])
const forbiddenBodyOverlayKeys = (body: Record<string, unknown>) =>
Object.keys(body).filter((key) => PROTOCOL_BODY_OVERLAY_DENYLIST.has(key))
const bodyWithOverlay = <Body>(body: Body, request: LLMRequest, encodeBody: (body: Body) => string) => const bodyWithOverlay = <Body>(body: Body, request: LLMRequest, encodeBody: (body: Body) => string) =>
Effect.gen(function* () { Effect.gen(function* () {
if (request.http?.body === undefined) return { jsonBody: body, bodyText: encodeBody(body) } if (request.http?.body === undefined) return { jsonBody: body, bodyText: encodeBody(body) }
const forbiddenKeys = forbiddenBodyOverlayKeys(request.http.body)
if (forbiddenKeys.length > 0)
return yield* ProviderShared.invalidRequest(
`http.body cannot overlay protocol-owned field(s): ${forbiddenKeys.join(", ")}`,
)
if (ProviderShared.isRecord(body)) { if (ProviderShared.isRecord(body)) {
const overlaid = mergeJsonRecords(body, request.http.body) ?? {} const overlaid = mergeJsonRecords(body, request.http.body) ?? {}
return { jsonBody: overlaid, bodyText: ProviderShared.encodeJson(overlaid) } return { jsonBody: overlaid, bodyText: ProviderShared.encodeJson(overlaid) }

View file

@ -167,9 +167,13 @@ export namespace ModelDefaults {
export const ModelToolSchemaCompatibility = Schema.Literals(["gemini", "moonshot"]) export const ModelToolSchemaCompatibility = Schema.Literals(["gemini", "moonshot"])
export type ModelToolSchemaCompatibility = Schema.Schema.Type<typeof ModelToolSchemaCompatibility> export type ModelToolSchemaCompatibility = Schema.Schema.Type<typeof ModelToolSchemaCompatibility>
export const ModelMaxTokensFieldCompatibility = Schema.Literals(["max_completion_tokens", "max_tokens"])
export type ModelMaxTokensFieldCompatibility = Schema.Schema.Type<typeof ModelMaxTokensFieldCompatibility>
export class ModelCompatibility extends Schema.Class<ModelCompatibility>("LLM.ModelCompatibility")({ export class ModelCompatibility extends Schema.Class<ModelCompatibility>("LLM.ModelCompatibility")({
toolSchema: Schema.optional(ModelToolSchemaCompatibility), toolSchema: Schema.optional(ModelToolSchemaCompatibility),
reasoningField: Schema.optional(Schema.String), reasoningField: Schema.optional(Schema.String),
maxTokensField: Schema.optional(ModelMaxTokensFieldCompatibility),
}) {} }) {}
export namespace ModelCompatibility { export namespace ModelCompatibility {

View file

@ -171,24 +171,27 @@ describe("request option precedence", () => {
), ),
) )
it.effect("rejects raw body overlays for protocol-owned roots", () => it.effect("applies raw body overlays after protocol lowering", () =>
Effect.gen(function* () { LLMClient.generate(
const model = OpenAIChat.route LLM.request({
.with({ endpoint: { baseURL: "https://api.openai.test/v1/" }, auth: Auth.bearer("test") }) model: OpenAIChat.route
.model({ id: "gpt-4o-mini" }) .with({ endpoint: { baseURL: "https://api.openai.test/v1/" }, auth: Auth.bearer("test") })
const error = yield* compileRequest( .model({ id: "gpt-4o-mini" }),
LLM.request({ prompt: "Say hello.",
model, http: { body: { model: "gpt-5", messages: [], tools: [] } },
prompt: "Say hello.", }),
http: { body: { model: "gpt-5", messages: [], tools: [] } }, ).pipe(
}), Effect.provide(
).pipe(Effect.flip) dynamicResponse((input) =>
Effect.gen(function* () {
expect(error.reason).toMatchObject({ expect(decodeJson(input.text)).toMatchObject({ model: "gpt-5", messages: [], tools: [] })
_tag: "InvalidRequest", return input.respond(sseEvents(deltaChunk({}, "stop")), {
message: "http.body cannot overlay protocol-owned field(s): model, messages, tools", headers: { "content-type": "text/event-stream" },
}) })
}), }),
),
),
),
) )
it.effect("uses model output limits after route limits and before call maxTokens", () => it.effect("uses model output limits after route limits and before call maxTokens", () =>

View file

@ -181,23 +181,6 @@ describe("Google Vertex providers", () => {
}), }),
) )
it.effect("protects the Vertex Messages API version from body overlays", () =>
Effect.gen(function* () {
const error = yield* compileRequest(
LLM.request({
model: GoogleVertexMessages.configure({
accessToken: "vertex-token",
http: { body: { anthropic_version: "wrong" } },
project: "vertex-project",
}).model("claude-sonnet-4-6"),
prompt: "Say hello.",
}),
).pipe(Effect.flip)
expect(error.message).toContain("http.body cannot overlay protocol-owned field(s): anthropic_version")
}),
)
it.effect("routes tuned Gemini models through their deployed endpoint", () => it.effect("routes tuned Gemini models through their deployed endpoint", () =>
Effect.gen(function* () { Effect.gen(function* () {
const response = yield* LLMClient.generate( const response = yield* LLMClient.generate(

View file

@ -144,6 +144,20 @@ describe("OpenAI-compatible Chat route", () => {
}), }),
) )
it.effect("configures the max tokens request field", () =>
Effect.gen(function* () {
const compatible = OpenAICompatibleChat.route
.with({ provider: "custom", endpoint: { baseURL: "https://api.custom.test/v1" } })
.model({ id: "custom-model", compatibility: { maxTokensField: "max_completion_tokens" } })
const prepared = yield* compileRequest(
LLM.request({ model: compatible, prompt: "Say hello.", generation: { maxTokens: 20 } }),
)
expect(prepared.body).toMatchObject({ max_completion_tokens: 20 })
expect(prepared.body).not.toHaveProperty("max_tokens")
}),
)
it.effect("matches AI SDK compatible tool request body fixture", () => it.effect("matches AI SDK compatible tool request body fixture", () =>
Effect.gen(function* () { Effect.gen(function* () {
const prepared = yield* compileRequest( const prepared = yield* compileRequest(