Python API reference

Reference for the river_client Python package, generated from river-client 0.6.1. The guide covers the same API as a walkthrough; this page lists every public class, method, and type.

pip install river-client

import river_client as river

Client

River API client.

Connects to the River API server over gRPC with automatic retry on transient failures and connection keepalive.

Client(
    api_key: str,
    endpoint: str = 'api.river.ai',
    port: int = 443,
    timeout: float = 86400.0,
    use_ssl: bool = True,
    enable_retries: bool = True,
)

Methods: session, sample, health_check, get_capabilities, get_streaming_replica, promote_streaming_replica, chat_complete_stream, chat_complete, chat_complete_from_checkpoint, chat_complete_from_training, close

Client.session

Client.session(*, timeout: float = 86400.0, **tags: str) -> SessionContext

Create a session context manager.

Returns: SessionContext — Context manager that yields a Session

Client.sample

Client.sample(
    prompts: str | list[str] | None = None,
    *,
    base_model: str,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    return_prompt_logprobs: bool = False,
    logprobs: int | None = None,
    images: list[bytes] | list[list[bytes]] | None = None,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    tokenizer: Any | None = None,
    metrics_type: str = '',
    timeout: float | None = None,
) -> list[Sample]

Sample from a base model (no session required).

Returns: list[Sample] — Flat list[Sample] — all samples across all prompts. For a single prompt with num_samples=1 (the default), this is a list with one element: result[0].text.

Client.health_check

Client.health_check() -> bool

Check API health.

Returns: bool — True if healthy

Client.get_capabilities

Client.get_capabilities() -> list[str]

Get supported models.

Returns: list[str] — List of supported model names

Client.get_streaming_replica

Client.get_streaming_replica(
    model: str,
    *,
    timeout: float | None = None,
) -> PromotedStreamingReplica | None

Return promoted streaming metadata for model when available.

This reads server-owned routing metadata from the control plane. Missing metadata returns None; other HTTP/auth failures raise a River client exception.

Client.promote_streaming_replica

Client.promote_streaming_replica(
    checkpoint: str | Checkpoint,
    model: str,
    *,
    timeout: float | None = None,
) -> PromotedStreamingReplica

Request beta promotion of a checkpoint to a streaming model alias.

The request is asynchronous: the returned metadata is usually status="provisioning". Poll get_streaming_replica until the status is "ready" or "degraded" before using chat_complete_stream.

Client.chat_complete_stream

Client.chat_complete_stream(
    messages: list[dict],
    *,
    model: str,
    timeout: float | None = None,
    on_not_ready: str = 'raise',
    **kwargs,
) -> Iterator[dict[str, Any]]

Stream OpenAI-compatible chat chunks from a promoted replica.

timeout is the socket read timeout for each blocking read, not an end-to-end generation deadline. Let the iterator finish, or call close() on it when breaking early, so the HTTP connection is closed.

Returns: Iterator[dict[str, Any]] — Iterator of decoded OpenAI streaming chunk dictionaries.

Client.chat_complete

Client.chat_complete(
    messages: list[dict],
    *,
    base_model: str,
    timeout: float | None = None,
    **kwargs,
) -> ChatCompleteResult | Iterator[dict[str, Any]]

Chat completion from a base model (no LoRA).

Builds an OpenAI-format request body and sends it through the gRPC ChatCompleteFromBase RPC.

Returns: ChatCompleteResult | Iterator[dict[str, Any]] — ChatCompleteResult with response_json and status_code.

Client.chat_complete_from_checkpoint

Client.chat_complete_from_checkpoint(
    messages: list[dict],
    *,
    checkpoint_path: str,
    base_model: str = '',
    timeout: float | None = None,
    **kwargs,
) -> ChatCompleteResult

Chat completion from a saved checkpoint.

Returns: ChatCompleteResult — ChatCompleteResult with response_json and status_code.

Client.chat_complete_from_training

