Refactor discriminated unions to ensure references to original discriminator are preserved

[bruno-f-cruz]

Consider the following example:

from pydantic import RootModel, Field
from aind_data_schema.base import AindModel
from typing import List, Union, Annotated, Literal


class Cat(AindModel):
    pet_type: Literal['cat'] = 'cat'
    name: str
    mean_cat: bool = Field(default=True)


class Dog(AindModel):
    pet_type: Literal['dog'] = 'dog'
    name: str
    good_dog: bool = Field(default=True)


Pet = Annotated[Union[Cat, Dog], Field(discriminator='pet_type')]


class MyFavoritePet(AindModel):
    pet: Pet
    dob: str


class AllPets(AindModel):
    my_favorite_pet: MyFavoritePet
    all_pets: List[Pet] = Field(default_factory=list)
    one_pet: Pet

Pet defines the discriminated union that distinguishes between dog and cat based on the discriminator pet_type. This reference is then passed onto both MyFavoritePet (used in the pet property) and AllPets (used in multiple properties). However, once a json-schema is generated using the AllPets.model_json_schema method, this association is lost since Pydantic generates the discriminators "inline" with the properties. i.e.:

"properties": {
  "my_favorite_pet": {
   "$ref": "#/$defs/MyFavoritePet"
  },
  "all_pets": {
   "items": {
    "discriminator": {
     "mapping": {
      "cat": "#/$defs/Cat",
      "dog": "#/$defs/Dog"
     },
     "propertyName": "pet_type"
    },
    "oneOf": [
     {
      "$ref": "#/$defs/Cat"
     },
     {
      "$ref": "#/$defs/Dog"
     }
    ]
   },
   "title": "All Pets",
   "type": "array"
  },
  "one_pet": {
   "discriminator": {
    "mapping": {
     "cat": "#/$defs/Cat",
     "dog": "#/$defs/Dog"
    },
    "propertyName": "pet_type"
   },
   "oneOf": [
    {
     "$ref": "#/$defs/Cat"
    },
    {
     "$ref": "#/$defs/Dog"
    }
   ],
   "title": "One Pet"
  }
 },
 "required": [
  "my_favorite_pet",
  "one_pet"
 ],
 "title": "AllPets",
 "type": "object"

This has a couple of problems:

1. The schema is unnecessarily large, especially for unions that are extensively used (as is the case with several from the aind-data-schemas)
2. Makes the job of tools that want to leverage the structure of the json incredibly hard since the relationship between classes must now be inferred instead of explicit. One would need to go through the full schema, find identical unions, and resolve them as the same type.

While both are valid under json-spec, if Pydantic were to construct the model using references to Pet instead of defining it "inline" every single time both of these would be solved.

Luckily, a simple change to the definition of the Pet discriminator should fix this issue.

class Pet(RootModel):
    root: Annotated[Union[Cat, Dog], Field(discriminator='pet_type')]

When generating the json-schema, we now got the correct references:

"properties": {
  "my_favorite_pet": {
   "$ref": "#/$defs/MyFavoritePet"
  },
  "all_pets": {
   "items": {
    "$ref": "#/$defs/Pet"
   },
   "title": "All Pets",
   "type": "array"
  },
  "one_pet": {
   "$ref": "#/$defs/Pet"
  }
 },
 "required": [
  "my_favorite_pet",
  "one_pet"

For interoperability with other tools, I would greatly appreciate the chance to propose this refactor to the current aind-data-schema. We could even try to come up with a way to add some syntactic sugar so users don't have to worry about these details. Happy to start working on a PR if people think it is worth it!

[bruno-f-cruz]

After some experimentation, this seems to be the code we need to wrap around:

Pet = create_model("Pet", root=(Annotated[Union[Cat, Dog], Field(discriminator='pet_type')], ...), __base__=RootModel)

This would make it trivial to go from the current implementation to this one with a single function applied on top of each discriminator.

[bruno-f-cruz]

@dyf

[dyf]

Looking at this now. I think maybe we should make a create_discriminated_union function that takes in the list of types and the discriminator field, and then it would be relatively straightforward to use throughout.

[saskiad]

@bruno-f-cruz is there a specific instance where the same union is being repeated that you are looking to improve?

[bruno-f-cruz]

Hi, thanks for looking into this!

When I first wrote the issue I was far more worried about the point 2. Even if a class is only duplicated once, it makes creating tools that leverage the json-schema harder than it needs to be. (Now that I think about it, even if a class is not duplicated, it makes it much harder to even know that since you cant ever know if the definition is defined in line because there is only one, or because you are not passing by reference)

One of the candidates would likely be this one:

registry: Annotated[Union[ResearchOrganizationRegistry], Field(default=Registry.ROR, discriminator="name")]

Hope it helps!

[bruno-f-cruz]

A pydantic maintainer has recently provided an alternative pattern that seems to address this issue with a more succinct syntax. I might tackle it in the Behavior.Services repository at some point but for now I will just leave it here as reference AllenNeuralDynamics/Aind.Behavior.Services#114

[jtyoung84]

Here is a link to the issue Bruno is referring to: pydantic/pydantic#10635 (comment)

I wonder how it de-serializes json?