Fix pyright type errors

[benjibc]

---

name: Pull Request
about: Propose changes to the codebase
title: "Fix: Improve type safety for reward functions and message content handling"
labels: ''
assignees: ''

---

Description

This PR addresses several Pyright type errors by improving type safety in reward functions and the reward_function decorator. The primary motivation was to resolve type inconsistencies related to message content and argument types, making the reward functions more robust to different message representations (Pydantic Message objects or raw dictionaries).

Changes include:

• Updating the reward_function decorator in eval_protocol/typed_interface.py to better handle List[Message] coercion and preserve function signatures, particularly when Union types are involved.
• Widening the messages and ground_truth parameters in affected reward functions to accept Union[List[Message], List[Dict[str, Any]]].
• Introducing a _to_text helper function in reward modules to consistently extract string content from various message content types (str or List[ChatCompletionContentPartTextParam]).
• Adjusting logic within reward functions to correctly parse message roles and content from both Message objects and dictionary representations.

Fixes # (issue)
Implements # (issue)

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update
• Refactoring/Code cleanup
• Build/CI/CD related changes
• Other (please describe):

How Has This Been Tested?

The changes were verified by running the project's type checker (basedpyright) locally.

• Ran basedpyright specifically on the eval_protocol/rewards directory.
• Confirmed that type errors in the modified files (multiple_choice_math_reward.py, reasoning_steps.py, repetition.py, tag_count.py, typed_interface.py) are resolved or significantly reduced. The four edited reward files are now type-clean, aside from an external optional import warning in repetition.py for jieba.

Test Configuration:

• Toolchain: basedpyright

Checklist:

• My code follows the style guidelines of this project (ran black ., isort ., flake8 .)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings (aside from pre-existing external dependency warnings)
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• Any dependent changes have been merged and published in downstream modules
• I have checked my code and corrected any misspellings

Screenshots (if applicable)

Additional context

The optional import warning for jieba in repetition.py remains, as it is an external dependency issue outside the scope of this PR.

---

[cursor]

Cursor Agent can help with this pull request. Just @cursor in comments and I'll start working on changes in this branch.
_{Learn more about Cursor Agents}