Skip to main content

​
Reward Kit API Reference

This API reference provides detailed documentation for the key classes, functions, and data models in the Reward Kit.

​
Core Components

​
Classes and Decorators

​
Data Models

  • Data Models: Documentation for Message, EvaluateResult, MetricResult, and other data models

​
Modules

​
reward_function Module

The reward_function module contains the core functionality for creating and using reward functions. 
from reward_kit.reward_function import RewardFunction, reward_function

​
evaluation Module

The evaluation module provides the Evaluator class for managing evaluation configurations and functions for creating and previewing evaluations. 
from reward_kit.evaluation import Evaluator, preview_evaluation, create_evaluation
Key components:
  • Evaluator class: Manages metric loading, sample loading, and evaluator creation on the platform.
  • preview_evaluation: Previews an evaluation with sample data before deployment.
  • create_evaluation: Creates and deploys an evaluator to the platform.

​
config Module

The config module handles loading and managing configurations for the Reward Kit, typically from a rewardkit.yaml file. 
from reward_kit.config import load_config, get_config, RewardKitConfig
Key functions and classes:
  • load_config() / get_config(): Load the global Reward Kit configuration.
  • RewardKitConfig: Pydantic model for the main configuration structure.
  • Other models like GCPCloudRunConfig, AWSLambdaConfig.

​
models Module

The models module contains data models used throughout the Reward Kit. 
from reward_kit.models import EvaluateResult, MetricResult, Message

​
rewards Module

The rewards module contains specialized reward functions for specific use cases. 
from reward_kit.rewards.function_calling import match_function_call

​
server Module

The server module provides the RewardServer class and serve function to host reward functions as a FastAPI application. 
from reward_kit.server import RewardServer, serve
Key components:
  • RewardServer class: A class to encapsulate a reward function and run it as a server.
  • serve() function: A utility to quickly serve a given reward function.

​
auth Module

The auth module provides utility functions to retrieve authentication credentials, primarily for Fireworks AI. 
from reward_kit.auth import get_fireworks_api_key, get_fireworks_account_id
Key functions:
  • get_fireworks_api_key(): Retrieves the Fireworks API key.
  • get_fireworks_account_id(): Retrieves the Fireworks account ID.

​
gcp_tools Module

The gcp_tools module offers utilities for working with Google Cloud Platform, such as building and pushing Docker images to Artifact Registry and deploying to Cloud Run. 
from reward_kit.gcp_tools import build_and_push_docker_image, deploy_to_cloud_run

​
packaging Module

The packaging module assists in preparing reward functions for deployment, for example, by generating Dockerfile content. 
from reward_kit.packaging import generate_dockerfile_content

​
platform_api Module

The platform_api module provides functions for direct interaction with the Fireworks AI platform API, such as managing secrets. 
from reward_kit.platform_api import create_or_update_fireworks_secret

​
rl_processing Module

The rl_processing module contains tools for processing data for Reinforcement Learning workflows, such as the RLDataAligner. 
from reward_kit.rl_processing import RLDataAligner

​
mcp Module (reward_kit.mcp)

This sub-package contains components related to the Model Context Protocol (MCP).
  • reward_kit.mcp.clients: Provides clients for interacting with MCP-compliant servers.

​
mcp_agent Module (reward_kit.mcp_agent)

This sub-package provides a framework for building and running agents that interact with MCP servers. It includes orchestration logic, various backend implementations, and a collection of pre-built MCP servers for common tasks (e.g., filesystem, git).

​
Command Line Interface

The Reward Kit provides a command-line interface for common operations: 
# Show help
reward-kit --help

# Preview an evaluator
reward-kit preview --metrics-folders "metric=./path" --samples ./samples.jsonl

# Deploy an evaluator
reward-kit deploy --id my-evaluator --metrics-folders "metric=./path" --force
For detailed CLI documentation, see the CLI Reference.

​
Common Patterns

​
Creating a Basic Reward Function

from reward_kit import reward_function, EvaluateResult, MetricResult

@reward_function
def my_reward_function(messages, original_messages=None, **kwargs):
    # Your evaluation logic here
    response = messages[-1].get("content", "")
    # Assume calculate_score returns a float between 0.0 and 1.0
    # and calculate_success returns a boolean
    score = calculate_score(response)
    success = calculate_success(response) # Assume calculate_success is defined

    return EvaluateResult(
        score=score,
        reason="Overall evaluation reason for my_reward_function", # Added top-level reason
        metrics={
            "my_metric": MetricResult(
                score=score,
                success=success, # Added success field
                reason="Explanation for the metric score"
            )
        }
    )

​
Using a Deployed Reward Function

from reward_kit import RewardFunction

# Create a reference to a deployed reward function
reward_fn = RewardFunction(
    name="my-deployed-evaluator",
    mode="remote"
)

# Call the reward function
result = reward_fn(messages=[
    {"role": "user", "content": "What is machine learning?"},
    {"role": "assistant", "content": "Machine learning is..."}
])

print(f"Score: {result.score}")

​
Next Steps