Optimize

[AndrewZuo01]

🔍 What type of PR is this?

👀 What this PR does / why we need it:

• My pull request adheres to the code style of this project
• My code requires changes to the documentation
• I have updated the documentation as required
• All the tests have passed

🅰 Which issue(s) this PR fixes:

Fixes #333

📝 Special notes for your reviewer:

🎯 Describe how to verify it

📑 Additional documentation e.g., RFC, notion, Google docs, usage docs, etc.:

[sweep-ai]

Apply Sweep Rules to your PR?

• Apply: All new business logic should have corresponding unit tests.
• Apply: Refactor large functions to be more modular.
• Apply: Add docstrings to all functions and file headers.