Explorar o código

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

Shoubhit Dash hai 2 meses
pai
achega
18466b8020

+ 8 - 3
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

+ 12 - 5
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<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)

+ 14 - 3
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<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

+ 11 - 3
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 },

+ 11 - 3
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,

+ 1 - 34
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 = <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

+ 1 - 3
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))
 

+ 86 - 0
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<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

+ 117 - 0
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<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,
+      })
+    }),
+  )
+})