An experiment on how to inject custom model definitions:

[bruno-f-cruz]

In our current experiment, we define task logic parameters under a custom json schema:

https://github.com/AllenNeuralDynamics/aind-vr-foraging/blob/main/src/DataSchemas/aind-vr-foraging-task.json

The goal of this experiment was to try to find a way to "inject" this schema definition in the current Session class provided by aind-data-schema. The obvious target would be the generic field:

aind-data-schema/src/aind_data_schema/stimulus.py

Line 107 in 8e1eaa3

input_parameters: Dict[str, Any] = Field(

Since the field already accepts a Dict[str, any] this should be compatible with any json schema definition, even without knowledge of that specific schema (i.e. it would just deserialize to the raw Dict type). However, what would be great is to be able to type this Dict to something more constrained. Here's what I propose:

1. Create the custom schema that defines the task logic. This can be done by inheriting from AindModel:

class TaskSchema(AindModel):
    param1: int = Field(..., description="This is a parameter")
    param2: str = Field(..., description="This is another parameter")

or even from raw json schemas using tools like https://docs.pydantic.dev/latest/integrations/datamodel_code_generator/

2. Create a custom BehaviorStimulation class by inheriting from the original base class. This guarantees that the children class will be compatible with the deserialization process of the parent, assuming we are only overloading the aforementioned property. Eg:

class CustomBehaviorStimulation(BehaviorStimulation):
    input_parameters: [TaskSchema] = Field(..., title = 'Task logic parameters', description="This is a custom input parameter")

3. Create your Session object and "inject" your custom object:

session.stimulus_epochs[0].stimulus.input_parameters = TaskSchema(param1=1, param2='hello')

4. Leverage the deserialization behavior:

session.write_standard_file()
this = Session.parse_file("session.json")
TaskLogic = TaskSchema(**this.stimulus_epochs[0].stimulus.input_parameters)

Notice how the Schema model definition never really gets updated, so everything should be backward and forward-compatible so long as the type of this field remains a Dict[str, any].

Alternatively, one could explore the idea of making the input to this property a Union[BaseModel, Dict[str,any]]

My current plan is to generate, on each update of the task logic schema, a new json schema that results from merging the untouched Session object with a new, automatically generated pydantic object (using the datamodel-code-generator tool), that gets injected into the full schema. As a result, for each experiment, I should be able to provide a fully compatible Session object but, at the same time, have a fully spec'ed json schema that I can use to automatically generate code, and validate task logic, in Bonsai or even during analysis.

[hanhou]

Is it possible to do something like this without creating CustomBehaviorStimulation?

task_schema = TaskSchema(param1=1, param2='hello')
behavior_simulation = BehaviorStimulation(input_parameters=task_schema.dict(), ...)

[bruno-f-cruz]

Ignore my previous comment, I misunderstood your question. Yes, you are right, I think you could pass it as a dict. However, while it would work for serialization and deserialization, you would have to do it manually, since your "full model" would not be aware that input_parameters are of type TaskSchema.

We could change input_parameters type to be of Union[Dict[str,any], BaseModel], as I had suggested. If we then have a discriminator (could just be the type of the class I guess...), , it might be possible for the model to know how to deserialize automatically, must test this approach.

[bruno-f-cruz]

After some experimentation I think you might be right. I just did not find a satisfying way of injecting the custom model into the original aind data schema. I guess it will be up to the experimenter to make sure it ships a valid schema (eg a json schema) with each experiment version. This is what I am already currently doing anyway. And it guarantees that from the experimental code you can access a valid schemas since they are tracked in the same repo.

[hanhou]

Thanks! I need to do some tests on my side especially for deserialization you mentioned. For dynamic foraging task, the task parameters could be quite "flat", in the sense that the Dict structure should be enough. But I agree putting a whole task schema into the full model is more solid.