We describe new efforts to model radio active galactic nuclei (AGN) in a cosmological context using the SAGE semi-analytic galaxy model. Our new method tracks the physical properties of radio jets in massive galaxies, including the evolution of radio lobes and their impact on the surrounding gas. This model also self consistently follows the gas cooling-heating cycle that significantly shapes star formation and the life and death of many galaxy types. Adding jet physics to SAGE adds new physical properties to the model output, which in turn allows us to make more detailed predictions for the radio AGN population. After calibrating the model to a set of core observations we analyse predictions for jet power, radio cocoon size, radio luminosity, and stellar mass. We find that the model is able to match the stellar mass--radio luminosity relation at $zsim0$, and the radio luminosity function out to $zsim1$. This updated model will make possible the construction of customised AGN-focused mock survey catalogues to be used for large-scale observing programs.