Reward Function API

Relevant source files

The following files were used as context for generating this wiki page:

• osmosis_ai/rubric_eval.py
• examples/reward_functions.py
• osmosis_ai/rubric_types.py
• osmosis_ai/utils.py
• README.md
• examples/README.md
• examples/rubric_functions.py

The Reward Function API within the Osmosis SDK provides a structured framework for evaluating Large Language Model (LLM) outputs. It facilitates two primary evaluation methods: local deterministic reward functions using string or numeric logic, and remote rubric-based evaluation delegating judgment to hosted LLM providers.

This API ensures consistency through strict type enforcement and signature validation via Python decorators. It integrates with LiteLLM to support over 100 model providers, allowing developers to implement complex scoring logic—ranging from simple exact matches to nuanced natural language rubric assessments—with built-in error handling and credential management.

Sources: README.md:4-8, osmosis_ai/rubric_eval.py:1-10

Core Decorators

The SDK uses decorators to enforce specific function signatures at definition time. This ensures that any reward or rubric function integrated into an Osmosis workflow adheres to the expected data contract.

@osmosis_reward

The @osmosis_reward decorator is used for local evaluation functions. It validates that the function accepts three specific parameters and returns a float.

The flow above illustrates how the decorator enforces constraints both at definition and execution time.
Sources: osmosis_ai/utils.py:12-28, examples/reward_functions.py:61-82

Parameter	Type	Required	Description
solution_str	str	Yes	The model output string to be evaluated.
ground_truth	str	Yes	The reference or expected answer.
extra_info	dict	No	Optional configuration or context (must default to None).

Sources: osmosis_ai/utils.py:38-53, examples/README.md:19-23

@osmosis_rubric

The @osmosis_rubric decorator is specifically for natural language evaluations. It allows ground_truth to be Optional[str] (accepting None if no reference exists) and mandates that extra_info be provided as a non-optional dictionary containing provider-specific configuration.

Sources: osmosis_ai/utils.py:100-115, README.md:129-140

Remote Rubric Evaluation

Remote evaluation is handled by the evaluate_rubric function, which delegates the judging process to external LLMs via LiteLLM.

Architecture and Data Flow

The SDK handles the complexity of prompt engineering (using sentinel markers to prevent prompt injection from model outputs) and response sanitization.
Sources: osmosis_ai/rubric_eval.py:408-440, README.md:46-55

Model Configuration (ModelInfo)

The model_info dictionary configures the remote judge.

Key	Type	Description
provider	str	The LLM provider (e.g., "openai", "anthropic").
model	str	Specific model identifier.
api_key	str	Explicit API key (optional).
api_key_env	str	Environment variable name to look up the key.
score_min	float	Minimum score bound (default 0.0).
score_max	float	Maximum score bound (default 1.0).
timeout	float	Request timeout in seconds.

Sources: osmosis_ai/rubric_types.py:6-17, osmosis_ai/rubric_eval.py:448-460

Response Handling and Data Structures

RewardRubricRunResult

When return_details=True is passed to evaluate_rubric, the API returns a structured dictionary instead of a raw float.

class RewardRubricRunResult(TypedDict):
    score: float
    explanation: str
    raw: Any

Sources: osmosis_ai/rubric_types.py:20-23

Error Handling

The API defines specific exceptions for common failure modes in remote evaluation:

• MissingAPIKeyError: Raised if credentials cannot be resolved from model_info or environment variables.
• ProviderRequestError: Raised when the remote call fails (e.g., rate limits, connection issues).
• ModelNotFoundError: A subclass of ProviderRequestError raised when the provider does not recognize the model identifier.

Sources: osmosis_ai/rubric_types.py:26-44, osmosis_ai/rubric_eval.py:345-360

Prompt Engineering and Sanitization

The API implements internal measures to ensure judge impartiality and output reliability:

1. Sentinel Markers: Inputs like CANDIDATE_OUTPUT and GROUND_TRUTH are wrapped in <<<BEGIN_...>>> and <<<END_...>>> tags to prevent models from following instructions contained within the data being evaluated.
2. JSON Sanitization: The SDK strips <think> blocks from reasoning models and markdown code fences before parsing the judge's response.
3. Score Clamping: Scores returned by the model are automatically clamped within the [score_min, score_max] range.

Sources: osmosis_ai/rubric_eval.py:115-132, osmosis_ai/rubric_eval.py:155-170, osmosis_ai/rubric_eval.py:397-400

Summary

The Reward Function API provides a robust bridge between deterministic logic and LLM-based qualitative evaluation. By combining signature enforcement via decorators with a unified interface for remote providers, it allows for scalable and reliable evaluation workflows in LLM training and testing.