MycLib/Src/AST/Myc.Ast.Json.Schema.pas

unit Myc.Ast.Json.Schema;

interface

uses
    system.sysutils,
    system.classes,
    system.generics.collections,
    system.generics.defaults,
    system.rtti,
    system.json,
    Myc.Ast.Nodes,
    Myc.Ast.Attributes,
    Myc.Ast.Types,
    Myc.Ast;

type
    { Generates JSON schema for Structured Output based on AST RTTI. }
    TAstSchema = class
    private
        class function GetJsonType(Kind: TFieldKind): TJSONObject; static;
    public
        { Generates a complete JSON schema for LLM structured output APIs. }
        class function GenerateFullSchema: TJSONObject; static;
    end;

implementation

{ TAstSchema }

class function TAstSchema.GetJsonType(Kind: TFieldKind): TJSONObject;
    {$region 'helpers'}
    function JRef(const APath: string): TJSONObject;
    begin
        Result := TJSONObject.Create;
        Result.AddPair('$ref', APath);
    end;

    function JDefRef(const ATag: string): TJSONObject;
    begin
        Result := JRef('#/$defs/' + ATag);
    end;

    function JType(const ATypeName: string): TJSONObject;
    begin
        Result := TJSONObject.Create;
        Result.AddPair('type', ATypeName);
    end;

    function JConst(const AVal: string): TJSONObject;
    begin
        Result := TJSONObject.Create;
        Result.AddPair('const', AVal);
    end;

    function JArray(AItems: TJSONValue): TJSONObject;
    begin
        Result := JType('array');
        Result.AddPair('items', AItems);
    end;

    // Combined function: Creates array from items and sets prefixItems + limits
    function JFixedArray(const AItems: array of TJSONValue): TJSONObject;
    var
        arr: TJSONArray;
        item: TJSONValue;
        count: Integer;
    begin
        Result := JType('array');

        arr := TJSONArray.Create;
        for item in AItems do
            arr.AddElement(item);

        count := arr.Count;
        Result.AddPair('prefixItems', arr);
        Result.AddPair('minItems', TJSONNumber.Create(count));
        Result.AddPair('maxItems', TJSONNumber.Create(count));
    end;

    function JTuple(AItems: TJSONValue): TJSONObject;
    begin
        Result := JFixedArray([JConst('Tuple'), AItems]);
    end;

    function JAnyOf(AOptions: array of TJSONValue): TJSONObject;
    var
        arr: TJSONArray;
        opt: TJSONValue;
    begin
        arr := TJSONArray.Create;
        for opt in AOptions do
            arr.AddElement(opt);
        Result := TJSONObject.Create;
        Result.AddPair('anyOf', arr);
    end;
{$endregion}

var
    pair, selItems, selTuple, entryContent, entryNode: TJSONObject;
begin
    case Kind of
        fkNode: Result := JDefRef('Node');
        fkTuple: Result := JDefRef('Tuple');
        fkKeyword: Result := JDefRef('Key');
        fkLambda: Result := JDefRef('Fn');
        fkIdentifier: Result := JDefRef('Id');
        fkString: Result := JType('string');
        fkValue: Result := JAnyOf([JType('number'), JType('string'), JType('boolean')]);
        fkNullableNode: Result := JAnyOf([JDefRef('Node'), JType('null')]);
        fkArrayOfNodes: Result := JArray(JDefRef('Node'));

        fkArrayOfPairs:
        begin
            // [[Cond, Branch], ...]
            pair := JFixedArray([JDefRef('Node'), JDefRef('Node')]);
            Result := JArray(pair);
        end;

        fkPipeInputs:
        begin
            // 1. Selector Tuple: ["Tuple", [Key, ...]]
            selTuple := JTuple(JArray(JDefRef('Key')));

            // 2. Entry Node: ["Tuple", [Id, SelectorTuple]]
            entryNode := JTuple(JFixedArray([JDefRef('Id'), selTuple]));

            // 3. Outer Inputs List: ["Tuple", [EntryNode, ...]]
            Result := JTuple(JArray(entryNode));
        end;
    else
        Result := TJSONObject.Create;
    end;
end;

class function TAstSchema.GenerateFullSchema: TJSONObject;
var
    ctx: TRttiContext;
    typ: TRttiType;
    attr: TCustomAttribute;
    defs, nodeDef, astNodeRef: TJSONObject;
    anyOfNodes: TJSONArray;
    tag: string;
    fields: TList<AstFieldAttribute>;
    minCount, maxCount: Integer;
begin
    ctx := TRttiContext.Create;
    Result := TJSONObject.Create;
    defs := TJSONObject.Create;
    anyOfNodes := TJSONArray.Create;

    try
        // 1. Scan all interfaces for AstTag metadata
        for typ in ctx.GetTypes do
        begin
            if typ.TypeKind <> tkInterface then
                continue;

            tag := '';
            fields := TList<AstFieldAttribute>.Create;
            try
                for attr in typ.GetAttributes do
                begin
                    if attr is AstTagAttribute then
                        tag := AstTagAttribute(attr).Tag;
                    if attr is AstFieldAttribute then
                        fields.Add(AstFieldAttribute(attr));
                end;

                if tag = '' then
                    continue;

                // Sort fields by index
                fields.Sort(
                    TComparer<AstFieldAttribute>
                        .Construct(function(const L, R: AstFieldAttribute): Integer begin Result := L.Index - R.Index; end)
                );

                nodeDef := TJSONObject.Create;
                nodeDef.AddPair('type', 'array');

                var items := TJSONArray.Create;
                items.AddElement(TJSONObject.Create.AddPair('const', tag));

                minCount := 1;
                maxCount := 1;

                for var f in fields do
                begin
                    inc(maxCount);
                    // Only non-nullable fields increase minCount
                    if f.Kind <> fkNullableNode then
                        minCount := maxCount;

                    items.AddElement(GetJsonType(f.Kind));
                end;

                nodeDef.AddPair('prefixItems', items);
                nodeDef.AddPair('minItems', TJSONNumber.Create(minCount));
                nodeDef.AddPair('maxItems', TJSONNumber.Create(maxCount));
                // additionalItems removed as maxItems guarantees strictness

                defs.AddPair(tag, nodeDef);

                // Add reference to the global AstNode union
                astNodeRef := TJSONObject.Create;
                astNodeRef.AddPair('$ref', '#/$defs/' + tag);
                anyOfNodes.AddElement(astNodeRef);

            finally
                fields.Free;
            end;
        end;

        // 2. Define the base AstNode anyOf union
        var astNodeBase := TJSONObject.Create;
        astNodeBase.AddPair('anyOf', anyOfNodes);
        defs.AddPair('Node', astNodeBase);

        // 3. Assemble final Root Object
        Result.AddPair('$ref', '#/$defs/Node');
        Result.AddPair('$defs', defs);

    finally
        ctx.Free;
    end;
end;

end.