feat: performance and code style updates for PR

feat: updating tests and improving generation_methods

bugfix: adjusting tests following vmap updates

feat: refactoring generation methods into generator - refactoring generation methods into one generator class and updating codebase to pass linting tests

feat: updating typing in Agent dataclass

docs: updating docstrings in generatory.py