Document parameters given to store_tokens callback

### Is your feature request related to a problem? Please describe.

I couldn't find this documented anywhere and had to dig into the source to find its arguments

### Describe the solution you'd like

Document it in the relevant docstrings and the mentions in the README

### Describe alternatives you've considered

n/a

### Additional context

n/a