Extend your trajectories

Adaptive Data

Generate a next user turn for each conversation to build multi-turn training data that depends on earlier context.

Extend your trajectories when you want multi-turn training data that tests whether a model remembers what was said earlier in a conversation. For each row’s conversation, datasets.extend_trajectories generates the user’s next message: a single question or request that can only be answered correctly by recalling details from earlier turns, written in the conversation’s language.

Each request creates a new dataset that holds the extended trajectories and leaves the source dataset unchanged. Only the next user turn is generated; no completion is generated for it.

The SDK examples assume you have configured the Python SDK and set ADAPTION_DATASET_ID to an imported dataset whose import has finished:

import os

from adaption import Adaption

client = Adaption()
dataset_id = os.environ["ADAPTION_DATASET_ID"]

How it works

Map the conversation. Point column_mapping at a chat column, or at a prompt and completion pair. See Map the conversation.
Estimate (optional). Price the exact request with estimate=True. Nothing is created or charged.
Extend. Call datasets.extend_trajectories to create the new dataset and start the run.
Wait and download. Wait for the new dataset, then download the rows.

Map the conversation

The trajectory is the conversation up to and including the assistant’s final reply, and the new user turn is generated after that reply. A row whose conversation has no user turn, or does not end on a non-empty assistant reply, is skipped and left out of the new dataset. Provide the trajectory with one of two mappings:

Mapping	Use when
`chat`	One column holds each row’s whole conversation as a list of turns, ending on the assistant’s reply. Send `chat` alone.
`prompt` + `completion`	`prompt` holds the conversation up to and including the user’s last turn. `completion` holds the assistant’s final reply as plain text. Send both.

A list of turns is an array of {"role": ..., "content": ...} objects, such as a column of arrays in JSONL or Parquet, or a JSON string encoding that array. prompt also accepts a plain string, read as a single user turn, so a single-turn prompt and completion dataset works as is:

{"prompt": [{"role": "user", "content": "I'm planning a trip to Kyoto in April."}, {"role": "assistant", "content": "April is cherry blossom season..."}, {"role": "user", "content": "Which temples should I visit?"}], "response": "Kiyomizu-dera and Fushimi Inari..."}
{"prompt": "What's a good first houseplant?", "response": "A pothos is hard to kill..."}

For a chat dataset:

response = client.datasets.extend_trajectories(
    dataset_id,
    column_mapping={"chat": "messages"},
)

For a prompt and completion dataset:

response = client.datasets.extend_trajectories(
    dataset_id,
    column_mapping={"prompt": "prompt", "completion": "response"},
)

Extend the trajectories

Set estimate=True to validate the mapping and price the request before it runs:

estimate = client.datasets.extend_trajectories(
    dataset_id,
    column_mapping={"chat": "messages"},
    estimate=True,
)
print(f"Estimated credits: {estimate.estimated_credits_consumed}")
print(f"Estimated minutes: {estimate.estimated_minutes}")

Then start the run. Use job_specification to cap the number of rows processed or to make retries safe. Row counts are subject to your plan’s per-launch row limit:

from uuid import uuid4

response = client.datasets.extend_trajectories(
    dataset_id,
    column_mapping={"chat": "messages"},
    job_specification={
        "max_rows": 1_000,
        "idempotency_key": str(uuid4()),
    },
)
print(response.dataset_id, response.run_id)

The response carries the new dataset’s dataset_id and the run_id of its run; both are None for an estimate. Use the new dataset_id, not the source dataset’s, for status and download:

from adaption import DatasetTimeout

try:
    final = client.datasets.wait_for_completion(response.dataset_id)
except DatasetTimeout as exc:
    print(exc.resource_id, exc.last_status)
    raise

print(final.status, final.row_count)

wait_for_completion backs off from 2 to 30 seconds and has a one-hour default timeout. A timeout stops waiting, not the run. Call it again with the same dataset ID to resume tracking.

The new dataset can have fewer rows than the source. Rows that are skipped or can’t be extended are left out of it.

Extend an adapted dataset

If the dataset has finished an Adaptive Data run, the extension builds on its adapted dataset instead of the original uploaded file. column_mapping then names columns of the adapted dataset, so you can extend the adapted replies with enhanced_completion:

response = client.datasets.extend_trajectories(
    dataset_id,
    column_mapping={
        "prompt": "prompt",
        "completion": "enhanced_completion",
    },
)

Map the original completion column instead, such as response, to extend the replies as they were uploaded.

Only the prompt + completion mapping applies to an adapted dataset. Adaptation splits a chat column into original_prompt and original_completion, so map original_prompt as prompt instead of sending chat.

Download the extended dataset

When you download the new dataset, each row carries its conversation plus an extended_prompt column holding the generated next user turn. The conversation columns follow the same naming as an adaptation run’s download:

Mapping	Conversation columns
`chat`	`original_prompt` holds the conversation up to the user’s last turn, and `original_completion` the final reply.
`prompt` + `completion`	Your own column names, such as `prompt` and `response`.

An extension of an adapted dataset also keeps the adapted columns, such as enhanced_completion, next to extended_prompt.

client.datasets.download(response.dataset_id, file_format="jsonl").write_to_file(
    "extended.jsonl"
)

See the extend_trajectories reference for the current request and response fields.