feat: support JSON-native union inputs

[markphelps]

What

Adds support for JSON-native union input types, e.g.:

def predict(self, value: str | float) -> str: ...
def predict(self, value: str | float | None = Input(default=None)) -> str: ...
def predict(self, nums: list[int] | list[float]) -> str: ...

Unions are restricted to JSON-native members (str, int, float, bool, dict/Any, list[T], None) so request validation happens at the HTTP edge against the OpenAPI schema. Unions involving Path, File, Secret, custom coders, and BaseModel are rejected at build/schema-generation time. Output unions remain unsupported.

How

• pkg/schema — recursive InputType model + resolver; emits OpenAPI anyOf for union inputs. Multi-variant nullable unions stay required when no default is supplied; non-union optionals are unchanged.
• pkg/predict — CLI -i parsing inspects anyOf members: numeric values parse as numbers for unions that accept one (schemaAcceptsNumber), and a failed numeric parse falls back to the string member for number-first unions like float | str (schemaAcceptsString).
• python/cog/_adt — deterministic union normalisation with strict per-member compatibility (bool is not int/float, scalars don't pass through dict/Any, list unions validate each element); anyOf json_type emission.
• docs — architecture/02-schema.md and docs/python.md document supported members, anyOf, and nullable semantics.

Behavioral notes

• No runtime primitive coercion. Input validation happens at the HTTP edge via the OpenAPI schema; the worker does not re-coerce primitives. So -i value=1 for str | float stays an int (consistent with how a plain float field already receives a Python int for -i num=10); only fractional values become floats.
• Nullable means "may be omitted", not "accepts explicit null". nullable: true lets an omitted optional fall back to its default. An explicit JSON null is still validated against the field type and rejected at the HTTP edge — identical to every other optional field today.

Testing

• Go unit tests: pkg/schema/..., pkg/predict/...
• Python regression tests: python/tests/test_adt.py (union ordering, bool/int/float overlap, list-element validation, optional/required None handling, empty list, mixed scalar+list)
• End-to-end txtar integration tests: union_input_http, union_input_cli (both member orderings), union_input_list_http

All passing locally; lint:go clean; ruff clean.

[ask-bonk]

I'm Bonk, and I've done a quick review of your PR.

Adds support for JSON-native union input types with recursive InputType modelling, OpenAPI anyOf emission, deterministic Python normalisation, and schema-directed CLI parsing.

Issues found (2 related bugs in CLI union parsing):

1. int | float fractional values error instead of coercing to float — When the resolved schema type is integer (e.g. int | float where int is first in anyOf), a fractional CLI value like -i value=1.5 hits the integer branch, ParseInt fails, and since schemaAcceptsString is false the function returns an error. It should fall back to the float member.

2. str | int non-integer numeric values incorrectly sent as float — When resolveSchemaType returns string (e.g. str | int), the switch falls through to schemaAcceptsNumber, which tries ParseFloat unconditionally. A value like 1.5 for str | int gets sent as float even though the union doesn't accept floats. It should fall back to string.

Both bugs stem from the schemaAcceptsNumber / integer branch not distinguishing between unions that accept any number vs. unions that only accept integers.

Suggested fixes:

Add a schemaAcceptsFloat helper (mirroring schemaAcceptsString / schemaAcceptsNumber):

func schemaAcceptsFloat(s *openapi3.Schema) bool {
	if s == nil {
		return false
	}
	if s.Type != nil && s.Type.Is("number") {
		return true
	}
	for _, ref := range s.AnyOf {
		if ref.Value != nil && schemaAcceptsFloat(ref.Value) {
			return true
		}
	}
	for _, ref := range s.AllOf {
		if ref.Value != nil && schemaAcceptsFloat(ref.Value) {
			return true
		}
	}
	return false
}

In the integer branch (around line 115), add a float fallback:

case propertySchema.Type.Is("integer"):
	value, err := strconv.ParseInt(val, 10, 32)
	if err != nil {
		// See the number case above: fall back to a string
		// member for unions such as `int | str`.
		if schemaAcceptsString(originalSchema) {
			break
		}
		// Also fall back to a float member for unions such as `int | float`.
		if schemaAcceptsFloat(originalSchema) {
			if value, err := strconv.ParseFloat(val, 32); err == nil {
				float := float32(value)
				input[key] = Input{Float: &float}
				continue
			}
		}
		return input, err
	}
	valueInt := int32(value)
	input[key] = Input{Int: &valueInt}
	continue

In the schemaAcceptsNumber branch (around line 127), gate ParseFloat:

case schemaAcceptsNumber(originalSchema):
	// Union input (anyOf) that includes a numeric member, e.g.
	// `str | float`. Parse numeric-looking values as numbers so
	// the runtime receives the intended type; otherwise fall
	// through to the string member below.
	if value, err := strconv.ParseInt(val, 10, 32); err == nil {
		valueInt := int32(value)
		input[key] = Input{Int: &valueInt}
		continue
	}
	if schemaAcceptsFloat(originalSchema) {
		if value, err := strconv.ParseFloat(val, 32); err == nil {
			float := float32(value)
			input[key] = Input{Float: &float}
			continue
		}
	}

Please also add unit tests for int | float and str | int unions in pkg/predict/input_test.go.

[ask-bonk]

Posted a review comment to PR #3048 identifying 2 related CLI parsing bugs for union inputs:

1. int | float with fractional values errors instead of falling back to the float member
2. str | int with non-integer numeric values incorrectly parses as float instead of falling back to string

Both stem from schemaAcceptsNumber not distinguishing integer-only unions from float-accepting unions. I suggested adding a schemaAcceptsFloat helper and gating ParseFloat behind it, plus adding unit tests for int | float and str | int unions.

github run

[ask-bonk]

@markphelps Bonk workflow was cancelled.

View workflow run · To retry, trigger Bonk again.

[dennnnnnnnnis]

so it's always better to include the default = None if an input filed is optional right? even if we have Optional[A] defined, otherwise the schema might not pick them up as optional?
I don't recall what I did with optional fields previously but I always had default=something, guess it's not a problem with me but just asking

[markphelps]

Good question, and you've actually hit on a real subtlety.

Short answer: Optional[T] / T | None does not need default=None. What it needs is some = Input(...) assignment. The Input(...) is what makes the field omittable; the None default is implicit for optionals.

BUT doing value: Optional[str] (bare, no Input) generates the openAPI schema correctly (not required, nullable: true), but it fails at runtime with no argument with TypeError: missing 1 required positional argument.. so that is actually a preexisting bug that we've been papering over with =Input() which we should fix!

But for now will update the docs so we are consistent everywhere and look into what it will take to actually fix or atleast make it more clear as to why it doesnt work as expected with a better error message or something

[ask-bonk]

LGTM

github run