| import type { Tool } from "../types.ts"; |
|
|
| interface JsonSchemaObject { |
| [key: string]: unknown; |
| type?: unknown; |
| properties?: Record<string, JsonSchemaObject | undefined>; |
| required?: unknown; |
| } |
|
|
| class UnsupportedStrictJsonSchemaError extends Error {} |
|
|
| const UNSUPPORTED_STRICT_SCHEMA_KEYS = [ |
| "$ref", |
| "$defs", |
| "definitions", |
| "allOf", |
| "oneOf", |
| "patternProperties", |
| "dependentSchemas", |
| "dependencies", |
| "unevaluatedProperties", |
| "propertyNames", |
| "contains", |
| "prefixItems", |
| "not", |
| "if", |
| "then", |
| "else", |
| ] as const; |
|
|
| function isJsonSchemaObject(value: unknown): value is JsonSchemaObject { |
| return typeof value === "object" && value !== null && !Array.isArray(value); |
| } |
|
|
| function isStructuredSchema(schema: unknown): boolean { |
| if (!isJsonSchemaObject(schema)) return false; |
| const types = typeof schema.type === "string" ? [schema.type] : Array.isArray(schema.type) ? schema.type : []; |
| return ( |
| types.includes("object") || |
| types.includes("array") || |
| schema.properties !== undefined || |
| schema.items !== undefined |
| ); |
| } |
|
|
| function schemaAllowsNull(schema: unknown): boolean { |
| if (!isJsonSchemaObject(schema)) return false; |
| if (schema.type === "null" || (Array.isArray(schema.type) && schema.type.includes("null"))) return true; |
| if (schema.const === null || (Array.isArray(schema.enum) && schema.enum.includes(null))) return true; |
| return Array.isArray(schema.anyOf) && schema.anyOf.some((variant) => schemaAllowsNull(variant)); |
| } |
|
|
| function makeJsonSchemaNodeStrict(schema: unknown): void { |
| if (!isJsonSchemaObject(schema)) { |
| throw new UnsupportedStrictJsonSchemaError("boolean schemas are unsupported"); |
| } |
| for (const key of UNSUPPORTED_STRICT_SCHEMA_KEYS) { |
| if (schema[key] !== undefined) { |
| throw new UnsupportedStrictJsonSchemaError(`${key} schemas are unsupported`); |
| } |
| } |
|
|
| if (schema.anyOf !== undefined) { |
| if (!Array.isArray(schema.anyOf) || schema.anyOf.length === 0) { |
| throw new UnsupportedStrictJsonSchemaError("anyOf must contain at least one schema"); |
| } |
| for (const variant of schema.anyOf) { |
| if (isStructuredSchema(variant)) { |
| throw new UnsupportedStrictJsonSchemaError("object and array unions are unsupported"); |
| } |
| makeJsonSchemaNodeStrict(variant); |
| } |
| } |
|
|
| if (schema.items !== undefined) { |
| if (Array.isArray(schema.items)) { |
| throw new UnsupportedStrictJsonSchemaError("tuple schemas are unsupported"); |
| } |
| makeJsonSchemaNodeStrict(schema.items); |
| } |
|
|
| const isObjectSchema = schema.type === "object"; |
| if (schema.properties !== undefined && !isObjectSchema) { |
| throw new UnsupportedStrictJsonSchemaError("properties require type object"); |
| } |
| if (!isObjectSchema) return; |
| if (schema.additionalProperties !== undefined && schema.additionalProperties !== false) { |
| throw new UnsupportedStrictJsonSchemaError("schema-valued or true additionalProperties is unsupported"); |
| } |
| if (schema.properties !== undefined && !isJsonSchemaObject(schema.properties)) { |
| throw new UnsupportedStrictJsonSchemaError("object properties must be a schema map"); |
| } |
| if ( |
| schema.required !== undefined && |
| (!Array.isArray(schema.required) || schema.required.some((key) => typeof key !== "string")) |
| ) { |
| throw new UnsupportedStrictJsonSchemaError("object required must be a string array"); |
| } |
|
|
| const properties = schema.properties ?? {}; |
| const propertyNames = Object.keys(properties); |
| const required = new Set(Array.isArray(schema.required) ? schema.required : []); |
| if ([...required].some((key) => !propertyNames.includes(key))) { |
| throw new UnsupportedStrictJsonSchemaError("required contains an unknown property"); |
| } |
| for (const [key, property] of Object.entries(properties)) { |
| makeJsonSchemaNodeStrict(property); |
| if (!required.has(key) && !schemaAllowsNull(property)) { |
| properties[key] = { anyOf: [property, { type: "null" }] }; |
| } |
| } |
| schema.required = propertyNames; |
| schema.additionalProperties = false; |
| } |
|
|
| |
| export function makeStrictJsonSchema(schema: Tool["parameters"]): Record<string, unknown> { |
| const cloned: unknown = structuredClone(schema); |
| if (!isJsonSchemaObject(cloned)) { |
| throw new UnsupportedStrictJsonSchemaError("root schema must have type object"); |
| } |
| makeJsonSchemaNodeStrict(cloned); |
| if (cloned.type !== "object") { |
| throw new UnsupportedStrictJsonSchemaError("root schema must have type object"); |
| } |
| return cloned; |
| } |
|
|
| export function getJsonSchemaToolParameters(tool: Tool, strict: boolean | undefined): Tool["parameters"] { |
| return (strict === true ? makeStrictJsonSchema(tool.parameters) : tool.parameters) as Tool["parameters"]; |
| } |
|
|
| export interface GrammarConstrainedSampling { |
| format: "lark" | "regex"; |
| definition: string; |
| inputProperty: string; |
| } |
|
|
| export interface GrammarToolInputJsonBuffer { |
| input: string; |
| started: boolean; |
| closed: boolean; |
| } |
|
|
| export function getGrammarToolInput( |
| toolName: string, |
| arguments_: Record<string, unknown>, |
| inputProperty: string, |
| ): string { |
| const input = arguments_[inputProperty]; |
| if (typeof input !== "string") { |
| throw new Error(`Grammar tool call "${toolName}" requires argument "${inputProperty}" to be a string.`); |
| } |
| return input; |
| } |
|
|
| export function appendGrammarToolInputJsonDelta( |
| buffer: GrammarToolInputJsonBuffer, |
| inputProperty: string, |
| nextInput: string, |
| close: boolean, |
| ): string | undefined { |
| if (buffer.closed) { |
| if (close && nextInput === buffer.input) return undefined; |
| throw new Error(`grammar tool input for property "${inputProperty}" changed after it was closed`); |
| } |
| if (!nextInput.startsWith(buffer.input)) { |
| throw new Error(`grammar tool input for property "${inputProperty}" changed non-monotonically`); |
| } |
|
|
| const inputDelta = nextInput.slice(buffer.input.length); |
| if (!close && inputDelta.length === 0) return undefined; |
|
|
| let delta = ""; |
| if (!buffer.started) { |
| delta += `{${JSON.stringify(inputProperty)}:"`; |
| buffer.started = true; |
| } |
| delta += JSON.stringify(inputDelta).slice(1, -1); |
| buffer.input = nextInput; |
|
|
| if (close) { |
| delta += '"}'; |
| buffer.closed = true; |
| } |
| return delta; |
| } |
|
|
| function inferGrammarInputProperty(tool: Tool): string { |
| const schema = tool.parameters as JsonSchemaObject; |
| if (schema.type !== "object") { |
| throw new Error("grammar constrained sampling requires an object parameter schema"); |
| } |
| if (!Array.isArray(schema.required) || schema.required.length !== 1 || typeof schema.required[0] !== "string") { |
| throw new Error("grammar constrained sampling requires exactly one required string property"); |
| } |
|
|
| const inputProperty = schema.required[0]; |
| if (!schema.properties?.[inputProperty]) { |
| throw new Error(`grammar constrained sampling requires a properties entry for ${inputProperty}`); |
| } |
| if (schema.properties[inputProperty]?.type !== "string") { |
| throw new Error(`grammar constrained sampling property ${inputProperty} must have type string`); |
| } |
| return inputProperty; |
| } |
|
|
| export function resolveJsonSchemaStrictSampling(tool: Tool, supportsStrictMode: boolean): boolean | undefined { |
| const config = tool.constrainedSampling; |
| if (!config || config.type !== "json_schema") return undefined; |
|
|
| if (supportsStrictMode) { |
| try { |
| makeStrictJsonSchema(tool.parameters); |
| return true; |
| } catch (error) { |
| if (!(error instanceof UnsupportedStrictJsonSchemaError)) throw error; |
| if (config.strict !== "require") return undefined; |
| throw new Error(`Tool "${tool.name}" requires JSON-schema constrained sampling, but ${error.message}.`); |
| } |
| } |
| if (config.strict === "require") { |
| throw new Error( |
| `Tool "${tool.name}" requires JSON-schema constrained sampling, but strict tools are unsupported.`, |
| ); |
| } |
| return undefined; |
| } |
|
|
| export function resolveGrammarConstrainedSampling( |
| tool: Tool, |
| supportsOpenAIGrammarTools: boolean, |
| ): GrammarConstrainedSampling | undefined { |
| const config = tool.constrainedSampling; |
| if (!config || config.type !== "grammar") { |
| return undefined; |
| } |
|
|
| if (!supportsOpenAIGrammarTools) { |
| return undefined; |
| } |
|
|
| const larkDefinition = config.variants.openai_lark; |
| const regexDefinition = config.variants.openai_regex; |
| const hasLarkDefinition = typeof larkDefinition === "string" && larkDefinition.trim().length > 0; |
| const hasRegexDefinition = typeof regexDefinition === "string" && regexDefinition.trim().length > 0; |
| if (!hasLarkDefinition && !hasRegexDefinition) { |
| throw new Error( |
| `Tool "${tool.name}" cannot use grammar constrained sampling: no supported grammar variant was provided.`, |
| ); |
| } |
|
|
| try { |
| return { |
| format: hasLarkDefinition ? "lark" : "regex", |
| definition: hasLarkDefinition ? larkDefinition : regexDefinition!, |
| inputProperty: inferGrammarInputProperty(tool), |
| }; |
| } catch (error) { |
| const message = error instanceof Error ? error.message : String(error); |
| throw new Error(`Tool "${tool.name}" cannot use grammar constrained sampling: ${message}.`); |
| } |
| } |
|
|
| export function createGrammarToolInputProperties( |
| tools: Tool[] | undefined, |
| supportsOpenAIGrammarTools: boolean, |
| ): ReadonlyMap<string, string> { |
| const properties = new Map<string, string>(); |
| for (const tool of tools ?? []) { |
| const grammar = resolveGrammarConstrainedSampling(tool, supportsOpenAIGrammarTools); |
| if (grammar) { |
| properties.set(tool.name, grammar.inputProperty); |
| } |
| } |
| return properties; |
| } |
|
|
