From 4995f0f8a7c2834ca2ce9388587c7e2b3ccb574a Mon Sep 17 00:00:00 2001 From: kywch Date: Sat, 5 Sep 2026 00:27:21 -0700 Subject: [PATCH 1/6] Handle Gemini passthrough trajectories in training exports --- .../scripts/validate_run_artifacts.py | 43 ++++- .../trajectories/export_prime_sft.py | 171 +++++++++++++++++- src/benchflow/trajectories/results.py | 8 +- ...t_benchflow_experiment_review_validator.py | 55 ++++-- tests/trajectories/test_export_prime_sft.py | 134 ++++++++++++++ 5 files changed, 396 insertions(+), 15 deletions(-) diff --git a/.agents/skills/benchflow-experiment-review/scripts/validate_run_artifacts.py b/.agents/skills/benchflow-experiment-review/scripts/validate_run_artifacts.py index b98f3d211..c6a127e21 100755 --- a/.agents/skills/benchflow-experiment-review/scripts/validate_run_artifacts.py +++ b/.agents/skills/benchflow-experiment-review/scripts/validate_run_artifacts.py @@ -376,12 +376,51 @@ def response_consumed_by_later_request( return bool( call_ids and any( - any(call_id in request_body for call_id in call_ids) - for request_body in request_bodies[exchange_idx + 1 :] + any(call_id in request_bodies[index] for call_id in call_ids) + or bool(call_ids & gemini_history_call_ids(rows[index])) + for index in range(exchange_idx + 1, len(rows)) ) ) +def gemini_history_call_ids(row: dict[str, Any]) -> set[str]: + """Match native IDs plus recorded signatures, never strip or guess call IDs.""" + metadata = row.get("metadata") + if ( + not isinstance(metadata, dict) + or metadata.get("call_type") != "pass_through_endpoint" + ): + return set() + request = row.get("request") + body = request.get("body") if isinstance(request, dict) else None + messages = body.get("messages") if isinstance(body, dict) else None + if ( + not isinstance(messages, list) + or len(messages) != 1 + or not isinstance(messages[0], dict) + ): + return set() + try: + native = json.loads(messages[0].get("content", "")) + except (ValueError, TypeError): + return set() + contents = native.get("contents") if isinstance(native, dict) else None + if not isinstance(contents, list): + return set() + ids = set() + for content in contents: + parts = content.get("parts") if isinstance(content, dict) else None + for part in parts if isinstance(parts, list) else []: + call = part.get("functionCall") if isinstance(part, dict) else None + if isinstance(call, dict) and isinstance(call.get("id"), str): + signature = part.get("thoughtSignature") + if signature is None or isinstance(signature, str): + ids.add( + call["id"] + (f"__thought__{signature}" if signature else "") + ) + return ids + + def response_call_ids(body: dict[str, Any]) -> set[str]: return {call_id for call_id, _ in response_tool_calls(body)} diff --git a/src/benchflow/trajectories/export_prime_sft.py b/src/benchflow/trajectories/export_prime_sft.py index 4ea47a56c..73e766d41 100644 --- a/src/benchflow/trajectories/export_prime_sft.py +++ b/src/benchflow/trajectories/export_prime_sft.py @@ -15,7 +15,7 @@ from copy import deepcopy from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Literal, cast +from typing import Any, Literal, NoReturn, cast from benchflow._utils.json_safe import dumps_finite, scrub_non_finite from benchflow.trajectories.types import redact_trajectory_obj @@ -537,11 +537,180 @@ def _assistant_from_responses_response( return message +def normalize_provider_exchange(exchange: dict[str, Any]) -> dict[str, Any]: + """Decode LiteLLM's Gemini passthrough envelope without changing raw evidence.""" + metadata = exchange.get("metadata") or {} + if ( + not isinstance(metadata, dict) + or metadata.get("call_type") != "pass_through_endpoint" + or metadata.get("training_input_format") == "gemini" + ): + return exchange + request = exchange.get("request") or {} + if not isinstance(request, dict): + return exchange + body = request.get("body") or {} + if not isinstance(body, dict): + return exchange + + def invalid(detail: str) -> NoReturn: + raise ValueError(f"Unsupported Gemini passthrough: {detail}") + + try: + envelope = body["messages"] + if not isinstance(envelope, list) or len(envelope) != 1: + raise ValueError("expected singleton envelope") + content = envelope[0]["content"] + if not isinstance(content, str): + raise ValueError("expected JSON text") + native = json.loads(content) + if not isinstance(native, dict) or "contents" not in native: + raise ValueError("expected native contents") + except (KeyError, TypeError, ValueError): + model = ( + metadata.get("provider_model") + or metadata.get("request_model") + or body.get("model") + or "" + ) + if str(model).rsplit("/", 1)[-1].startswith(("gemini-", "gemma-")): + invalid("malformed native request envelope") + return exchange + + contents = native["contents"] + if not isinstance(contents, list): + invalid("contents must be a list") + messages: list[dict[str, Any]] = [] + call_ids: dict[str, str] = {} + system = native.get("systemInstruction") + if system is not None: + if not isinstance(system, dict): + invalid("systemInstruction must be an object") + contents = [{**system, "role": "system"}, *contents] + for turn in contents: + if not isinstance(turn, dict) or turn.get("role") not in { + "system", + "user", + "model", + }: + invalid("unknown content role") + parts = turn.get("parts") + if not isinstance(parts, list): + invalid("parts must be a list") + role = "assistant" if turn["role"] == "model" else turn["role"] + message: dict[str, Any] = {"role": role, "content": ""} + for part in parts: + if not isinstance(part, dict) or set(part) - { + "text", + "thought", + "thoughtSignature", + "functionCall", + "functionResponse", + }: + invalid("unsupported content part") + if len(set(part) & {"text", "functionCall", "functionResponse"}) != 1: + invalid("content part must have exactly one payload") + if "text" in part: + if not isinstance(part["text"], str): + invalid("text must be a string") + if not part.get("thought"): + message["content"] += part["text"] + elif "functionCall" in part: + call = part["functionCall"] + if ( + role != "assistant" + or not isinstance(call, dict) + or not isinstance(call.get("id"), str) + or not call["id"] + or not isinstance(call.get("name"), str) + or not call["name"] + or not isinstance(call.get("args", {}), dict) + ): + invalid("function call requires model role, id and name") + call_id = call["id"] + signature = part.get("thoughtSignature") + if signature is not None and not isinstance(signature, str): + invalid("thought signature must be a string") + decorated = call_id + (f"__thought__{signature}" if signature else "") + if call_id in call_ids and call_ids[call_id] != decorated: + invalid("conflicting signatures for function call id") + call_ids[call_id] = decorated + message.setdefault("tool_calls", []).append( + { + "id": decorated, + "name": call["name"], + "arguments": call.get("args", {}), + } + ) + elif "functionResponse" in part: + response = part["functionResponse"] + if ( + role != "user" + or not isinstance(response, dict) + or not isinstance(response.get("id"), str) + or not response["id"] + or not isinstance(response.get("response"), dict) + ): + invalid("function response requires user role and id") + if message["content"]: + messages.append(message) + message = {"role": role, "content": ""} + response_id = response["id"] + if response_id not in call_ids: + invalid("function response references unknown call id") + messages.append( + { + "role": "tool", + "tool_call_id": call_ids[response_id], + "content": json.dumps(response.get("response")), + } + ) + if message["content"] or message.get("tool_calls"): + messages.append(message) + tools = [] + declarations = native.get("tools", []) + if not isinstance(declarations, list): + invalid("tools must be a list") + for tool in declarations: + if ( + not isinstance(tool, dict) + or set(tool) != {"functionDeclarations"} + or not isinstance(tool["functionDeclarations"], list) + ): + invalid("unsupported tool declaration") + for declaration in tool["functionDeclarations"]: + if not isinstance(declaration, dict) or not declaration.get("name"): + invalid("function declaration requires a name") + parameters = declaration.get( + "parametersJsonSchema", + declaration.get("parameters", {"type": "object", "properties": {}}), + ) + if not isinstance(parameters, dict): + invalid("function declaration requires an object schema") + tools.append({**declaration, "parameters": parameters}) + return { + **exchange, + "metadata": {**metadata, "training_input_format": "gemini"}, + "request": { + **request, + "body": { + **body, + "messages": messages, + "tools": _tool_defs_from_body({"tools": tools}), + }, + }, + } + + def _exchange_to_messages_and_tools( exchange: dict[str, Any], *, redact: bool = True, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], str | None]: + try: + exchange = normalize_provider_exchange(exchange) + except ValueError as exc: + return [], [], str(exc) request = ( cast(dict[str, Any], exchange.get("request")) if isinstance(exchange.get("request"), dict) diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index 700093fb3..1b3dff92b 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -26,6 +26,7 @@ PrimeSftTrajectoryJsonlError, load_llm_trajectory_jsonl, normalize_prime_sft_exchange, + normalize_provider_exchange, prime_sft_last_user_training_window, validate_prime_sft_row, ) @@ -162,9 +163,14 @@ def _llm_steps_from_trajectory( if not path.exists(): return steps, tool_defs, None try: - exchanges = load_llm_trajectory_jsonl(path, strict=True) + exchanges = [ + normalize_provider_exchange(exchange) + for exchange in load_llm_trajectory_jsonl(path, strict=True) + ] except PrimeSftTrajectoryJsonlError as exc: return [], [], f"Invalid LLM trajectory JSONL: {exc}" + except ValueError as exc: + return [], [], str(exc) training_success_indices = _training_success_exchange_indices(exchanges) skipped_successful: list[str] = [] for exchange_idx, exchange in enumerate(exchanges): diff --git a/tests/test_benchflow_experiment_review_validator.py b/tests/test_benchflow_experiment_review_validator.py index d1c20ec5d..2060e865a 100644 --- a/tests/test_benchflow_experiment_review_validator.py +++ b/tests/test_benchflow_experiment_review_validator.py @@ -4,6 +4,8 @@ import json from pathlib import Path +import pytest + SCRIPT = ( Path(__file__).resolve().parents[1] / ".agents" @@ -476,10 +478,12 @@ def test_validator_deduplicates_completed_retry_race(tmp_path: Path) -> None: assert report["artifacts"]["llm"]["deduplicated_completed_responses"] == 1 +@pytest.mark.parametrize("native_gemini", [False, True]) def test_validator_prefers_duplicate_consumed_by_later_request( tmp_path: Path, + native_gemini: bool, ) -> None: - """Guards PR #921 against keeping a late abandoned response.""" + """Guards PR #921 and a0b16985's native Gemini signature-consumption gap.""" validator = _load_validator() rollout = _rollout(tmp_path) llm_path = rollout / "trajectory" / "llm_trajectory.jsonl" @@ -489,7 +493,9 @@ def test_validator_prefers_duplicate_consumed_by_later_request( "content": "", "tool_calls": [ { - "id": "call_consumed", + "id": "call_consumed__thought__signature" + if native_gemini + else "call_consumed", "type": "function", "function": {"name": "finish", "arguments": "{}"}, } @@ -500,16 +506,39 @@ def test_validator_prefers_duplicate_consumed_by_later_request( "call_abandoned" ) followup = json.loads(llm_path.read_text()) - followup["request"]["body"]["messages"].extend( - [ - consumed["response"]["body"]["choices"][0]["message"], - { - "role": "tool", - "tool_call_id": "call_consumed", - "content": "done", - }, + if native_gemini: + call = { + "functionCall": {"id": "call_consumed", "name": "finish", "args": {}}, + "thoughtSignature": "signature", + } + output = { + "functionResponse": { + "id": "call_consumed", + "name": "finish", + "response": {}, + } + } + native = { + "contents": [ + {"role": "model", "parts": [call]}, + {"role": "user", "parts": [output]}, + ] + } + followup["metadata"] = {"call_type": "pass_through_endpoint"} + followup["request"]["body"]["messages"] = [ + {"role": "user", "content": json.dumps(native)} ] - ) + else: + followup["request"]["body"]["messages"].extend( + [ + consumed["response"]["body"]["choices"][0]["message"], + { + "role": "tool", + "tool_call_id": "call_consumed", + "content": "done", + }, + ] + ) _write_jsonl(llm_path, [consumed, abandoned, followup]) row = json.loads((rollout / "results.jsonl").read_text()) row["trajectory"] = [ @@ -528,6 +557,10 @@ def test_validator_prefers_duplicate_consumed_by_later_request( assert report["healthy"] is True assert report["artifacts"]["llm"]["successful_exchange_indices"] == [0, 2] + if native_gemini: + row["trajectory"] = row["trajectory"][1:] + _write_jsonl(rollout / "results.jsonl", [row]) + assert validator.validate_rollout(rollout)["healthy"] is False def test_validator_excludes_unique_late_unconsumed_nonterminal_retry( diff --git a/tests/trajectories/test_export_prime_sft.py b/tests/trajectories/test_export_prime_sft.py index 3d8c12981..d305e402f 100644 --- a/tests/trajectories/test_export_prime_sft.py +++ b/tests/trajectories/test_export_prime_sft.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from copy import deepcopy from pathlib import Path import pytest @@ -13,10 +14,143 @@ export_prime_sft_jsonl, load_llm_trajectory_jsonl, normalize_prime_sft_exchange, + normalize_provider_exchange, validate_prime_sft_jsonl, ) +@pytest.mark.parametrize("second_signature", ["", "other-signature"]) +def test_gemini_passthrough_export_preserves_tools_and_signatures( + tmp_path, second_signature +): + """Guards Gemini export on a0b16985: nested schemas and thought-decorated IDs.""" + from benchflow.trajectories.results import _llm_steps_from_trajectory + + first = _exchange(final=False) + function = first["request"]["body"]["tools"][0]["function"] + declaration = { + "name": function["name"], + "parametersJsonSchema": function["parameters"], + } + native = { + "systemInstruction": {"parts": [{"text": "Use tools."}]}, + "contents": [{"role": "user", "parts": [{"text": "List twice."}]}], + "tools": [{"functionDeclarations": [declaration]}], + } + first["metadata"] = {"call_type": "pass_through_endpoint"} + first["request"]["body"] = { + "messages": [{"role": "user", "content": json.dumps(native)}] + } + calls, parts, outputs = [], [], [] + for index, signature in enumerate(("signature", second_signature)): + call_id = f"call_{index}" + decorated_id = call_id + (f"__thought__{signature}" if signature else "") + calls.append( + { + "id": decorated_id, + "function": {"name": "bash", "arguments": '{"command":"ls"}'}, + } + ) + part = { + "functionCall": {"id": call_id, "name": "bash", "args": {"command": "ls"}} + } + if signature: + part["thoughtSignature"] = signature + parts.append(part) + outputs.append( + { + "functionResponse": { + "id": call_id, + "name": "bash", + "response": {"output": "README.md"}, + } + } + ) + first["response"]["body"]["choices"][0]["message"] = { + "role": "assistant", + "tool_calls": calls, + } + final = deepcopy(first) + native["contents"].extend( + [{"role": "model", "parts": parts}, {"role": "user", "parts": outputs}] + ) + final["request"]["body"]["messages"][0]["content"] = json.dumps(native) + final["response"]["body"]["choices"][0]["message"] = { + "role": "assistant", + "content": "Done.", + } + original = deepcopy(final) + normalized, reason = normalize_prime_sft_exchange(final, redact=False) + assert reason is None + assert normalized.tool_defs[0]["function"]["parameters"] == function["parameters"] + assert normalized.messages[0] == {"role": "system", "content": "Use tools."} + results = [m for m in normalized.messages if m["role"] == "tool"] + assert [m["tool_call_id"] for m in results] == [c["id"] for c in calls] + assert all("README.md" in m["content"] for m in results) + assert final == original + orphan = deepcopy(final) + orphan_native = deepcopy(native) + orphan_native["contents"][-1]["parts"][0]["functionResponse"]["id"] = "unknown" + orphan["request"]["body"]["messages"][0]["content"] = json.dumps(orphan_native) + rejected, error = normalize_prime_sft_exchange(orphan) + assert rejected is None + assert "unknown call id" in error + rollout = tmp_path / "rollout" + _write_rollout(rollout, exchanges=[first, final]) + steps, tools, error = _llm_steps_from_trajectory( + rollout, reward=1, is_truncated=False, trajectory_id_prefix="test" + ) + assert error is None + assert len(steps) == 2 + assert tools[0]["function"]["parameters"] == function["parameters"] + + +@pytest.mark.parametrize( + "native", + [ + {"contents": None}, + {"contents": [{"role": "user", "parts": [{"inlineData": {}}]}]}, + {"contents": [], "tools": None}, + {"contents": [], "tools": [{"functionDeclarations": "invalid"}]}, + { + "contents": [], + "tools": [{"functionDeclarations": [{"name": "bad", "parameters": None}]}], + }, + ], +) +def test_gemini_unrepresentable_history_blocks_training(native): + """Guards a0b16985 against silently accepting malformed native history.""" + exchange = _exchange(final=False) + exchange["metadata"] = {"call_type": "pass_through_endpoint"} + exchange["request"]["body"]["messages"] = [ + {"role": "user", "content": json.dumps(native)} + ] + normalized, reason = normalize_prime_sft_exchange(exchange) + assert normalized is None + assert reason.startswith("Unsupported Gemini passthrough:") + + +def test_gemini_envelope_guard_preserves_ordinary_json_and_no_arg_schema(): + """Guards a0b16985: JSON user text is not native provider history.""" + exchange = _exchange(final=False) + native = { + "contents": [{"role": "user", "parts": [{"text": '{"contents": []}'}]}], + "tools": [{"functionDeclarations": [{"name": "no_args"}]}], + } + exchange["request"]["body"]["messages"] = [ + {"role": "user", "content": json.dumps(native)} + ] + assert normalize_provider_exchange(exchange) is exchange + exchange["metadata"] = {"call_type": "pass_through_endpoint"} + normalized = normalize_provider_exchange(exchange) + assert normalize_provider_exchange(normalized) == normalized + assert normalized["request"]["body"]["messages"][0]["content"] == '{"contents": []}' + assert normalized["request"]["body"]["tools"][0]["function"]["parameters"] == { + "type": "object", + "properties": {}, + } + + def _write_rollout( rollout_dir: Path, *, From 213bbd5c5b28936464dbdff2bbc10b76920ff559 Mon Sep 17 00:00:00 2001 From: kywch Date: Sat, 5 Sep 2026 00:33:19 -0700 Subject: [PATCH 2/6] Extract shared trajectory message contract --- src/benchflow/eval_artifacts.py | 2 +- .../trajectories/export_prime_sft.py | 878 +----------------- src/benchflow/trajectories/export_trl_sft.py | 2 + .../trajectories/message_contract.py | 867 +++++++++++++++++ src/benchflow/trajectories/results.py | 2 +- tests/trajectories/test_export_prime_sft.py | 2 +- 6 files changed, 897 insertions(+), 856 deletions(-) create mode 100644 src/benchflow/trajectories/message_contract.py diff --git a/src/benchflow/eval_artifacts.py b/src/benchflow/eval_artifacts.py index f3169aa2a..1e3b3fc1a 100644 --- a/src/benchflow/eval_artifacts.py +++ b/src/benchflow/eval_artifacts.py @@ -11,7 +11,7 @@ from benchflow._utils.task_authoring import task_digest from benchflow._utils.text import truncate_end from benchflow.task.discovery import is_task_dir, resolve_task_collection_root -from benchflow.trajectories.export_prime_sft import ( +from benchflow.trajectories.message_contract import ( PrimeSftTrajectoryJsonlError, load_llm_trajectory_jsonl, ) diff --git a/src/benchflow/trajectories/export_prime_sft.py b/src/benchflow/trajectories/export_prime_sft.py index 73e766d41..af624be88 100644 --- a/src/benchflow/trajectories/export_prime_sft.py +++ b/src/benchflow/trajectories/export_prime_sft.py @@ -12,37 +12,39 @@ import json from collections.abc import Iterator -from copy import deepcopy from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Literal, NoReturn, cast +from typing import Any, Literal, cast from benchflow._utils.json_safe import dumps_finite, scrub_non_finite +from benchflow.trajectories.message_contract import ( + PrimeSftExchangeData as PrimeSftExchangeData, +) +from benchflow.trajectories.message_contract import ( + PrimeSftTrajectoryJsonlError as PrimeSftTrajectoryJsonlError, +) +from benchflow.trajectories.message_contract import ( + _align_legacy_tool_call_ids, + _has_tool_calls, + _normalize_tool_call, + _row_messages, +) +from benchflow.trajectories.message_contract import ( + load_llm_trajectory_jsonl as load_llm_trajectory_jsonl, +) +from benchflow.trajectories.message_contract import ( + normalize_prime_sft_exchange as normalize_prime_sft_exchange, +) +from benchflow.trajectories.message_contract import ( + prime_sft_last_user_training_window as prime_sft_last_user_training_window, +) +from benchflow.trajectories.message_contract import ( + validate_prime_sft_row as validate_prime_sft_row, +) from benchflow.trajectories.types import redact_trajectory_obj PrimeSftRowMode = Literal["rollout", "exchange"] -ALLOWED_ROLES = {"system", "user", "assistant", "tool"} -BANNED_ROW_KEYS = { - "gold", - "gold_solution", - "verify_source", - "tools_py", - "initial_db", - "db_json", - "target_constants", - "private_reasoning", - "reasoning_content", - "thinking_blocks", -} -BANNED_MESSAGE_KEYS = { - "reasoning_content", - "thinking_blocks", - "private_reasoning", - "provider_specific_fields", - "function_call", -} - @dataclass class PrimeSftExportStats: @@ -89,16 +91,6 @@ def as_dict(self) -> dict[str, Any]: } -@dataclass(frozen=True) -class PrimeSftExchangeData: - messages: list[dict[str, Any]] - tool_defs: list[dict[str, Any]] - - -class PrimeSftTrajectoryJsonlError(ValueError): - """Raised when an LLM trajectory JSONL file is not parseable.""" - - def _json_line(record: dict[str, Any], *, redact: bool = True) -> str: # Redact secrets in the record's string values BEFORE serializing so the # emitted SFT row is always valid JSON; redacting the serialized text could @@ -118,40 +110,6 @@ def _load_json(path: Path) -> dict[str, Any] | None: return data if isinstance(data, dict) else None -def load_llm_trajectory_jsonl( - path: Path, - *, - strict: bool = False, -) -> list[dict[str, Any]]: - records: list[dict[str, Any]] = [] - try: - lines = path.read_text().splitlines() - except OSError as exc: - if strict: - raise PrimeSftTrajectoryJsonlError( - f"{path}: cannot read LLM trajectory JSONL: {exc}" - ) from exc - return records - for line_num, line in enumerate(lines, start=1): - if not line.strip(): - continue - try: - record = json.loads(line) - except json.JSONDecodeError as exc: - if strict: - raise PrimeSftTrajectoryJsonlError( - f"{path}: line {line_num}: invalid JSON: {exc}" - ) from exc - continue - if isinstance(record, dict): - records.append(record) - elif strict: - raise PrimeSftTrajectoryJsonlError( - f"{path}: line {line_num}: top-level record must be an object" - ) - return records - - def _iter_rollout_dirs(root: str | Path) -> list[Path]: path = Path(root) if (path / "result.json").is_file(): @@ -219,585 +177,6 @@ def _reward_from_result(result: dict[str, Any] | None) -> float | None: return None -def _content_to_text(content: Any) -> str: - if content is None: - return "" - if isinstance(content, str): - return content - if isinstance(content, list): - parts: list[str] = [] - for item in content: - if isinstance(item, str): - parts.append(item) - elif isinstance(item, dict): - text = item.get("text") - if text is None: - text = item.get("content") - if isinstance(text, str): - parts.append(text) - return "\n".join(parts) - return str(content) - - -def _normalize_role(role: Any) -> str: - if role == "developer": - return "system" - if role == "model": - return "assistant" - return str(role or "user") - - -def _json_tool_call_arguments(arguments: Any, *, redact: bool = True) -> str: - if isinstance(arguments, str): - try: - parsed = json.loads(arguments) - except json.JSONDecodeError: - parsed = {"_malformed_json_arguments": arguments} - else: - if not isinstance(parsed, dict): - parsed = {"_non_object_json_arguments": parsed} - else: - clean = redact_trajectory_obj(parsed) if redact else parsed - if clean == parsed: - return arguments - return dumps_finite(clean, sort_keys=False, default=str) - elif isinstance(arguments, dict): - parsed = arguments - elif arguments is None: - parsed = {} - else: - parsed = {"_non_object_arguments": arguments} - clean = redact_trajectory_obj(parsed) if redact else parsed - return dumps_finite(clean, sort_keys=False, default=str) - - -def _normalize_tool_call( - call: dict[str, Any], index: int = 0, *, redact: bool = True -) -> dict[str, Any]: - function = call.get("function") - if not isinstance(function, dict): - function = {} - name = function.get("name") or call.get("name") or "tool" - arguments = function.get("arguments", call.get("arguments", {})) - return { - "id": str(call.get("id") or call.get("tool_call_id") or f"call_{index:06d}"), - "type": "function", - "function": { - "name": str(name), - "arguments": _json_tool_call_arguments(arguments, redact=redact), - }, - } - - -def _normalize_message( - message: dict[str, Any], index: int, *, redact: bool = True -) -> dict[str, Any]: - message_type = message.get("type") - if message_type == "function_call": - return { - "role": "assistant", - "content": "", - "tool_calls": [ - _normalize_tool_call( - { - "id": message.get("call_id") or message.get("id"), - "type": "function", - "function": { - "name": message.get("name"), - "arguments": message.get("arguments", {}), - }, - }, - index, - redact=redact, - ) - ], - } - if message_type == "function_call_output": - return { - "role": "tool", - "tool_call_id": str(message.get("call_id") or message.get("id") or ""), - "content": _content_to_text(message.get("output")), - } - role = _normalize_role(message.get("role")) - out: dict[str, Any] = {"role": role} - if role == "tool": - tool_call_id = message.get("tool_call_id") - if tool_call_id is not None: - out["tool_call_id"] = str(tool_call_id) - content = message.get("content") - out["content"] = _content_to_text(content) - tool_calls = message.get("tool_calls") - if tool_calls is None and isinstance(message.get("function_call"), dict): - tool_calls = [message["function_call"]] - if isinstance(tool_calls, list) and tool_calls: - out["tool_calls"] = [ - _normalize_tool_call(call, i, redact=redact) - for i, call in enumerate(tool_calls) - if isinstance(call, dict) - ] - return out - - -def _normalize_system_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - # Prime-RL SFT allows a system message only at index 0 (see - # validate_prime_sft_row). Any system message after the first position — - # including a *second consecutive* leading system message — is remapped to - # "user" so the whole row isn't silently dropped into skipped_invalid. - normalized: list[dict[str, Any]] = [] - for idx, message in enumerate(messages): - out = dict(message) - if out.get("role") == "system" and idx != 0: - out["role"] = "user" - normalized.append(out) - return normalized - - -def prime_sft_last_user_training_window( - messages: list[dict[str, Any]], -) -> tuple[list[dict[str, Any]], list[dict[str, Any]]] | None: - """Return a compact prompt/completion window anchored at the last user turn. - - OpenHands-style system prompts are large enough that a full conversation - prefix can push the first trainable assistant token beyond an 8k SFT - sequence. Prime-RL can then skip the row even though the JSONL is valid. - Keeping the latest user instruction plus the following assistant/tool turns - preserves the supervised action trace while moving trainable tokens into the - loaded context window. - """ - for idx in range(len(messages) - 2, -1, -1): - message = messages[idx] - if message.get("role") != "user": - continue - completion = messages[idx + 1 :] - if any(item.get("role") == "assistant" for item in completion): - return [message], completion - return None - - -def _messages_from_chat_request( - body: dict[str, Any], *, redact: bool = True -) -> list[dict[str, Any]]: - messages = body.get("messages") - if not isinstance(messages, list): - return [] - normalized: list[dict[str, Any]] = [] - for idx, message in enumerate(messages): - if not isinstance(message, dict): - continue - message = cast(dict[str, Any], message) - if message.get("type") == "reasoning": - continue - normalized.append(_normalize_message(message, idx, redact=redact)) - return normalized - - -def _messages_from_responses_request( - body: dict[str, Any], *, redact: bool = True -) -> list[dict[str, Any]]: - messages: list[dict[str, Any]] = [] - instructions = body.get("instructions") - if instructions: - messages.append({"role": "system", "content": _content_to_text(instructions)}) - raw_input = body.get("input") - if isinstance(raw_input, str): - messages.append({"role": "user", "content": raw_input}) - elif isinstance(raw_input, list): - for idx, item in enumerate(raw_input): - if not isinstance(item, dict): - continue - item = cast(dict[str, Any], item) - if item.get("type") != "reasoning": - messages.append(_normalize_message(item, idx, redact=redact)) - return messages - - -def _tool_defs_from_body(body: dict[str, Any]) -> list[dict[str, Any]]: - raw_tools = body.get("tools") or body.get("tool_defs") or [] - if not isinstance(raw_tools, list): - return [] - tools: list[dict[str, Any]] = [] - for item in raw_tools: - if not isinstance(item, dict): - continue - if isinstance(item.get("function"), dict): - function = dict(item["function"]) - else: - function = { - "name": item.get("name"), - "description": item.get("description", ""), - "parameters": item.get( - "parameters", {"type": "object", "properties": {}} - ), - } - if not function.get("name"): - continue - function.setdefault("description", "") - function.setdefault("parameters", {"type": "object", "properties": {}}) - tools.append({"type": "function", "function": function}) - return tools - - -def _assistant_from_anthropic_content( - content: Any, *, redact: bool = True -) -> dict[str, Any] | None: - """Build an assistant row from Anthropic ``/v1/messages`` content blocks. - - Anthropic responses carry a list of typed blocks: ``text`` blocks hold the - visible reply and ``tool_use`` blocks hold tool calls. The previous fallback - flattened the whole list to text, silently dropping the tool calls and - turning a tool-using assistant turn into corrupted SFT data. Preserve - ``tool_use`` blocks as OpenAI-shaped ``tool_calls`` instead. Returns ``None`` - when ``content`` is not a block list, so the caller can fall back to text. - """ - if not isinstance(content, list): - return None - raw_tool_calls = [ - { - "id": item.get("id"), - "type": "function", - "function": { - "name": item.get("name"), - "arguments": item.get("input", {}), - }, - } - for item in content - if isinstance(item, dict) and item.get("type") == "tool_use" - ] - message: dict[str, Any] = { - "role": "assistant", - "content": _content_to_text(content), - } - if raw_tool_calls: - message["tool_calls"] = [ - _normalize_tool_call(call, i, redact=redact) - for i, call in enumerate(raw_tool_calls) - ] - return message - - -def _assistant_from_chat_response( - body: dict[str, Any], *, redact: bool = True -) -> dict[str, Any] | None: - choices = body.get("choices") - if isinstance(choices, list) and choices: - first = choices[0] - if isinstance(first, dict) and isinstance(first.get("message"), dict): - return _normalize_message(first["message"], 0, redact=redact) - message = body.get("message") - if isinstance(message, dict): - return _normalize_message(message, 0, redact=redact) - content = body.get("content") - if content: - assistant = _assistant_from_anthropic_content(content, redact=redact) - if assistant is not None: - return assistant - return {"role": "assistant", "content": _content_to_text(content)} - assistant = _assistant_from_responses_response(body, redact=redact) - if assistant is not None: - return assistant - return None - - -def _assistant_from_responses_response( - body: dict[str, Any], *, redact: bool = True -) -> dict[str, Any] | None: - output = body.get("output") - if not isinstance(output, list): - return None - texts: list[str] = [] - tool_calls: list[dict[str, Any]] = [] - for item in output: - if not isinstance(item, dict): - continue - item_type = item.get("type") - if item_type == "message": - texts.append(_content_to_text(item.get("content"))) - elif item_type in {"function_call", "tool_call"}: - tool_calls.append( - { - "id": item.get("call_id") or item.get("id"), - "type": "function", - "function": { - "name": item.get("name"), - "arguments": item.get("arguments", {}), - }, - } - ) - if not texts and not tool_calls: - return None - message: dict[str, Any] = { - "role": "assistant", - "content": "\n".join(t for t in texts if t), - } - if tool_calls: - message["tool_calls"] = [ - _normalize_tool_call(call, i, redact=redact) - for i, call in enumerate(tool_calls) - ] - return message - - -def normalize_provider_exchange(exchange: dict[str, Any]) -> dict[str, Any]: - """Decode LiteLLM's Gemini passthrough envelope without changing raw evidence.""" - metadata = exchange.get("metadata") or {} - if ( - not isinstance(metadata, dict) - or metadata.get("call_type") != "pass_through_endpoint" - or metadata.get("training_input_format") == "gemini" - ): - return exchange - request = exchange.get("request") or {} - if not isinstance(request, dict): - return exchange - body = request.get("body") or {} - if not isinstance(body, dict): - return exchange - - def invalid(detail: str) -> NoReturn: - raise ValueError(f"Unsupported Gemini passthrough: {detail}") - - try: - envelope = body["messages"] - if not isinstance(envelope, list) or len(envelope) != 1: - raise ValueError("expected singleton envelope") - content = envelope[0]["content"] - if not isinstance(content, str): - raise ValueError("expected JSON text") - native = json.loads(content) - if not isinstance(native, dict) or "contents" not in native: - raise ValueError("expected native contents") - except (KeyError, TypeError, ValueError): - model = ( - metadata.get("provider_model") - or metadata.get("request_model") - or body.get("model") - or "" - ) - if str(model).rsplit("/", 1)[-1].startswith(("gemini-", "gemma-")): - invalid("malformed native request envelope") - return exchange - - contents = native["contents"] - if not isinstance(contents, list): - invalid("contents must be a list") - messages: list[dict[str, Any]] = [] - call_ids: dict[str, str] = {} - system = native.get("systemInstruction") - if system is not None: - if not isinstance(system, dict): - invalid("systemInstruction must be an object") - contents = [{**system, "role": "system"}, *contents] - for turn in contents: - if not isinstance(turn, dict) or turn.get("role") not in { - "system", - "user", - "model", - }: - invalid("unknown content role") - parts = turn.get("parts") - if not isinstance(parts, list): - invalid("parts must be a list") - role = "assistant" if turn["role"] == "model" else turn["role"] - message: dict[str, Any] = {"role": role, "content": ""} - for part in parts: - if not isinstance(part, dict) or set(part) - { - "text", - "thought", - "thoughtSignature", - "functionCall", - "functionResponse", - }: - invalid("unsupported content part") - if len(set(part) & {"text", "functionCall", "functionResponse"}) != 1: - invalid("content part must have exactly one payload") - if "text" in part: - if not isinstance(part["text"], str): - invalid("text must be a string") - if not part.get("thought"): - message["content"] += part["text"] - elif "functionCall" in part: - call = part["functionCall"] - if ( - role != "assistant" - or not isinstance(call, dict) - or not isinstance(call.get("id"), str) - or not call["id"] - or not isinstance(call.get("name"), str) - or not call["name"] - or not isinstance(call.get("args", {}), dict) - ): - invalid("function call requires model role, id and name") - call_id = call["id"] - signature = part.get("thoughtSignature") - if signature is not None and not isinstance(signature, str): - invalid("thought signature must be a string") - decorated = call_id + (f"__thought__{signature}" if signature else "") - if call_id in call_ids and call_ids[call_id] != decorated: - invalid("conflicting signatures for function call id") - call_ids[call_id] = decorated - message.setdefault("tool_calls", []).append( - { - "id": decorated, - "name": call["name"], - "arguments": call.get("args", {}), - } - ) - elif "functionResponse" in part: - response = part["functionResponse"] - if ( - role != "user" - or not isinstance(response, dict) - or not isinstance(response.get("id"), str) - or not response["id"] - or not isinstance(response.get("response"), dict) - ): - invalid("function response requires user role and id") - if message["content"]: - messages.append(message) - message = {"role": role, "content": ""} - response_id = response["id"] - if response_id not in call_ids: - invalid("function response references unknown call id") - messages.append( - { - "role": "tool", - "tool_call_id": call_ids[response_id], - "content": json.dumps(response.get("response")), - } - ) - if message["content"] or message.get("tool_calls"): - messages.append(message) - tools = [] - declarations = native.get("tools", []) - if not isinstance(declarations, list): - invalid("tools must be a list") - for tool in declarations: - if ( - not isinstance(tool, dict) - or set(tool) != {"functionDeclarations"} - or not isinstance(tool["functionDeclarations"], list) - ): - invalid("unsupported tool declaration") - for declaration in tool["functionDeclarations"]: - if not isinstance(declaration, dict) or not declaration.get("name"): - invalid("function declaration requires a name") - parameters = declaration.get( - "parametersJsonSchema", - declaration.get("parameters", {"type": "object", "properties": {}}), - ) - if not isinstance(parameters, dict): - invalid("function declaration requires an object schema") - tools.append({**declaration, "parameters": parameters}) - return { - **exchange, - "metadata": {**metadata, "training_input_format": "gemini"}, - "request": { - **request, - "body": { - **body, - "messages": messages, - "tools": _tool_defs_from_body({"tools": tools}), - }, - }, - } - - -def _exchange_to_messages_and_tools( - exchange: dict[str, Any], - *, - redact: bool = True, -) -> tuple[list[dict[str, Any]], list[dict[str, Any]], str | None]: - try: - exchange = normalize_provider_exchange(exchange) - except ValueError as exc: - return [], [], str(exc) - request = ( - cast(dict[str, Any], exchange.get("request")) - if isinstance(exchange.get("request"), dict) - else {} - ) - response = ( - cast(dict[str, Any], exchange.get("response")) - if isinstance(exchange.get("response"), dict) - else {} - ) - request_body = ( - cast(dict[str, Any], request.get("body")) - if isinstance(request.get("body"), dict) - else {} - ) - response_body = ( - cast(dict[str, Any], response.get("body")) - if isinstance(response.get("body"), dict) - else {} - ) - - if "messages" in request_body: - messages = _messages_from_chat_request(request_body, redact=redact) - assistant = _assistant_from_chat_response(response_body, redact=redact) - else: - messages = _messages_from_responses_request(request_body, redact=redact) - assistant = _assistant_from_responses_response(response_body, redact=redact) - - if assistant is None: - return [], [], "no_assistant" - messages.append(assistant) - return ( - _normalize_system_messages(messages), - _tool_defs_from_body(request_body), - None, - ) - - -def _has_tool_calls(messages: list[dict[str, Any]]) -> bool: - return any(bool(message.get("tool_calls")) for message in messages) - - -def _normalize_tools_for_validation( - row: dict[str, Any], row_num: int -) -> list[Any] | None: - tools = row.get("tool_defs", row.get("tools")) - if tools is None: - return None - if isinstance(tools, str): - try: - tools = json.loads(tools) - except json.JSONDecodeError as exc: - raise ValueError( - f"row {row_num}: tool_defs/tools is not valid JSON: {exc}" - ) from exc - if not isinstance(tools, list): - raise ValueError(f"row {row_num}: tool_defs/tools must be a list") - return tools - - -def _tool_names_for_validation(tools: list[Any] | None) -> set[str]: - names: set[str] = set() - for tool in tools or []: - if not isinstance(tool, dict): - continue - function = tool.get("function") - name = function.get("name") if isinstance(function, dict) else tool.get("name") - if isinstance(name, str) and name: - names.add(name) - return names - - -def _row_messages(row: dict[str, Any], row_num: int) -> list[Any]: - messages = row.get("messages") - if isinstance(messages, list) and messages: - return messages - prompt = row.get("prompt") - completion = row.get("completion") - if isinstance(prompt, list) and isinstance(completion, list): - combined = prompt + completion - if combined: - return combined - raise ValueError( - f"row {row_num}: expected non-empty messages or prompt+completion lists" - ) - - def _sanitize_message_tool_call_arguments( messages: list[Any], *, redact: bool = True ) -> list[Any]: @@ -853,77 +232,6 @@ def _row_message_segments( return [] -def _content_join(left: Any, right: Any) -> str: - return "\n".join( - part for part in (_content_to_text(left), _content_to_text(right)) if part - ) - - -def _align_legacy_tool_call_ids( - segments: list[tuple[Literal["messages", "prompt", "completion"], Any]], -) -> tuple[ - list[tuple[Literal["messages", "prompt", "completion"], Any]], dict[str, int] -]: - """Repair legacy BenchFlow rows whose provider call ids drifted. - - Some historical ``results.jsonl`` artifacts preserved assistant tool-call ids - from one provider layer (``fc_*``) while the following tool messages used the - OpenAI-compatible ``call_*`` ids that the runtime sent back on the next turn. - Pair by message order and rewrite only when a pending assistant call exists; - true orphan tool outputs still fail validation. - """ - out: list[tuple[Literal["messages", "prompt", "completion"], Any]] = [] - pending: list[dict[str, Any]] = [] - stats = {"tool_call_ids_rewritten": 0, "tool_messages_merged": 0} - - for segment, raw_message in segments: - message = deepcopy(raw_message) - if not isinstance(message, dict): - out.append((segment, message)) - continue - - tool_calls = message.get("tool_calls") - if message.get("role") == "assistant" and isinstance(tool_calls, list): - pending.extend( - tool_call for tool_call in tool_calls if isinstance(tool_call, dict) - ) - - if message.get("role") == "tool": - tool_call_id = message.get("tool_call_id") - if ( - out - and isinstance(out[-1][1], dict) - and out[-1][1].get("role") == "tool" - and out[-1][1].get("tool_call_id") == tool_call_id - ): - out[-1][1]["content"] = _content_join( - out[-1][1].get("content"), - message.get("content"), - ) - stats["tool_messages_merged"] += 1 - continue - - match_index = next( - ( - idx - for idx, tool_call in enumerate(pending) - if tool_call.get("id") == tool_call_id - ), - None, - ) - if match_index is not None: - pending.pop(match_index) - elif pending and tool_call_id: - tool_call = pending.pop(0) - if tool_call.get("id") != tool_call_id: - tool_call["id"] = tool_call_id - stats["tool_call_ids_rewritten"] += 1 - - out.append((segment, message)) - - return out, stats - - def _canonicalize_existing_prime_sft_row( row: dict[str, Any], row_num: int, @@ -957,118 +265,6 @@ def _canonicalize_existing_prime_sft_row( return out, stats -def validate_prime_sft_row(row: dict[str, Any], row_num: int = 1) -> None: - leaked = sorted(BANNED_ROW_KEYS.intersection(row)) - if leaked: - raise ValueError( - f"row {row_num}: banned leakage keys present: {', '.join(leaked)}" - ) - - messages = _row_messages(row, row_num) - tools = _normalize_tools_for_validation(row, row_num) - known_tool_names = _tool_names_for_validation(tools) - pending_tool_call_ids: set[str] = set() - - for idx, message in enumerate(messages): - if not isinstance(message, dict): - raise ValueError(f"row {row_num}: messages[{idx}] must be object") - message = cast(dict[str, Any], message) - leaked_message = sorted(BANNED_MESSAGE_KEYS.intersection(message)) - if leaked_message: - raise ValueError( - f"row {row_num}: messages[{idx}] has banned keys: {', '.join(leaked_message)}" - ) - role = message.get("role") - if role not in ALLOWED_ROLES: - raise ValueError(f"row {row_num}: messages[{idx}].role invalid: {role!r}") - if role == "system" and idx != 0: - raise ValueError( - f"row {row_num}: system message must be at index 0, got index {idx}" - ) - if "content" not in message and "tool_calls" not in message: - raise ValueError( - f"row {row_num}: messages[{idx}] needs content or tool_calls" - ) - tool_calls = message.get("tool_calls") - if tool_calls and role != "assistant": - raise ValueError( - f"row {row_num}: only assistant messages may contain tool_calls" - ) - if role == "tool" and not message.get("tool_call_id"): - raise ValueError(f"row {row_num}: tool message requires tool_call_id") - if role == "tool" and message.get("tool_call_id") not in pending_tool_call_ids: - raise ValueError( - f"row {row_num}: tool message references unknown tool_call_id" - ) - if role == "tool": - pending_tool_call_ids.discard(cast(str, message.get("tool_call_id"))) - if tool_calls is not None and not isinstance(tool_calls, list): - raise ValueError( - f"row {row_num}: messages[{idx}].tool_calls must be a list" - ) - if isinstance(tool_calls, list): - for tool_call_idx, tool_call in enumerate(tool_calls): - prefix = f"row {row_num}: messages[{idx}].tool_calls[{tool_call_idx}]" - if not isinstance(tool_call, dict): - raise ValueError(f"{prefix} must be object") - tool_call = cast(dict[str, Any], tool_call) - function = tool_call.get("function") - if not isinstance(function, dict): - raise ValueError(f"{prefix}.function must be object") - tool_call_id = tool_call.get("id") - if not isinstance(tool_call_id, str) or not tool_call_id: - raise ValueError(f"{prefix}.id must be a non-empty string") - if tool_call.get("type") != "function": - raise ValueError(f"{prefix}.type must be 'function'") - name = function.get("name") - if not isinstance(name, str) or not name: - raise ValueError( - f"{prefix}.function.name must be a non-empty string" - ) - if known_tool_names and name not in known_tool_names: - raise ValueError( - f"{prefix}.function.name {name!r} not found in tool_defs/tools" - ) - arguments = function.get("arguments") - if isinstance(arguments, str): - try: - parsed_arguments = json.loads(arguments) - except json.JSONDecodeError as exc: - raise ValueError( - f"{prefix}.function.arguments is not valid JSON: {exc}" - ) from exc - if not isinstance(parsed_arguments, dict): - raise ValueError( - f"{prefix}.function.arguments must be a JSON object" - ) - elif not isinstance(arguments, dict): - raise ValueError( - f"{prefix}.function.arguments must be a JSON object or JSON-encoded object" - ) - pending_tool_call_ids.add(tool_call_id) - - if not any(isinstance(m, dict) and m.get("role") == "assistant" for m in messages): - raise ValueError(f"row {row_num}: no assistant message") - - typed_messages = [m for m in messages if isinstance(m, dict)] - if _has_tool_calls(typed_messages) and not tools: - raise ValueError( - f"row {row_num}: assistant tool_calls require non-empty tool_defs/tools" - ) - if tools is not None: - for tool_idx, tool in enumerate(tools): - if not isinstance(tool, dict): - raise ValueError(f"row {row_num}: tool_defs[{tool_idx}] must be object") - function = tool.get("function") - name = ( - function.get("name") if isinstance(function, dict) else tool.get("name") - ) - if not isinstance(name, str) or not name: - raise ValueError( - f"row {row_num}: tool_defs[{tool_idx}] missing function name" - ) - - def validate_prime_sft_jsonl( jsonl: str | Path, *, @@ -1280,30 +476,6 @@ def _copy_existing_prime_sft_jsonl( handle.write(_json_line(compact, redact=redact) + "\n") -def normalize_prime_sft_exchange( - exchange: dict[str, Any], - *, - redact: bool = True, -) -> tuple[PrimeSftExchangeData | None, str | None]: - """Normalize one raw LLM exchange through the Prime-SFT validator path.""" - messages, tool_defs, skip_reason = _exchange_to_messages_and_tools( - exchange, redact=redact - ) - if skip_reason: - return None, skip_reason - repaired, _ = _align_legacy_tool_call_ids( - [("messages", message) for message in messages] - ) - messages = [message for _, message in repaired] - if _has_tool_calls(messages) and not tool_defs: - return None, "missing_tool_defs" - try: - validate_prime_sft_row({"messages": messages, "tool_defs": tool_defs}, 1) - except ValueError as exc: - return None, f"invalid_prime_sft_row: {exc}" - return PrimeSftExchangeData(messages=messages, tool_defs=tool_defs), None - - def _row_from_exchange( *, exchange: dict[str, Any], diff --git a/src/benchflow/trajectories/export_trl_sft.py b/src/benchflow/trajectories/export_trl_sft.py index 8a51fd65e..252ee9c81 100644 --- a/src/benchflow/trajectories/export_trl_sft.py +++ b/src/benchflow/trajectories/export_trl_sft.py @@ -20,6 +20,8 @@ _result_training_skip_reason, _reward_from_result, _row_reward, +) +from benchflow.trajectories.message_contract import ( load_llm_trajectory_jsonl, normalize_prime_sft_exchange, validate_prime_sft_row, diff --git a/src/benchflow/trajectories/message_contract.py b/src/benchflow/trajectories/message_contract.py new file mode 100644 index 000000000..faa64df1e --- /dev/null +++ b/src/benchflow/trajectories/message_contract.py @@ -0,0 +1,867 @@ +"""Decode captured LLM exchanges into validated message records.""" + +from __future__ import annotations + +import json +from copy import deepcopy +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, NoReturn, cast + +from benchflow._utils.json_safe import dumps_finite +from benchflow.trajectories.types import redact_trajectory_obj + +ALLOWED_ROLES = {"system", "user", "assistant", "tool"} + + +BANNED_ROW_KEYS = { + "gold", + "gold_solution", + "verify_source", + "tools_py", + "initial_db", + "db_json", + "target_constants", + "private_reasoning", + "reasoning_content", + "thinking_blocks", +} + + +BANNED_MESSAGE_KEYS = { + "reasoning_content", + "thinking_blocks", + "private_reasoning", + "provider_specific_fields", + "function_call", +} + + +@dataclass(frozen=True) +class PrimeSftExchangeData: + messages: list[dict[str, Any]] + tool_defs: list[dict[str, Any]] + + +class PrimeSftTrajectoryJsonlError(ValueError): + """Raised when an LLM trajectory JSONL file is not parseable.""" + + +def load_llm_trajectory_jsonl( + path: Path, + *, + strict: bool = False, +) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + try: + lines = path.read_text().splitlines() + except OSError as exc: + if strict: + raise PrimeSftTrajectoryJsonlError( + f"{path}: cannot read LLM trajectory JSONL: {exc}" + ) from exc + return records + for line_num, line in enumerate(lines, start=1): + if not line.strip(): + continue + try: + record = json.loads(line) + except json.JSONDecodeError as exc: + if strict: + raise PrimeSftTrajectoryJsonlError( + f"{path}: line {line_num}: invalid JSON: {exc}" + ) from exc + continue + if isinstance(record, dict): + records.append(record) + elif strict: + raise PrimeSftTrajectoryJsonlError( + f"{path}: line {line_num}: top-level record must be an object" + ) + return records + + +def _content_to_text(content: Any) -> str: + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for item in content: + if isinstance(item, str): + parts.append(item) + elif isinstance(item, dict): + text = item.get("text") + if text is None: + text = item.get("content") + if isinstance(text, str): + parts.append(text) + return "\n".join(parts) + return str(content) + + +def _normalize_role(role: Any) -> str: + if role == "developer": + return "system" + if role == "model": + return "assistant" + return str(role or "user") + + +def _json_tool_call_arguments(arguments: Any, *, redact: bool = True) -> str: + if isinstance(arguments, str): + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + parsed = {"_malformed_json_arguments": arguments} + else: + if not isinstance(parsed, dict): + parsed = {"_non_object_json_arguments": parsed} + else: + clean = redact_trajectory_obj(parsed) if redact else parsed + if clean == parsed: + return arguments + return dumps_finite(clean, sort_keys=False, default=str) + elif isinstance(arguments, dict): + parsed = arguments + elif arguments is None: + parsed = {} + else: + parsed = {"_non_object_arguments": arguments} + clean = redact_trajectory_obj(parsed) if redact else parsed + return dumps_finite(clean, sort_keys=False, default=str) + + +def _normalize_tool_call( + call: dict[str, Any], index: int = 0, *, redact: bool = True +) -> dict[str, Any]: + function = call.get("function") + if not isinstance(function, dict): + function = {} + name = function.get("name") or call.get("name") or "tool" + arguments = function.get("arguments", call.get("arguments", {})) + return { + "id": str(call.get("id") or call.get("tool_call_id") or f"call_{index:06d}"), + "type": "function", + "function": { + "name": str(name), + "arguments": _json_tool_call_arguments(arguments, redact=redact), + }, + } + + +def _normalize_message( + message: dict[str, Any], index: int, *, redact: bool = True +) -> dict[str, Any]: + message_type = message.get("type") + if message_type == "function_call": + return { + "role": "assistant", + "content": "", + "tool_calls": [ + _normalize_tool_call( + { + "id": message.get("call_id") or message.get("id"), + "type": "function", + "function": { + "name": message.get("name"), + "arguments": message.get("arguments", {}), + }, + }, + index, + redact=redact, + ) + ], + } + if message_type == "function_call_output": + return { + "role": "tool", + "tool_call_id": str(message.get("call_id") or message.get("id") or ""), + "content": _content_to_text(message.get("output")), + } + role = _normalize_role(message.get("role")) + out: dict[str, Any] = {"role": role} + if role == "tool": + tool_call_id = message.get("tool_call_id") + if tool_call_id is not None: + out["tool_call_id"] = str(tool_call_id) + content = message.get("content") + out["content"] = _content_to_text(content) + tool_calls = message.get("tool_calls") + if tool_calls is None and isinstance(message.get("function_call"), dict): + tool_calls = [message["function_call"]] + if isinstance(tool_calls, list) and tool_calls: + out["tool_calls"] = [ + _normalize_tool_call(call, i, redact=redact) + for i, call in enumerate(tool_calls) + if isinstance(call, dict) + ] + return out + + +def _normalize_system_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + # Prime-RL SFT allows a system message only at index 0 (see + # validate_prime_sft_row). Any system message after the first position — + # including a *second consecutive* leading system message — is remapped to + # "user" so the whole row isn't silently dropped into skipped_invalid. + normalized: list[dict[str, Any]] = [] + for idx, message in enumerate(messages): + out = dict(message) + if out.get("role") == "system" and idx != 0: + out["role"] = "user" + normalized.append(out) + return normalized + + +def prime_sft_last_user_training_window( + messages: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]] | None: + """Return a compact prompt/completion window anchored at the last user turn. + + OpenHands-style system prompts are large enough that a full conversation + prefix can push the first trainable assistant token beyond an 8k SFT + sequence. Prime-RL can then skip the row even though the JSONL is valid. + Keeping the latest user instruction plus the following assistant/tool turns + preserves the supervised action trace while moving trainable tokens into the + loaded context window. + """ + for idx in range(len(messages) - 2, -1, -1): + message = messages[idx] + if message.get("role") != "user": + continue + completion = messages[idx + 1 :] + if any(item.get("role") == "assistant" for item in completion): + return [message], completion + return None + + +def _messages_from_chat_request( + body: dict[str, Any], *, redact: bool = True +) -> list[dict[str, Any]]: + messages = body.get("messages") + if not isinstance(messages, list): + return [] + normalized: list[dict[str, Any]] = [] + for idx, message in enumerate(messages): + if not isinstance(message, dict): + continue + message = cast(dict[str, Any], message) + if message.get("type") == "reasoning": + continue + normalized.append(_normalize_message(message, idx, redact=redact)) + return normalized + + +def _messages_from_responses_request( + body: dict[str, Any], *, redact: bool = True +) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [] + instructions = body.get("instructions") + if instructions: + messages.append({"role": "system", "content": _content_to_text(instructions)}) + raw_input = body.get("input") + if isinstance(raw_input, str): + messages.append({"role": "user", "content": raw_input}) + elif isinstance(raw_input, list): + for idx, item in enumerate(raw_input): + if not isinstance(item, dict): + continue + item = cast(dict[str, Any], item) + if item.get("type") != "reasoning": + messages.append(_normalize_message(item, idx, redact=redact)) + return messages + + +def _tool_defs_from_body(body: dict[str, Any]) -> list[dict[str, Any]]: + raw_tools = body.get("tools") or body.get("tool_defs") or [] + if not isinstance(raw_tools, list): + return [] + tools: list[dict[str, Any]] = [] + for item in raw_tools: + if not isinstance(item, dict): + continue + if isinstance(item.get("function"), dict): + function = dict(item["function"]) + else: + function = { + "name": item.get("name"), + "description": item.get("description", ""), + "parameters": item.get( + "parameters", {"type": "object", "properties": {}} + ), + } + if not function.get("name"): + continue + function.setdefault("description", "") + function.setdefault("parameters", {"type": "object", "properties": {}}) + tools.append({"type": "function", "function": function}) + return tools + + +def _assistant_from_anthropic_content( + content: Any, *, redact: bool = True +) -> dict[str, Any] | None: + """Build an assistant row from Anthropic ``/v1/messages`` content blocks. + + Anthropic responses carry a list of typed blocks: ``text`` blocks hold the + visible reply and ``tool_use`` blocks hold tool calls. The previous fallback + flattened the whole list to text, silently dropping the tool calls and + turning a tool-using assistant turn into corrupted SFT data. Preserve + ``tool_use`` blocks as OpenAI-shaped ``tool_calls`` instead. Returns ``None`` + when ``content`` is not a block list, so the caller can fall back to text. + """ + if not isinstance(content, list): + return None + raw_tool_calls = [ + { + "id": item.get("id"), + "type": "function", + "function": { + "name": item.get("name"), + "arguments": item.get("input", {}), + }, + } + for item in content + if isinstance(item, dict) and item.get("type") == "tool_use" + ] + message: dict[str, Any] = { + "role": "assistant", + "content": _content_to_text(content), + } + if raw_tool_calls: + message["tool_calls"] = [ + _normalize_tool_call(call, i, redact=redact) + for i, call in enumerate(raw_tool_calls) + ] + return message + + +def _assistant_from_chat_response( + body: dict[str, Any], *, redact: bool = True +) -> dict[str, Any] | None: + choices = body.get("choices") + if isinstance(choices, list) and choices: + first = choices[0] + if isinstance(first, dict) and isinstance(first.get("message"), dict): + return _normalize_message(first["message"], 0, redact=redact) + message = body.get("message") + if isinstance(message, dict): + return _normalize_message(message, 0, redact=redact) + content = body.get("content") + if content: + assistant = _assistant_from_anthropic_content(content, redact=redact) + if assistant is not None: + return assistant + return {"role": "assistant", "content": _content_to_text(content)} + assistant = _assistant_from_responses_response(body, redact=redact) + if assistant is not None: + return assistant + return None + + +def _assistant_from_responses_response( + body: dict[str, Any], *, redact: bool = True +) -> dict[str, Any] | None: + output = body.get("output") + if not isinstance(output, list): + return None + texts: list[str] = [] + tool_calls: list[dict[str, Any]] = [] + for item in output: + if not isinstance(item, dict): + continue + item_type = item.get("type") + if item_type == "message": + texts.append(_content_to_text(item.get("content"))) + elif item_type in {"function_call", "tool_call"}: + tool_calls.append( + { + "id": item.get("call_id") or item.get("id"), + "type": "function", + "function": { + "name": item.get("name"), + "arguments": item.get("arguments", {}), + }, + } + ) + if not texts and not tool_calls: + return None + message: dict[str, Any] = { + "role": "assistant", + "content": "\n".join(t for t in texts if t), + } + if tool_calls: + message["tool_calls"] = [ + _normalize_tool_call(call, i, redact=redact) + for i, call in enumerate(tool_calls) + ] + return message + + +def normalize_provider_exchange(exchange: dict[str, Any]) -> dict[str, Any]: + """Decode LiteLLM's Gemini passthrough envelope without changing raw evidence.""" + metadata = exchange.get("metadata") or {} + if ( + not isinstance(metadata, dict) + or metadata.get("call_type") != "pass_through_endpoint" + or metadata.get("training_input_format") == "gemini" + ): + return exchange + request = exchange.get("request") or {} + if not isinstance(request, dict): + return exchange + body = request.get("body") or {} + if not isinstance(body, dict): + return exchange + + def invalid(detail: str) -> NoReturn: + raise ValueError(f"Unsupported Gemini passthrough: {detail}") + + try: + envelope = body["messages"] + if not isinstance(envelope, list) or len(envelope) != 1: + raise ValueError("expected singleton envelope") + content = envelope[0]["content"] + if not isinstance(content, str): + raise ValueError("expected JSON text") + native = json.loads(content) + if not isinstance(native, dict) or "contents" not in native: + raise ValueError("expected native contents") + except (KeyError, TypeError, ValueError): + model = ( + metadata.get("provider_model") + or metadata.get("request_model") + or body.get("model") + or "" + ) + if str(model).rsplit("/", 1)[-1].startswith(("gemini-", "gemma-")): + invalid("malformed native request envelope") + return exchange + + contents = native["contents"] + if not isinstance(contents, list): + invalid("contents must be a list") + messages: list[dict[str, Any]] = [] + call_ids: dict[str, str] = {} + system = native.get("systemInstruction") + if system is not None: + if not isinstance(system, dict): + invalid("systemInstruction must be an object") + contents = [{**system, "role": "system"}, *contents] + for turn in contents: + if not isinstance(turn, dict) or turn.get("role") not in { + "system", + "user", + "model", + }: + invalid("unknown content role") + parts = turn.get("parts") + if not isinstance(parts, list): + invalid("parts must be a list") + role = "assistant" if turn["role"] == "model" else turn["role"] + message: dict[str, Any] = {"role": role, "content": ""} + for part in parts: + if not isinstance(part, dict) or set(part) - { + "text", + "thought", + "thoughtSignature", + "functionCall", + "functionResponse", + }: + invalid("unsupported content part") + if len(set(part) & {"text", "functionCall", "functionResponse"}) != 1: + invalid("content part must have exactly one payload") + if "text" in part: + if not isinstance(part["text"], str): + invalid("text must be a string") + if not part.get("thought"): + message["content"] += part["text"] + elif "functionCall" in part: + call = part["functionCall"] + if ( + role != "assistant" + or not isinstance(call, dict) + or not isinstance(call.get("id"), str) + or not call["id"] + or not isinstance(call.get("name"), str) + or not call["name"] + or not isinstance(call.get("args", {}), dict) + ): + invalid("function call requires model role, id and name") + call_id = call["id"] + signature = part.get("thoughtSignature") + if signature is not None and not isinstance(signature, str): + invalid("thought signature must be a string") + decorated = call_id + (f"__thought__{signature}" if signature else "") + if call_id in call_ids and call_ids[call_id] != decorated: + invalid("conflicting signatures for function call id") + call_ids[call_id] = decorated + message.setdefault("tool_calls", []).append( + { + "id": decorated, + "name": call["name"], + "arguments": call.get("args", {}), + } + ) + elif "functionResponse" in part: + response = part["functionResponse"] + if ( + role != "user" + or not isinstance(response, dict) + or not isinstance(response.get("id"), str) + or not response["id"] + or not isinstance(response.get("response"), dict) + ): + invalid("function response requires user role and id") + if message["content"]: + messages.append(message) + message = {"role": role, "content": ""} + response_id = response["id"] + if response_id not in call_ids: + invalid("function response references unknown call id") + messages.append( + { + "role": "tool", + "tool_call_id": call_ids[response_id], + "content": json.dumps(response.get("response")), + } + ) + if message["content"] or message.get("tool_calls"): + messages.append(message) + tools = [] + declarations = native.get("tools", []) + if not isinstance(declarations, list): + invalid("tools must be a list") + for tool in declarations: + if ( + not isinstance(tool, dict) + or set(tool) != {"functionDeclarations"} + or not isinstance(tool["functionDeclarations"], list) + ): + invalid("unsupported tool declaration") + for declaration in tool["functionDeclarations"]: + if not isinstance(declaration, dict) or not declaration.get("name"): + invalid("function declaration requires a name") + parameters = declaration.get( + "parametersJsonSchema", + declaration.get("parameters", {"type": "object", "properties": {}}), + ) + if not isinstance(parameters, dict): + invalid("function declaration requires an object schema") + tools.append({**declaration, "parameters": parameters}) + return { + **exchange, + "metadata": {**metadata, "training_input_format": "gemini"}, + "request": { + **request, + "body": { + **body, + "messages": messages, + "tools": _tool_defs_from_body({"tools": tools}), + }, + }, + } + + +def _exchange_to_messages_and_tools( + exchange: dict[str, Any], + *, + redact: bool = True, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]], str | None]: + try: + exchange = normalize_provider_exchange(exchange) + except ValueError as exc: + return [], [], str(exc) + request = ( + cast(dict[str, Any], exchange.get("request")) + if isinstance(exchange.get("request"), dict) + else {} + ) + response = ( + cast(dict[str, Any], exchange.get("response")) + if isinstance(exchange.get("response"), dict) + else {} + ) + request_body = ( + cast(dict[str, Any], request.get("body")) + if isinstance(request.get("body"), dict) + else {} + ) + response_body = ( + cast(dict[str, Any], response.get("body")) + if isinstance(response.get("body"), dict) + else {} + ) + + if "messages" in request_body: + messages = _messages_from_chat_request(request_body, redact=redact) + assistant = _assistant_from_chat_response(response_body, redact=redact) + else: + messages = _messages_from_responses_request(request_body, redact=redact) + assistant = _assistant_from_responses_response(response_body, redact=redact) + + if assistant is None: + return [], [], "no_assistant" + messages.append(assistant) + return ( + _normalize_system_messages(messages), + _tool_defs_from_body(request_body), + None, + ) + + +def _has_tool_calls(messages: list[dict[str, Any]]) -> bool: + return any(bool(message.get("tool_calls")) for message in messages) + + +def _normalize_tools_for_validation( + row: dict[str, Any], row_num: int +) -> list[Any] | None: + tools = row.get("tool_defs", row.get("tools")) + if tools is None: + return None + if isinstance(tools, str): + try: + tools = json.loads(tools) + except json.JSONDecodeError as exc: + raise ValueError( + f"row {row_num}: tool_defs/tools is not valid JSON: {exc}" + ) from exc + if not isinstance(tools, list): + raise ValueError(f"row {row_num}: tool_defs/tools must be a list") + return tools + + +def _tool_names_for_validation(tools: list[Any] | None) -> set[str]: + names: set[str] = set() + for tool in tools or []: + if not isinstance(tool, dict): + continue + function = tool.get("function") + name = function.get("name") if isinstance(function, dict) else tool.get("name") + if isinstance(name, str) and name: + names.add(name) + return names + + +def _row_messages(row: dict[str, Any], row_num: int) -> list[Any]: + messages = row.get("messages") + if isinstance(messages, list) and messages: + return messages + prompt = row.get("prompt") + completion = row.get("completion") + if isinstance(prompt, list) and isinstance(completion, list): + combined = prompt + completion + if combined: + return combined + raise ValueError( + f"row {row_num}: expected non-empty messages or prompt+completion lists" + ) + + +def _content_join(left: Any, right: Any) -> str: + return "\n".join( + part for part in (_content_to_text(left), _content_to_text(right)) if part + ) + + +def _align_legacy_tool_call_ids( + segments: list[tuple[Literal["messages", "prompt", "completion"], Any]], +) -> tuple[ + list[tuple[Literal["messages", "prompt", "completion"], Any]], dict[str, int] +]: + """Repair legacy BenchFlow rows whose provider call ids drifted. + + Some historical ``results.jsonl`` artifacts preserved assistant tool-call ids + from one provider layer (``fc_*``) while the following tool messages used the + OpenAI-compatible ``call_*`` ids that the runtime sent back on the next turn. + Pair by message order and rewrite only when a pending assistant call exists; + true orphan tool outputs still fail validation. + """ + out: list[tuple[Literal["messages", "prompt", "completion"], Any]] = [] + pending: list[dict[str, Any]] = [] + stats = {"tool_call_ids_rewritten": 0, "tool_messages_merged": 0} + + for segment, raw_message in segments: + message = deepcopy(raw_message) + if not isinstance(message, dict): + out.append((segment, message)) + continue + + tool_calls = message.get("tool_calls") + if message.get("role") == "assistant" and isinstance(tool_calls, list): + pending.extend( + tool_call for tool_call in tool_calls if isinstance(tool_call, dict) + ) + + if message.get("role") == "tool": + tool_call_id = message.get("tool_call_id") + if ( + out + and isinstance(out[-1][1], dict) + and out[-1][1].get("role") == "tool" + and out[-1][1].get("tool_call_id") == tool_call_id + ): + out[-1][1]["content"] = _content_join( + out[-1][1].get("content"), + message.get("content"), + ) + stats["tool_messages_merged"] += 1 + continue + + match_index = next( + ( + idx + for idx, tool_call in enumerate(pending) + if tool_call.get("id") == tool_call_id + ), + None, + ) + if match_index is not None: + pending.pop(match_index) + elif pending and tool_call_id: + tool_call = pending.pop(0) + if tool_call.get("id") != tool_call_id: + tool_call["id"] = tool_call_id + stats["tool_call_ids_rewritten"] += 1 + + out.append((segment, message)) + + return out, stats + + +def validate_prime_sft_row(row: dict[str, Any], row_num: int = 1) -> None: + leaked = sorted(BANNED_ROW_KEYS.intersection(row)) + if leaked: + raise ValueError( + f"row {row_num}: banned leakage keys present: {', '.join(leaked)}" + ) + + messages = _row_messages(row, row_num) + tools = _normalize_tools_for_validation(row, row_num) + known_tool_names = _tool_names_for_validation(tools) + pending_tool_call_ids: set[str] = set() + + for idx, message in enumerate(messages): + if not isinstance(message, dict): + raise ValueError(f"row {row_num}: messages[{idx}] must be object") + message = cast(dict[str, Any], message) + leaked_message = sorted(BANNED_MESSAGE_KEYS.intersection(message)) + if leaked_message: + raise ValueError( + f"row {row_num}: messages[{idx}] has banned keys: {', '.join(leaked_message)}" + ) + role = message.get("role") + if role not in ALLOWED_ROLES: + raise ValueError(f"row {row_num}: messages[{idx}].role invalid: {role!r}") + if role == "system" and idx != 0: + raise ValueError( + f"row {row_num}: system message must be at index 0, got index {idx}" + ) + if "content" not in message and "tool_calls" not in message: + raise ValueError( + f"row {row_num}: messages[{idx}] needs content or tool_calls" + ) + tool_calls = message.get("tool_calls") + if tool_calls and role != "assistant": + raise ValueError( + f"row {row_num}: only assistant messages may contain tool_calls" + ) + if role == "tool" and not message.get("tool_call_id"): + raise ValueError(f"row {row_num}: tool message requires tool_call_id") + if role == "tool" and message.get("tool_call_id") not in pending_tool_call_ids: + raise ValueError( + f"row {row_num}: tool message references unknown tool_call_id" + ) + if role == "tool": + pending_tool_call_ids.discard(cast(str, message.get("tool_call_id"))) + if tool_calls is not None and not isinstance(tool_calls, list): + raise ValueError( + f"row {row_num}: messages[{idx}].tool_calls must be a list" + ) + if isinstance(tool_calls, list): + for tool_call_idx, tool_call in enumerate(tool_calls): + prefix = f"row {row_num}: messages[{idx}].tool_calls[{tool_call_idx}]" + if not isinstance(tool_call, dict): + raise ValueError(f"{prefix} must be object") + tool_call = cast(dict[str, Any], tool_call) + function = tool_call.get("function") + if not isinstance(function, dict): + raise ValueError(f"{prefix}.function must be object") + tool_call_id = tool_call.get("id") + if not isinstance(tool_call_id, str) or not tool_call_id: + raise ValueError(f"{prefix}.id must be a non-empty string") + if tool_call.get("type") != "function": + raise ValueError(f"{prefix}.type must be 'function'") + name = function.get("name") + if not isinstance(name, str) or not name: + raise ValueError( + f"{prefix}.function.name must be a non-empty string" + ) + if known_tool_names and name not in known_tool_names: + raise ValueError( + f"{prefix}.function.name {name!r} not found in tool_defs/tools" + ) + arguments = function.get("arguments") + if isinstance(arguments, str): + try: + parsed_arguments = json.loads(arguments) + except json.JSONDecodeError as exc: + raise ValueError( + f"{prefix}.function.arguments is not valid JSON: {exc}" + ) from exc + if not isinstance(parsed_arguments, dict): + raise ValueError( + f"{prefix}.function.arguments must be a JSON object" + ) + elif not isinstance(arguments, dict): + raise ValueError( + f"{prefix}.function.arguments must be a JSON object or JSON-encoded object" + ) + pending_tool_call_ids.add(tool_call_id) + + if not any(isinstance(m, dict) and m.get("role") == "assistant" for m in messages): + raise ValueError(f"row {row_num}: no assistant message") + + typed_messages = [m for m in messages if isinstance(m, dict)] + if _has_tool_calls(typed_messages) and not tools: + raise ValueError( + f"row {row_num}: assistant tool_calls require non-empty tool_defs/tools" + ) + if tools is not None: + for tool_idx, tool in enumerate(tools): + if not isinstance(tool, dict): + raise ValueError(f"row {row_num}: tool_defs[{tool_idx}] must be object") + function = tool.get("function") + name = ( + function.get("name") if isinstance(function, dict) else tool.get("name") + ) + if not isinstance(name, str) or not name: + raise ValueError( + f"row {row_num}: tool_defs[{tool_idx}] missing function name" + ) + + +def normalize_prime_sft_exchange( + exchange: dict[str, Any], + *, + redact: bool = True, +) -> tuple[PrimeSftExchangeData | None, str | None]: + """Normalize one raw LLM exchange through the Prime-SFT validator path.""" + messages, tool_defs, skip_reason = _exchange_to_messages_and_tools( + exchange, redact=redact + ) + if skip_reason: + return None, skip_reason + repaired, _ = _align_legacy_tool_call_ids( + [("messages", message) for message in messages] + ) + messages = [message for _, message in repaired] + if _has_tool_calls(messages) and not tool_defs: + return None, "missing_tool_defs" + try: + validate_prime_sft_row({"messages": messages, "tool_defs": tool_defs}, 1) + except ValueError as exc: + return None, f"invalid_prime_sft_row: {exc}" + return PrimeSftExchangeData(messages=messages, tool_defs=tool_defs), None diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index 1b3dff92b..5a78f635e 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -22,7 +22,7 @@ from typing import Any, cast from benchflow._utils.json_safe import scrub_non_finite -from benchflow.trajectories.export_prime_sft import ( +from benchflow.trajectories.message_contract import ( PrimeSftTrajectoryJsonlError, load_llm_trajectory_jsonl, normalize_prime_sft_exchange, diff --git a/tests/trajectories/test_export_prime_sft.py b/tests/trajectories/test_export_prime_sft.py index d305e402f..c17a2aafd 100644 --- a/tests/trajectories/test_export_prime_sft.py +++ b/tests/trajectories/test_export_prime_sft.py @@ -14,9 +14,9 @@ export_prime_sft_jsonl, load_llm_trajectory_jsonl, normalize_prime_sft_exchange, - normalize_provider_exchange, validate_prime_sft_jsonl, ) +from benchflow.trajectories.message_contract import normalize_provider_exchange @pytest.mark.parametrize("second_signature", ["", "other-signature"]) From d340440241bba0960e33fb0a0c1259dfd185f87c Mon Sep 17 00:00:00 2001 From: kywch Date: Sat, 5 Sep 2026 00:37:56 -0700 Subject: [PATCH 3/6] Use neutral names for trajectory message contracts --- src/benchflow/eval_artifacts.py | 4 +- .../trajectories/export_prime_sft.py | 25 ++++---- src/benchflow/trajectories/export_trl_sft.py | 8 +-- .../trajectories/message_contract.py | 58 ++++++++----------- src/benchflow/trajectories/results.py | 16 ++--- tests/trajectories/test_export_prime_sft.py | 23 +++++++- 6 files changed, 69 insertions(+), 65 deletions(-) diff --git a/src/benchflow/eval_artifacts.py b/src/benchflow/eval_artifacts.py index 1e3b3fc1a..6c88d351d 100644 --- a/src/benchflow/eval_artifacts.py +++ b/src/benchflow/eval_artifacts.py @@ -12,7 +12,7 @@ from benchflow._utils.text import truncate_end from benchflow.task.discovery import is_task_dir, resolve_task_collection_root from benchflow.trajectories.message_contract import ( - PrimeSftTrajectoryJsonlError, + TrajectoryJsonlError, load_llm_trajectory_jsonl, ) @@ -200,7 +200,7 @@ def _llm_trajectory_status(rollout_dir: Path) -> tuple[bool, bool, int]: return False, False, 0 try: rows = load_llm_trajectory_jsonl(path, strict=True) - except PrimeSftTrajectoryJsonlError: + except TrajectoryJsonlError: return True, False, 0 return True, True, len(rows) diff --git a/src/benchflow/trajectories/export_prime_sft.py b/src/benchflow/trajectories/export_prime_sft.py index af624be88..724288b78 100644 --- a/src/benchflow/trajectories/export_prime_sft.py +++ b/src/benchflow/trajectories/export_prime_sft.py @@ -18,32 +18,27 @@ from benchflow._utils.json_safe import dumps_finite, scrub_non_finite from benchflow.trajectories.message_contract import ( - PrimeSftExchangeData as PrimeSftExchangeData, -) -from benchflow.trajectories.message_contract import ( - PrimeSftTrajectoryJsonlError as PrimeSftTrajectoryJsonlError, -) -from benchflow.trajectories.message_contract import ( + NormalizedExchange, + TrajectoryJsonlError, _align_legacy_tool_call_ids, _has_tool_calls, _normalize_tool_call, _row_messages, + last_user_training_window, + normalize_exchange, + validate_message_record, ) from benchflow.trajectories.message_contract import ( load_llm_trajectory_jsonl as load_llm_trajectory_jsonl, ) -from benchflow.trajectories.message_contract import ( - normalize_prime_sft_exchange as normalize_prime_sft_exchange, -) -from benchflow.trajectories.message_contract import ( - prime_sft_last_user_training_window as prime_sft_last_user_training_window, -) -from benchflow.trajectories.message_contract import ( - validate_prime_sft_row as validate_prime_sft_row, -) from benchflow.trajectories.types import redact_trajectory_obj PrimeSftRowMode = Literal["rollout", "exchange"] +PrimeSftExchangeData = NormalizedExchange +PrimeSftTrajectoryJsonlError = TrajectoryJsonlError +normalize_prime_sft_exchange = normalize_exchange +prime_sft_last_user_training_window = last_user_training_window +validate_prime_sft_row = validate_message_record @dataclass diff --git a/src/benchflow/trajectories/export_trl_sft.py b/src/benchflow/trajectories/export_trl_sft.py index 252ee9c81..56b633ba2 100644 --- a/src/benchflow/trajectories/export_trl_sft.py +++ b/src/benchflow/trajectories/export_trl_sft.py @@ -23,8 +23,8 @@ ) from benchflow.trajectories.message_contract import ( load_llm_trajectory_jsonl, - normalize_prime_sft_exchange, - validate_prime_sft_row, + normalize_exchange, + validate_message_record, ) from benchflow.trajectories.types import redact_trajectory_obj @@ -226,7 +226,7 @@ def validate_trl_sft_row(row: dict[str, Any], row_num: int = 1) -> None: ) if "tool_defs" in row: raise ValueError(f"row {row_num}: TRL rows must use tools, not tool_defs") - validate_prime_sft_row( + validate_message_record( {"prompt": prompt, "completion": completion, "tools": tools}, row_num, ) @@ -256,7 +256,7 @@ def _row_from_exchange( exchange_idx: int, redact: bool, ) -> tuple[dict[str, Any] | None, str | None]: - normalized, skip_reason = normalize_prime_sft_exchange( + normalized, skip_reason = normalize_exchange( exchange, redact=redact, ) diff --git a/src/benchflow/trajectories/message_contract.py b/src/benchflow/trajectories/message_contract.py index faa64df1e..0fc4487cf 100644 --- a/src/benchflow/trajectories/message_contract.py +++ b/src/benchflow/trajectories/message_contract.py @@ -38,12 +38,12 @@ @dataclass(frozen=True) -class PrimeSftExchangeData: +class NormalizedExchange: messages: list[dict[str, Any]] tool_defs: list[dict[str, Any]] -class PrimeSftTrajectoryJsonlError(ValueError): +class TrajectoryJsonlError(ValueError): """Raised when an LLM trajectory JSONL file is not parseable.""" @@ -57,7 +57,7 @@ def load_llm_trajectory_jsonl( lines = path.read_text().splitlines() except OSError as exc: if strict: - raise PrimeSftTrajectoryJsonlError( + raise TrajectoryJsonlError( f"{path}: cannot read LLM trajectory JSONL: {exc}" ) from exc return records @@ -68,14 +68,14 @@ def load_llm_trajectory_jsonl( record = json.loads(line) except json.JSONDecodeError as exc: if strict: - raise PrimeSftTrajectoryJsonlError( + raise TrajectoryJsonlError( f"{path}: line {line_num}: invalid JSON: {exc}" ) from exc continue if isinstance(record, dict): records.append(record) elif strict: - raise PrimeSftTrajectoryJsonlError( + raise TrajectoryJsonlError( f"{path}: line {line_num}: top-level record must be an object" ) return records @@ -163,11 +163,8 @@ def _normalize_message( _normalize_tool_call( { "id": message.get("call_id") or message.get("id"), - "type": "function", - "function": { - "name": message.get("name"), - "arguments": message.get("arguments", {}), - }, + "name": message.get("name"), + "arguments": message.get("arguments", {}), }, index, redact=redact, @@ -201,8 +198,8 @@ def _normalize_message( def _normalize_system_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - # Prime-RL SFT allows a system message only at index 0 (see - # validate_prime_sft_row). Any system message after the first position — + # The message contract allows a system message only at index 0 (see + # validate_message_record). Any system message after the first position — # including a *second consecutive* leading system message — is remapped to # "user" so the whole row isn't silently dropped into skipped_invalid. normalized: list[dict[str, Any]] = [] @@ -214,7 +211,7 @@ def _normalize_system_messages(messages: list[dict[str, Any]]) -> list[dict[str, return normalized -def prime_sft_last_user_training_window( +def last_user_training_window( messages: list[dict[str, Any]], ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]] | None: """Return a compact prompt/completion window anchored at the last user turn. @@ -264,12 +261,9 @@ def _messages_from_responses_request( if isinstance(raw_input, str): messages.append({"role": "user", "content": raw_input}) elif isinstance(raw_input, list): - for idx, item in enumerate(raw_input): - if not isinstance(item, dict): - continue - item = cast(dict[str, Any], item) - if item.get("type") != "reasoning": - messages.append(_normalize_message(item, idx, redact=redact)) + messages.extend( + _messages_from_chat_request({"messages": raw_input}, redact=redact) + ) return messages @@ -316,11 +310,8 @@ def _assistant_from_anthropic_content( raw_tool_calls = [ { "id": item.get("id"), - "type": "function", - "function": { - "name": item.get("name"), - "arguments": item.get("input", {}), - }, + "name": item.get("name"), + "arguments": item.get("input", {}), } for item in content if isinstance(item, dict) and item.get("type") == "tool_use" @@ -378,11 +369,8 @@ def _assistant_from_responses_response( tool_calls.append( { "id": item.get("call_id") or item.get("id"), - "type": "function", - "function": { - "name": item.get("name"), - "arguments": item.get("arguments", {}), - }, + "name": item.get("name"), + "arguments": item.get("arguments", {}), } ) if not texts and not tool_calls: @@ -731,7 +719,7 @@ def _align_legacy_tool_call_ids( return out, stats -def validate_prime_sft_row(row: dict[str, Any], row_num: int = 1) -> None: +def validate_message_record(row: dict[str, Any], row_num: int = 1) -> None: leaked = sorted(BANNED_ROW_KEYS.intersection(row)) if leaked: raise ValueError( @@ -843,12 +831,12 @@ def validate_prime_sft_row(row: dict[str, Any], row_num: int = 1) -> None: ) -def normalize_prime_sft_exchange( +def normalize_exchange( exchange: dict[str, Any], *, redact: bool = True, -) -> tuple[PrimeSftExchangeData | None, str | None]: - """Normalize one raw LLM exchange through the Prime-SFT validator path.""" +) -> tuple[NormalizedExchange | None, str | None]: + """Normalize one raw LLM exchange through the shared message contract.""" messages, tool_defs, skip_reason = _exchange_to_messages_and_tools( exchange, redact=redact ) @@ -861,7 +849,7 @@ def normalize_prime_sft_exchange( if _has_tool_calls(messages) and not tool_defs: return None, "missing_tool_defs" try: - validate_prime_sft_row({"messages": messages, "tool_defs": tool_defs}, 1) + validate_message_record({"messages": messages, "tool_defs": tool_defs}, 1) except ValueError as exc: return None, f"invalid_prime_sft_row: {exc}" - return PrimeSftExchangeData(messages=messages, tool_defs=tool_defs), None + return NormalizedExchange(messages=messages, tool_defs=tool_defs), None diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index 5a78f635e..ab46524fd 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -23,12 +23,12 @@ from benchflow._utils.json_safe import scrub_non_finite from benchflow.trajectories.message_contract import ( - PrimeSftTrajectoryJsonlError, + TrajectoryJsonlError, + last_user_training_window, load_llm_trajectory_jsonl, - normalize_prime_sft_exchange, + normalize_exchange, normalize_provider_exchange, - prime_sft_last_user_training_window, - validate_prime_sft_row, + validate_message_record, ) from benchflow.trajectories.types import redact_trajectory_obj from benchflow.usage_tracking import USAGE_SOURCE_AGENT_NATIVE_ACP @@ -167,7 +167,7 @@ def _llm_steps_from_trajectory( normalize_provider_exchange(exchange) for exchange in load_llm_trajectory_jsonl(path, strict=True) ] - except PrimeSftTrajectoryJsonlError as exc: + except TrajectoryJsonlError as exc: return [], [], f"Invalid LLM trajectory JSONL: {exc}" except ValueError as exc: return [], [], str(exc) @@ -182,7 +182,7 @@ def _llm_steps_from_trajectory( f"exchange {exchange_idx}: selected response is not an object" ) continue - normalized, skip_reason = normalize_prime_sft_exchange(exchange) + normalized, skip_reason = normalize_exchange(exchange) if normalized is None: skipped_successful.append( f"exchange {exchange_idx}: {skip_reason or 'normalization failed'}" @@ -417,7 +417,7 @@ def _top_level_prompt_completion( typed_full_messages = [ message for message in full_messages if isinstance(message, dict) ] - window = prime_sft_last_user_training_window(typed_full_messages) + window = last_user_training_window(typed_full_messages) if window is not None: return window if prompt and len(full_messages) >= len(prompt): @@ -584,7 +584,7 @@ def _prime_sft_validation_error( if not completion: return None try: - validate_prime_sft_row( + validate_message_record( { "prompt": prompt, "completion": completion, diff --git a/tests/trajectories/test_export_prime_sft.py b/tests/trajectories/test_export_prime_sft.py index c17a2aafd..788491561 100644 --- a/tests/trajectories/test_export_prime_sft.py +++ b/tests/trajectories/test_export_prime_sft.py @@ -9,14 +9,35 @@ import pytest from benchflow.trajectories.export_prime_sft import ( + PrimeSftExchangeData, PrimeSftTrajectoryJsonlError, convert_benchflow_rollouts_to_prime_sft_rows, export_prime_sft_jsonl, load_llm_trajectory_jsonl, normalize_prime_sft_exchange, + prime_sft_last_user_training_window, validate_prime_sft_jsonl, + validate_prime_sft_row, ) -from benchflow.trajectories.message_contract import normalize_provider_exchange +from benchflow.trajectories.message_contract import ( + NormalizedExchange, + TrajectoryJsonlError, + last_user_training_window, + normalize_exchange, + normalize_provider_exchange, + validate_message_record, +) + + +def test_message_contract_preserves_prime_compatibility_aliases(): + assert PrimeSftExchangeData is NormalizedExchange + assert PrimeSftTrajectoryJsonlError is TrajectoryJsonlError + assert normalize_prime_sft_exchange is normalize_exchange + assert prime_sft_last_user_training_window is last_user_training_window + assert validate_prime_sft_row is validate_message_record + assert repr(NormalizedExchange(messages=[], tool_defs=[])).startswith( + "NormalizedExchange(" + ) @pytest.mark.parametrize("second_signature", ["", "other-signature"]) From c6a8cec8cf0d10e2c8ebc34756f91640a9e3f1e2 Mon Sep 17 00:00:00 2001 From: kywch Date: Sat, 5 Sep 2026 00:47:25 -0700 Subject: [PATCH 4/6] Preserve raw Gemini request identity during result selection --- src/benchflow/trajectories/results.py | 50 +++++++++++++-------- tests/test_train_mode_artifact_emission.py | 31 +++++++++++++ tests/trajectories/test_export_prime_sft.py | 8 +++- 3 files changed, 69 insertions(+), 20 deletions(-) diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index ab46524fd..e528e235c 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -163,15 +163,17 @@ def _llm_steps_from_trajectory( if not path.exists(): return steps, tool_defs, None try: + raw_exchanges = load_llm_trajectory_jsonl(path, strict=True) exchanges = [ - normalize_provider_exchange(exchange) - for exchange in load_llm_trajectory_jsonl(path, strict=True) + normalize_provider_exchange(exchange) for exchange in raw_exchanges ] except TrajectoryJsonlError as exc: return [], [], f"Invalid LLM trajectory JSONL: {exc}" except ValueError as exc: return [], [], str(exc) - training_success_indices = _training_success_exchange_indices(exchanges) + training_success_indices = _training_success_exchange_indices( + exchanges, request_identity_exchanges=raw_exchanges + ) skipped_successful: list[str] = [] for exchange_idx, exchange in enumerate(exchanges): response = exchange.get("response") @@ -268,26 +270,36 @@ def _response_is_training_success(response: Any) -> bool: def _training_success_exchange_indices( exchanges: list[dict[str, Any]], + *, + request_identity_exchanges: list[dict[str, Any]] | None = None, ) -> set[int]: - request_bodies = [ - json.dumps( - ( - request.get("body") - if isinstance(request := exchange.get("request"), dict) - and isinstance(request.get("body"), dict) - else {} - ), - sort_keys=True, - separators=(",", ":"), - default=str, - ) - for exchange in exchanges - ] + def request_bodies(rows: list[dict[str, Any]]) -> list[str]: + return [ + json.dumps( + ( + request.get("body") + if isinstance(request := exchange.get("request"), dict) + and isinstance(request.get("body"), dict) + else {} + ), + sort_keys=True, + separators=(",", ":"), + default=str, + ) + for exchange in rows + ] + + normalized_request_bodies = request_bodies(exchanges) + request_identities = ( + request_bodies(request_identity_exchanges) + if request_identity_exchanges is not None + else normalized_request_bodies + ) candidates_by_request: dict[str, list[int]] = {} for exchange_idx, exchange in enumerate(exchanges): if not _response_is_training_success(exchange.get("response")): continue - candidates_by_request.setdefault(request_bodies[exchange_idx], []).append( + candidates_by_request.setdefault(request_identities[exchange_idx], []).append( exchange_idx ) selected: set[int] = set() @@ -296,7 +308,7 @@ def _training_success_exchange_indices( exchange_idx for exchange_idx in candidates if _response_consumed_by_later_request( - exchanges, request_bodies, exchange_idx + exchanges, normalized_request_bodies, exchange_idx ) ] if consumed: diff --git a/tests/test_train_mode_artifact_emission.py b/tests/test_train_mode_artifact_emission.py index c5da44596..ad9143a27 100644 --- a/tests/test_train_mode_artifact_emission.py +++ b/tests/test_train_mode_artifact_emission.py @@ -28,6 +28,7 @@ from benchflow.trajectories.export_prime_sft import validate_prime_sft_jsonl from benchflow.trajectories.results import ( JOB_RESULTS_ERRORS_FILENAME, + _llm_steps_from_trajectory, _training_success_exchange_indices, write_job_results_jsonl, ) @@ -939,6 +940,36 @@ def test_results_drop_unique_late_unconsumed_nonterminal_retry(): } +def test_results_retry_grouping_preserves_raw_request_identity(tmp_path): + exchanges = [] + for content, temperature in (("first", 0.1), ("second", 0.9)): + exchange = _llm_exchange(assistant={"role": "assistant", "content": content}) + native = { + "contents": [{"role": "user", "parts": [{"text": "same prompt"}]}], + "generationConfig": {"temperature": temperature}, + } + exchange["metadata"] = { + "call_type": "pass_through_endpoint", + "provider_model": "gemini-test", + } + exchange["request"]["body"] = { + "model": "gemini-test", + "messages": [{"role": "user", "content": json.dumps(native)}], + } + exchanges.append(exchange) + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + (trajectory_dir / "llm_trajectory.jsonl").write_text( + "".join(json.dumps(exchange) + "\n" for exchange in exchanges) + ) + + steps, _, error = _llm_steps_from_trajectory( + tmp_path, reward=1, is_truncated=False, trajectory_id_prefix="test" + ) + assert error is None + assert [step["completion"][0]["content"] for step in steps] == ["first", "second"] + + def test_results_jsonl_preserves_llm_exchange_metadata(tmp_path): """Guards PR #925: trainer conversion retains LLM call purpose metadata.""" rollout_dir = tmp_path / "rollout-call-purpose" diff --git a/tests/trajectories/test_export_prime_sft.py b/tests/trajectories/test_export_prime_sft.py index 788491561..da9a2b25f 100644 --- a/tests/trajectories/test_export_prime_sft.py +++ b/tests/trajectories/test_export_prime_sft.py @@ -137,14 +137,20 @@ def test_gemini_passthrough_export_preserves_tools_and_signatures( "contents": [], "tools": [{"functionDeclarations": [{"name": "bad", "parameters": None}]}], }, + "{malformed", ], ) def test_gemini_unrepresentable_history_blocks_training(native): """Guards a0b16985 against silently accepting malformed native history.""" exchange = _exchange(final=False) exchange["metadata"] = {"call_type": "pass_through_endpoint"} + if isinstance(native, str): + exchange["request"]["body"]["model"] = "gemini-test" exchange["request"]["body"]["messages"] = [ - {"role": "user", "content": json.dumps(native)} + { + "role": "user", + "content": native if isinstance(native, str) else json.dumps(native), + } ] normalized, reason = normalize_prime_sft_exchange(exchange) assert normalized is None From 5dc2e7f494bb8ffb73b62cfd6c8fdc080278585a Mon Sep 17 00:00:00 2001 From: kywch Date: Sat, 5 Sep 2026 01:14:02 -0700 Subject: [PATCH 5/6] Support Gemini tool histories without call IDs --- .../trajectories/message_contract.py | 83 +++++++++++++++---- tests/trajectories/test_export_prime_sft.py | 78 +++++++++++++++++ 2 files changed, 145 insertions(+), 16 deletions(-) diff --git a/src/benchflow/trajectories/message_contract.py b/src/benchflow/trajectories/message_contract.py index 0fc4487cf..12c1b7ba8 100644 --- a/src/benchflow/trajectories/message_contract.py +++ b/src/benchflow/trajectories/message_contract.py @@ -431,7 +431,19 @@ def invalid(detail: str) -> NoReturn: if not isinstance(contents, list): invalid("contents must be a list") messages: list[dict[str, Any]] = [] - call_ids: dict[str, str] = {} + pending_calls: list[tuple[str | None, str, str]] = [] + explicit_ids: set[str] = set() + for turn in contents: + parts = turn.get("parts") if isinstance(turn, dict) else None + for part in parts if isinstance(parts, list) else []: + call = part.get("functionCall") if isinstance(part, dict) else None + if isinstance(call, dict) and isinstance(call.get("id"), str): + signature = part.get("thoughtSignature") + explicit_ids.add( + call["id"] + + (f"__thought__{signature}" if isinstance(signature, str) else "") + ) + generated_id = 0 system = native.get("systemInstruction") if system is not None: if not isinstance(system, dict): @@ -470,21 +482,38 @@ def invalid(detail: str) -> NoReturn: if ( role != "assistant" or not isinstance(call, dict) - or not isinstance(call.get("id"), str) - or not call["id"] or not isinstance(call.get("name"), str) or not call["name"] or not isinstance(call.get("args", {}), dict) ): - invalid("function call requires model role, id and name") - call_id = call["id"] + invalid("function call requires model role and name") + call_id = call.get("id") + if call_id is not None and ( + not isinstance(call_id, str) or not call_id + ): + invalid("function call id must be a non-empty string") signature = part.get("thoughtSignature") if signature is not None and not isinstance(signature, str): invalid("thought signature must be a string") - decorated = call_id + (f"__thought__{signature}" if signature else "") - if call_id in call_ids and call_ids[call_id] != decorated: - invalid("conflicting signatures for function call id") - call_ids[call_id] = decorated + if call_id is None: + while True: + candidate = f"gemini_call_{generated_id}" + generated_id += 1 + decorated = candidate + ( + f"__thought__{signature}" if signature else "" + ) + if ( + candidate not in explicit_ids + and decorated not in explicit_ids + ): + break + normalized_id = candidate + else: + normalized_id = call_id + decorated = normalized_id + ( + f"__thought__{signature}" if signature else "" + ) + pending_calls.append((call_id, call["name"], decorated)) message.setdefault("tool_calls", []).append( { "id": decorated, @@ -497,21 +526,43 @@ def invalid(detail: str) -> NoReturn: if ( role != "user" or not isinstance(response, dict) - or not isinstance(response.get("id"), str) - or not response["id"] + or not isinstance(response.get("name"), str) + or not response["name"] or not isinstance(response.get("response"), dict) ): - invalid("function response requires user role and id") + invalid("function response requires user role and name") + response_id = response.get("id") + if response_id is not None and ( + not isinstance(response_id, str) or not response_id + ): + invalid("function response id must be a non-empty string") if message["content"]: messages.append(message) message = {"role": role, "content": ""} - response_id = response["id"] - if response_id not in call_ids: - invalid("function response references unknown call id") + matches = [ + index + for index, (source_id, name, _) in enumerate(pending_calls) + if ( + source_id == response_id + if response_id is not None + else name == response["name"] + ) + ] + if not matches: + invalid( + "function response references unknown call id" + if response_id is not None + else "function response references unknown call name" + ) + if len(matches) > 1: + invalid("function response is ambiguous") + _, name, normalized_id = pending_calls.pop(matches[0]) + if response_id is not None and name != response["name"]: + invalid("function response name does not match call") messages.append( { "role": "tool", - "tool_call_id": call_ids[response_id], + "tool_call_id": normalized_id, "content": json.dumps(response.get("response")), } ) diff --git a/tests/trajectories/test_export_prime_sft.py b/tests/trajectories/test_export_prime_sft.py index da9a2b25f..f59fb9561 100644 --- a/tests/trajectories/test_export_prime_sft.py +++ b/tests/trajectories/test_export_prime_sft.py @@ -126,6 +126,84 @@ def test_gemini_passthrough_export_preserves_tools_and_signatures( assert tools[0]["function"]["parameters"] == function["parameters"] +def test_gemini_passthrough_pairs_sequential_idless_calls_by_name(): + exchange = _exchange(final=False) + function = exchange["request"]["body"]["tools"][0]["function"] + exchange["metadata"] = {"call_type": "pass_through_endpoint"} + native = { + "contents": [{"role": "user", "parts": [{"text": "Run twice."}]}], + "tools": [ + { + "functionDeclarations": [ + { + "name": function["name"], + "parametersJsonSchema": function["parameters"], + } + ] + } + ], + } + for command in ("pwd", "ls"): + native["contents"].extend( + [ + { + "role": "model", + "parts": [ + {"functionCall": {"name": "bash", "args": {"command": command}}} + ], + }, + { + "role": "user", + "parts": [ + { + "functionResponse": { + "name": "bash", + "response": {"output": command}, + } + } + ], + }, + ] + ) + exchange["request"]["body"] = { + "messages": [{"role": "user", "content": json.dumps(native)}] + } + + normalized, reason = normalize_exchange(exchange, redact=False) + + assert reason is None + calls = [m for m in normalized.messages if m.get("tool_calls")] + results = [m for m in normalized.messages if m["role"] == "tool"] + assert [m["tool_calls"][0]["id"] for m in calls] == [ + "gemini_call_0", + "gemini_call_1", + ] + assert [m["tool_call_id"] for m in results] == [ + "gemini_call_0", + "gemini_call_1", + ] + + native["contents"][1:] = [ + { + "role": "model", + "parts": [ + {"functionCall": {"name": "bash"}}, + {"functionCall": {"name": "bash"}}, + ], + }, + { + "role": "user", + "parts": [ + {"functionResponse": {"name": "bash", "response": {"output": "x"}}} + ], + }, + ] + exchange["request"]["body"]["messages"][0]["content"] = json.dumps(native) + rejected, reason = normalize_exchange(exchange) + assert rejected is None + assert "ambiguous" in reason + + @pytest.mark.parametrize( "native", [ From e0d7d9d64384734fc31fbf91cc8839a3c46dd264 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 5 Sep 2026 04:38:44 -0700 Subject: [PATCH 6/6] fix(trajectories): retain idless Gemini tool steps --- .../scripts/validate_run_artifacts.py | 171 +++++++++++++----- src/benchflow/trajectories/results.py | 151 +++++++++++----- ...t_benchflow_experiment_review_validator.py | 66 +++++++ tests/test_litellm_config.py | 4 +- tests/test_train_mode_artifact_emission.py | 83 +++++++++ 5 files changed, 386 insertions(+), 89 deletions(-) diff --git a/.agents/skills/benchflow-experiment-review/scripts/validate_run_artifacts.py b/.agents/skills/benchflow-experiment-review/scripts/validate_run_artifacts.py index c6a127e21..9102eeb59 100755 --- a/.agents/skills/benchflow-experiment-review/scripts/validate_run_artifacts.py +++ b/.agents/skills/benchflow-experiment-review/scripts/validate_run_artifacts.py @@ -27,6 +27,7 @@ import argparse import json +from contextlib import suppress from pathlib import Path from typing import Any @@ -373,11 +374,17 @@ def response_consumed_by_later_request( response = rows[exchange_idx].get("response") body = response.get("body") if isinstance(response, dict) else {} call_ids = response_call_ids(body if isinstance(body, dict) else {}) + call_fingerprints = response_tool_call_fingerprints( + body if isinstance(body, dict) else {} + ) return bool( - call_ids + (call_ids or call_fingerprints) and any( any(call_id in request_bodies[index] for call_id in call_ids) or bool(call_ids & gemini_history_call_ids(rows[index])) + or bool( + call_fingerprints & gemini_history_tool_call_fingerprints(rows[index]) + ) for index in range(exchange_idx + 1, len(rows)) ) ) @@ -421,6 +428,96 @@ def gemini_history_call_ids(row: dict[str, Any]) -> set[str]: return ids +def canonical_tool_arguments(arguments: Any) -> str: + if isinstance(arguments, str): + with suppress(json.JSONDecodeError): + arguments = json.loads(arguments) + return json.dumps(arguments, sort_keys=True, separators=(",", ":"), default=str) + + +def thought_signature(call_id: Any) -> str | None: + if not isinstance(call_id, str) or "__thought__" not in call_id: + return None + return call_id.split("__thought__", 1)[1] + + +def tool_call_fingerprint(call: dict[str, Any]) -> tuple[str, str, str | None] | None: + function = call.get("function") + name = function.get("name") if isinstance(function, dict) else call.get("name") + if not isinstance(name, str) or not name: + return None + arguments = ( + function.get("arguments", {}) + if isinstance(function, dict) + else call.get("arguments", call.get("args", {})) + ) + return ( + name, + canonical_tool_arguments(arguments), + thought_signature(call.get("id") or call.get("tool_call_id")), + ) + + +def response_tool_call_fingerprints( + body: dict[str, Any], +) -> set[tuple[str, str, str | None]]: + return { + fingerprint + for call in response_tool_call_objects(body) + if (fingerprint := tool_call_fingerprint(call)) is not None + } + + +def gemini_history_tool_call_fingerprints( + row: dict[str, Any], +) -> set[tuple[str, str, str | None]]: + metadata = row.get("metadata") + if ( + not isinstance(metadata, dict) + or metadata.get("call_type") != "pass_through_endpoint" + ): + return set() + request = row.get("request") + body = request.get("body") if isinstance(request, dict) else None + messages = body.get("messages") if isinstance(body, dict) else None + if ( + not isinstance(messages, list) + or len(messages) != 1 + or not isinstance(messages[0], dict) + ): + return set() + try: + native = json.loads(messages[0].get("content", "")) + except (ValueError, TypeError): + return set() + contents = native.get("contents") if isinstance(native, dict) else None + if not isinstance(contents, list): + return set() + fingerprints: set[tuple[str, str, str | None]] = set() + for content in contents: + parts = content.get("parts") if isinstance(content, dict) else None + for part in parts if isinstance(parts, list) else []: + call = part.get("functionCall") if isinstance(part, dict) else None + if not isinstance(call, dict): + continue + fingerprint = tool_call_fingerprint( + { + **call, + "id": ( + str(call.get("id") or "") + + ( + f"__thought__{part['thoughtSignature']}" + if isinstance(part.get("thoughtSignature"), str) + else "" + ) + ), + } + ) + if fingerprint is not None: + fingerprints.add(fingerprint) + return fingerprints + + def response_call_ids(body: dict[str, Any]) -> set[str]: return {call_id for call_id, _ in response_tool_calls(body)} @@ -442,55 +539,39 @@ def response_safe_without_later_consumption( def response_tool_calls(body: dict[str, Any]) -> list[tuple[str, str | None]]: + calls = [] + for call in response_tool_call_objects(body): + call_id = call.get("call_id") or call.get("id") + if not call_id: + continue + function = call.get("function") + name = function.get("name") if isinstance(function, dict) else call.get("name") + calls.append((str(call_id), str(name) if name else None)) + return calls + + +def response_tool_call_objects(body: dict[str, Any]) -> list[dict[str, Any]]: calls = [ - ( - str(item.get("call_id") or item.get("id")), - str(item.get("name")) if item.get("name") else None, - ) + item for item in body.get("output") or [] - if isinstance(item, dict) - and item.get("type") in {"function_call", "tool_call"} - and (item.get("call_id") or item.get("id")) + if isinstance(item, dict) and item.get("type") in {"function_call", "tool_call"} ] + messages = [] choices = body.get("choices") if isinstance(choices, list): - for choice in choices: - message = choice.get("message") if isinstance(choice, dict) else None - if not isinstance(message, dict): - continue - calls.extend( - ( - str(call.get("id") or call.get("tool_call_id")), - ( - str(function.get("name")) - if isinstance(function := call.get("function"), dict) - and function.get("name") - else str(call.get("name")) - if call.get("name") - else None - ), - ) - for call in message.get("tool_calls") or [] - if isinstance(call, dict) - and (call.get("id") or call.get("tool_call_id")) - ) + messages.extend( + message + for choice in choices + if isinstance(choice, dict) + and isinstance(message := choice.get("message"), dict) + ) message = body.get("message") if isinstance(message, dict): - calls.extend( - ( - str(call.get("id") or call.get("tool_call_id")), - ( - str(function.get("name")) - if isinstance(function := call.get("function"), dict) - and function.get("name") - else str(call.get("name")) - if call.get("name") - else None - ), - ) - for call in message.get("tool_calls") or [] - if isinstance(call, dict) and (call.get("id") or call.get("tool_call_id")) - ) + messages.append(message) + for message in messages: + raw_calls = message.get("tool_calls") + if isinstance(raw_calls, list): + calls.extend(call for call in raw_calls if isinstance(call, dict)) return calls @@ -646,7 +727,9 @@ def validate_results_row( if native_subscription_without_llm: expected_reason = "missing_healthy_structured_llm_trajectory" if training_ready is not False: - issues.append(f"{row_path}: native subscription row must not be training-ready") + issues.append( + f"{row_path}: native subscription row must not be training-ready" + ) if info.get("training_ready_reason") != expected_reason: issues.append( f"{row_path}: native subscription row has unexpected training_ready_reason" diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index e528e235c..fa9b365ad 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -332,11 +332,20 @@ def _response_consumed_by_later_request( response = exchanges[exchange_idx].get("response") body = response.get("body") if isinstance(response, dict) else {} call_ids = _response_call_ids(body if isinstance(body, dict) else {}) + call_fingerprints = _response_tool_call_fingerprints( + body if isinstance(body, dict) else {} + ) return bool( - call_ids + (call_ids or call_fingerprints) and any( any(call_id in request_body for call_id in call_ids) - for request_body in request_bodies[exchange_idx + 1 :] + or bool( + call_fingerprints + & _gemini_history_tool_call_fingerprints(exchanges[later_idx]) + ) + for later_idx, request_body in enumerate( + request_bodies[exchange_idx + 1 :], start=exchange_idx + 1 + ) ) ) @@ -345,6 +354,76 @@ def _response_call_ids(body: dict[str, Any]) -> set[str]: return {call_id for call_id, _ in _response_tool_calls(body)} +def _canonical_tool_arguments(arguments: Any) -> str: + if isinstance(arguments, str): + with suppress(json.JSONDecodeError): + arguments = json.loads(arguments) + return json.dumps(arguments, sort_keys=True, separators=(",", ":"), default=str) + + +def _thought_signature(call_id: Any) -> str | None: + if not isinstance(call_id, str) or "__thought__" not in call_id: + return None + return call_id.split("__thought__", 1)[1] + + +def _tool_call_fingerprint(call: dict[str, Any]) -> tuple[str, str, str | None] | None: + function = call.get("function") + name = function.get("name") if isinstance(function, dict) else call.get("name") + if not isinstance(name, str) or not name: + return None + arguments = ( + function.get("arguments", {}) + if isinstance(function, dict) + else call.get("arguments", call.get("args", {})) + ) + return ( + name, + _canonical_tool_arguments(arguments), + _thought_signature(call.get("id") or call.get("tool_call_id")), + ) + + +def _response_tool_call_fingerprints( + body: dict[str, Any], +) -> set[tuple[str, str, str | None]]: + return { + fingerprint + for call in _response_tool_call_objects(body) + if (fingerprint := _tool_call_fingerprint(call)) is not None + } + + +def _gemini_history_tool_call_fingerprints( + exchange: dict[str, Any], +) -> set[tuple[str, str, str | None]]: + metadata = exchange.get("metadata") + if ( + not isinstance(metadata, dict) + or metadata.get("training_input_format") != "gemini" + ): + return set() + request = exchange.get("request") + body = request.get("body") if isinstance(request, dict) else None + messages = body.get("messages") if isinstance(body, dict) else None + if not isinstance(messages, list): + return set() + calls = [ + call + for message in messages + if isinstance(message, dict) + and message.get("role") == "assistant" + and isinstance(message.get("tool_calls"), list) + for call in message["tool_calls"] + if isinstance(call, dict) + ] + return { + fingerprint + for call in calls + if (fingerprint := _tool_call_fingerprint(call)) is not None + } + + def _response_safe_without_later_consumption( exchanges: list[dict[str, Any]], exchange_idx: int ) -> bool: @@ -362,55 +441,39 @@ def _response_safe_without_later_consumption( def _response_tool_calls(body: dict[str, Any]) -> list[tuple[str, str | None]]: + calls = [] + for call in _response_tool_call_objects(body): + call_id = call.get("call_id") or call.get("id") + if not call_id: + continue + function = call.get("function") + name = function.get("name") if isinstance(function, dict) else call.get("name") + calls.append((str(call_id), str(name) if name else None)) + return calls + + +def _response_tool_call_objects(body: dict[str, Any]) -> list[dict[str, Any]]: calls = [ - ( - str(item.get("call_id") or item.get("id")), - str(item.get("name")) if item.get("name") else None, - ) + item for item in body.get("output") or [] - if isinstance(item, dict) - and item.get("type") in {"function_call", "tool_call"} - and (item.get("call_id") or item.get("id")) + if isinstance(item, dict) and item.get("type") in {"function_call", "tool_call"} ] + messages = [] choices = body.get("choices") if isinstance(choices, list): - for choice in choices: - message = choice.get("message") if isinstance(choice, dict) else None - if not isinstance(message, dict): - continue - calls.extend( - ( - str(call.get("id") or call.get("tool_call_id")), - ( - str(function.get("name")) - if isinstance(function := call.get("function"), dict) - and function.get("name") - else str(call.get("name")) - if call.get("name") - else None - ), - ) - for call in message.get("tool_calls") or [] - if isinstance(call, dict) - and (call.get("id") or call.get("tool_call_id")) - ) + messages.extend( + message + for choice in choices + if isinstance(choice, dict) + and isinstance(message := choice.get("message"), dict) + ) message = body.get("message") if isinstance(message, dict): - calls.extend( - ( - str(call.get("id") or call.get("tool_call_id")), - ( - str(function.get("name")) - if isinstance(function := call.get("function"), dict) - and function.get("name") - else str(call.get("name")) - if call.get("name") - else None - ), - ) - for call in message.get("tool_calls") or [] - if isinstance(call, dict) and (call.get("id") or call.get("tool_call_id")) - ) + messages.append(message) + for message in messages: + raw_calls = message.get("tool_calls") + if isinstance(raw_calls, list): + calls.extend(call for call in raw_calls if isinstance(call, dict)) return calls diff --git a/tests/test_benchflow_experiment_review_validator.py b/tests/test_benchflow_experiment_review_validator.py index 2060e865a..dd9ccfc6b 100644 --- a/tests/test_benchflow_experiment_review_validator.py +++ b/tests/test_benchflow_experiment_review_validator.py @@ -563,6 +563,72 @@ def test_validator_prefers_duplicate_consumed_by_later_request( assert validator.validate_rollout(rollout)["healthy"] is False +def test_validator_counts_consumed_gemini_call_without_provider_id( + tmp_path: Path, +) -> None: + """Guards PR #1106: validator must match idless native Gemini history.""" + validator = _load_validator() + rollout = _rollout(tmp_path) + llm_path = rollout / "trajectory" / "llm_trajectory.jsonl" + first = json.loads(llm_path.read_text()) + first["response"]["body"]["choices"][0]["message"] = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_callback_123", + "type": "function", + "function": { + "name": "bash", + "arguments": '{"command":"pwd"}', + }, + } + ], + } + final = json.loads(json.dumps(first)) + final["metadata"] = {"call_type": "pass_through_endpoint"} + native = { + "contents": [ + {"role": "user", "parts": [{"text": "Run pwd."}]}, + { + "role": "model", + "parts": [ + {"functionCall": {"name": "bash", "args": {"command": "pwd"}}} + ], + }, + { + "role": "user", + "parts": [ + { + "functionResponse": { + "name": "bash", + "response": {"output": "/workspace"}, + } + } + ], + }, + ] + } + final["request"]["body"] = { + "messages": [{"role": "user", "content": json.dumps(native)}] + } + final["response"]["body"]["choices"][0]["message"] = { + "role": "assistant", + "content": "Done.", + } + _write_jsonl(llm_path, [first, final]) + row = json.loads((rollout / "results.jsonl").read_text()) + row["trajectory"] = [ + {**row["trajectory"][0], "extras": {"exchange_index": 0}}, + {**row["trajectory"][0], "extras": {"exchange_index": 1}}, + ] + _write_jsonl(rollout / "results.jsonl", [row]) + + report = validator.validate_rollout(rollout) + + assert report["artifacts"]["llm"]["successful_exchange_indices"] == [0, 1] + + def test_validator_excludes_unique_late_unconsumed_nonterminal_retry( tmp_path: Path, ) -> None: diff --git a/tests/test_litellm_config.py b/tests/test_litellm_config.py index 1bd8a19d4..7dbbb6072 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -13,7 +13,9 @@ @pytest.mark.parametrize( "model", ["gemini-3.1-pro-preview", "gemini-3.8-flash", "gemini-3.5-flash-lite"] ) -async def test_gemini_vertex_passthrough_exchanges_gateway_auth_for_adc(monkeypatch, model): +async def test_gemini_vertex_passthrough_exchanges_gateway_auth_for_adc( + monkeypatch, model +): """Guards Vertex proxy auth against the regression in commit 28b82e33.""" from unittest.mock import AsyncMock, Mock diff --git a/tests/test_train_mode_artifact_emission.py b/tests/test_train_mode_artifact_emission.py index ad9143a27..e811c185d 100644 --- a/tests/test_train_mode_artifact_emission.py +++ b/tests/test_train_mode_artifact_emission.py @@ -970,6 +970,89 @@ def test_results_retry_grouping_preserves_raw_request_identity(tmp_path): assert [step["completion"][0]["content"] for step in steps] == ["first", "second"] +def test_results_preserve_consumed_gemini_call_without_provider_id(tmp_path): + """Guards PR #1106: idless Gemini history must retain the consumed tool step.""" + first = _llm_exchange( + assistant={ + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_callback_123", + "type": "function", + "function": { + "name": "bash", + "arguments": '{"command":"pwd"}', + }, + } + ], + } + ) + native = { + "contents": [{"role": "user", "parts": [{"text": "Run pwd."}]}], + "tools": [ + { + "functionDeclarations": [ + { + "name": "bash", + "parametersJsonSchema": { + "type": "object", + "properties": {"command": {"type": "string"}}, + }, + } + ] + } + ], + } + first["metadata"] = { + "call_type": "pass_through_endpoint", + "provider_model": "gemini-test", + } + first["request"]["body"] = { + "model": "gemini-test", + "messages": [{"role": "user", "content": json.dumps(native)}], + } + final = json.loads(json.dumps(first)) + native["contents"].extend( + [ + { + "role": "model", + "parts": [ + {"functionCall": {"name": "bash", "args": {"command": "pwd"}}} + ], + }, + { + "role": "user", + "parts": [ + { + "functionResponse": { + "name": "bash", + "response": {"output": "/workspace"}, + } + } + ], + }, + ] + ) + final["request"]["body"]["messages"][0]["content"] = json.dumps(native) + final["response"]["body"]["choices"][0]["message"] = { + "role": "assistant", + "content": "Done.", + } + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + (trajectory_dir / "llm_trajectory.jsonl").write_text( + json.dumps(first) + "\n" + json.dumps(final) + "\n" + ) + + steps, _, error = _llm_steps_from_trajectory( + tmp_path, reward=1, is_truncated=False, trajectory_id_prefix="test" + ) + + assert error is None + assert [step["extras"]["exchange_index"] for step in steps] == [0, 1] + + def test_results_jsonl_preserves_llm_exchange_metadata(tmp_path): """Guards PR #925: trainer conversion retains LLM call purpose metadata.""" rollout_dir = tmp_path / "rollout-call-purpose"