Simplify implementation per feedback - use helper function instead of wrapper

Co-authored-by: pwwang <[email protected]>

Author: Copilot