From 18466b8020aa6775d2d72c23c22241c2fe553cab Mon Sep 17 00:00:00 2001 From: Shoubhit Dash Date: Mon, 29 Jun 2026 21:45:42 +0530 Subject: [PATCH] feat(llm): add tool schema projections (#34454) --- .../llm/src/protocols/anthropic-messages.ts | 11 +- .../llm/src/protocols/bedrock-converse.ts | 17 ++- packages/llm/src/protocols/gemini.ts | 17 ++- packages/llm/src/protocols/openai-chat.ts | 14 ++- .../llm/src/protocols/openai-responses.ts | 14 ++- packages/llm/src/protocols/shared.ts | 35 +----- .../src/protocols/utils/gemini-tool-schema.ts | 4 +- .../llm/src/protocols/utils/tool-schema.ts | 86 +++++++++++++ .../llm/test/tool-schema-projection.test.ts | 117 ++++++++++++++++++ 9 files changed, 261 insertions(+), 54 deletions(-) create mode 100644 packages/llm/src/protocols/utils/tool-schema.ts create mode 100644 packages/llm/test/tool-schema-projection.test.ts diff --git a/packages/llm/src/protocols/anthropic-messages.ts b/packages/llm/src/protocols/anthropic-messages.ts index d51da9851d6..26990b0063b 100644 --- a/packages/llm/src/protocols/anthropic-messages.ts +++ b/packages/llm/src/protocols/anthropic-messages.ts @@ -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 diff --git a/packages/llm/src/protocols/bedrock-converse.ts b/packages/llm/src/protocols/bedrock-converse.ts index 80412ca9ddb..c447a1a39d3 100644 --- a/packages/llm/src/protocols/bedrock-converse.ts +++ b/packages/llm/src/protocols/bedrock-converse.ts @@ -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 // ============================================================================= // 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): BedrockTool[] => { +const lowerTools = ( + compatibility: ModelToolSchemaCompatibility | undefined, + breakpoints: BedrockCache.Breakpoints, + tools: ReadonlyArray, +): 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) diff --git a/packages/llm/src/protocols/gemini.ts b/packages/llm/src/protocols/gemini.ts index c7324608200..c4bb9476a49 100644 --- a/packages/llm/src/protocols/gemini.ts +++ b/packages/llm/src/protocols/gemini.ts @@ -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(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) => @@ -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 diff --git a/packages/llm/src/protocols/openai-chat.ts b/packages/llm/src/protocols/openai-chat.ts index 4e41a34ce4a..9ac85b07b13 100644 --- a/packages/llm/src/protocols/openai-chat.ts +++ b/packages/llm/src/protocols/openai-chat.ts @@ -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 }, diff --git a/packages/llm/src/protocols/openai-responses.ts b/packages/llm/src/protocols/openai-responses.ts index 00100c1a0a6..4936d31c921 100644 --- a/packages/llm/src/protocols/openai-responses.ts +++ b/packages/llm/src/protocols/openai-responses.ts @@ -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, diff --git a/packages/llm/src/protocols/shared.ts b/packages/llm/src/protocols/shared.ts index c5b6003fd28..173dc511bb0 100644 --- a/packages/llm/src/protocols/shared.ts +++ b/packages/llm/src/protocols/shared.ts @@ -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 = (schema: S) => Schema.optional(Schema.Array(schema)) export const optionalNull = (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 diff --git a/packages/llm/src/protocols/utils/gemini-tool-schema.ts b/packages/llm/src/protocols/utils/gemini-tool-schema.ts index 7690b2e6001..efdbe3f6ec6 100644 --- a/packages/llm/src/protocols/utils/gemini-tool-schema.ts +++ b/packages/llm/src/protocols/utils/gemini-tool-schema.ts @@ -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)) diff --git a/packages/llm/src/protocols/utils/tool-schema.ts b/packages/llm/src/protocols/utils/tool-schema.ts new file mode 100644 index 00000000000..3a311eb34ce --- /dev/null +++ b/packages/llm/src/protocols/utils/tool-schema.ts @@ -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) => { + 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 diff --git a/packages/llm/test/tool-schema-projection.test.ts b/packages/llm/test/tool-schema-projection.test.ts new file mode 100644 index 00000000000..a9df815dafa --- /dev/null +++ b/packages/llm/test/tool-schema-projection.test.ts @@ -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( + 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, + }) + }), + ) +})