Client.chat_complete_from_training(
    messages: list[dict],
    *,
    model_id: str,
    timeout: float | None = None,
    **kwargs,
) -> ChatCompleteResult

Chat completion from in-memory training weights.

Returns: ChatCompleteResult — ChatCompleteResult with response_json and status_code.

Client.close

Client.close() -> None

Close the client connection.

Session

A training session with GPU allocation.

Entered through Client.session; it owns the models you train.

Methods: attest_training_data, create_model, sample, submit_sample

Session.attest_training_data

Session.attest_training_data(
    artifacts: list[TrainingDataArtifact],
    timeout: float = 86400.0,
) -> TrainingDataAttestation

Ask the API to hash source artifacts and retain their manifest.

The source bytes are discarded by the API after hashing. Pass the result to create_model to make the server fail closed before every forward/backward request if that manifest disappears or no longer belongs to the model's session.

Session.create_model

Session.create_model(
    base_model: str,
    lora: LoraConfig | None = None,
    tokenizer: str | Any | None = None,
    checkpoint: str | Checkpoint | None = None,
    timeout: float = 86400.0,
    training_data_attestation: TrainingDataAttestation | str | None = None,
) -> Model

Create a new model for training.

Returns: Model — Model object for training

Session.sample

Session.sample(
    prompts: str | list[str] | None = None,
    *,
    base_model: str,
    checkpoint: str | Checkpoint | None = None,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    return_prompt_logprobs: bool = False,
    logprobs: int | None = None,
    return_expert_routing: bool = False,
    images: list[bytes] | list[list[bytes]] | None = None,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    tokenizer: Any | None = None,
    metrics_type: str = '',
    timeout: float = 86400.0,
) -> list[list[Sample]]

Sample from a base model or checkpoint.

When checkpoint is provided, the server loads the LoRA from the saved checkpoint, generates text, then unloads.

Returns: list[list[Sample]] — list[list[Sample]] — outer list is per-prompt, inner list is per-sample.

Session.submit_sample

Session.submit_sample(
    prompts: str | list[str] | None = None,
    *,
    base_model: str,
    checkpoint: str | Checkpoint | None = None,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    return_prompt_logprobs: bool = False,
    logprobs: int | None = None,
    return_expert_routing: bool = False,
    images: list[bytes] | list[list[bytes]] | None = None,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    tokenizer: Any | None = None,
    metrics_type: str = '',
    timeout: float = 86400.0,
) -> PendingSample

Submit sampling from a base model or checkpoint without waiting.

See sample for the full kwarg reference.

Session.session_id

Session.session_id: str

SessionContext

Context manager for Session with auto-heartbeat.

The context manager returned by Client.session.

Model

A training model with mutable in-memory weights.

Created by Session.create_model.

Methods: forward, forward_backward, optim_step, train_step, submit_forward_backward, submit_optim_step, submit_train_step, sample, submit_sample, chat_complete, save_weights, promote_to_streaming, load_weights

Model.forward

Model.forward(
    data: list[dict],
    loss_fn: str = 'cross_entropy',
    timeout: float = 86400.0,
    **loss_config: float,
) -> ForwardResult

Forward pass only (compute loss, no gradients).

Returns: ForwardResult — ForwardResult with metrics

Model.forward_backward

Model.forward_backward(
    data: list[dict],
    loss_fn: str = 'cross_entropy',
    timeout: float = 86400.0,
    return_logprobs: bool = False,
    zero_out: bool = True,
    compute_expert_flip_metric: bool = False,
    force_routing_replay: bool = False,
    **loss_config: float,
) -> ForwardResult

Forward + backward pass (compute gradients).

Returns: ForwardResult — ForwardResult with metrics and logprobs when returned by the worker.

Model.optim_step

Model.optim_step(
    lr: float,
    beta1: float = 0.9,
    beta2: float = 0.999,
    eps: float = 1e-08,
    weight_decay: float = 0.0,
    grad_clip_norm: float | None = None,
    timeout: float = 86400.0,
) -> OptimStepResult

