In the code the movement is taken from the player's animation in this line:
root_motion = animation_tree.get_root_motion_transform()
This line is called two times, both when the player is on floor. In my opinion this line should be out from those IFs, just before this line:
orientation *= root_motion
This way you could modify the animation to have a displacement and then use it from the root motion, but my recommendation is to use another vector to "push" the player, this way if you need to add other forces (like wind) this vector can be used.
Use the velocity to set your push vector, do it in the jump:
if not on_air and Input.is_action_just_pressed("jump"):
add it to the character velocity before move and slide. Then you have to decide if the velocity is constant while the player is on the air or if it will decrease over time, just remember to set it to zero when player is not on the air.