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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: eval-protocol
-Version: 0.3.29
+Version: 0.3.31
 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.31}/eval_protocol/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-04-28T20:17:42-0700",
+ "date": "2026-06-19T13:02:26-0700",
  "dirty": false,
  "error": null,
- "full-revisionid": "6b9bea9979400c7e0d2eec9f4f167a98d4a2a057",
- "version": "0.3.29"
+ "full-revisionid": "08e9b85f0315f420f5977c04549b56a406a826d5",
+ "version": "0.3.31"
 }
 '''  # END VERSION_JSON
 

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

@@ -15,6 +15,7 @@ from typing import Any, Dict, List, Optional, Protocol
 import os
 
 from eval_protocol.models import EvaluationRow, InputMetadata, ExecutionMetadata, Message
+from eval_protocol.tracing import PayloadType, decode_payloads
 from .base import BaseAdapter
 from .utils import extract_messages_from_data
 from ..common_utils import get_user_agent
@@ -100,13 +101,50 @@ def convert_trace_dict_to_evaluation_row(
                 ):
                     break  # Break early if we've found all the metadata we need
 
+        # Decoding lives in eval_protocol.tracing; here we only map results onto the row.
+        payloads = trace.get("payloads")
+        if isinstance(payloads, dict):
+            decoded = decode_payloads(
+                payloads,
+                on_error=lambda pt, e: logger.warning(
+                    "Failed to decode %s payload for trace %s: %s", pt.value, trace.get("id"), e
+                ),
+            )
+            if decoded and execution_metadata.extra is None:
+                execution_metadata.extra = {}
+
+            if (dp := decoded.get(PayloadType.ROUTER_REPLAY)) is not None:
+                execution_metadata.extra["routing_matrices"] = dp.value
+                execution_metadata.extra["routing_metadata"] = dp.metadata
+
+            if (dp := decoded.get(PayloadType.LOGPROBS)) is not None:
+                logprobs = dp.value
+                token_ids = dp.token_ids
+                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"] = dp.metadata
+
+                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
+
+            if (dp := decoded.get(PayloadType.PROMPT_TOKEN_IDS)) is not None:
+                execution_metadata.extra["prompt_token_ids"] = dp.value
+                execution_metadata.extra["prompt_token_ids_metadata"] = dp.metadata
+
         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 +413,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 +464,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 +488,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 +520,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.31/eval_protocol/adapters/lp_deserializer.py

@@ -0,0 +1,41 @@
+"""Deprecated compatibility shim for ``eval_protocol.tracing.logprobs``.
+
+Import from ``eval_protocol.tracing.logprobs`` (or ``decode_payloads`` from
+``eval_protocol.tracing``) instead. This module re-exports the LP/v1 helpers
+that lived here before the tracing package refactor.
+"""
+
+from __future__ import annotations
+
+import warnings
+
+warnings.warn(
+    "eval_protocol.adapters.lp_deserializer is deprecated; "
+    "import from eval_protocol.tracing.logprobs instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+from eval_protocol.tracing.logprobs import (  # noqa: E402
+    ENTRY_FORMAT,
+    ENTRY_SIZE,
+    HEADER_FORMAT,
+    HEADER_SIZE,
+    HEADER_VERSION,
+    MAGIC,
+    MISSING_TOKEN_ID,
+    decompress_and_parse_lp,
+    parse_logprobs,
+)
+
+__all__ = [
+    "ENTRY_FORMAT",
+    "ENTRY_SIZE",
+    "HEADER_FORMAT",
+    "HEADER_SIZE",
+    "HEADER_VERSION",
+    "MAGIC",
+    "MISSING_TOKEN_ID",
+    "decompress_and_parse_lp",
+    "parse_logprobs",
+]

eval_protocol-0.3.31/eval_protocol/adapters/r3_deserializer.py

@@ -0,0 +1,41 @@
+"""Deprecated compatibility shim for ``eval_protocol.tracing.router_replay``.
+
+Import from ``eval_protocol.tracing.router_replay`` (or ``decode_payloads`` from
+``eval_protocol.tracing``) instead. This module re-exports the R3/v1 helpers
+that lived here before the tracing package refactor.
+"""
+
+from __future__ import annotations
+
+import warnings
+
+warnings.warn(
+    "eval_protocol.adapters.r3_deserializer is deprecated; "
+    "import from eval_protocol.tracing.router_replay instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+from eval_protocol.tracing.router_replay import (  # noqa: E402
+    BITS_PER_BYTE,
+    HEADER_FORMAT,
+    HEADER_SIZE,
+    MAGIC,
+    _RoutingDtype,
+    _SelectorMode,
+    _parse_header,
+    _read_bitmap_positions,
+    decompress_and_parse_r3,
+)
+
+__all__ = [
+    "BITS_PER_BYTE",
+    "HEADER_FORMAT",
+    "HEADER_SIZE",
+    "MAGIC",
+    "_RoutingDtype",
+    "_SelectorMode",
+    "_parse_header",
+    "_read_bitmap_positions",
+    "decompress_and_parse_r3",
+]

{eval_protocol-0.3.29 → eval_protocol-0.3.31}/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.31}/eval_protocol/pytest/tracing_utils.py

@@ -22,9 +22,63 @@ 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",
+                "prompt_token_ids",
+                "prompt_token_ids_metadata",
+                "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 +153,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 +183,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 +202,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]

eval_protocol-0.3.31/eval_protocol/tracing/__init__.py

@@ -0,0 +1,31 @@
+"""Decode Fireworks tracing-gateway payloads.
+
+Standalone, dependency-light helpers (stdlib + ``zstandard`` only) for turning
+the binary/JSON ``payloads`` returned by the Fireworks tracing gateway
+(``GET /traces?include_payloads=true``) into Python values. No EvaluationRow or
+rollout machinery required -- usable on its own.
+
+Typical use::
+
+    from eval_protocol.tracing import decode_payloads, PayloadType
+
+    decoded = decode_payloads(trace["payloads"])
+    decoded[PayloadType.PROMPT_TOKEN_IDS].value   # List[int]
+    decoded[PayloadType.LOGPROBS].value           # List[float]
+    decoded[PayloadType.ROUTER_REPLAY].value      # List[Optional[str]]
+
+See ``README.md`` in this package for details.
+"""
+
+from __future__ import annotations
+
+from .registry import decode_payload, decode_payloads, decode_trace
+from .types import DecodedPayload, PayloadType
+
+__all__ = [
+    "PayloadType",
+    "DecodedPayload",
+    "decode_payloads",
+    "decode_payload",
+    "decode_trace",
+]

eval_protocol-0.3.31/eval_protocol/tracing/logprobs.py

@@ -0,0 +1,126 @@
+"""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
+
+from .types import DecodedPayload, PayloadType
+
+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)
+
+
+def decode_logprobs(data_b64: str) -> DecodedPayload:
+    """Decode a gateway ``payloads.logprobs.data`` blob into a ``DecodedPayload``.
+
+    ``value`` is the per-completion-token logprob list; per-token ids (when all
+    valid) are available under ``token_ids``.
+    """
+    logprobs, token_ids, metadata = decompress_and_parse_lp(data_b64)
+    return DecodedPayload(
+        payload_type=PayloadType.LOGPROBS,
+        value=logprobs,
+        metadata=metadata,
+        token_ids=token_ids,
+    )

eval_protocol-0.3.31/eval_protocol/tracing/prompt_token_ids.py

@@ -0,0 +1,34 @@
+"""``pti/v1`` decoder for prompt token ID payloads.
+
+Inverse of the tracing gateway's ``serialize_prompt_token_ids``: the gateway
+stores prompt token IDs as ``base64(zstd(json.dumps(token_ids)))`` -- a compact
+JSON int array, no bespoke binary header.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+from typing import Any, Dict, List, Tuple
+
+import zstandard as zstd
+
+from .types import DecodedPayload, PayloadType
+
+
+def parse_prompt_token_ids(raw: bytes) -> Tuple[List[int], Dict[str, Any]]:
+    """Parse uncompressed ``pti/v1`` bytes (a JSON int array) into ids + metadata."""
+    token_ids = json.loads(raw)
+    metadata: Dict[str, Any] = {"scope": "prompt_only", "token_count": len(token_ids)}
+    return token_ids, metadata
+
+
+def decode_prompt_token_ids(data_b64: str) -> DecodedPayload:
+    """Decode a gateway ``payloads.prompt_token_ids.data`` blob."""
+    raw = zstd.ZstdDecompressor().decompress(base64.b64decode(data_b64))
+    token_ids, metadata = parse_prompt_token_ids(raw)
+    return DecodedPayload(
+        payload_type=PayloadType.PROMPT_TOKEN_IDS,
+        value=token_ids,
+        metadata=metadata,
+    )