Apply gradients with Adam optimizer.

Returns: OptimStepResult — OptimStepResult with metrics (step, lr, grad_norm, grad_norm_finite)

Model.train_step

Model.train_step(
    data: list[dict],
    lr: float,
    *,
    loss_fn: str = 'cross_entropy',
    beta1: float = 0.9,
    beta2: float = 0.999,
    eps: float = 1e-08,
    weight_decay: float = 0.0,
    grad_clip_norm: float | None = None,
    compute_expert_flip_metric: bool = False,
    force_routing_replay: bool = False,
    timeout: float = 86400.0,
    **loss_config: float,
) -> tuple[ForwardResult, OptimStepResult]

Complete training step: forward+backward plus optimizer update.

Submits forward+backward and the optimizer step back-to-back — the server runs them in order as one pipelined unit, without a client round trip in between — then waits for both results. See forward_backward and optim_step for the full parameter reference.

The error path differs from calling forward_backward() then optim_step(): the optimizer step is already submitted when forward-backward resolves, so if forward-backward fails, this call raises its error while the optimizer step still runs server-side and Model.step has already advanced. Callers that need to inspect both outcomes should use submit_train_step.

A train step is a complete step: gradients are always cleared first. For micro-batch gradient accumulation, use submit_forward_backward(zero_out=...) and submit_optim_step directly.

Returns: tuple[ForwardResult, OptimStepResult] — Tuple of (ForwardResult, OptimStepResult).

Model.submit_forward_backward

Model.submit_forward_backward(
    data: list[dict],
    loss_fn: str = 'cross_entropy',
    timeout: float = 86400.0,
    return_logprobs: bool = False,
    zero_out: bool = True,
    compute_expert_flip_metric: bool = False,
    force_routing_replay: bool = False,
    **loss_config: float,
) -> PendingOp

Submit forward+backward without blocking. Returns a PendingOp.

Call pending.result() later to get the ForwardResult. This enables pipelining: submit step N+1 while step N is still running. See forward_backward for the full kwarg reference.

When submitting multiple micro-batches before submit_optim_step, use zero_out=True for the first one and False for later submissions.

timeout bounds the submit RPC and, separately, the wait inside pending.result().

Model.submit_optim_step

Model.submit_optim_step(
    lr: float,
    beta1: float = 0.9,
    beta2: float = 0.999,
    eps: float = 1e-08,
    weight_decay: float = 0.0,
    grad_clip_norm: float | None = None,
    timeout: float = 86400.0,
) -> PendingOp

Submit optimizer step without blocking. Returns a PendingOp.

Call pending.result() later to get the OptimStepResult. Model.step advances at submit time, even if the operation later fails.

Model.submit_train_step

Model.submit_train_step(
    data: list[dict],
    lr: float,
    *,
    loss_fn: str = 'cross_entropy',
    beta1: float = 0.9,
    beta2: float = 0.999,
    eps: float = 1e-08,
    weight_decay: float = 0.0,
    grad_clip_norm: float | None = None,
    compute_expert_flip_metric: bool = False,
    force_routing_replay: bool = False,
    timeout: float = 86400.0,
    **loss_config: float,
) -> tuple[PendingOp, PendingOp]

Submit a complete training step without blocking.

Fires forward+backward and the optimizer step back-to-back; the server runs them in submission order per model. Returns the two PendingOps as (forward_backward, optim_step).

Because both are submitted up front, a failed forward-backward does not cancel the already-submitted optimizer step. Model.step advances at submit time. See train_step for the full kwarg reference.

Model.sample

Model.sample(
    prompts: str | list[str] | None = None,
    *,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    return_prompt_logprobs: bool = False,
    logprobs: int | None = None,
    images: list[bytes] | list[list[bytes]] | None = None,
    return_expert_routing: bool = False,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    metrics_type: str = '',
    timeout: float = 86400.0,
    poll_interval: float = 1.0,
) -> list[list[Sample]]

