Import from the right module
Improve logic and fix string format