feat: Explicitly export IO schema field order (#595)

#### Description of changes

Export IO schemas' field declaration order in JSON. This is needed for
downstream clients to reliably reproduce the order of inputs/outputs as
declared in the schema (e.g. Postgres will scramble JSON field order).

#### Testing done

Local, CI

---------

Co-authored-by: Dion Häfner <dion.haefner@simulation.science>

Author: linusseelinger

tesseract_core/runtime/schema_generation.py

@@ -214,6 +214,19 @@ def _recurse_over_model_tree(treeobj: Any, path: list[str]) -> Any:
     return _recurse_over_model_tree(Schema, [])
 
 
+def _path_tuple_to_str(pathtuple: tuple) -> str:
+    """Convert a path tuple to a dotted string, replacing sentinels with [] / {} indexing."""
+    str_parts = []
+    for part in pathtuple:
+        if part is SEQ_INDEX_SENTINEL:
+            str_parts.append("[]")
+        elif part is DICT_INDEX_SENTINEL:
+            str_parts.append("{}")
+        else:
+            str_parts.append(part)
+    return ".".join(str_parts)
+
+
 def _serialize_diffable_arrays(
     obj: dict[tuple, Any],
 ) -> dict[str, dict[str, Any]]:
@@ -233,24 +246,19 @@ def _serialize_diffable_arrays(
         else:
             shape = tuple(shape)
 
-        # Replace sentinel values with indexing syntax
-        str_parts = []
-        for part in pathtuple:
-            if part is SEQ_INDEX_SENTINEL:
-                str_parts.append("[]")
-            elif part is DICT_INDEX_SENTINEL:
-                str_parts.append("{}")
-            else:
-                str_parts.append(part)
-
-        serialized[".".join(str_parts)] = {
+        serialized[_path_tuple_to_str(pathtuple)] = {
             "shape": shape,
             "dtype": dtype,
         }
 
     return serialized
 
 
+def _serialize_field_order(obj: dict[tuple, list[str]]) -> dict[str, list[str]]:
+    """Convert a dict {path_tuple: [field_names]} to a dict {path_str: [field_names]}."""
+    return {_path_tuple_to_str(p): fields for p, fields in obj.items()}
+
+
 def create_apply_schema(
     InputSchema: type[BaseModel], OutputSchema: type[BaseModel]
 ) -> tuple[type[BaseModel], type[BaseModel]]:
@@ -266,6 +274,11 @@ def create_apply_schema(
         OutputSchema, filter_fn=is_differentiable
     )
 
+    # Snapshot field declaration order for every (nested) model, since JSON object key order is
+    # not guaranteed to be preserved by downstream clients (e.g. postgres).
+    input_field_order = get_field_order(InputSchema)
+    output_field_order = get_field_order(OutputSchema)
+
     InputSchema = apply_function_to_model_tree(
         InputSchema,
         lambda x, _: x,
@@ -287,9 +300,15 @@ class ApplyInputSchema(BaseModel):
         differentiable_arrays: ClassVar[dict[str, dict[str, Any]]] = (
             _serialize_diffable_arrays(diffable_input_paths)
         )
+        field_order: ClassVar[dict[str, list[str]]] = _serialize_field_order(
+            input_field_order
+        )
         model_config = ConfigDict(
             extra="forbid",
-            json_schema_extra={"differentiable_arrays": differentiable_arrays},
+            json_schema_extra={
+                "differentiable_arrays": differentiable_arrays,
+                "field_order": field_order,
+            },
         )
 
     class ApplyOutputSchema(RootModel):
@@ -299,8 +318,14 @@ class ApplyOutputSchema(RootModel):
         differentiable_arrays: ClassVar[dict[str, dict[str, Any]]] = (
             _serialize_diffable_arrays(diffable_output_paths)
         )
+        field_order: ClassVar[dict[str, list[str]]] = _serialize_field_order(
+            output_field_order
+        )
         model_config = ConfigDict(
-            json_schema_extra={"differentiable_arrays": differentiable_arrays}
+            json_schema_extra={
+                "differentiable_arrays": differentiable_arrays,
+                "field_order": field_order,
+            },
         )
 
     return ApplyInputSchema, ApplyOutputSchema
@@ -382,6 +407,25 @@ def add_path(obj: T, path: tuple) -> T:
     return path_to_type
 
 
+def get_field_order(schema: type[BaseModel]) -> dict[tuple, list[str]]:
+    """Return a mapping {model_path: [field_names_in_declaration_order]} for every BaseModel in the tree.
+
+    The root model is keyed by the empty tuple. Field names are collected as a side effect while
+    traversing leaf paths: any string at position `i` in a leaf path is a field of the model at
+    prefix `path[:i]`.
+    """
+    field_order_by_path: dict[tuple, dict[str, None]] = {}
+
+    def add_field(obj: T, path: tuple) -> T:
+        for i, part in enumerate(path):
+            if isinstance(part, str):
+                field_order_by_path.setdefault(path[:i], {})[part] = None
+        return obj
+
+    apply_function_to_model_tree(schema, add_field)
+    return {p: list(fields) for p, fields in field_order_by_path.items()}
+
+
 def _path_to_pattern(path: Sequence[str | object]) -> str:
     """Return a type describing valid paths for all passed paths."""
     # Check if path includes sequence or dict indexing -- in this case, we can't use the

tests/runtime_tests/test_schema_generation.py

@@ -119,6 +119,19 @@ def test_create_apply_schema():
         OutputSchema.model_validate({**testinput, "foo": 1})
 
 
+def test_field_order_exported():
+    InputSchema, OutputSchema = create_apply_schema(NestedModel, NestedModel)
+
+    expected = {
+        "": list(NestedModel.model_fields.keys()),
+        "testfoo.[]": list(SubModel.model_fields.keys()),
+        "testbar.{}": list(SubModel.model_fields.keys()),
+        "testrootmodel": ["root"],
+    }
+    assert InputSchema.model_json_schema()["field_order"] == expected
+    assert OutputSchema.model_json_schema()["field_order"] == expected
+
+
 def test_create_abstract_eval_schema():
     testinput_abstract = replace_arrays(
         testinput, lambda arr: {"shape": arr["shape"], "dtype": arr["dtype"]}