Small cleanup of updated api