river-client

Python client for the River ML training API — sampling, LoRA and full-weight fine-tuning, and reinforcement learning against River-hosted models.

Installation

pip install river-client

Requires Python 3.12+.

Quick start

River provides a low-level API for sampling, training, and checkpointing, plus an integrated RL library that manages rollouts and training from your environment and reward function. The example below uses the low-level API.

import river_client as river

client = river.Client(api_key="your-key", endpoint="api.river.ai")

# Stateless sampling from a base model
samples = client.sample(
    "What is 2+2?",
    base_model="Qwen/Qwen3.6-35B-A3B-FP8",
    max_tokens=50,
)
print(samples[0].text)

# Training with a session
with client.session() as session:
    model = session.create_model(
        base_model="Qwen/Qwen3.6-35B-A3B-FP8",
        lora=river.LoraConfig(rank=16),
    )

    # Forward + backward, then an optimizer step
    result = model.forward_backward(data, loss_fn="cross_entropy")
    model.optim_step(lr=1e-4)

    # Sample from the current weights
    sample_groups = model.sample("Continue:", max_tokens=100)

Reinforcement learning

river_client.rl manages rollouts and training around your environment, tools, and reward function. It supports multi-turn conversations, images, KV-cache reuse, and synchronous or bounded asynchronous training.

import operator
import os
from typing import Literal

import river_client as river
from river_client import rl
from river_client.renderers import get_renderer, get_text_content

# Tool schemas come from type hints and docstrings; execution stays in your code.
@rl.tool
async def calculator(a: float, operation: Literal["+", "-", "*", "/"], b: float) -> str:
    """Calculate the result of an arithmetic operation on two numbers."""
    operations = {"+": operator.add, "-": operator.sub, "*": operator.mul, "/": operator.truediv}
    return f"{operations[operation](a, b):g}"

class AnswerEnv(rl.Env):
    tools = [calculator]  # The engine runs tool calls and feeds results back.

    async def reset(self, row):
        return [{"role": "user", "content": (
            row["question"] + " Use the calculator if helpful. Return only the answer."
        )}]

    async def reward(self, traj, row):
        # Score the final answer, excluding the model's reasoning.
        answer = get_text_content(traj.messages[-1]).strip()
        return float(answer == row["answer"])

# Replace this toy dataset and reward with your task.
dataset = [
    {"question": "What is 17 * 23?", "answer": "391"},
    {"question": "What is 144 / 12?", "answer": "12"},
]
base_model = "zai-org/GLM-5.3-Flash"
renderer = get_renderer(base_model)
client = river.Client(api_key=os.environ["RIVER_API_KEY"], endpoint="api.river.ai")

with client.session() as session:
    model = session.create_model(
        base_model=base_model,
        tokenizer=renderer.tokenizer,
        lora=river.LoraConfig(rank=16),
    )
    trainer = rl.AsyncTrainer(
        engine=rl.RolloutEngine(
            model, env=AnswerEnv, renderer=renderer,
            budget=rl.Budget(max_turns=4, max_generated_tokens=2048),
            # Sampling stays on one policy for the whole trajectory by default.
        ),
        optimizer=rl.Adam(lr=1e-5),
        advantage=rl.GroupCentered(),  # Center rewards within each prompt group.
        completion=rl.GroupCompletion(mode="wait"),
        loss=rl.CISPO(normalize="token"),
        groups_per_step=2, group_size=8,  # Two prompts, eight rollouts each.
        max_staleness=0,  # Synchronous; >0 allows bounded sampling ahead.
    )
    rl.run(trainer, dataset, steps=10, on_step=lambda step: print(step.n, step.metrics))

See the bundled agent skill for image observations, cache policies, checkpoint/resume, and evaluation.

Dedicated streaming inference

Gated feature — disabled by default. Contact River to enable dedicated deployments for your team and the checkpoint's base model before using these APIs. Use a team API key with that access; personal API keys cannot create deployments.

client.create_deployment(checkpoint, ...) provisions capacity for a checkpoint and returns a base URL that the standard OpenAI client streams from unchanged; list_deployments, get_deployment_usage, scale_on_target and delete_deployment manage it from there. See the deployment walkthrough for a complete example. The bundled agent skill carries the full workflow: replica roles, scale-to-zero and resume, streaming error handling, and usage accounting.

Export checkpoints to Hugging Face

Install the hf-export extra to convert a downloaded LoRA checkpoint into a Hugging Face PEFT adapter:

pip install -U "river-client[hf-export]"
python -m river_client.hf_export ./river-checkpoint ./hf-adapter

In Console, Download HF export script on an inference checkpoint downloads and converts in one step. The bundled agent skill covers the converter's options, Python API, supported layouts, and loader caveats.

AI agent skill

The package bundles an agent skill — a SKILL.md that teaches AI coding agents (Claude Code and compatible tools) the current training API: train_step semantics, data formats, RL/SFT/distillation loop patterns, image uploads and handles, checkpoints, full-weight training, capacity status, and fault-tolerant auto-recovery. Because it ships inside the wheel, the skill always matches the installed client version.

Install it into your agent's skills directory:

python -m river_client.skill --install

This copies the skill into ~/.claude/skills/; pass --dest for a different location (e.g. a project's .claude/skills/). Run without --install to print the bundled skill's path instead.

License

Apache-2.0. See LICENSE.