Fix: Build errors
Missing imports and wrong package name

Signed-off-by: Chintan Shah <cgintan1998@gmail.com>