Would be helpful to have the token count on streamed responses.
Maybe add it to the last response from the output?