Function backed data for Vega only supported via Typescript (not Python)

Are function-backed data inputs for Vega plots only supported via Typescript functions? I tried using python, but the function is not being recognised. AIP Assist says only Typescript functions will work, just checking here.
Tried both
def WorkflowWaterfall() -> ThreeDimensionalAggregation[String, String, Double]:

and

def WorkflowWaterfall() -> List[Dict[str, Optional[str]]]:

which both seem to give the right outputs but function is not showing in dropdown.

ah I just needed this
@dataclass
class SummaryData:
parent: Optional[str] = None
id: Optional[str] = None
total_count: Optional[int] = None

and then
@function
def StormWorkflowWaterfall() → list[SummaryData]:

and then usage after aggregation code:
for (parent, child), cnt in counts.items():
rows.append(SummaryData(
parent=parent,
id=child,
total_count=cnt
))

return rows

And then workshop recognises the function :partying_face: