-
Notifications
You must be signed in to change notification settings - Fork 1
test(reindex): add whale fixture bounds #3891
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
14f4465
8ecaa1c
f2d2f36
cf68f69
14dea39
db6cbfd
cb05930
6b5de09
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Large diffs are not rendered by default.
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4,12 +4,15 @@ | |
|
|
||
| import hashlib | ||
| import json | ||
| import pickle | ||
| import re | ||
| import shlex | ||
| import tempfile | ||
| from collections import defaultdict | ||
| from collections.abc import Iterable, Sequence | ||
| from collections.abc import Iterable, Iterator, Sequence | ||
| from dataclasses import dataclass, replace | ||
| from datetime import datetime | ||
| from typing import BinaryIO | ||
|
|
||
| from pydantic import ValidationError | ||
|
|
||
|
|
@@ -78,6 +81,22 @@ | |
| _STRUCTURAL_BYTE_KEYS = frozenset({"bytes", "byte_count", "bytes_written", "size_bytes", "written_bytes"}) | ||
|
|
||
|
|
||
| @dataclass(slots=True) | ||
| class _PickleRecordReplay: | ||
| """Re-iterable disk spool for preserving parser lookahead without retaining records.""" | ||
|
|
||
| handle: BinaryIO | ||
|
|
||
| def __iter__(self) -> Iterator[object]: | ||
| self.handle.seek(0) | ||
| while True: | ||
| try: | ||
| # This reads only the private spool written immediately above. | ||
| yield pickle.load(self.handle) | ||
| except EOFError: | ||
| return | ||
|
|
||
|
|
||
| @dataclass(frozen=True, slots=True) | ||
| class _CodexExecChildType: | ||
| kind: str | ||
|
|
@@ -1555,7 +1574,7 @@ def _response_inner_record(item: object) -> dict[str, object] | None: | |
| return inner if inner is not None and not _is_message(inner) else None | ||
|
|
||
|
|
||
| def _code_mode_exec_envelopes(records: Sequence[object]) -> dict[int, _CodexExecEnvelope]: | ||
| def _code_mode_exec_envelopes(records: Iterable[object]) -> dict[int, _CodexExecEnvelope]: | ||
| call_occurrences: dict[str, list[tuple[int, _CodexExecEnvelope]]] = defaultdict(list) | ||
| output_occurrences: dict[str, list[tuple[int, dict[str, object]]]] = defaultdict(list) | ||
| envelopes_by_record: dict[int, _CodexExecEnvelope] = {} | ||
|
|
@@ -2160,11 +2179,34 @@ def _sanitize_codex_data_url(value: str) -> str: | |
| return value | ||
| header, encoded = value.split(",", 1) | ||
| mime = header.removeprefix("data:").split(";", 1)[0] or "image/unknown" | ||
| digest = hashlib.sha256(encoded.encode("ascii", errors="ignore")).hexdigest() | ||
| approx_bytes = (len(encoded.rstrip("=")) * 3) // 4 | ||
| digest_builder = hashlib.sha256() | ||
| for offset in range(0, len(encoded), 1024 * 1024): | ||
| digest_builder.update(encoded[offset : offset + 1024 * 1024].encode("ascii", errors="ignore")) | ||
| digest = digest_builder.hexdigest() | ||
| padding = 2 if encoded.endswith("==") else 1 if encoded.endswith("=") else 0 | ||
| approx_bytes = max(0, (len(encoded) * 3) // 4 - padding) | ||
| return f"<inline image omitted; mime={mime}; approx_bytes={approx_bytes}; sha256_base64={digest}>" | ||
|
|
||
|
|
||
| def _codex_inline_image_blocks(content: object) -> tuple[ParsedContentBlock, ...]: | ||
| """Return typed, bounded evidence for inline images without authored prose inflation.""" | ||
| if not isinstance(content, list): | ||
| return () | ||
| blocks: list[ParsedContentBlock] = [] | ||
| for item in content: | ||
| if not isinstance(item, dict) or item.get("type") not in {"input_image", "image"}: | ||
| continue | ||
| image_url = item.get("image_url") | ||
| if not isinstance(image_url, str): | ||
| continue | ||
| summary = _sanitize_codex_data_url(image_url) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
For every ordinary inline image, this new call reaches Useful? React with 👍 / 👎. |
||
| if summary != image_url: | ||
| header = image_url.split(",", 1)[0] | ||
| mime = header.removeprefix("data:").split(";", 1)[0] or "image/unknown" | ||
| blocks.append(ParsedContentBlock(type=BlockType.IMAGE, text=summary, media_type=mime)) | ||
| return tuple(blocks) | ||
|
|
||
|
|
||
| def _message_signature(role: Role | str, text: str | None) -> tuple[str, str]: | ||
| role_value = role.value if isinstance(role, Role) else str(role) | ||
| return (role_value, " ".join((text or "").split())) | ||
|
|
@@ -2337,7 +2379,7 @@ def is_supported_session_stream(payload: Sequence[object]) -> bool: | |
| return has_session_header or has_direct_record or has_envelope_record | ||
|
|
||
|
|
||
| def _parse_records(records: Iterable[object], fallback_id: str) -> ParsedSession: | ||
| def _parse_records(records: Iterable[object], fallback_id: str, *, _reiterable: bool = False) -> ParsedSession: | ||
| """Parse Codex JSONL session file using typed CodexRecord model. | ||
|
|
||
| Supports two format generations via CodexRecord.format_type: | ||
|
|
@@ -2350,9 +2392,17 @@ def _parse_records(records: Iterable[object], fallback_id: str) -> ParsedSession | |
| - text_content: Extracted text from any format | ||
| - format_type: Detected format generation | ||
| """ | ||
| record_list = list(records) | ||
| code_mode_envelopes = _code_mode_exec_envelopes(record_list) | ||
| response_signatures = _response_message_signatures(record_list) | ||
| if not isinstance(records, Sequence) and not _reiterable: | ||
| # The parser needs two lookahead-derived indexes before the materializing | ||
| # pass. Persist a private replay spool so a multi-million-record JSONL | ||
| # stream stays bounded by one decoded record instead of list(records). | ||
| with tempfile.TemporaryFile(mode="w+b") as spool: | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When production ingest supplies a one-shot Codex JSONL iterator, every decoded record is now pickled into Useful? React with 👍 / 👎. |
||
| for item in records: | ||
| pickle.dump(item, spool, protocol=pickle.HIGHEST_PROTOCOL) | ||
| return _parse_records(_PickleRecordReplay(spool), fallback_id, _reiterable=True) | ||
|
|
||
| code_mode_envelopes = _code_mode_exec_envelopes(records) | ||
| response_signatures = _response_message_signatures(records) | ||
| messages: list[ParsedMessage] = [] | ||
| session_events: list[ParsedSessionEvent] = [] | ||
| session_id = fallback_id | ||
|
|
@@ -2393,7 +2443,7 @@ def _parse_records(records: Iterable[object], fallback_id: str) -> ParsedSession | |
| session_model_provider: str | None = None | ||
| session_developer_instructions: str | None = None | ||
|
|
||
| for idx, item in enumerate(record_list, start=1): | ||
| for idx, item in enumerate(records, start=1): | ||
| record = _dict_record(item) | ||
| if record is None: | ||
| continue | ||
|
|
@@ -2750,12 +2800,15 @@ def _parse_records(records: Iterable[object], fallback_id: str) -> ParsedSession | |
| raw_role = _effective_role(message_record) | ||
| content = _effective_content(message_record) | ||
| text = extract_codex_text(content) | ||
| inline_image_blocks = _codex_inline_image_blocks(content) | ||
| timestamp_pair = parse_timestamp_pair(_message_timestamp(record, message_record)) | ||
| timestamp = timestamp_pair[1] if timestamp_pair is not None else None | ||
|
|
||
| content_blocks = content_blocks_from_segments(content) | ||
| content_blocks.extend(inline_image_blocks) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When a Codex message contains an Useful? React with 👍 / 👎.
coderabbitai[bot] marked this conversation as resolved.
|
||
| has_structured = any( | ||
| cb.type in (BlockType.TOOL_USE, BlockType.TOOL_RESULT, BlockType.THINKING) for cb in content_blocks | ||
| cb.type in (BlockType.TOOL_USE, BlockType.TOOL_RESULT, BlockType.THINKING, BlockType.IMAGE) | ||
| for cb in content_blocks | ||
| ) | ||
| if not raw_role or raw_role == "unknown": | ||
| continue | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,23 @@ | ||
| """Generate the private-data-free Codex whale fixture pack on disk.""" | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import argparse | ||
| from collections.abc import Sequence | ||
| from pathlib import Path | ||
|
|
||
| from tests.infra.whale_fixtures import write_codex_whale_fixture_pack | ||
|
|
||
|
|
||
| def main(argv: Sequence[str] | None = None) -> int: | ||
| parser = argparse.ArgumentParser(description=__doc__) | ||
| parser.add_argument("output_dir", type=Path, help="directory for codex-whale.jsonl and manifest.json") | ||
| args = parser.parse_args(argv) | ||
| source_path, manifest_path = write_codex_whale_fixture_pack(args.output_dir) | ||
| print(source_path) | ||
| print(manifest_path) | ||
| return 0 | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| raise SystemExit(main()) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
For a
type: "image"segment carrying a data URL inimage_url,content_blocks_from_segments()already emits anIMAGEblock, and this helper emits a second block for the same segment. Such messages are persisted and rendered with duplicate attachment blocks, inflating block counts and attachment evidence; restrict this additional conversion to the unsupportedinput_imageshape or enrich the existing image block instead.Useful? React with 👍 / 👎.