Trajectory Logging¶

class TrajectoryEventType(str, Enum):
    # Run lifecycle
    RUN_START = "run_start"
    RUN_END = "run_end"

    # REPL iterations
    ITERATION_START = "iteration_start"
    ITERATION_REASONING = "iteration_reasoning"
    ITERATION_CODE = "iteration_code"
    ITERATION_OUTPUT = "iteration_output"
    ITERATION_END = "iteration_end"

    # LLM calls
    LLM_REQUEST = "llm_request"
    LLM_RESPONSE = "llm_response"

    # Sub-LLM (llm_query from code)
    SUB_LLM_REQUEST = "sub_llm_request"
    SUB_LLM_RESPONSE = "sub_llm_response"

    # Child agents
    CHILD_SPAWN = "child_spawn"
    CHILD_RESULT = "child_result"

    # Termination
    FINAL_DETECTED = "final_detected"

    # Context
    CONTEXT_LOAD = "context_load"
    CONTEXT_UPDATE = "context_update"

    # Memory
    MEMORY_COMPACT = "memory_compact"

    # Errors
    ERROR = "error"

@dataclass
class TrajectoryEvent:
    event_type: TrajectoryEventType          # Event category
    timestamp: float = field(...)            # Unix timestamp (time.time())
    run_id: str = ""                         # Run identifier
    iteration: int | None = None             # Iteration number
    depth: int = 0                           # Recursion depth (0 = root)
    parent_id: str | None = None             # Parent agent ID

    # Event-specific data
    data: dict[str, Any] = field(...)        # Arbitrary event payload

    # Metrics
    tokens_in: int | None = None             # Input tokens
    tokens_out: int | None = None            # Output tokens
    duration_ms: float | None = None         # Duration in milliseconds

event = TrajectoryEvent(
    event_type=TrajectoryEventType.ITERATION_CODE,
    run_id="run_abc123",
    iteration=2,
    data={"code": "print(len(context))"},
)

d = event.to_dict()
# {"event_type": "iteration_code", "timestamp": 1706400001.5,
#  "run_id": "run_abc123", "iteration": 2,
#  "data": {"code": "print(len(context))"}}

event = TrajectoryEvent.from_dict({
    "event_type": "iteration_output",
    "timestamp": 1706400002.1,
    "run_id": "run_abc123",
    "iteration": 2,
    "data": {"output": "45230"},
    "duration_ms": 15.3,
})

class TrajectoryLogger:
    def __init__(
        self,
        output_path: str | Path,
        run_id: str | None = None,
        metadata: dict[str, Any] | None = None,
    ):

with TrajectoryLogger("traces/run_001.jsonl") as logger:
    logger.log_run_start(task="Analyze data")
    # ... log events ...
    logger.log_run_end(success=True, answer="42")

logger.log_run_start(
    task="Analyze sentiment of customer reviews",
    context_length=45230,
    model="gpt-4o",
)

logger.log_run_end(
    success=True,
    answer="Overall sentiment is positive (87% confidence)",
    total_tokens=8523,
)

logger.log_iteration(
    iteration=1,
    reasoning="First, explore the data structure",
    code='print(f"Length: {len(context)}")',
    output="Length: 45230",
    duration_ms=15.3,
)

# Root LLM call
logger.log_llm_call(
    prompt="Analyze this data...",
    response="The data shows...",
    tokens_in=500,
    tokens_out=200,
    duration_ms=1200,
)

# Sub-LLM call from llm_query()
logger.log_llm_call(
    prompt="Summarize this chunk...",
    response="This chunk discusses...",
    is_sub_llm=True,
)

logger.push_depth("child_agent_001")  # depth becomes 1
# ... log child events ...
logger.pop_depth()                     # depth returns to 0

class TrajectoryViewer:
    def __init__(self, trajectory_path: str | Path):

viewer = TrajectoryViewer("traces/run_001.jsonl")
for iteration_events in viewer.iterations():
    print(f"Iteration with {len(iteration_events)} events")

summary = viewer.summary()

print(viewer.format_tree())

viewer.export_html("traces/run_001.html")

from rlm_code.rlm.trajectory import load_trajectory

viewer = load_trajectory("traces/run_001.jsonl")
print(viewer.summary())

from rlm_code.rlm.trajectory import compare_trajectories

comparison = compare_trajectories([
    "traces/run_001.jsonl",
    "traces/run_002.jsonl",
    "traces/run_003.jsonl",
])

{
    "trajectories": [
        {
            "path": "traces/run_001.jsonl",
            "run_id": "run_001",
            "task": "Analyze sentiment...",
            "success": True,
            "iterations": 3,
            "tokens": 8523,
            "duration_ms": 4500,
        },
        # ... more trajectories ...
    ],
    "comparison": {
        "avg_iterations": 3.7,
        "avg_tokens": 9100,
        "avg_duration_ms": 5200,
        "success_rate": 0.67,  # 2 out of 3 succeeded
    },
}

from rlm_code.rlm.trajectory import TrajectoryLogger, TrajectoryViewer

# --- Logging ---
with TrajectoryLogger("traces/my_run.jsonl", metadata={"model": "gpt-4o"}) as logger:
    logger.log_run_start(task="Analyze data", context_length=50000)

    logger.log_iteration(
        iteration=1,
        reasoning="Explore data structure",
        code='print(type(context), len(context))',
        output="<class 'str'> 50000",
        duration_ms=12,
    )

    logger.log_llm_call(
        prompt="Summarize the first section...",
        response="The first section covers...",
        tokens_in=600,
        tokens_out=150,
        duration_ms=1100,
        is_sub_llm=True,
    )

    logger.log_final("Analysis complete: found 3 key themes")
    logger.log_run_end(success=True, answer="3 key themes", total_tokens=2500)

# --- Viewing ---
viewer = TrajectoryViewer("traces/my_run.jsonl")
print(viewer.format_tree())
viewer.export_html("traces/my_run.html")

summary = viewer.summary()
print(f"Run {summary['run_id']}: {summary['total_iterations']} iterations, "
      f"{summary['total_tokens']} tokens")