JSON to BAML TypeBuilder Parser

parser.go

package parser

import (
	"fmt"
	"strings"

	"example.com/baml_client/type_builder"
	baml "github.com/boundaryml/baml/engine/language_client_go/pkg"
)

type SchemaAdder struct {
	tb       *type_builder.TypeBuilder
	schema   map[string]interface{}
	refCache map[string]baml.Type
}

func NewSchemaAdder(tb *type_builder.TypeBuilder, schema map[string]interface{}) *SchemaAdder {
	return &SchemaAdder{
		tb:       tb,
		schema:   schema,
		refCache: make(map[string]baml.Type),
	}
}

func (s *SchemaAdder) parseObject(jsonSchema map[string]interface{}) (baml.Type, error) {
	schemaType, ok := jsonSchema["type"].(string)
	if !ok || schemaType != "object" {
		return nil, fmt.Errorf("expected object type, got %v", schemaType)
	}

	title, ok := jsonSchema["title"].(string)
	if !ok || title == "" {
		return nil, fmt.Errorf("title is required in JSON schema for object type")
	}

	requiredFields := []string{}
	if required, ok := jsonSchema["required"].([]interface{}); ok {
		for _, field := range required {
			if fieldStr, ok := field.(string); ok {
				requiredFields = append(requiredFields, fieldStr)
			}
		}
	}

	newClass, err := s.tb.AddClass(title)
	if err != nil {
		return nil, err
	}

	if properties, ok := jsonSchema["properties"].(map[string]interface{}); ok {
		for fieldName, fieldSchema := range properties {
			fieldSchemaMap, ok := fieldSchema.(map[string]interface{})
			if !ok {
				continue
			}

			var defaultValue interface{}
			if defVal, exists := fieldSchemaMap["default"]; exists {
				defaultValue = defVal
			}

			fieldType, err := s.Parse(fieldSchemaMap)
			if err != nil {
				return nil, err
			}

			// Check if field is optional
			isRequired := false
			for _, req := range requiredFields {
				if req == fieldName {
					isRequired = true
					break
				}
			}

			if !isRequired && defaultValue == nil {
				fieldType, err = s.tb.Optional(fieldType)
				if err != nil {
					return nil, err
				}
			}

			property, err := newClass.AddProperty(fieldName, fieldType)
			if err != nil {
				return nil, err
			}

			// Handle description
			if description, ok := fieldSchemaMap["description"].(string); ok {
				if defaultValue != nil {
					description = strings.TrimSpace(description) + "\n" + fmt.Sprintf("Default: %v", defaultValue)
					description = strings.TrimSpace(description)
				}
				if len(description) > 0 {
					if err := property.SetDescription(description); err != nil {
						return nil, err
					}
				}
			}
		}
	}

	return newClass.Type()
}

func (s *SchemaAdder) parseString(jsonSchema map[string]interface{}) (baml.Type, error) {
	schemaType, ok := jsonSchema["type"].(string)
	if !ok || schemaType != "string" {
		return nil, fmt.Errorf("expected string type, got %v", schemaType)
	}

	title, _ := jsonSchema["title"].(string)

	if enumValues, ok := jsonSchema["enum"].([]interface{}); ok {
		if title == "" {
			// Treat as a union of literals
			types := []baml.Type{}
			for _, value := range enumValues {
				if strValue, ok := value.(string); ok {
					literalType, err := s.tb.LiteralString(strValue)
					if err != nil {
						return nil, err
					}
					types = append(types, literalType)
				}
			}
			return s.tb.Union(types)
		}

		// Create an enum
		newEnum, err := s.tb.AddEnum(title)
		if err != nil {
			return nil, err
		}

		for _, value := range enumValues {
			if strValue, ok := value.(string); ok {
				if _, err := newEnum.AddValue(strValue); err != nil {
					return nil, err
				}
			}
		}

		return newEnum.Type()
	}

	return s.tb.String()
}

func (s *SchemaAdder) loadRef(ref string) (baml.Type, error) {
	if !strings.HasPrefix(ref, "#/") {
		return nil, fmt.Errorf("only local references are supported: %s", ref)
	}

	// Check cache first
	if cachedType, exists := s.refCache[ref]; exists {
		return cachedType, nil
	}

	parts := strings.SplitN(ref, "/", 3)
	if len(parts) != 3 {
		return nil, fmt.Errorf("invalid reference format: %s", ref)
	}

	left := parts[1]
	right := parts[2]

	refs, ok := s.schema[left].(map[string]interface{})
	if !ok {
		return nil, fmt.Errorf("reference section %s not found in schema", left)
	}

	refSchema, ok := refs[right].(map[string]interface{})
	if !ok {
		return nil, fmt.Errorf("reference %s not found in schema", ref)
	}

	// Parse and cache the result
	parsedType, err := s.Parse(refSchema)
	if err != nil {
		return nil, err
	}

	s.refCache[ref] = parsedType
	return parsedType, nil
}

func (s *SchemaAdder) Parse(jsonSchema map[string]interface{}) (baml.Type, error) {
	// Handle anyOf
	if anyOf, ok := jsonSchema["anyOf"].([]interface{}); ok {
		types := []baml.Type{}
		for _, subSchema := range anyOf {
			if subSchemaMap, ok := subSchema.(map[string]interface{}); ok {
				fieldType, err := s.Parse(subSchemaMap)
				if err != nil {
					return nil, err
				}
				types = append(types, fieldType)
			}
		}
		return s.tb.Union(types)
	}

	// Handle $ref
	if ref, ok := jsonSchema["$ref"].(string); ok {
		return s.loadRef(ref)
	}

	// Handle type field
	schemaType, ok := jsonSchema["type"].(string)
	if !ok {
		return nil, fmt.Errorf("type is required in JSON schema: %v", jsonSchema)
	}

	switch schemaType {
	case "string":
		return s.parseString(jsonSchema)
	case "number":
		return s.tb.Float()
	case "integer":
		return s.tb.Int()
	case "object":
		return s.parseObject(jsonSchema)
	case "array":
		items, ok := jsonSchema["items"].(map[string]interface{})
		if !ok {
			return nil, fmt.Errorf("array type requires items field")
		}
		itemType, err := s.Parse(items)
		if err != nil {
			return nil, err
		}
		return s.tb.List(itemType)
	case "boolean":
		return s.tb.Bool()
	case "null":
		return s.tb.Null()
	default:
		return nil, fmt.Errorf("unsupported type: %s", schemaType)
	}
}

func ParseJSONSchema(jsonSchema map[string]interface{}, tb *type_builder.TypeBuilder) (baml.Type, error) {
	parser := NewSchemaAdder(tb, jsonSchema)
	return parser.Parse(jsonSchema)
}