eval-protocol 0.3.29__tar.gz → 0.3.30__tar.gz

{eval_protocol-0.3.29/eval_protocol.egg-info → eval_protocol-0.3.30}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: eval-protocol
-Version: 0.3.29
+Version: 0.3.30
 Summary: The official Python SDK for Eval Protocol (EP.) EP is an open protocol that standardizes how developers author evals for large language model (LLM) applications.
 Author-email: Fireworks AI <info@fireworks.ai>
 License-Expression: MIT
@@ -40,6 +40,7 @@ Requires-Dist: addict>=2.4.0
 Requires-Dist: deepdiff>=6.0.0
 Requires-Dist: websockets>=15.0.1
 Requires-Dist: fastapi>=0.116.1
+Requires-Dist: zstandard>=0.19.0
 Provides-Extra: dev
 Requires-Dist: build; extra == "dev"
 Requires-Dist: twine; extra == "dev"

{eval_protocol-0.3.29 → eval_protocol-0.3.30}/eval_protocol/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-04-28T20:17:42-0700",
+ "date": "2026-05-29T16:09:24-0700",
  "dirty": false,
  "error": null,
- "full-revisionid": "6b9bea9979400c7e0d2eec9f4f167a98d4a2a057",
- "version": "0.3.29"
+ "full-revisionid": "1bd5447a3afbca3b71e0f0d205ed7cff6c3afe5d",
+ "version": "0.3.30"
 }
 '''  # END VERSION_JSON
 

{eval_protocol-0.3.29 → eval_protocol-0.3.30}/eval_protocol/adapters/fireworks_tracing.py

@@ -16,6 +16,8 @@ import os
 
 from eval_protocol.models import EvaluationRow, InputMetadata, ExecutionMetadata, Message
 from .base import BaseAdapter
+from .lp_deserializer import decompress_and_parse_lp
+from .r3_deserializer import decompress_and_parse_r3
 from .utils import extract_messages_from_data
 from ..common_utils import get_user_agent
 
@@ -100,13 +102,53 @@ def convert_trace_dict_to_evaluation_row(
                 ):
                     break  # Break early if we've found all the metadata we need
 
+        # Extract router replay payloads when present
+        payloads = trace.get("payloads")
+        if isinstance(payloads, dict):
+            router_replay = payloads.get("router_replay")
+            if isinstance(router_replay, dict) and router_replay.get("data"):
+                try:
+                    matrices, r3_meta = decompress_and_parse_r3(router_replay["data"])
+                    if execution_metadata.extra is None:
+                        execution_metadata.extra = {}
+                    execution_metadata.extra["routing_matrices"] = matrices
+                    execution_metadata.extra["routing_metadata"] = r3_meta
+                except Exception as e:
+                    logger.warning("Failed to decompress R3 payload for trace %s: %s", trace.get("id"), e)
+
+            logprobs_payload = payloads.get("logprobs")
+            if isinstance(logprobs_payload, dict) and logprobs_payload.get("data"):
+                try:
+                    logprobs, token_ids, lp_meta = decompress_and_parse_lp(logprobs_payload["data"])
+                    if execution_metadata.extra is None:
+                        execution_metadata.extra = {}
+                    execution_metadata.extra["completion_logprobs"] = logprobs
+                    if token_ids is not None:
+                        execution_metadata.extra["completion_token_ids"] = token_ids
+                    execution_metadata.extra["logprobs_metadata"] = lp_meta
+
+                    for i in range(len(messages) - 1, -1, -1):
+                        if messages[i].role == "assistant":
+                            content_entries = [{"logprob": lp} for lp in logprobs]
+                            if token_ids is not None:
+                                for entry, tid in zip(content_entries, token_ids):
+                                    entry["token_id"] = tid
+                            messages[i].logprobs = {"content": content_entries}
+                            break
+                except Exception as e:
+                    logger.warning(
+                        "Failed to decompress logprobs payload for trace %s: %s",
+                        trace.get("id"),
+                        e,
+                    )
+
         return EvaluationRow(
             messages=messages,
             tools=tools,
             input_metadata=InputMetadata(
                 row_id=row_id,
                 session_data={
-                    "langfuse_trace_id": trace.get("id"),  # Store the trace ID here
+                    "langfuse_trace_id": trace.get("id"),
                 },
             ),
             execution_metadata=execution_metadata,
@@ -375,6 +417,37 @@ class FireworksTracingAdapter(BaseAdapter):
             )
         return results
 
+    async def async_get_status(self, session: aiohttp.ClientSession, rollout_id: str) -> Optional[Dict[str, Any]]:
+        """Fetch rollout status from the lightweight /status endpoint.
+
+        Returns the parsed JSON response or None if the status is not yet available.
+        Response shape: {"rollout_id": "...", "status": {"code": ...} | null, "extras": {...} | null}
+        """
+        headers = {
+            "Authorization": f"Bearer {self._get_api_key()}",
+            "User-Agent": get_user_agent(),
+        }
+        params: Dict[str, Any] = {"rollout_id": rollout_id}
+        timeout = aiohttp.ClientTimeout(total=self.timeout)
+
+        urls_to_try = [f"{self.base_url}/v1/status", f"{self.base_url}/status"]
+        last_error: Optional[str] = None
+        for url in urls_to_try:
+            try:
+                async with session.get(url, params=params, headers=headers, timeout=timeout) as resp:
+                    if resp.status == 404:
+                        last_error = f"404 for {url}"
+                        continue
+                    resp.raise_for_status()
+                    return (await resp.json(content_type=None)) or {}
+            except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError) as e:
+                last_error = str(e)
+                continue
+
+        if last_error:
+            logger.error("Failed to fetch status from Fireworks (tried %s): %s", urls_to_try, last_error)
+        return None
+
     def get_evaluation_rows(
         self,
         tags: List[str],
@@ -395,6 +468,7 @@ class FireworksTracingAdapter(BaseAdapter):
         max_retries: int = 3,
         span_name: Optional[str] = None,
         converter: Optional[TraceDictConverter] = None,
+        include_payloads: bool = False,
     ) -> List[EvaluationRow]:
         """Pull traces from Langfuse via proxy and convert to EvaluationRow format.
 
@@ -418,6 +492,8 @@ class FireworksTracingAdapter(BaseAdapter):
             max_retries: Max retry attempts used by proxy (default: 3)
             converter: Optional custom converter implementing TraceDictConverter protocol.
                 If provided, this will be used instead of the default conversion logic.
+            include_payloads: If True, request payload data (e.g., router replay)
+                from the gateway and decompress it into the returned EvaluationRows.
 
         Returns:
             List[EvaluationRow]: Converted evaluation rows
@@ -448,6 +524,7 @@ class FireworksTracingAdapter(BaseAdapter):
             "to_timestamp": to_timestamp.isoformat() if to_timestamp else None,
             "sleep_between_gets": sleep_between_gets,
             "max_retries": max_retries,
+            "include_payloads": include_payloads if include_payloads else None,
         }
 
         # Remove None values

eval_protocol-0.3.30/eval_protocol/adapters/lp_deserializer.py

@@ -0,0 +1,109 @@
+"""LP/v1 binary deserializer for per-token logprobs payloads.
+
+Implements the inverse of the tracing gateway's ``logprobs_serializer.serialize_logprobs``.
+See that module for the full header specification.
+"""
+
+from __future__ import annotations
+
+import base64
+import struct
+from typing import Any, Dict, List, Optional, Tuple
+
+import zstandard as zstd
+
+MAGIC = b"LP01"
+HEADER_VERSION = 1
+MISSING_TOKEN_ID = -1
+ENTRY_FORMAT = "<if"
+ENTRY_SIZE = struct.calcsize(ENTRY_FORMAT)  # 8 bytes
+HEADER_FORMAT = "<4sBBHIIQ"
+HEADER_SIZE = struct.calcsize(HEADER_FORMAT)  # 24 bytes
+
+
+def _parse_header(raw: bytes) -> Dict[str, Any]:
+    if len(raw) < HEADER_SIZE:
+        raise ValueError(f"Payload too short for lp/v1 header: {len(raw)} < {HEADER_SIZE}")
+
+    (
+        magic,
+        version,
+        flags,
+        reserved_u16,
+        token_count,
+        body_byte_length,
+        reserved_u64,
+    ) = struct.unpack(HEADER_FORMAT, raw[:HEADER_SIZE])
+
+    if magic != MAGIC:
+        raise ValueError(f"Bad LP/v1 magic: {magic!r}")
+    if version != HEADER_VERSION:
+        raise ValueError(f"Unsupported lp/v1 header version: {version}")
+
+    return {
+        "flags": flags,
+        "reserved_u16": reserved_u16,
+        "token_count": token_count,
+        "body_byte_length": body_byte_length,
+        "reserved_u64": reserved_u64,
+    }
+
+
+def parse_logprobs(raw: bytes) -> Tuple[List[float], Optional[List[int]], Dict[str, Any]]:
+    """Parse uncompressed LP/v1 bytes into logprobs, optional token ids, and metadata."""
+    header = _parse_header(raw)
+    token_count = header["token_count"]
+    body_byte_length = header["body_byte_length"]
+
+    if token_count == 0:
+        raise ValueError("LP/v1 token_count must be > 0")
+    if body_byte_length != token_count * ENTRY_SIZE:
+        raise ValueError(
+            f"body_byte_length ({body_byte_length}) != token_count * {ENTRY_SIZE} "
+            f"({token_count * ENTRY_SIZE})"
+        )
+
+    expected_len = HEADER_SIZE + body_byte_length
+    if len(raw) != expected_len:
+        raise ValueError(f"LP/v1 payload length mismatch: {len(raw)} != {expected_len}")
+
+    logprobs: List[float] = []
+    token_ids: List[int] = []
+    all_token_ids_valid = True
+    offset = HEADER_SIZE
+    for _ in range(token_count):
+        wire_id, logprob = struct.unpack(ENTRY_FORMAT, raw[offset : offset + ENTRY_SIZE])
+        offset += ENTRY_SIZE
+        logprobs.append(logprob)
+        if wire_id == MISSING_TOKEN_ID:
+            all_token_ids_valid = False
+            token_ids.append(wire_id)
+        else:
+            token_ids.append(wire_id)
+
+    metadata: Dict[str, Any] = {
+        "scope": "completion_only",
+        "completion_token_count": token_count,
+        "all_token_ids_valid": all_token_ids_valid,
+    }
+    header.update(metadata)
+    ids_out: Optional[List[int]] = token_ids if all_token_ids_valid else None
+    return logprobs, ids_out, header
+
+
+def decompress_and_parse_lp(data_b64: str) -> Tuple[List[float], Optional[List[int]], Dict[str, Any]]:
+    """Decompress and unpack an LP/v1 payload into completion logprobs and token ids.
+
+    Args:
+        data_b64: Base64-encoded zstd-compressed LP binary blob from
+            ``payloads.logprobs.data``.
+
+    Returns:
+        ``(logprobs, token_ids, metadata)`` where ``logprobs`` is per-completion-token
+        scalars, ``token_ids`` is ``None`` if any wire id was ``MISSING_TOKEN_ID``,
+        and ``metadata`` includes ``all_token_ids_valid`` and ``completion_token_count``.
+    """
+    compressed = base64.b64decode(data_b64)
+    decompressor = zstd.ZstdDecompressor()
+    raw = decompressor.decompress(compressed)
+    return parse_logprobs(raw)

eval_protocol-0.3.30/eval_protocol/adapters/r3_deserializer.py

@@ -0,0 +1,187 @@
+"""R3/v1 binary deserializer for router-replay payloads.
+
+Implements the inverse of the packed binary format produced by the tracing
+gateway's ``r3_serializer.serialize_r3``.  See that module for the full
+header specification.
+
+The main entry point is :func:`decompress_and_parse_r3`, which accepts the
+base64-encoded compressed blob returned by the gateway's
+``/v1/traces/pointwise?include_payloads=true`` endpoint and produces
+per-token routing matrices in the same ``List[Optional[str]]`` format used
+by the direct inference path (``DeploymentSampler.sample_with_tokens()``).
+"""
+
+from __future__ import annotations
+
+import base64
+import struct
+from enum import IntEnum
+from typing import Any, Dict, List, Optional, Tuple
+
+import zstandard as zstd
+
+MAGIC = b"R3V1"
+HEADER_FORMAT = "<4sBBBBIIIIQ"
+HEADER_SIZE = struct.calcsize(HEADER_FORMAT)  # 32 bytes
+BITS_PER_BYTE = 8
+
+
+class _SelectorMode(IntEnum):
+    ALL = 0
+    SUFFIX = 1
+    BITMAP = 2
+
+
+class _RoutingDtype(IntEnum):
+    UINT8 = 1
+    UINT16 = 2
+
+
+_SELECTOR_MODE_NAMES = {v: v.name.lower() for v in _SelectorMode}
+_ROUTING_DTYPE_NAMES = {v: v.name.lower() for v in _RoutingDtype}
+
+
+def _parse_header(raw: bytes) -> Dict[str, Any]:
+    if len(raw) < HEADER_SIZE:
+        raise ValueError(
+            f"Payload too short for r3/v1 header: {len(raw)} < {HEADER_SIZE}"
+        )
+
+    (
+        magic,
+        version,
+        selector_mode,
+        routing_dtype,
+        flags,
+        total_token_count,
+        replayed_token_count,
+        replay_start_token,
+        selector_byte_length,
+        matrix_byte_length,
+    ) = struct.unpack(HEADER_FORMAT, raw[:HEADER_SIZE])
+
+    if magic != MAGIC:
+        raise ValueError(f"Bad R3 magic: {magic!r}")
+    if version != 1:
+        raise ValueError(f"Unsupported R3 header version: {version}")
+
+    return {
+        "selector_mode": selector_mode,
+        "routing_dtype": routing_dtype,
+        "flags": flags,
+        "total_token_count": total_token_count,
+        "replayed_token_count": replayed_token_count,
+        "replay_start_token": replay_start_token,
+        "selector_byte_length": selector_byte_length,
+        "matrix_byte_length": matrix_byte_length,
+    }
+
+
+def _read_bitmap_positions(
+    selector_bytes: bytes, total_token_count: int
+) -> List[int]:
+    """Return sorted token indices where the bitmap bit is set."""
+    positions: List[int] = []
+    for i in range(total_token_count):
+        byte_idx = i // BITS_PER_BYTE
+        bit_idx = i % BITS_PER_BYTE
+        if byte_idx < len(selector_bytes) and (selector_bytes[byte_idx] >> bit_idx) & 1:
+            positions.append(i)
+    return positions
+
+
+def decompress_and_parse_r3(
+    data_b64: str,
+) -> Tuple[List[Optional[str]], Dict[str, Any]]:
+    """Decompress and unpack an R3/v1 payload into per-token routing matrices.
+
+    Args:
+        data_b64: Base64-encoded zstd-compressed R3 binary blob, as returned
+            by the tracing gateway in ``payloads.router_replay.data``.
+
+    Returns:
+        A tuple of ``(routing_matrices, metadata)`` where:
+
+        - ``routing_matrices`` is a ``List[Optional[str]]`` of length
+          ``total_token_count``.  Each present position contains a
+          base64-encoded routing matrix (matching the format returned by
+          the direct inference path); absent positions are ``None``.
+        - ``metadata`` is a dict with keys ``routing_dtype``,
+          ``selector_mode``, ``total_token_count``, ``replayed_token_count``,
+          ``replay_start_token``.
+    """
+    compressed = base64.b64decode(data_b64)
+
+    # ZstdCompressor.compress() embeds the uncompressed size in the frame
+    # header by default, so the library can auto-allocate the output buffer.
+    decompressor = zstd.ZstdDecompressor()
+    raw = decompressor.decompress(compressed)
+
+    header = _parse_header(raw)
+
+    selector_mode = header["selector_mode"]
+    routing_dtype = header["routing_dtype"]
+    total_token_count = header["total_token_count"]
+    replayed_token_count = header["replayed_token_count"]
+    replay_start_token = header["replay_start_token"]
+    selector_byte_length = header["selector_byte_length"]
+    matrix_byte_length = header["matrix_byte_length"]
+
+    metadata: Dict[str, Any] = {
+        "routing_dtype": _ROUTING_DTYPE_NAMES.get(routing_dtype, str(routing_dtype)),
+        "selector_mode": _SELECTOR_MODE_NAMES.get(selector_mode, str(selector_mode)),
+        "total_token_count": total_token_count,
+        "replayed_token_count": replayed_token_count,
+        "replay_start_token": replay_start_token,
+    }
+
+    if replayed_token_count == 0:
+        return [None] * total_token_count, metadata
+
+    # Per-token matrix byte size is implicit in the payload: all replayed
+    # tokens share the same matrix length, so we can recover it from the
+    # matrix section total length divided by the replayed-token count.
+    if matrix_byte_length % replayed_token_count != 0:
+        raise ValueError(
+            f"matrix_byte_length ({matrix_byte_length}) is not a multiple of "
+            f"replayed_token_count ({replayed_token_count}); cannot split "
+            "into per-token matrices"
+        )
+    matrix_elem_size = matrix_byte_length // replayed_token_count
+
+    body = raw[HEADER_SIZE:]
+    expected_body_length = selector_byte_length + matrix_byte_length
+    if len(body) < expected_body_length:
+        raise ValueError(
+            f"Payload body too short for selector and matrix sections: "
+            f"{len(body)} < {expected_body_length}"
+        )
+
+    selector_bytes = body[:selector_byte_length]
+    matrix_bytes = body[selector_byte_length : selector_byte_length + matrix_byte_length]
+
+    if selector_mode == _SelectorMode.ALL:
+        replayed_positions = list(range(total_token_count))
+    elif selector_mode == _SelectorMode.SUFFIX:
+        replayed_positions = list(
+            range(replay_start_token, replay_start_token + replayed_token_count)
+        )
+    elif selector_mode == _SelectorMode.BITMAP:
+        replayed_positions = _read_bitmap_positions(selector_bytes, total_token_count)
+    else:
+        raise ValueError(f"Unknown selector_mode: {selector_mode}")
+
+    if len(replayed_positions) != replayed_token_count:
+        raise ValueError(
+            f"Selector produced {len(replayed_positions)} replayed positions, "
+            f"but header replayed_token_count is {replayed_token_count}"
+        )
+
+    # Split matrix bytes into per-token chunks and base64-encode each one
+    matrices: List[Optional[str]] = [None] * total_token_count
+    for idx, pos in enumerate(replayed_positions):
+        start = idx * matrix_elem_size
+        end = start + matrix_elem_size
+        matrices[pos] = base64.b64encode(matrix_bytes[start:end]).decode("ascii")
+
+    return matrices, metadata

{eval_protocol-0.3.29 → eval_protocol-0.3.30}/eval_protocol/pytest/remote_rollout_processor.py

@@ -35,11 +35,13 @@ class RemoteRolloutProcessor(RolloutProcessor):
         model_base_url: str = "https://tracing.fireworks.ai",
         poll_interval: float = 1.0,
         timeout_seconds: float = 120.0,
+        include_payloads: bool = False,
     ):
         # Prefer constructor-provided configuration. These can be overridden via
         # config.kwargs at call time for backward compatibility.
         self._remote_base_url = remote_base_url
         self._model_base_url = model_base_url
+        self._include_payloads = include_payloads
         if os.getenv("EP_REMOTE_ROLLOUT_PROCESSOR_BASE_URL"):
             self._remote_base_url = os.getenv("EP_REMOTE_ROLLOUT_PROCESSOR_BASE_URL")
         _ep_model_base_url = os.getenv("EP_MODEL_BASE_URL")
@@ -122,45 +124,46 @@ class RemoteRolloutProcessor(RolloutProcessor):
 
             while time.time() < deadline:
                 session = self._get_or_create_session()
-                completed_logs = await self._tracing_adapter.async_search_logs(
-                    session, tags=[f"rollout_id:{row.execution_metadata.rollout_id}"]
+                status_result = await self._tracing_adapter.async_get_status(
+                    session,
+                    rollout_id=row.execution_metadata.rollout_id,
                 )
-                # Filter for logs that actually have status information
-                status_logs = []
-                for log in completed_logs:
-                    status_dict = log.get("status")
-                    if status_dict and isinstance(status_dict, dict) and "code" in status_dict:
-                        status_logs.append(log)
-
-                if status_logs:
-                    if len(status_logs) > 1:
-                        logger.warning(
-                            "Found %s status logs for rollout %s; expected at most 1. Using the first one: %s",
-                            len(status_logs),
-                            row.execution_metadata.rollout_id,
-                            status_logs[0],
-                        )
-                    # Use the first log with status information
-                    status_log = status_logs[0]
-                    status_dict = status_log.get("status")
-                    raw_extras = status_log.get("extras") or {}
-                    status_extras = {
-                        k: v for k, v in raw_extras.items() if k not in ("logger_name", "level", "timestamp")
-                    }
+                status = (status_result or {}).get("status")
+                if isinstance(status, dict) and "code" in status:
+                    status_code = status["code"]
+                    if status_code == Status.Code.RUNNING:
+                        await asyncio.sleep(poll_interval)
+                        continue
 
                     logger.info(
-                        f"Found status log for rollout {row.execution_metadata.rollout_id}: {status_log.get('message', '')}"
+                        "Found status for rollout %s with code %s",
+                        row.execution_metadata.rollout_id,
+                        status_code,
                     )
 
-                    status_code = status_dict.get("code")
-                    status_message = status_dict.get("message", "")
-                    status_details = status_dict.get("details", [])
-
-                    logger.info(
-                        f"Found Fireworks log for rollout {row.execution_metadata.rollout_id} with status code {status_code}"
+                    # /status only returns the code; backfill message/details/extras from Logs once.
+                    status_message: str = ""
+                    status_details: list = []
+                    status_extras: dict = {}
+                    completed_logs = await self._tracing_adapter.async_search_logs(
+                        session, tags=[f"rollout_id:{row.execution_metadata.rollout_id}"]
                     )
+                    # Pick the log row whose status code matches the terminal
+                    # code from /status, so intermediate RUNNING checkpoints
+                    # don't poison the backfill.
+                    for log in completed_logs:
+                        sd = log.get("status")
+                        if isinstance(sd, dict) and sd.get("code") == status_code:
+                            status_message = sd.get("message", "") or ""
+                            status_details = sd.get("details", []) or []
+                            raw_extras = log.get("extras") or {}
+                            status_extras = {
+                                k: v
+                                for k, v in raw_extras.items()
+                                if k not in ("logger_name", "level", "timestamp")
+                            }
+                            break
 
-                    # Create and raise exception if appropriate, preserving original message
                     exception = exception_for_status_code(status_code, status_message)
                     if exception is not None:
                         raise exception
@@ -171,10 +174,11 @@ class RemoteRolloutProcessor(RolloutProcessor):
                         details=status_details,
                     )
 
-                    if row.execution_metadata.extra:
-                        row.execution_metadata.extra.update(status_extras)
-                    else:
-                        row.execution_metadata.extra = status_extras
+                    if status_extras:
+                        if row.execution_metadata.extra:
+                            row.execution_metadata.extra.update(status_extras)
+                        else:
+                            row.execution_metadata.extra = status_extras
 
                     logger.info("Stopping polling for rollout %s", row.execution_metadata.rollout_id)
                     break
@@ -192,7 +196,10 @@ class RemoteRolloutProcessor(RolloutProcessor):
             row.execution_metadata.rollout_duration_seconds = time.perf_counter() - start_time
 
             def _update_with_trace() -> None:
-                return update_row_with_remote_trace(row, default_fireworks_output_data_loader, model_base_url)
+                return update_row_with_remote_trace(
+                    row, default_fireworks_output_data_loader, model_base_url,
+                    include_payloads=self._include_payloads,
+                )
 
             await asyncio.to_thread(_update_with_trace)  # Update row with remote trace in-place
             return row

{eval_protocol-0.3.29 → eval_protocol-0.3.30}/eval_protocol/pytest/tracing_utils.py

@@ -22,9 +22,61 @@ def default_fireworks_output_data_loader(config: DataLoaderConfig) -> DynamicDat
         # Use EP_REMOTE_API_KEY for fetching remote traces, falling back to FIREWORKS_API_KEY
         api_key = os.environ.get("EP_REMOTE_API_KEY") or os.environ.get("FIREWORKS_API_KEY")
         adapter = FireworksTracingAdapter(base_url=base_url, api_key=api_key)
-        return adapter.get_evaluation_rows(tags=[f"rollout_id:{config.rollout_id}"], max_retries=5)
+        return adapter.get_evaluation_rows(
+            tags=[f"rollout_id:{config.rollout_id}"],
+            max_retries=5,
+            include_payloads=config.include_payloads,
+        )
 
-    return DynamicDataLoader(generators=[fetch_traces], preprocess_fn=filter_longest_conversation)
+    def preprocess_traces(rows: List[EvaluationRow]) -> List[EvaluationRow]:
+        filtered_rows = filter_longest_conversation(rows)
+        if config.include_payloads and filtered_rows:
+            _merge_payloads_into_longest_row(filtered_rows[0], rows)
+        return filtered_rows
+
+    return DynamicDataLoader(generators=[fetch_traces], preprocess_fn=preprocess_traces)
+
+
+def _merge_payloads_into_longest_row(longest_row: EvaluationRow, rows: List[EvaluationRow]) -> None:
+    """
+    Preserve per-turn payload-derived metadata after selecting the longest trace row.
+
+    Each trace row carries payloads for its final assistant turn. The longest row
+    keeps the full conversation, while its top-level execution metadata remains
+    the payload metadata for the final completion for backward compatibility.
+    """
+    target_assistants = longest_row.get_assistant_messages()
+    assistant_turn_payloads = []
+
+    for row in sorted(rows, key=lambda item: len(item.messages)):
+        source = row.last_assistant_message()
+        source_turn_index = len(row.get_assistant_messages()) - 1
+        if source_turn_index < 0 or source_turn_index >= len(target_assistants):
+            continue
+
+        if source and source.logprobs and not target_assistants[source_turn_index].logprobs:
+            target_assistants[source_turn_index].logprobs = source.logprobs
+
+        extra = row.execution_metadata.extra or {}
+        turn_payload = {
+            key: extra[key]
+            for key in (
+                "completion_logprobs",
+                "completion_token_ids",
+                "logprobs_metadata",
+                "routing_matrices",
+                "routing_metadata",
+            )
+            if key in extra
+        }
+        if turn_payload:
+            turn_payload["assistant_turn_index"] = source_turn_index
+            assistant_turn_payloads.append(turn_payload)
+
+    if assistant_turn_payloads:
+        if longest_row.execution_metadata.extra is None:
+            longest_row.execution_metadata.extra = {}
+        longest_row.execution_metadata.extra["assistant_turn_payloads"] = assistant_turn_payloads
 
 
 def build_fireworks_tracing_url(
@@ -99,7 +151,7 @@ def build_init_request(
     if not completion_params_dict.get("model"):
         raise ValueError("Model must be provided in completion_params")
 
-    # Extract base_url from completion_params
+    # Extract base_url from completion_params for tracing-gateway URL encoding
     completion_params_base_url: Optional[str] = completion_params_dict.get("base_url")
 
     # Strip non-OpenAI fields from messages
@@ -129,7 +181,7 @@ def build_init_request(
 
     # Build final model base URL with tracing metadata
     final_model_base_url = model_base_url
-    if model_base_url and ("tracing.fireworks.ai" in model_base_url or model_base_url.startswith("http://localhost")):
+    if model_base_url and ("tracing.fireworks.ai" in model_base_url or model_base_url.startswith("http://localhost") or "litellm-gateway" in model_base_url):
         final_model_base_url = build_fireworks_tracing_url(model_base_url, meta, completion_params_base_url)
 
     # Extract API key from environment or completion_params
@@ -148,13 +200,20 @@ def build_init_request(
 
 
 def update_row_with_remote_trace(
-    row: EvaluationRow, output_data_loader: Callable[[DataLoaderConfig], DynamicDataLoader], model_base_url: str
+    row: EvaluationRow,
+    output_data_loader: Callable[[DataLoaderConfig], DynamicDataLoader],
+    model_base_url: str,
+    include_payloads: bool = False,
 ) -> None:
     """Update row with remote trace data using output_data_loader (shared logic)."""
     if not row.execution_metadata.rollout_id:
         return None
 
-    loader_config = DataLoaderConfig(rollout_id=row.execution_metadata.rollout_id, model_base_url=model_base_url)
+    loader_config = DataLoaderConfig(
+        rollout_id=row.execution_metadata.rollout_id,
+        model_base_url=model_base_url,
+        include_payloads=include_payloads,
+    )
     data_loader = output_data_loader(loader_config)
     results = data_loader.load()
     output_rows: List[EvaluationRow] = [r for result in results for r in result.rows]