Updated library to work better with the existing CI library