Sample from the current in-memory training weights.

Generates text using the model's current weights, without needing to save a checkpoint first. Per-token logprobs are always returned.

Returns: list[list[Sample]] — list[list[Sample]] — outer list is per-prompt, inner list is per-sample. Each Sample has .tokens, .text, .logprobs, and .stop_reason.

Model.submit_sample

Model.submit_sample(
    prompts: str | list[str] | None = None,
    *,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    return_prompt_logprobs: bool = False,
    logprobs: int | None = None,
    images: list[bytes] | list[list[bytes]] | None = None,
    return_expert_routing: bool = False,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    metrics_type: str = '',
    timeout: float = 86400.0,
    poll_interval: float = 1.0,
) -> PendingSample

Submit sampling from current training weights without waiting.

See sample for the full kwarg reference.

Model.chat_complete

Model.chat_complete(
    messages: list[dict],
    *,
    timeout: float | None = None,
    **kwargs,
) -> ChatCompleteResult

OpenAI-compatible chat completion from current training weights.

Like model.sample() but using the OpenAI chat-completions format instead of raw prompts.

Returns: ChatCompleteResult — ChatCompleteResult with response_json (full OpenAI-format JSON string) and status_code.

Model.save_weights

Model.save_weights(
    name: str,
    mode: str = 'training',
    timeout: float = 86400.0,
    ttl: datetime.timedelta | None = None,
) -> Checkpoint

Save a checkpoint of the current model weights.

Returns: Checkpoint — Checkpoint object with river:// path.

Model.promote_to_streaming

Model.promote_to_streaming(
    model: str,
    checkpoint: str | Checkpoint | None = None,
    checkpoint_name: str | None = None,
    timeout: float | None = 86400.0,
) -> PromotedStreamingReplica

Promote this model's current or saved checkpoint to a stream alias.

When checkpoint is omitted, this saves the current weights with mode="inference" using checkpoint_name or a generated name, then asks the control plane to promote that checkpoint asynchronously. The timeout is applied separately to the save and promotion calls, not as one end-to-end deadline. Omitted checkpoint_name values create a new server-side inference checkpoint for each call.

Model.load_weights

Model.load_weights(
    checkpoint: str | Checkpoint,
    load_optimizer: bool = True,
    timeout: float = 86400.0,
) -> None

Load weights from a checkpoint.

Model.model_id

Model.model_id: str

Model.training_run_id

Model.training_run_id: str

Model.step

Model.step: int

Current training step.

Advances when an optimizer step is submitted (so pipelined submissions observe a consistent value), not when it resolves — a failed or timed-out optimizer step leaves it advanced.

Data types

Values returned by the client and the configuration objects you pass to it.

LoraConfig

LoRA adapter configuration.

Sample

A generated sample.

tokens contains generated token IDs when the response includes them; older responses may fall back to retokenizing text client-side.

prompt_token_ids / top_logprobs / prompt_top_logprobs are None when the server did not provide them or the feature was not requested.

expert_routing is populated when return_expert_routing=True was passed on the request.

metrics carries per-result scalar metrics; empty unless the originating request passed a recognized metrics_type token.

Sample.routing_datum_keys

Sample.routing_datum_keys(*, required: bool = False) -> dict[str, bytes | str]

Return the per-datum keys to splat into a forward_backward datum when enabling force_routing_replay or compute_expert_flip_metric.

Returns an empty dict when this sample carries no routing capture and required is false. Raises ValueError when required is true and no routing handle is available.

Example:

data = []
for sample in samples:
    datum = {
        "input_ids": ...,
        "advantages": ...,
        **sample.routing_datum_keys(required=True),
    }
    data.append(datum)

ChatCompleteResult

Result of a chat completion request.

ForwardResult

Result of forward or forward_backward pass.

OptimStepResult

Result of an optimizer step.

Checkpoint

A saved model checkpoint.

TopLogprob

One top-K candidate token at a single position.

