Plugin architecture

[tomdurrant]

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/

[rafa-guedes]

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))