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

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))
 

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