PendingOp

A submitted but not-yet-resolved async operation.

Returned by Model.submit_forward_backward() and Model.submit_optim_step(). Call .result() to block until complete.

PendingOp.result

PendingOp.result() -> ForwardResult | OptimStepResult

Block until the operation completes and return the result.

PendingSample

A submitted but not-yet-resolved sampling operation.

PendingSample.result

PendingSample.result() -> list[list[Sample]]

Block until sampling completes and return grouped samples.

ExpertRouting

Captured MoE expert routing for one sample.

Populated when the caller passed return_expert_routing=True on the sample request and routing capture was available.

To enable router replay / flip-rate metrics, splat the canonical per-datum keys via Sample.routing_datum_keys() rather than building them by hand:

datum = {..., **sample.routing_datum_keys(required=True)}

Current servers return routing captures as an opaque handle; the splat yields expert_routing_handle. topk_ids is retained for inspection and proto round-trip compatibility only; replay always recomputes the routing weights in the trainer.

The shape header fields (num_decoder_layers / top_k / layer_indices) are informational only and let the user inspect capture metadata.

Byte layout (when ids are present): [num_tokens, num_decoder_layers, top_k] row-major int16 little-endian. num_tokens = seqlen - 1 because routing capture omits the trailing position.

PromotedStreamingReplica

Server-owned routing metadata for a promoted streaming replica.

TrainingDataArtifact

Source bytes and expected digest for server-side integrity verification.

The API hashes content itself and retains only the resulting manifest. Callers should verify their source files locally before constructing these objects, then bind the returned TrainingDataAttestation to the model they create.

AttestedTrainingDataArtifact

Server-observed source-artifact digest and size.

TrainingDataAttestation

A server-verified source-artifact manifest bound to one session.

Errors

Every server and transport failure derives from RiverError, so a single except river.RiverError catches all of them. Invalid arguments still raise the standard ValueError and TypeError.

RiverError

Inherits Exception.

Base exception for River client errors.

AuthenticationError

Inherits RiverError.

Authentication failed.

CapacityError

Inherits RiverError.

No capacity available for the operation.

ModelNotFoundError

Inherits RiverError.

Model not found.

RiverConnectionError

Inherits RiverError.

Connection or communication error with the River API server.

This wraps gRPC errors with a more user-friendly message while preserving the original error details for debugging.

RiverConnectionError.from_grpc_error

RiverConnectionError.from_grpc_error(
    error: Exception,
    context: str = 'API call',
) -> 'RiverConnectionError'

Create a RiverConnectionError from a gRPC RpcError.

RiverTimeoutError

Inherits RiverError.

Operation timed out while retaining its recoverable future ID.

SessionHeartbeatError

Inherits RiverConnectionError.

Session heartbeat was lost or rejected while a session was active.

Tokenizers

Helpers for tokenizing prompts yourself, for example when you build prompt_token_ids for training data.

load_tokenizer

load_tokenizer(
    tokenizer: str | Any | None = None,
    *,
    base_model: str | None = None,
    revision: str | None = None,
    local_files_only: bool = False,
    resolve_aliases: bool = True,
)

Load or return a tokenizer for River client result parsing.

base_model remains the public River model name used for routing. When it is a hyphen-suffixed deployment alias of a known canonical model, this helper resolves it to the underlying Hugging Face tokenizer id before loading. local_files_only keeps sealed jobs from reaching Hugging Face after their tokenizer revision has been frozen. resolve_aliases=False retains a deployment's own tokenizer source for unpinned compatibility paths.

resolve_tokenizer_name

resolve_tokenizer_name(model_name: str) -> str

Return the tokenizer for a canonical River model or one of its aliases.

Deployment aliases must append a hyphen-delimited suffix to a canonical model name. Longest-root matching keeps the result deterministic if a future canonical model name extends another root.

MODEL_TOKENIZER_ALIASES

MODEL_TOKENIZER_ALIASES: dict[str, str]