Get a tool class by name. Args: name: The name of the tool to retrieve Returns: The tool class if f

Get a trace/signal by ID. Args: trace_id: Trace/signal ID Returns: TraceContext object or None if n

Load a dataset from the registry. Args: name: Name of the dataset to load split: Split name to load (e.g., 'train', '

Load a dataset from a JSON file based on the dataset enum. This function takes a dataset enum value and loads the corresponding JSON file fro

Register a dataset by saving it to disk and updating the registry. Args: name: Name of the dataset data: List of dict

Get current metadata from baggage.

Get the dataset data.

Extract tool calls from the model response.

Encode queries or passages into embeddings.

Create an OpenTelemetry session context manager.

Start the connection manager thread.

Get current metadata from ContextVar. Returns: Metadata dict or empty dict if no metadata is set.

Select a subset of the dataset (HuggingFace-compatible API). Args: indices: List or range of indices to select Returns:

Get the output from the trajectory (backward compatibility).

Sleep all rollout replica instances asynchronously.

Apply proxy URL rewriting and extra headers to client.

Reset the environment and return initial observations.

Get the absolute path of the dataset file. Returns: Optional[str]: The absolute path of the dataset file, or None if the dataset

Execute the workflow on a single task Args: task: The task to execute. uid: The unique identifier for the task.

Take a step in the environment based on the action. Args: action: Action from the agent (tool calls or final response)

Reset the environment and return initial observations.

Run a function in seperate thread pool executor. Args: fn: The function to run. *args: The arguments to pass to the f

Take a step in the environment based on the action. Args: actions: List containing a single action string from the agent

Get the forwarded environment variables. The `RLLM_EXCLUDE` environment variable can be used to exclude specific environment variables or all

Extract completion token IDs from response payload. Looks for token IDs in various locations within the response: 1. choice.output_token_ids

Generates a random valid map (one that has a path from start to goal) Args: size: size of each side of the grid p: probability th

Repeat the dataset n times, keeping repeated entries adjacent. Args: n: Number of times to repeat the dataset Returns:

- Map custom action to gymnasium FrozenLakeEnv action and take the step - Check if the action is effective (whether player moves in t

Log LLM call to tracer.

Return the metadata dict that should be routed through the proxy slug.

Close the environment and clean up resources.

Merge positional and keyword arguments into a single dict. Supports passing a single dict as the first positional argument, which will be mer

Load the dataset registry from the registry file.

Get async OpenAI chat client with automatic session tracking. Async version of get_chat_client(). See get_chat_client() for details. Returns

Process environment feedback and update internal state.

Create a configured async SQLite connection.

Close tracer and flush pending traces. This method should clean up resources and ensure all traces are persisted before retu

Factory method to get the appropriate tool parser based on a string identifier. Args: tokenizer: The tokenizer to use with the pa

Execute the main PPO training workflow. This method sets up the distributed training environment, initializes workers, datasets, and

Create session context for automatic trace tracking with auto-generated name. Session name is auto-generated. Nested sessions inherit parent meta

Shuffle the dataset (HuggingFace-compatible API). Args: seed: Random seed for reproducibility Returns: Datas

Update the agent's state based on environment feedback. Args: observation: Environment observation (task instruction or

Verify that parsing messages together is equivalent to parsing them individually. Args: messages (list): List of message dictiona

Ensure database is initialized (lazy initialization).

Extract Python code from the model's response. Supported formats: 1. ```python ... ``` 2. Code: ... 3. Whole

Formats a token string with the given click style.

Commit a trajectory for training. Args: name: The name of the trajectory. agent: The agent that generated the traject

Create a sampling client from a checkpoint path. Args: sampler_path: Path to sampler checkpoint Returns:

Get the unified proxy endpoint URL.

Dense retrieval using multi-GPU sharded FAISS.

Stop the connection manager thread.

Check if the agent has reacched the goal (G) or hole (H)

Wrapper for environment calls with timing. Args: func: The function to call (typically env.reset or env.step) *args:

Updates the agent's state based on environment feedback. Formats observation and updates the trajectory.

Updates the agent's internal state based on the model's response.

Initialize the agent with task information from the first observation. Args: observation: Initial observation containing

Run asynchronous workflow execution with retry logic for multiple tasks. Args: tasks: List of task dictionaries to process.

Filter trajectories based on reward

Synchronous version of generate that wraps the async method.

Async version of generate that can be called from async contexts.

Fast lookup of all traces for a session context. Uses the composite indexed junction table for optimal performance. Args:

Returns the agent's current state as a dictionary. This method provides access to the agent's internal state at the current step,

Get the tool prompt for the model.

Initialize semaphore for controlling concurrent task execution. Creates asyncio semaphore to limit parallel execution. Idempotent - s

Override request() to add tracing and proxy routing to ALL API calls.

Save checkpoint and return paths. Args: batch_idx: Current batch index kind: Checkpoint kind ("state", "samp

Set the sampling client for inference. Args: sampling_client: Tinker SamplingClient instance

Updates the agent's internal state after an environment step. Includes logic to check if the observation changed from the previous st

Log LLM call to tracer.

Read session context from W3C baggage.