Open tomdurrant opened 1 month ago
Something to potentially work from:
from typing import Dict, Type, Literal
from pydantic import BaseModel, create_model, Field
# 1. First, let's create a type registry and a function to register new types:
class TypeRegistry:
_types: Dict[str, Type[BaseModel]] = {}
@classmethod
def register(cls, model_type: str, model_class: Type[BaseModel]):
cls._types[model_type] = model_class
@classmethod
def get_types(cls):
return cls._types
def register_type(model_type: str):
def decorator(cls):
TypeRegistry.register(model_type, cls)
return cls
return decorator
# 2. Now let's define our base types using the registry
@register_type('foo')
class Foo(BaseModel):
model_type: Literal['foo']
@register_type('bar')
class Bar(BaseModel):
model_type: Literal['bar']
# 3. Create a function to dynamically generate the MyModel class:
from typing import Union
def create_my_model():
types = TypeRegistry.get_types()
foobar_type = Union[tuple(types.values())]
return create_model(
'MyModel',
foobar=(foobar_type, Field(..., discriminator='model_type'))
)
# 4. Use the dynamically created model:
MyModel = create_my_model()
kwargs = {"foobar": {"model_type": "foo"}}
print(MyModel(**kwargs))
# 5. If you want to add a new type later, you can simply define it and register it
@register_type('baz')
class Baz(BaseModel):
model_type: Literal['baz']
# Recreate the model to include the new type
MyModel = create_my_model()
kwargs = {"foobar": {"model_type": "baz"}}
print(MyModel(**kwargs))
Explore options for a plug in architecture. Ideally, this would allow users to import any config object, and register it as a option in available to ModelRun
We are currently using literal discriminators to distinguish which options are instantiated.
https://stackoverflow.com/questions/77947027/how-can-i-use-a-classvar-literal-as-a-discriminator-for-pydantic-fields
For this to work, we need to hard code available option to pydantic fields as inputs, limiting a new users ability to add new classes without changing the whole parent classes as well.
pydantic's native plugin architecture is worth a look https://docs.pydantic.dev/latest/concepts/plugins/