Fix agent_demo

- Missing imports
- Forgot class follow