feat(llm): add tool schema projections (#34454)

This commit is contained in:
Shoubhit Dash 2026-06-29 21:45:42 +05:30 committed by GitHub
parent 71ec022b47
commit 18466b8020
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 261 additions and 54 deletions

View file

@ -9,6 +9,7 @@ import {
Usage,
type CacheHint,
type FinishReason,
type JsonSchema,
type LLMRequest,
type MediaPart,
type ProviderMetadata,
@ -21,6 +22,7 @@ import { JsonObject, optionalArray, optionalNull, ProviderShared } from "./share
import { isContextOverflow } from "../provider-error"
import * as Cache from "./utils/cache"
import { Lifecycle } from "./utils/lifecycle"
import { ToolSchemaProjection } from "./utils/tool-schema"
import { ToolStream } from "./utils/tool-stream"
const ADAPTER = "anthropic-messages"
@ -256,10 +258,10 @@ const signatureFromMetadata = (metadata: ProviderMetadata | undefined): string |
return typeof anthropic.signature === "string" ? anthropic.signature : undefined
}
const lowerTool = (breakpoints: Cache.Breakpoints, tool: ToolDefinition): AnthropicTool => ({
const lowerTool = (breakpoints: Cache.Breakpoints, tool: ToolDefinition, inputSchema: JsonSchema): AnthropicTool => ({
name: tool.name,
description: tool.description,
input_schema: tool.inputSchema,
input_schema: inputSchema,
cache_control: cacheControl(breakpoints, tool.cache),
})
@ -504,6 +506,7 @@ const lowerThinking = Effect.fn("AnthropicMessages.lowerThinking")(function* (re
const fromRequest = Effect.fn("AnthropicMessages.fromRequest")(function* (request: LLMRequest) {
const toolChoice = request.toolChoice ? yield* lowerToolChoice(request.toolChoice) : undefined
const generation = request.generation
const toolSchemaCompatibility = request.model.compatibility?.toolSchema
const outputLimit = request.model.defaults?.limits?.output ?? request.model.route.defaults.limits?.output ?? 4096
// Allocate the 4-breakpoint budget in invalidation order: tools → system →
// messages. Tools live highest in the cache hierarchy, so when callers
@ -512,7 +515,9 @@ const fromRequest = Effect.fn("AnthropicMessages.fromRequest")(function* (reques
const tools =
request.tools.length === 0 || request.toolChoice?.type === "none"
? undefined
: request.tools.map((tool) => lowerTool(breakpoints, tool))
: request.tools.map((tool) =>
lowerTool(breakpoints, tool, ToolSchemaProjection.modelCompatibility(tool.inputSchema, toolSchemaCompatibility)),
)
const system =
request.system.length === 0
? undefined

View file

@ -7,7 +7,9 @@ import {
Usage,
type CacheHint,
type FinishReason,
type JsonSchema,
type LLMRequest,
type ModelToolSchemaCompatibility,
type ProviderMetadata,
type ReasoningPart,
type ToolCallPart,
@ -21,6 +23,7 @@ import { BedrockAuth } from "./utils/bedrock-auth"
import { BedrockCache } from "./utils/bedrock-cache"
import { BedrockMedia } from "./utils/bedrock-media"
import { Lifecycle } from "./utils/lifecycle"
import { ToolSchemaProjection } from "./utils/tool-schema"
import { ToolStream } from "./utils/tool-stream"
const ADAPTER = "bedrock-converse"
@ -205,18 +208,22 @@ type BedrockEvent = Schema.Schema.Type<typeof BedrockEvent>
// =============================================================================
// Request Lowering
// =============================================================================
const lowerToolSpec = (tool: ToolDefinition): BedrockToolSpec => ({
const lowerToolSpec = (tool: ToolDefinition, inputSchema: JsonSchema): BedrockToolSpec => ({
toolSpec: {
name: tool.name,
description: tool.description,
inputSchema: { json: tool.inputSchema },
inputSchema: { json: inputSchema },
},
})
const lowerTools = (breakpoints: BedrockCache.Breakpoints, tools: ReadonlyArray<ToolDefinition>): BedrockTool[] => {
const lowerTools = (
compatibility: ModelToolSchemaCompatibility | undefined,
breakpoints: BedrockCache.Breakpoints,
tools: ReadonlyArray<ToolDefinition>,
): BedrockTool[] => {
const result: BedrockTool[] = []
for (const tool of tools) {
result.push(lowerToolSpec(tool))
result.push(lowerToolSpec(tool, ToolSchemaProjection.modelCompatibility(tool.inputSchema, compatibility)))
const cachePoint = BedrockCache.block(breakpoints, tool.cache)
if (cachePoint) result.push(cachePoint)
}
@ -386,7 +393,7 @@ const fromRequest = Effect.fn("BedrockConverse.fromRequest")(function* (request:
const breakpoints = BedrockCache.breakpoints()
const toolConfig =
request.tools.length > 0 && request.toolChoice?.type !== "none"
? { tools: lowerTools(breakpoints, request.tools), toolChoice }
? { tools: lowerTools(request.model.compatibility?.toolSchema, breakpoints, request.tools), toolChoice }
: undefined
const system = request.system.length === 0 ? undefined : lowerSystem(breakpoints, request.system)
const messages = yield* lowerMessages(request, breakpoints)

View file

@ -8,6 +8,7 @@ import {
LLMEvent,
Usage,
type FinishReason,
type JsonSchema,
type LLMRequest,
type MediaPart,
type ProviderMetadata,
@ -19,6 +20,7 @@ import {
import { JsonObject, optionalArray, ProviderShared } from "./shared"
import { GeminiToolSchema } from "./utils/gemini-tool-schema"
import { Lifecycle } from "./utils/lifecycle"
import { ToolSchemaProjection } from "./utils/tool-schema"
const ADAPTER = "gemini"
const MEDIA_MIMES = new Set<string>(ProviderShared.MEDIA_MIMES)
@ -166,10 +168,10 @@ interface ParserState {
// =============================================================================
// Request Lowering
// =============================================================================
const lowerTool = (tool: ToolDefinition) => ({
const lowerTool = (tool: ToolDefinition, inputSchema: JsonSchema) => ({
name: tool.name,
description: tool.description,
parameters: GeminiToolSchema.convert(tool.inputSchema),
parameters: GeminiToolSchema.convert(inputSchema),
})
const lowerToolConfig = (toolChoice: NonNullable<LLMRequest["toolChoice"]>) =>
@ -300,6 +302,7 @@ const thinkingConfig = (request: LLMRequest) => {
const fromRequest = Effect.fn("Gemini.fromRequest")(function* (request: LLMRequest) {
const toolsEnabled = request.tools.length > 0 && request.toolChoice?.type !== "none"
const generation = request.generation
const toolSchemaCompatibility = request.model.compatibility?.toolSchema
const generationConfig = {
maxOutputTokens: generation?.maxTokens,
temperature: generation?.temperature,
@ -313,7 +316,15 @@ const fromRequest = Effect.fn("Gemini.fromRequest")(function* (request: LLMReque
contents: yield* lowerMessages(request),
systemInstruction:
request.system.length === 0 ? undefined : { parts: [{ text: ProviderShared.joinText(request.system) }] },
tools: toolsEnabled ? [{ functionDeclarations: request.tools.map(lowerTool) }] : undefined,
tools: toolsEnabled
? [
{
functionDeclarations: request.tools.map((tool) =>
lowerTool(tool, ToolSchemaProjection.modelCompatibility(tool.inputSchema, toolSchemaCompatibility)),
),
},
]
: undefined,
toolConfig: toolsEnabled && request.toolChoice ? yield* lowerToolConfig(request.toolChoice) : undefined,
generationConfig: Object.values(generationConfig).some((value) => value !== undefined)
? generationConfig

View file

@ -8,6 +8,7 @@ import {
LLMEvent,
Usage,
type FinishReason,
type JsonSchema,
type LLMRequest,
type MediaPart,
type ReasoningPart,
@ -19,6 +20,7 @@ import {
import { isRecord, JsonObject, optionalArray, optionalNull, ProviderShared } from "./shared"
import { OpenAIOptions } from "./utils/openai-options"
import { Lifecycle } from "./utils/lifecycle"
import { ToolSchemaProjection } from "./utils/tool-schema"
import { ToolStream } from "./utils/tool-stream"
const ADAPTER = "openai-chat"
@ -174,12 +176,12 @@ const invalid = ProviderShared.invalidRequest
// Lowering is the only place that knows how common LLM messages map onto the
// OpenAI Chat wire format. Keep provider quirks here instead of leaking native
// fields into `LLMRequest`.
const lowerTool = (tool: ToolDefinition): OpenAIChatTool => ({
const lowerTool = (tool: ToolDefinition, inputSchema: JsonSchema): OpenAIChatTool => ({
type: "function",
function: {
name: tool.name,
description: tool.description,
parameters: ProviderShared.openAiToolInputSchema(tool.inputSchema),
parameters: ToolSchemaProjection.openAI(inputSchema),
},
})
@ -343,10 +345,16 @@ const fromRequest = Effect.fn("OpenAIChat.fromRequest")(function* (request: LLMR
// `fromRequest` returns the provider body only. Endpoint, auth, framing,
// validation, and HTTP execution are composed by `Route.make`.
const generation = request.generation
const toolSchemaCompatibility = request.model.compatibility?.toolSchema
return {
model: request.model.id,
messages: yield* lowerMessages(request),
tools: request.tools.length === 0 ? undefined : request.tools.map(lowerTool),
tools:
request.tools.length === 0
? undefined
: request.tools.map((tool) =>
lowerTool(tool, ToolSchemaProjection.modelCompatibility(tool.inputSchema, toolSchemaCompatibility)),
),
tool_choice: request.toolChoice ? yield* lowerToolChoice(request.toolChoice) : undefined,
stream: true as const,
stream_options: { include_usage: true },

View file

@ -8,6 +8,7 @@ import {
LLMEvent,
Usage,
type FinishReason,
type JsonSchema,
type LLMRequest,
type ProviderMetadata,
type ReasoningPart,
@ -21,6 +22,7 @@ import { JsonObject, optionalArray, optionalNull, ProviderShared } from "./share
import { isContextOverflow } from "../provider-error"
import { OpenAIOptions } from "./utils/openai-options"
import { Lifecycle } from "./utils/lifecycle"
import { ToolSchemaProjection } from "./utils/tool-schema"
import { ToolStream } from "./utils/tool-stream"
const ADAPTER = "openai-responses"
@ -254,11 +256,11 @@ const invalid = ProviderShared.invalidRequest
// =============================================================================
// Request Lowering
// =============================================================================
const lowerTool = (tool: ToolDefinition): OpenAIResponsesTool => ({
const lowerTool = (tool: ToolDefinition, inputSchema: JsonSchema): OpenAIResponsesTool => ({
type: "function",
name: tool.name,
description: tool.description,
parameters: ProviderShared.openAiToolInputSchema(tool.inputSchema),
parameters: ToolSchemaProjection.openAI(inputSchema),
// TODO: Read this from OpenAI-specific tool options so direct LLM callers can opt into strict schemas.
strict: false,
})
@ -476,10 +478,16 @@ const lowerOptions = Effect.fn("OpenAIResponses.lowerOptions")(function* (reques
const fromRequest = Effect.fn("OpenAIResponses.fromRequest")(function* (request: LLMRequest) {
const generation = request.generation
const options = yield* lowerOptions(request)
const toolSchemaCompatibility = request.model.compatibility?.toolSchema
return {
model: request.model.id,
input: yield* lowerMessages(request),
tools: request.tools.length === 0 ? undefined : request.tools.map(lowerTool),
tools:
request.tools.length === 0
? undefined
: request.tools.map((tool) =>
lowerTool(tool, ToolSchemaProjection.modelCompatibility(tool.inputSchema, toolSchemaCompatibility)),
),
tool_choice: request.toolChoice ? yield* lowerToolChoice(request.toolChoice) : undefined,
stream: true as const,
max_output_tokens: generation?.maxTokens,

View file

@ -1,5 +1,5 @@
import { Buffer } from "node:buffer"
import { Effect, JsonSchema, Schema, Stream } from "effect"
import { Effect, Schema, Stream } from "effect"
import * as Sse from "effect/unstable/encoding/Sse"
import { Headers, HttpClientRequest } from "effect/unstable/http"
import {
@ -24,39 +24,6 @@ export const JsonObject = Schema.Record(Schema.String, Schema.Unknown)
export const optionalArray = <const S extends Schema.Top>(schema: S) => Schema.optional(Schema.Array(schema))
export const optionalNull = <const S extends Schema.Top>(schema: S) => Schema.optional(Schema.NullOr(schema))
/** OpenAI function schemas require one flat object at the top level. */
export const openAiToolInputSchema = (schema: JsonSchema.JsonSchema): JsonSchema.JsonSchema => {
const variants = Array.isArray(schema.anyOf) ? schema.anyOf.filter(isRecord) : []
const flattened =
variants.length === 0
? { ...schema, type: "object" }
: {
...Object.fromEntries(Object.entries(schema).filter(([key]) => key !== "anyOf")),
type: "object",
properties: variants.reduce(
(properties, variant) => ({ ...(isRecord(variant.properties) ? variant.properties : {}), ...properties }),
{},
),
additionalProperties: false,
}
const normalized = removeNullSchemas(flattened)
return isRecord(normalized) ? normalized : { type: "object" }
}
const removeNullSchemas = (value: unknown): unknown => {
if (Array.isArray(value)) return value.map(removeNullSchemas)
if (!isRecord(value)) return value
const fields = Object.fromEntries(
Object.entries(value)
.filter(([key]) => key !== "anyOf")
.map(([key, field]) => [key, removeNullSchemas(field)]),
)
if (!Array.isArray(value.anyOf)) return fields
const variants = value.anyOf.filter((variant) => !isRecord(variant) || variant.type !== "null").map(removeNullSchemas)
if (variants.length === 1 && isRecord(variants[0])) return { ...fields, ...variants[0] }
return { ...fields, anyOf: variants }
}
/**
* Streaming tool-call accumulator. Adapters that build a tool call across
* multiple `tool-input-delta` chunks store the partial JSON input string here

View file

@ -1,4 +1,4 @@
import { ProviderShared } from "../shared"
import { isRecord } from "../../utils/record"
// Gemini accepts a JSON Schema-like dialect for tool parameters, but rejects a
// handful of common JSON Schema shapes. Keep this projection isolated so the
@ -20,8 +20,6 @@ const SCHEMA_INTENT_KEYS = [
"else",
]
const isRecord = ProviderShared.isRecord
const hasCombiner = (schema: unknown) =>
isRecord(schema) && (Array.isArray(schema.anyOf) || Array.isArray(schema.oneOf) || Array.isArray(schema.allOf))

View file

@ -0,0 +1,86 @@
import type { JsonSchema, ModelToolSchemaCompatibility } from "../../schema"
import { isRecord } from "../../utils/record"
import { GeminiToolSchema } from "./gemini-tool-schema"
const removeNullSchemas = (value: unknown): unknown => {
if (Array.isArray(value)) return value.map(removeNullSchemas)
if (!isRecord(value)) return value
const fields = Object.fromEntries(
Object.entries(value)
.filter(([key]) => key !== "anyOf")
.map(([key, field]) => [key, removeNullSchemas(field)]),
)
if (!Array.isArray(value.anyOf)) return fields
const variants = value.anyOf.filter((variant) => !isRecord(variant) || variant.type !== "null").map(removeNullSchemas)
if (variants.length === 1 && isRecord(variants[0])) return { ...fields, ...variants[0] }
return { ...fields, anyOf: variants }
}
const tupleItemsSchema = (items: ReadonlyArray<unknown>) => {
const projected = items.map(moonshotNode)
if (projected.length === 0) return {}
if (projected.length === 1) return projected[0]
return { anyOf: projected }
}
const moonshotNode = (schema: unknown): unknown => {
if (Array.isArray(schema)) return schema.map(moonshotNode)
if (!isRecord(schema)) return schema
if (typeof schema.$ref === "string") return { $ref: schema.$ref }
return Object.fromEntries(
Object.entries(schema).flatMap(([key, value]) => {
if (key === "items" && Array.isArray(value)) return [[key, tupleItemsSchema(value)]]
if (key === "prefixItems") {
if ("items" in schema) return []
return [["items", tupleItemsSchema(Array.isArray(value) ? value : [])]]
}
if (key === "unevaluatedItems") return []
return [[key, moonshotNode(value)]]
}),
)
}
const moonshot = (schema: JsonSchema): JsonSchema => {
const projected = moonshotNode(schema)
return isRecord(projected) ? projected : {}
}
const openAI = (schema: JsonSchema): JsonSchema => {
const variants = Array.isArray(schema.anyOf) ? schema.anyOf.filter(isRecord) : []
const flattened =
variants.length === 0
? { ...schema, type: "object" }
: {
...Object.fromEntries(Object.entries(schema).filter(([key]) => key !== "anyOf")),
type: "object",
properties: variants.reduce(
(properties, variant) => ({ ...(isRecord(variant.properties) ? variant.properties : {}), ...properties }),
{},
),
additionalProperties: false,
}
const normalized = removeNullSchemas(flattened)
return isRecord(normalized) ? normalized : { type: "object" }
}
const gemini = (schema: JsonSchema): JsonSchema => GeminiToolSchema.convert(schema) ?? {}
const modelCompatibility = (
schema: JsonSchema,
compatibility: ModelToolSchemaCompatibility | undefined,
): JsonSchema => {
if (compatibility === undefined) return schema
switch (compatibility) {
case "gemini":
return gemini(schema)
case "moonshot":
return moonshot(schema)
}
}
export const ToolSchemaProjection = {
gemini,
modelCompatibility,
moonshot,
openAI,
} as const

View file

@ -0,0 +1,117 @@
import { describe, expect, test } from "bun:test"
import { Effect } from "effect"
import { LLM } from "../src"
import { OpenAIChat } from "../src/protocols"
import { ToolSchemaProjection } from "../src/protocols/utils/tool-schema"
import { Auth, LLMClient } from "../src/route"
import { it } from "./lib/effect"
describe("tool schema projections", () => {
test("moonshot strips $ref siblings and converts tuple arrays to a schema object", () => {
expect(
ToolSchemaProjection.moonshot({
type: "object",
properties: {
linked: { $ref: "#/$defs/Linked", description: "drop me" },
tuple: { type: "array", items: [{ type: "string" }, { type: "number" }] },
prefixTuple: { type: "array", prefixItems: [{ type: "boolean" }, { type: "string" }] },
},
}),
).toEqual({
type: "object",
properties: {
linked: { $ref: "#/$defs/Linked" },
tuple: { type: "array", items: { anyOf: [{ type: "string" }, { type: "number" }] } },
prefixTuple: { type: "array", items: { anyOf: [{ type: "boolean" }, { type: "string" }] } },
},
})
})
test("gemini handles numeric enums, dangling required fields, untyped arrays, and scalar object keys", () => {
expect(
ToolSchemaProjection.gemini({
type: "object",
required: ["status", "missing"],
properties: {
status: { type: "integer", enum: [1, 2] },
tags: { type: "array" },
name: { type: "string", properties: { ignored: { type: "string" } }, required: ["ignored"] },
},
}),
).toEqual({
type: "object",
required: ["status"],
properties: {
status: { type: "string", enum: ["1", "2"] },
tags: { type: "array", items: { type: "string" } },
name: { type: "string" },
},
})
})
test("openai keeps one flat object top-level schema", () => {
expect(
ToolSchemaProjection.openAI({
anyOf: [
{
type: "object",
properties: {
path: { type: "string" },
maybe: { anyOf: [{ type: "string" }, { type: "null" }] },
},
},
{ type: "object", properties: { resource: { type: "string" } } },
],
}),
).toEqual({
type: "object",
properties: {
path: { type: "string" },
maybe: { type: "string" },
resource: { type: "string" },
},
additionalProperties: false,
})
})
it.effect("applies model compatibility before protocol projection", () =>
Effect.gen(function* () {
const model = OpenAIChat.route
.with({ endpoint: { baseURL: "https://api.openai.test/v1/" }, auth: Auth.bearer("test") })
.model({ id: "kimi-k2", compatibility: { toolSchema: "moonshot" } })
const prepared = yield* LLMClient.prepare<OpenAIChat.OpenAIChatBody>(
LLM.request({
model,
prompt: "Use the tool.",
tools: [
{
name: "lookup",
description: "Lookup data.",
inputSchema: {
type: "object",
anyOf: [
{
type: "object",
properties: {
tuple: { type: "array", items: [{ type: "string" }, { type: "number" }] },
linked: { $ref: "#/$defs/Linked", description: "drop me" },
},
},
],
},
},
],
}),
)
expect(prepared.body.tools?.[0]?.function.parameters).toEqual({
type: "object",
properties: {
tuple: { type: "array", items: { anyOf: [{ type: "string" }, { type: "number" }] } },
linked: { $ref: "#/$defs/Linked" },
},
additionalProperties: false,
})
}),
)
})