fix: handle $defs in a $ref object

giladgd · giladgd · commit a31c6c5f7931 · 2025-06-09T04:03:57.000+03:00
diff --git a/src/evaluator/LlamaChat/utils/FunctionCallParamsGrammar.ts b/src/evaluator/LlamaChat/utils/FunctionCallParamsGrammar.ts
@@ -66,7 +66,7 @@ export class FunctionCallParamsGrammar<const Functions extends ChatModelFunction
 function getGbnfGrammarForFunctionParams(paramsSchema: GbnfJsonSchema): string {
     const grammarGenerator = new GbnfGrammarGenerator();
     const rootTerminal = getGbnfJsonTerminalForGbnfJsonSchema(paramsSchema, grammarGenerator);
-    const rootGrammar = rootTerminal.getGrammar(grammarGenerator);
+    const rootGrammar = rootTerminal.resolve(grammarGenerator, true);
 
     return grammarGenerator.generateGbnfFile(rootGrammar + ` "${"\\n".repeat(4)}"`);
 }
diff --git a/src/types.ts b/src/types.ts
@@ -193,8 +193,8 @@ export type ChatSessionModelFunctions = {
 
 export type ChatSessionModelFunction<Params extends GbnfJsonSchema | undefined = GbnfJsonSchema | undefined> = {
     readonly description?: string,
-    readonly params?: Readonly<Params>,
-    readonly handler: (params: GbnfJsonSchemaToType<Params>) => any
+    readonly params?: Params,
+    readonly handler: (params: GbnfJsonSchemaToType<NoInfer<Params>>) => any
 };
 
 export function isChatModelResponseFunctionCall(item: ChatModelResponse["response"][number] | undefined): item is ChatModelFunctionCall {
diff --git a/src/utils/gbnfJson/GbnfGrammarGenerator.ts b/src/utils/gbnfJson/GbnfGrammarGenerator.ts
@@ -2,11 +2,12 @@ import {MultiKeyMap} from "lifecycle-utils";
 import {GbnfJsonSchema} from "./types.js";
 
 export class GbnfGrammarGenerator {
-    public rules = new Map<string, string | null>();
+    public rules = new Map<string, string>();
     public ruleContentToRuleName = new Map<string, string>();
     public literalValueRuleNames = new Map<string | number, string>();
-    public defRuleNames = new MultiKeyMap<[string, GbnfJsonSchema], string>();
+    public defRuleNames = new MultiKeyMap<[string, GbnfJsonSchema], string | null>();
     public defScopeDefs = new MultiKeyMap<[string, GbnfJsonSchema], Record<string, GbnfJsonSchema>>();
+    public usedRootRuleName: boolean = false;
     private ruleId: number = 0;
     private valueRuleId: number = 0;
     private defRuleId: number = 0;
@@ -31,17 +32,17 @@ export class GbnfGrammarGenerator {
         return ruleName;
     }
 
-    public generateRuleNameForDef(defName: string, def: GbnfJsonSchema): [created: boolean, ruleName: string] {
+    public generateRuleNameForDef(defName: string, def: GbnfJsonSchema): string {
         const existingRuleName = this.defRuleNames.get([defName, def]);
         if (existingRuleName != null)
-            return [false, existingRuleName];
+            return existingRuleName;
 
         const ruleName = `def${this.defRuleId}`;
         this.defRuleId++;
 
         this.defRuleNames.set([defName, def], ruleName);
 
-        return [true, ruleName];
+        return ruleName;
     }
 
     public registerDefs(scopeDefs: Record<string, GbnfJsonSchema>) {
diff --git a/src/utils/gbnfJson/GbnfTerminal.ts b/src/utils/gbnfJson/GbnfTerminal.ts
@@ -25,7 +25,22 @@ export abstract class GbnfTerminal {
         return this.getGrammar(grammarGenerator);
     }
 
-    public resolve(grammarGenerator: GbnfGrammarGenerator): string {
+    private _getRootRuleName(grammarGenerator: GbnfGrammarGenerator) {
+        if (this._ruleName != null)
+            return this._ruleName;
+
+        const ruleName = grammarGenerator.usedRootRuleName
+            ? this.getRuleName(grammarGenerator)
+            : "root";
+        this._ruleName = ruleName;
+
+        if (ruleName === "root")
+            grammarGenerator.usedRootRuleName = true;
+
+        return ruleName;
+    }
+
+    public resolve(grammarGenerator: GbnfGrammarGenerator, resolveAsRootGrammar: boolean = false): string {
         if (this._ruleName != null)
             return this._ruleName;
 
@@ -37,7 +52,12 @@ export abstract class GbnfTerminal {
             return existingRuleName;
         }
 
-        const ruleName = this.getRuleName(grammarGenerator);
+        const ruleName = resolveAsRootGrammar
+            ? this._getRootRuleName(grammarGenerator)
+            : this.getRuleName(grammarGenerator);
+
+        if (resolveAsRootGrammar)
+            return grammar;
 
         if (grammar === ruleName) {
             this._ruleName = ruleName;
diff --git a/src/utils/gbnfJson/getGbnfGrammarForGbnfJsonSchema.ts b/src/utils/gbnfJson/getGbnfGrammarForGbnfJsonSchema.ts
@@ -14,7 +14,7 @@ export function getGbnfGrammarForGbnfJsonSchema(schema: Readonly<GbnfJsonSchema>
     const grammarGenerator = new GbnfGrammarGenerator();
     const scopeState = new GbnfJsonScopeState({allowNewLines, scopePadSpaces});
     const rootTerminal = getGbnfJsonTerminalForGbnfJsonSchema(schema, grammarGenerator, scopeState);
-    const rootGrammar = rootTerminal.getGrammar(grammarGenerator);
+    const rootGrammar = rootTerminal.resolve(grammarGenerator, true);
 
     return grammarGenerator.generateGbnfFile(rootGrammar + ` "${"\\n".repeat(4)}"` + " [\\n]*");
 }
diff --git a/src/utils/gbnfJson/terminals/GbnfGrammar.ts b/src/utils/gbnfJson/terminals/GbnfGrammar.ts
@@ -21,10 +21,10 @@ export class GbnfGrammar extends GbnfTerminal {
         return this.grammar;
     }
 
-    public override resolve(grammarGenerator: GbnfGrammarGenerator): string {
+    public override resolve(grammarGenerator: GbnfGrammarGenerator, resolveAsRootGrammar: boolean = false): string {
         if (this.resolveToRawGrammar)
             return this.getGrammar();
 
-        return super.resolve(grammarGenerator);
+        return super.resolve(grammarGenerator, resolveAsRootGrammar);
     }
 }
diff --git a/src/utils/gbnfJson/terminals/GbnfNumberValue.ts b/src/utils/gbnfJson/terminals/GbnfNumberValue.ts
@@ -14,12 +14,12 @@ export class GbnfNumberValue extends GbnfTerminal {
         return '"' + JSON.stringify(this.value) + '"';
     }
 
-    public override resolve(grammarGenerator: GbnfGrammarGenerator): string {
+    public override resolve(grammarGenerator: GbnfGrammarGenerator, resolveAsRootGrammar: boolean = false): string {
         const grammar = this.getGrammar();
         if (grammar.length <= grammarGenerator.getProposedLiteralValueRuleNameLength())
             return grammar;
 
-        return super.resolve(grammarGenerator);
+        return super.resolve(grammarGenerator, resolveAsRootGrammar);
     }
 
     protected override generateRuleName(grammarGenerator: GbnfGrammarGenerator): string {
diff --git a/src/utils/gbnfJson/terminals/GbnfOr.ts b/src/utils/gbnfJson/terminals/GbnfOr.ts
@@ -30,7 +30,7 @@ export class GbnfOr extends GbnfTerminal {
         return "( " + mappedValues.join(" | ") + " )";
     }
 
-    public override resolve(grammarGenerator: GbnfGrammarGenerator): string {
+    public override resolve(grammarGenerator: GbnfGrammarGenerator, resolveAsRootGrammar: boolean = false): string {
         const mappedValues = this.values
             .map((v) => v.resolve(grammarGenerator))
             .filter((value) => value !== "" && value !== grammarNoValue);
@@ -40,6 +40,6 @@ export class GbnfOr extends GbnfTerminal {
         else if (mappedValues.length === 1)
             return mappedValues[0]!;
 
-        return super.resolve(grammarGenerator);
+        return super.resolve(grammarGenerator, resolveAsRootGrammar);
     }
 }
diff --git a/src/utils/gbnfJson/terminals/GbnfRef.ts b/src/utils/gbnfJson/terminals/GbnfRef.ts
@@ -7,8 +7,6 @@ export class GbnfRef extends GbnfTerminal {
     public readonly getValueTerminal: () => GbnfTerminal;
     public readonly defName: string;
     public readonly def: GbnfJsonSchema;
-    private _valueTerminal?: GbnfTerminal;
-    private _grammar?: string;
 
     public constructor({
         getValueTerminal,
@@ -26,28 +24,30 @@ export class GbnfRef extends GbnfTerminal {
     }
 
     public override getGrammar(grammarGenerator: GbnfGrammarGenerator): string {
-        this._createRule(grammarGenerator);
-
-        if (this._valueTerminal != null)
-            return this._valueTerminal.getGrammar(grammarGenerator);
-        else if (this._grammar != null)
-            return this._grammar;
-
-        return this.getValueTerminal().getGrammar(grammarGenerator);
+        return this.generateRuleName(grammarGenerator);
     }
 
     protected override generateRuleName(grammarGenerator: GbnfGrammarGenerator): string {
-        return this._createRule(grammarGenerator);
-    }
+        if (!grammarGenerator.defRuleNames.has([this.defName, this.def])) {
+            const alreadyGeneratingGrammarForThisRef = grammarGenerator.defRuleNames.get([this.defName, this.def]) === null;
+            if (alreadyGeneratingGrammarForThisRef)
+                return grammarGenerator.generateRuleNameForDef(this.defName, this.def);
+
+            grammarGenerator.defRuleNames.set([this.defName, this.def], null);
+            const grammar = this.getValueTerminal().resolve(grammarGenerator);
+
+            if (grammarGenerator.rules.has(grammar) && grammarGenerator.defRuleNames.get([this.defName, this.def]) === null) {
+                grammarGenerator.defRuleNames.set([this.defName, this.def], grammar);
+                return grammar;
+            }
+
+            const ruleName = grammarGenerator.generateRuleNameForDef(this.defName, this.def);
+            grammarGenerator.rules.set(ruleName, grammar);
+            grammarGenerator.ruleContentToRuleName.set(grammar, ruleName);
 
-    private _createRule(grammarGenerator: GbnfGrammarGenerator) {
-        const [isNew, ruleName] = grammarGenerator.generateRuleNameForDef(this.defName, this.def);
-        if (!isNew) {
-            this._grammar = ruleName;
             return ruleName;
         }
 
-        this._valueTerminal = this.getValueTerminal();
-        return ruleName;
+        return grammarGenerator.generateRuleNameForDef(this.defName, this.def);
     }
 }
diff --git a/src/utils/gbnfJson/types.ts b/src/utils/gbnfJson/types.ts
@@ -180,7 +180,9 @@ export type GbnfJsonRefSchema<Defs extends GbnfJsonDefList<NoInfer<Defs>> = {}>
      *
      * Only passed to the model when using function calling, and has no effect when using JSON Schema grammar directly.
      */
-    readonly description?: string
+    readonly description?: string,
+
+    readonly $defs?: Defs
 };
 
 
@@ -211,7 +213,7 @@ export type GbnfJsonSchemaToTSType<T, Defs extends GbnfJsonDefList<NoInfer<Defs>
                                     : T extends GbnfJsonArraySchema<Record<any, any>>
                                         ? ArrayTypeToType<T, CombineDefs<NoInfer<Defs>, T["$defs"]>>
                                         : T extends GbnfJsonRefSchema<any>
-                                            ? GbnfJsonRefSchemaToType<T, NoInfer<Defs>>
+                                            ? GbnfJsonRefSchemaToType<T, CombineDefs<NoInfer<Defs>, T["$defs"]>>
                                             : undefined;
 
 type GbnfJsonBasicStringSchemaToType<T extends GbnfJsonBasicStringSchema> =
diff --git a/src/utils/gbnfJson/utils/getGbnfJsonTerminalForGbnfJsonSchema.ts b/src/utils/gbnfJson/utils/getGbnfJsonTerminalForGbnfJsonSchema.ts
@@ -32,6 +32,9 @@ export function getGbnfJsonTerminalForGbnfJsonSchema(
         throw new Error("Maximum nesting scope exceeded. Ensure that your schema does not have circular references or excessive nesting.");
 
     if (isGbnfJsonRefSchema(schema)) {
+        const currentDefs = joinDefs(defs, schema.$defs);
+        grammarGenerator.registerDefs(currentDefs);
+
         const ref = schema?.$ref;
         const referencePrefix = "#/$defs/";
         if (ref == null || !ref.startsWith(referencePrefix)) {
@@ -44,7 +47,7 @@ export function getGbnfJsonTerminalForGbnfJsonSchema(
         }
 
         const defName = ref.slice(referencePrefix.length);
-        const def = defs[defName];
+        const def = currentDefs[defName];
         if (def == null) {
             console.warn(
                 getConsoleLogPrefix(true, false),
diff --git a/src/utils/gbnfJson/utils/validateObjectAgainstGbnfSchema.ts b/src/utils/gbnfJson/utils/validateObjectAgainstGbnfSchema.ts
@@ -84,6 +84,9 @@ function validateRef<T extends GbnfJsonRefSchema<Record<any, any>>>(
     defs: Record<string, GbnfJsonSchema> = {},
     defScopeDefs: DefScopeDefs = new DefScopeDefs()
 ): object is GbnfJsonSchemaToType<T> {
+    const currentDefs = joinDefs(defs, schema.$defs);
+    defScopeDefs.registerDefs(currentDefs);
+
     const ref = schema.$ref;
     const referencePrefix = "#/$defs/";
 
@@ -94,15 +97,15 @@ function validateRef<T extends GbnfJsonRefSchema<Record<any, any>>>(
     }
 
     const defName = ref.slice(referencePrefix.length);
-    const def = defs[defName];
+    const def = currentDefs[defName];
     if (def == null) {
         // if the $ref points to a non-existing def, a warning was already shows when the grammar was generated,
         // so we don't perform validation on the object as it's considered an "any" type
         return true;
     }
 
-    const currentDefs = defScopeDefs.defScopeDefs.get([defName, def]);
-    return validateObjectWithGbnfSchema(object, def, currentDefs ?? {}, defScopeDefs);
+    const scopeDefs = defScopeDefs.defScopeDefs.get([defName, def]);
+    return validateObjectWithGbnfSchema(object, def, scopeDefs ?? {}, defScopeDefs);
 }
 
 function validateArray<T extends GbnfJsonArraySchema>(
diff --git a/src/utils/getTypeScriptTypeStringForGbnfJsonSchema.ts b/src/utils/getTypeScriptTypeStringForGbnfJsonSchema.ts
@@ -17,13 +17,16 @@ function _getTypeScriptTypeStringForGbnfJsonSchema(
     defScopeDefs: DefScopeDefs = new DefScopeDefs()
 ): string {
     if (isGbnfJsonRefSchema(schema)) {
+        const currentDefs = joinDefs(defs, schema.$defs);
+        defScopeDefs.registerDefs(currentDefs);
+
         const ref = schema?.$ref;
         const referencePrefix = "#/$defs/";
         if (ref == null || !ref.startsWith(referencePrefix))
             return "any";
 
         const defName = ref.slice(referencePrefix.length);
-        const def = defs[defName];
+        const def = currentDefs[defName];
         if (def == null)
             return "any";
         else if (printedDefs.has(def)) {
@@ -36,8 +39,8 @@ function _getTypeScriptTypeStringForGbnfJsonSchema(
             ].join("");
         }
 
-        const currentDefs = defScopeDefs.defScopeDefs.get([defName, def]);
-        if (currentDefs == null)
+        const scopeDefs = defScopeDefs.defScopeDefs.get([defName, def]);
+        if (scopeDefs == null)
             return "any";
 
         printedDefs.add(def);
@@ -47,7 +50,7 @@ function _getTypeScriptTypeStringForGbnfJsonSchema(
                 .replaceAll("\n", " ")
                 .replaceAll("*/", "* /"),
             " */ ",
-            _getTypeScriptTypeStringForGbnfJsonSchema(def, printedDefs, currentDefs, defScopeDefs)
+            _getTypeScriptTypeStringForGbnfJsonSchema(def, printedDefs, scopeDefs, defScopeDefs)
         ].join("");
     } else if (isGbnfJsonOneOfSchema(schema)) {
         const currentDefs = joinDefs(defs, schema.$defs);

Original file line number	Diff line number	Diff line change
`@@ -66,7 +66,7 @@ export class FunctionCallParamsGrammar<const Functions extends ChatModelFunction`
`66`	`66`	`function getGbnfGrammarForFunctionParams(paramsSchema: GbnfJsonSchema): string {`
`67`	`67`	`const grammarGenerator = new GbnfGrammarGenerator();`
`68`	`68`	`const rootTerminal = getGbnfJsonTerminalForGbnfJsonSchema(paramsSchema, grammarGenerator);`
`69`		`- const rootGrammar = rootTerminal.getGrammar(grammarGenerator);`
	`69`	`+ const rootGrammar = rootTerminal.resolve(grammarGenerator, true);`
`70`	`70`
`71`	`71`	return grammarGenerator.generateGbnfFile(rootGrammar + ` "${"\\n".repeat(4)}"`);
`72`	`72`	`}`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ export function getGbnfGrammarForGbnfJsonSchema(schema: Readonly<GbnfJsonSchema>`
`14`	`14`	`const grammarGenerator = new GbnfGrammarGenerator();`
`15`	`15`	`const scopeState = new GbnfJsonScopeState({allowNewLines, scopePadSpaces});`
`16`	`16`	`const rootTerminal = getGbnfJsonTerminalForGbnfJsonSchema(schema, grammarGenerator, scopeState);`
`17`		`- const rootGrammar = rootTerminal.getGrammar(grammarGenerator);`
	`17`	`+ const rootGrammar = rootTerminal.resolve(grammarGenerator, true);`
`18`	`18`
`19`	`19`	return grammarGenerator.generateGbnfFile(rootGrammar + ` "${"\\n".repeat(4)}"` + " [\\n]*");
`20`	`20`	`}`
Original file line number	Diff line number	Diff line change
`@@ -21,10 +21,10 @@ export class GbnfGrammar extends GbnfTerminal {`
`21`	`21`	`return this.grammar;`
`22`	`22`	`}`
`23`	`23`
`24`		`- public override resolve(grammarGenerator: GbnfGrammarGenerator): string {`
	`24`	`+ public override resolve(grammarGenerator: GbnfGrammarGenerator, resolveAsRootGrammar: boolean = false): string {`
`25`	`25`	`if (this.resolveToRawGrammar)`
`26`	`26`	`return this.getGrammar();`
`27`	`27`
`28`		`- return super.resolve(grammarGenerator);`
	`28`	`+ return super.resolve(grammarGenerator, resolveAsRootGrammar);`
`29`	`29`	`}`
`30`	`30`	`}`
Original file line number	Diff line number	Diff line change
`@@ -14,12 +14,12 @@ export class GbnfNumberValue extends GbnfTerminal {`
`14`	`14`	`return '"' + JSON.stringify(this.value) + '"';`
`15`	`15`	`}`
`16`	`16`
`17`		`- public override resolve(grammarGenerator: GbnfGrammarGenerator): string {`
	`17`	`+ public override resolve(grammarGenerator: GbnfGrammarGenerator, resolveAsRootGrammar: boolean = false): string {`
`18`	`18`	`const grammar = this.getGrammar();`
`19`	`19`	`if (grammar.length <= grammarGenerator.getProposedLiteralValueRuleNameLength())`
`20`	`20`	`return grammar;`
`21`	`21`
`22`		`- return super.resolve(grammarGenerator);`
	`22`	`+ return super.resolve(grammarGenerator, resolveAsRootGrammar);`
`23`	`23`	`}`
`24`	`24`
`25`	`25`	`protected override generateRuleName(grammarGenerator: GbnfGrammarGenerator): string {`
Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ export class GbnfOr extends GbnfTerminal {`
`30`	`30`	`return "( " + mappedValues.join(" \| ") + " )";`
`31`	`31`	`}`
`32`	`32`
`33`		`- public override resolve(grammarGenerator: GbnfGrammarGenerator): string {`
	`33`	`+ public override resolve(grammarGenerator: GbnfGrammarGenerator, resolveAsRootGrammar: boolean = false): string {`
`34`	`34`	`const mappedValues = this.values`
`35`	`35`	`.map((v) => v.resolve(grammarGenerator))`
`36`	`36`	`.filter((value) => value !== "" && value !== grammarNoValue);`
`@@ -40,6 +40,6 @@ export class GbnfOr extends GbnfTerminal {`
`40`	`40`	`else if (mappedValues.length === 1)`
`41`	`41`	`return mappedValues[0]!;`
`42`	`42`
`43`		`- return super.resolve(grammarGenerator);`
	`43`	`+ return super.resolve(grammarGenerator, resolveAsRootGrammar);`
`44`	`44`	`}`
`45`	`45`	`}`