Agent Rollout Ingestion

1
import data_designer.config as dd
2
from data_designer.interface import DataDesigner
3
4
data_designer = DataDesigner()
5
config_builder = dd.DataDesignerConfigBuilder(
6
    model_configs=[
7
        dd.ModelConfig(
8
            alias="trace-writer",
9
            model="nvidia/nemotron-3-nano-30b-a3b",
10
            provider="nvidia",
11
        )
12
    ]
13
)
14
15
config_builder.with_seed_dataset(
16
    dd.AgentRolloutSeedSource(
17
        format=dd.AgentRolloutFormat.CLAUDE_CODE,
18
    )
19
)
20
21
config_builder.add_column(
22
    dd.ExpressionColumnConfig(
23
        name="sampled_turn",
24
        expr="{{ messages | random }}",
25
    )
26
)
27
28
config_builder.add_column(
29
    dd.LLMTextColumnConfig(
30
        name="turn_summary",
31
        model_alias="trace-writer",
32
        prompt="""\
33
Summarize this randomly sampled rollout turn in one sentence.
34
The turn may come from the user, assistant, or a tool result.
35
36
Trace: {{ trace_id }}
37
Turn:
38
{{ sampled_turn }}
39
""",
40
    )
41
)
42
43
preview = data_designer.preview(config_builder, num_records=3)
44
preview.display_sample_record()

1
import data_designer.config as dd
2
from data_designer.interface import DataDesigner
3
from pydantic import BaseModel, Field
4
from typing import Literal
5
6
7
@dd.custom_column_generator(
8
    required_columns=["messages"],
9
    side_effect_columns=["tool_call", "tool_response", "tool_name"],
10
)
11
def explode_tool_interactions(row: dict) -> list[dict]:
12
    rows = []
13
    tool_calls_by_id = {}
14
    context_messages = []
15
16
    for message in row["messages"]:
17
        context_messages.append(message)
18
19
        for tool_call in message.get("tool_calls") or []:
20
            tool_call_id = tool_call.get("id")
21
            if tool_call_id:
22
                tool_calls_by_id[tool_call_id] = tool_call
23
24
        if message.get("role") != "tool":
25
            continue
26
27
        tool_call = tool_calls_by_id.get(
28
            message.get("tool_call_id"),
29
            {
30
                "id": message.get("tool_call_id"),
31
                "type": "function",
32
                "function": {"name": "unknown", "arguments": "{}"},
33
            },
34
        )
35
        tool_name = tool_call.get("function", {}).get("name", "unknown")
36
37
        rows.append(
38
            {
39
                **row,
40
                "tool_interaction_context": list(context_messages),
41
                "tool_call": tool_call,
42
                "tool_response": message,
43
                "tool_name": tool_name,
44
            }
45
        )
46
47
    return rows
48
49
50
class ToolInteractionAnalysis(BaseModel):
51
    outcome: Literal["success", "failure", "unclear"] = Field(
52
        description="Whether the tool interaction appears to have succeeded, failed, or is ambiguous."
53
    )
54
    summary: str = Field(
55
        description="One or two sentences summarizing what the tool was asked to do and what the response indicates."
56
    )
57
58
59
data_designer = DataDesigner()
60
config_builder = dd.DataDesignerConfigBuilder(
61
    model_configs=[
62
        dd.ModelConfig(
63
            alias="tool-analyst",
64
            model="nvidia/nemotron-3-nano-30b-a3b",
65
            provider="nvidia",
66
        )
67
    ]
68
)
69
70
config_builder.with_seed_dataset(
71
    dd.AgentRolloutSeedSource(
72
        format=dd.AgentRolloutFormat.CLAUDE_CODE,
73
    )
74
)
75
76
config_builder.add_column(
77
    dd.CustomColumnConfig(
78
        name="tool_interaction_context",
79
        generator_function=explode_tool_interactions,
80
        allow_resize=True,
81
    )
82
)
83
84
config_builder.add_column(
85
    dd.LLMStructuredColumnConfig(
86
        name="tool_interaction_analysis",
87
        model_alias="tool-analyst",
88
        output_format=ToolInteractionAnalysis,
89
        prompt="""\
90
You are analyzing one tool interaction from an imported agent rollout.
91
92
Context up to the tool response:
93
{{ tool_interaction_context }}
94
95
Tool name: {{ tool_name }}
96
97
Tool call:
98
{{ tool_call }}
99
100
Tool response:
101
{{ tool_response }}
102
103
Decide whether this interaction is a success, failure, or unclear outcome.
104
Then summarize what the tool was asked to do and what happened.
105
Base your answer on the tool call arguments, the tool response, and the immediate context.
106
""",
107
    )
108
)
109
110
preview = data_designer.preview(config_builder, num_records=5)
111
preview.display_sample_record()