atif-sql 0.2.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- atif_analytics/__init__.py +9 -0
- atif_analytics/application/__init__.py +3 -0
- atif_analytics/application/analyze.py +158 -0
- atif_analytics/application/prompts.py +750 -0
- atif_analytics/application/use_cases/__init__.py +3 -0
- atif_analytics/application/use_cases/_shared.py +264 -0
- atif_analytics/application/use_cases/classify.py +430 -0
- atif_analytics/application/use_cases/conflicts.py +448 -0
- atif_analytics/application/use_cases/friction.py +634 -0
- atif_analytics/application/use_cases/perceived.py +451 -0
- atif_analytics/domain/__init__.py +3 -0
- atif_analytics/domain/authorship.py +166 -0
- atif_analytics/domain/config.py +30 -0
- atif_analytics/domain/costs.py +50 -0
- atif_analytics/domain/friction.py +117 -0
- atif_analytics/domain/layout.py +94 -0
- atif_analytics/domain/models.py +385 -0
- atif_analytics/domain/ports.py +100 -0
- atif_analytics/domain/transcript.py +359 -0
- atif_analytics/infrastructure/__init__.py +3 -0
- atif_analytics/infrastructure/corpus_reader.py +546 -0
- atif_analytics/infrastructure/parquet_cache.py +186 -0
- atif_analytics/infrastructure/settings.py +114 -0
- atif_analytics/infrastructure/sqlite_state/__init__.py +3 -0
- atif_analytics/infrastructure/sqlite_state/checkpointer.py +326 -0
- atif_analytics/infrastructure/sqlite_state/retry_queue.py +296 -0
- atif_analytics/py.typed +0 -0
- atif_cli/__init__.py +3 -0
- atif_cli/__main__.py +8 -0
- atif_cli/app.py +2459 -0
- atif_cli/converter_adapter.py +151 -0
- atif_cli/corpus.py +409 -0
- atif_cli/cron.py +311 -0
- atif_cli/duck_errors.py +97 -0
- atif_cli/embed_lake.py +101 -0
- atif_cli/errors.py +89 -0
- atif_cli/lake.py +373 -0
- atif_cli/lake_sessions.py +270 -0
- atif_cli/output.py +297 -0
- atif_cli/py.typed +0 -0
- atif_converter/__init__.py +3 -0
- atif_converter/application/__init__.py +3 -0
- atif_converter/application/convert_and_audit.py +294 -0
- atif_converter/application/convert_codex.py +210 -0
- atif_converter/domain/__init__.py +9 -0
- atif_converter/domain/agents.py +40 -0
- atif_converter/domain/atif/__init__.py +55 -0
- atif_converter/domain/atif/agent.py +37 -0
- atif_converter/domain/atif/content.py +161 -0
- atif_converter/domain/atif/final_metrics.py +45 -0
- atif_converter/domain/atif/metrics.py +49 -0
- atif_converter/domain/atif/observation.py +21 -0
- atif_converter/domain/atif/observation_result.py +47 -0
- atif_converter/domain/atif/step.py +144 -0
- atif_converter/domain/atif/subagent_trajectory_ref.py +90 -0
- atif_converter/domain/atif/tool_call.py +36 -0
- atif_converter/domain/atif/trajectory.py +211 -0
- atif_converter/domain/atif/trajectory_validator.py +299 -0
- atif_converter/domain/blobs.py +510 -0
- atif_converter/domain/claude_code_conversion.py +1077 -0
- atif_converter/domain/codex_conversion.py +1036 -0
- atif_converter/domain/codex_edges.py +192 -0
- atif_converter/domain/codex_enrichment.py +419 -0
- atif_converter/domain/codex_fidelity.py +124 -0
- atif_converter/domain/codex_nested_calls.py +165 -0
- atif_converter/domain/edges.py +113 -0
- atif_converter/domain/enrichment.py +412 -0
- atif_converter/domain/errors.py +60 -0
- atif_converter/domain/fidelity.py +171 -0
- atif_converter/domain/model_prices.json +8538 -0
- atif_converter/domain/pricing.py +790 -0
- atif_converter/domain/result_signals.py +774 -0
- atif_converter/domain/schema_version.py +90 -0
- atif_converter/domain/session_events.py +408 -0
- atif_converter/infrastructure/__init__.py +3 -0
- atif_converter/infrastructure/census.py +89 -0
- atif_converter/infrastructure/claude_code_converter.py +127 -0
- atif_converter/infrastructure/codex_adapter.py +132 -0
- atif_converter/infrastructure/codex_census.py +100 -0
- atif_converter/infrastructure/codex_converter.py +82 -0
- atif_converter/infrastructure/harbor_adapter.py +187 -0
- atif_converter/infrastructure/raw_records.py +426 -0
- atif_converter/infrastructure/source_archive.py +154 -0
- atif_converter/py.typed +0 -0
- atif_corpus/__init__.py +13 -0
- atif_corpus/application/__init__.py +7 -0
- atif_corpus/application/materialize.py +1826 -0
- atif_corpus/domain/__init__.py +9 -0
- atif_corpus/domain/agents.py +33 -0
- atif_corpus/domain/generation.py +34 -0
- atif_corpus/domain/layout.py +168 -0
- atif_corpus/domain/ports.py +175 -0
- atif_corpus/domain/session_id.py +68 -0
- atif_corpus/domain/sessions.py +228 -0
- atif_corpus/domain/slug.py +72 -0
- atif_corpus/domain/source_layout.py +143 -0
- atif_corpus/domain/watermark.py +76 -0
- atif_corpus/infrastructure/__init__.py +7 -0
- atif_corpus/infrastructure/atomic.py +266 -0
- atif_corpus/infrastructure/compress_artifacts.py +161 -0
- atif_corpus/infrastructure/fake_converter.py +174 -0
- atif_corpus/infrastructure/scanner.py +316 -0
- atif_corpus/infrastructure/settings.py +130 -0
- atif_corpus/infrastructure/source_archive.py +202 -0
- atif_corpus/py.typed +0 -0
- atif_duck/__init__.py +43 -0
- atif_duck/domain/__init__.py +32 -0
- atif_duck/domain/artifacts.py +66 -0
- atif_duck/domain/authorship.py +165 -0
- atif_duck/domain/catalog.py +626 -0
- atif_duck/domain/columnar.py +149 -0
- atif_duck/domain/embedding_guard.py +88 -0
- atif_duck/domain/examples.py +265 -0
- atif_duck/domain/lake.py +427 -0
- atif_duck/domain/lake_sessions.py +97 -0
- atif_duck/domain/raw_readers.py +140 -0
- atif_duck/domain/session_id.py +69 -0
- atif_duck/domain/sql_literal.py +50 -0
- atif_duck/infrastructure/__init__.py +34 -0
- atif_duck/infrastructure/analytics.py +469 -0
- atif_duck/infrastructure/authorship.py +139 -0
- atif_duck/infrastructure/columnar.py +641 -0
- atif_duck/infrastructure/lake.py +1383 -0
- atif_duck/infrastructure/lake_sessions.py +194 -0
- atif_duck/infrastructure/lake_settings.py +63 -0
- atif_duck/infrastructure/lake_steps.py +225 -0
- atif_duck/infrastructure/projections.py +221 -0
- atif_duck/infrastructure/registry.py +2194 -0
- atif_duck/infrastructure/stored_artifacts.py +115 -0
- atif_duck/py.typed +0 -0
- atif_embed/__init__.py +12 -0
- atif_embed/application/__init__.py +3 -0
- atif_embed/application/embed.py +293 -0
- atif_embed/application/prune.py +80 -0
- atif_embed/domain/__init__.py +3 -0
- atif_embed/domain/discovery.py +70 -0
- atif_embed/domain/embedding_guard.py +83 -0
- atif_embed/domain/errors.py +94 -0
- atif_embed/domain/ports.py +223 -0
- atif_embed/domain/sql_literal.py +49 -0
- atif_embed/domain/text_stamp.py +71 -0
- atif_embed/infrastructure/__init__.py +3 -0
- atif_embed/infrastructure/cohere_bedrock.py +433 -0
- atif_embed/infrastructure/corpus_text_rows.py +273 -0
- atif_embed/infrastructure/gemma_local.py +355 -0
- atif_embed/infrastructure/lake_text_rows.py +203 -0
- atif_embed/infrastructure/lance_store.py +440 -0
- atif_embed/infrastructure/providers.py +30 -0
- atif_embed/infrastructure/settings.py +154 -0
- atif_embed/py.typed +0 -0
- atif_models/__init__.py +14 -0
- atif_models/domain/__init__.py +9 -0
- atif_models/domain/ports.py +143 -0
- atif_models/domain/registry.py +151 -0
- atif_models/domain/schema.py +109 -0
- atif_models/infrastructure/__init__.py +8 -0
- atif_models/infrastructure/openai_bedrock.py +314 -0
- atif_models/infrastructure/settings.py +85 -0
- atif_models/py.typed +0 -0
- atif_sql-0.2.0.dist-info/METADATA +272 -0
- atif_sql-0.2.0.dist-info/RECORD +164 -0
- atif_sql-0.2.0.dist-info/WHEEL +4 -0
- atif_sql-0.2.0.dist-info/entry_points.txt +2 -0
- atif_sql-0.2.0.dist-info/licenses/LICENSE +202 -0
|
@@ -0,0 +1,192 @@
|
|
|
1
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
2
|
+
|
|
3
|
+
"""edges.jsonl for Codex rollouts — the id-keyed skeleton of a rollout.
|
|
4
|
+
|
|
5
|
+
Pure functions over already-parsed record dicts; no harbor, no I/O. The line
|
|
6
|
+
shape is the SAME nine contract fields the Claude Code emitter produces
|
|
7
|
+
(:mod:`atif_converter.domain.edges`), so atif-duck's ``messages`` view reads
|
|
8
|
+
one projection for every agent:
|
|
9
|
+
|
|
10
|
+
{uuid, parent_uuid, message_id, type, ts, is_sidechain,
|
|
11
|
+
is_compact_summary, source_file, tool_use_ids: [..]}
|
|
12
|
+
|
|
13
|
+
Three fields mean something different for a rollout than for a Claude Code
|
|
14
|
+
transcript, and the difference is the point of this module rather than a
|
|
15
|
+
shortcut:
|
|
16
|
+
|
|
17
|
+
``uuid``
|
|
18
|
+
A rollout record has no session-wide uuid. ``response_item`` payloads
|
|
19
|
+
usually carry an ``id`` (``msg_*`` / ``call_*`` / ``rs_*``), and newer
|
|
20
|
+
Codex builds add a per-file ``ordinal`` to every record — but a
|
|
21
|
+
``function_call_output`` carries neither, and rollouts written before the
|
|
22
|
+
ordinal landed carry it nowhere. :func:`record_id` therefore falls back to
|
|
23
|
+
a SYNTHETIC ``<type>:L<line>`` id, deterministic in the file's bytes and
|
|
24
|
+
visibly synthetic so a reader cannot mistake it for something Codex
|
|
25
|
+
assigned.
|
|
26
|
+
|
|
27
|
+
``parent_uuid``
|
|
28
|
+
Always ``None``. A rollout is a flat append-only log with no parent
|
|
29
|
+
pointer: the ordering IS the structure. Inventing a chain by pairing each
|
|
30
|
+
record with its predecessor would produce a tree that looks authoritative
|
|
31
|
+
and describes nothing, so the column stays null and
|
|
32
|
+
``turn_id``-derived grouping is left to ``codex_turn_id`` on the steps.
|
|
33
|
+
|
|
34
|
+
``type``
|
|
35
|
+
The raw ROLE where the record has one (``user`` / ``assistant`` /
|
|
36
|
+
``developer``), and the item kind otherwise (``function_call``,
|
|
37
|
+
``reasoning``, ``compacted``, …). This is what preserves the
|
|
38
|
+
developer-vs-system distinction that harbor's step conversion flattens
|
|
39
|
+
(:attr:`~atif_converter.domain.codex_fidelity.CodexFidelityGap.DEVELOPER_ROLE_FLATTENED_TO_SYSTEM`).
|
|
40
|
+
"""
|
|
41
|
+
|
|
42
|
+
from __future__ import annotations
|
|
43
|
+
|
|
44
|
+
from typing import Any
|
|
45
|
+
|
|
46
|
+
from atif_converter.domain.edges import EDGE_FIELDS, edge_to_jsonl_line
|
|
47
|
+
|
|
48
|
+
#: Payload types whose ``call_id`` is a tool-call join key back to ATIF
|
|
49
|
+
#: ``tool_calls[].tool_call_id`` / ``observation.results[].source_call_id``.
|
|
50
|
+
_CALL_ID_ITEM_TYPES: frozenset[str] = frozenset(
|
|
51
|
+
{
|
|
52
|
+
"function_call",
|
|
53
|
+
"function_call_output",
|
|
54
|
+
"custom_tool_call",
|
|
55
|
+
"custom_tool_call_output",
|
|
56
|
+
"tool_search_call",
|
|
57
|
+
"tool_search_output",
|
|
58
|
+
}
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def _payload(record: dict[str, Any]) -> dict[str, Any]:
|
|
63
|
+
payload = record.get("payload")
|
|
64
|
+
return payload if isinstance(payload, dict) else {}
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def record_id(record: dict[str, Any], line_number: int) -> str:
|
|
68
|
+
"""A stable id for one rollout record: its own, else a synthetic one.
|
|
69
|
+
|
|
70
|
+
Preference order, and each rung exists because the one above it is absent
|
|
71
|
+
on real rollouts (measured 2026-09-11 over 78 local rollouts): the
|
|
72
|
+
payload's own ``id``, then the record's per-file ``ordinal``, then the
|
|
73
|
+
1-based ``line_number``. The synthetic forms carry the record type so a
|
|
74
|
+
reader can see at a glance that Codex did not assign the id.
|
|
75
|
+
"""
|
|
76
|
+
payload = _payload(record)
|
|
77
|
+
own_id = payload.get("id")
|
|
78
|
+
if isinstance(own_id, str) and own_id:
|
|
79
|
+
return own_id
|
|
80
|
+
record_type = record.get("type")
|
|
81
|
+
kind = record_type if isinstance(record_type, str) and record_type else "record"
|
|
82
|
+
ordinal = record.get("ordinal")
|
|
83
|
+
if isinstance(ordinal, int):
|
|
84
|
+
return f"{kind}:{ordinal}"
|
|
85
|
+
return f"{kind}:L{line_number}"
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
def _edge_type(record: dict[str, Any]) -> str | None:
|
|
89
|
+
"""The raw role for a message, else the item/record kind."""
|
|
90
|
+
record_type = record.get("type")
|
|
91
|
+
payload = _payload(record)
|
|
92
|
+
if record_type == "response_item":
|
|
93
|
+
item_type = payload.get("type")
|
|
94
|
+
if item_type == "message":
|
|
95
|
+
role = payload.get("role")
|
|
96
|
+
return role if isinstance(role, str) and role else "message"
|
|
97
|
+
return item_type if isinstance(item_type, str) else None
|
|
98
|
+
if record_type == "event_msg":
|
|
99
|
+
item_type = payload.get("type")
|
|
100
|
+
return f"event:{item_type}" if isinstance(item_type, str) else "event"
|
|
101
|
+
return record_type if isinstance(record_type, str) else None
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
def _tool_use_ids(record: dict[str, Any]) -> list[str]:
|
|
105
|
+
"""The record's ``call_id``, when it is a tool call or a tool output."""
|
|
106
|
+
payload = _payload(record)
|
|
107
|
+
item_type = payload.get("type")
|
|
108
|
+
# A `type` that is an object or a list is unhashable, and the membership test raised
|
|
109
|
+
# `TypeError` on it (found by the fuzzer). Only a string can name a call item.
|
|
110
|
+
if not isinstance(item_type, str) or item_type not in _CALL_ID_ITEM_TYPES:
|
|
111
|
+
return []
|
|
112
|
+
call_id = payload.get("call_id")
|
|
113
|
+
return [call_id] if isinstance(call_id, str) and call_id else []
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
def codex_edge_from_record(
|
|
117
|
+
record: dict[str, Any],
|
|
118
|
+
source_file: str,
|
|
119
|
+
line_number: int,
|
|
120
|
+
) -> dict[str, Any]:
|
|
121
|
+
"""Project one rollout record onto the contract's edge shape."""
|
|
122
|
+
payload = _payload(record)
|
|
123
|
+
message_id = payload.get("id") if payload.get("type") == "message" else None
|
|
124
|
+
return {
|
|
125
|
+
"uuid": record_id(record, line_number),
|
|
126
|
+
# Flat log, no parent pointer — see the module docstring.
|
|
127
|
+
"parent_uuid": None,
|
|
128
|
+
"message_id": message_id if isinstance(message_id, str) else None,
|
|
129
|
+
"type": _edge_type(record),
|
|
130
|
+
"ts": record.get("timestamp"),
|
|
131
|
+
# Codex sub-agents get their own rollout file rather than an inlined
|
|
132
|
+
# sidechain, so no record of THIS session is a sidechain record.
|
|
133
|
+
"is_sidechain": False,
|
|
134
|
+
"is_compact_summary": record.get("type") == "compacted",
|
|
135
|
+
"source_file": source_file,
|
|
136
|
+
"tool_use_ids": _tool_use_ids(record),
|
|
137
|
+
}
|
|
138
|
+
|
|
139
|
+
|
|
140
|
+
def build_codex_edges(records: list[tuple[dict[str, Any], str]]) -> list[dict[str, Any]]:
|
|
141
|
+
"""Build every edge from ``(record, source_file)`` pairs, in FILE order.
|
|
142
|
+
|
|
143
|
+
File order, not timestamp order: a rollout is append-only, so its line
|
|
144
|
+
order is the causal order, and two records inside one millisecond share a
|
|
145
|
+
timestamp often enough that sorting by it would shuffle a tool call past
|
|
146
|
+
its own output. The line number is also what
|
|
147
|
+
:func:`record_id` falls back to, so the two must be counted the same way —
|
|
148
|
+
per source file, 1-based.
|
|
149
|
+
"""
|
|
150
|
+
line_numbers: dict[str, int] = {}
|
|
151
|
+
used_ids: set[str] = set()
|
|
152
|
+
edges: list[dict[str, Any]] = []
|
|
153
|
+
for record, source_file in records:
|
|
154
|
+
line_numbers[source_file] = line_numbers.get(source_file, 0) + 1
|
|
155
|
+
edge = codex_edge_from_record(record, source_file, line_numbers[source_file])
|
|
156
|
+
# UNIQUENESS GUARD. `messages` is a uuid-keyed view, so two edges
|
|
157
|
+
# sharing a uuid would silently double a record. No duplicate payload
|
|
158
|
+
# id appears across the 74 local rollouts measured 2026-09-11, but the
|
|
159
|
+
# ids come from a source this workspace does not control, so a repeat
|
|
160
|
+
# falls back to the synthetic form rather than shadowing the first.
|
|
161
|
+
if str(edge["uuid"]) in used_ids:
|
|
162
|
+
kind = record.get("type")
|
|
163
|
+
prefix = kind if isinstance(kind, str) and kind else "record"
|
|
164
|
+
edge["uuid"] = f"{prefix}:L{line_numbers[source_file]}"
|
|
165
|
+
used_ids.add(str(edge["uuid"]))
|
|
166
|
+
edges.append(edge)
|
|
167
|
+
return edges
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
def codex_edges_jsonl_lines(records: list[tuple[dict[str, Any], str]]) -> list[str]:
|
|
171
|
+
"""``build_codex_edges`` + serialization: the ready-to-write edges.jsonl body."""
|
|
172
|
+
return [edge_to_jsonl_line(edge) for edge in build_codex_edges(records)]
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
def codex_edge_ids(records: list[tuple[dict[str, Any], str]]) -> list[str]:
|
|
176
|
+
"""Every record's edge id, in the same order :func:`build_codex_edges` uses.
|
|
177
|
+
|
|
178
|
+
The enrichment pass attributes steps to these ids, so it must derive them
|
|
179
|
+
from the one function the edges are built with rather than recomputing the
|
|
180
|
+
fallback rules and drifting from them.
|
|
181
|
+
"""
|
|
182
|
+
return [str(edge["uuid"]) for edge in build_codex_edges(records)]
|
|
183
|
+
|
|
184
|
+
|
|
185
|
+
__all__ = [
|
|
186
|
+
"EDGE_FIELDS",
|
|
187
|
+
"build_codex_edges",
|
|
188
|
+
"codex_edge_from_record",
|
|
189
|
+
"codex_edge_ids",
|
|
190
|
+
"codex_edges_jsonl_lines",
|
|
191
|
+
"record_id",
|
|
192
|
+
]
|
|
@@ -0,0 +1,419 @@
|
|
|
1
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
2
|
+
|
|
3
|
+
"""Post-conversion enrichment for Codex: restore what harbor 0.22.0 drops.
|
|
4
|
+
|
|
5
|
+
Pure function over a trajectory dict plus the rollout's raw records — no
|
|
6
|
+
harbor import, no I/O. The Claude Code counterpart is
|
|
7
|
+
:mod:`atif_converter.domain.enrichment`; this module repairs the Codex gaps
|
|
8
|
+
that are repairable wrap-locally and, exactly like that one, REFUSES rather
|
|
9
|
+
than guesses when its re-derivation disagrees with harbor.
|
|
10
|
+
|
|
11
|
+
What it repairs, by the gap it answers
|
|
12
|
+
--------------------------------------
|
|
13
|
+
``ITEM_IDS_NOT_PRESERVED``
|
|
14
|
+
``step.extra["source_uuids"]`` lists the rollout record ids that
|
|
15
|
+
contributed to the step, using the SAME key the Claude Code pass writes so
|
|
16
|
+
atif-duck's ``steps.source_uuids`` column means one thing for every agent.
|
|
17
|
+
|
|
18
|
+
Cache split
|
|
19
|
+
``trajectory.extra["cache_creation_total"]`` is surfaced from
|
|
20
|
+
``final_metrics.extra.total_cache_write_input_tokens``, again matching the
|
|
21
|
+
Claude Code key so one query answers cache-creation for both agents.
|
|
22
|
+
|
|
23
|
+
``COMPACTION_UNHANDLED``
|
|
24
|
+
``trajectory.extra["codex_compaction_count"]`` records how many
|
|
25
|
+
``compacted`` records the rollout carried. The summaries themselves are
|
|
26
|
+
not reconstructable into steps — harbor emitted none — so the count is
|
|
27
|
+
the honest repair: it tells a reader the step list has a seam without
|
|
28
|
+
inventing a step to mark it.
|
|
29
|
+
|
|
30
|
+
ATTRIBUTION STRATEGY: identity first, and position only where the format
|
|
31
|
+
leaves no id
|
|
32
|
+
-----------------------------------------------------------------------------
|
|
33
|
+
Three passes, ordered by how much they can be trusted.
|
|
34
|
+
|
|
35
|
+
1. **Tool records, by call id.** harbor keeps each Codex ``call_id`` verbatim
|
|
36
|
+
in ``tool_calls[].tool_call_id``, so a step's tool records are an exact
|
|
37
|
+
lookup: no ordering assumption, no off-by-one to make, and a step whose
|
|
38
|
+
call ids match nothing simply gets no tool attribution.
|
|
39
|
+
|
|
40
|
+
2. **Agent message records, by API call id.** harbor bundles the assistant
|
|
41
|
+
messages of one model request into one step and stamps that request's id
|
|
42
|
+
into ``step.extra.api_call_id``. This pass re-derives the same ids from the
|
|
43
|
+
raw records — the counter advances on a ``token_count`` event that closed a
|
|
44
|
+
request which produced output, which is harbor's own rule — and attributes
|
|
45
|
+
each assistant ``message`` record to the step bearing its request id. That
|
|
46
|
+
is what makes an assistant record with EMPTY text attributable at all: it
|
|
47
|
+
contributes no text to its step (harbor filters empty parts), so no
|
|
48
|
+
text-matching walk could ever place it, while its request id places it
|
|
49
|
+
exactly.
|
|
50
|
+
|
|
51
|
+
3. **User and system message steps, by position.** These carry no id anywhere:
|
|
52
|
+
harbor emits one step per non-assistant message record, in file order, with
|
|
53
|
+
no bundling and nothing filtered. The walk therefore pairs them 1:1 and
|
|
54
|
+
cross-checks each pairing against harbor's own text extraction. The FIRST
|
|
55
|
+
pair whose text or source disagrees stops this pass and records
|
|
56
|
+
``trajectory.extra["enrichment_truncated_at_step"]``: once positions
|
|
57
|
+
disagree, every later pairing is suspect.
|
|
58
|
+
|
|
59
|
+
Reasoning records are attributed only to a step that actually carries
|
|
60
|
+
``reasoning_content``. A reasoning record belongs to the API call that
|
|
61
|
+
consumed it, but recent Codex builds ship the content encrypted and harbor
|
|
62
|
+
reads only the plaintext summary
|
|
63
|
+
(:attr:`~atif_converter.domain.codex_fidelity.CodexFidelityGap.REASONING_ENCRYPTED_DROPPED`),
|
|
64
|
+
so claiming a dropped record "contributed to" a step would overstate what the
|
|
65
|
+
step holds.
|
|
66
|
+
|
|
67
|
+
After all three passes, any step left with no ``source_uuids`` and any
|
|
68
|
+
non-assistant message record never consumed are counted into
|
|
69
|
+
``trajectory.extra["enrichment_unattributed_steps"]`` and
|
|
70
|
+
``trajectory.extra["enrichment_leftover_messages"]`` (present only when
|
|
71
|
+
nonzero). ``extra`` is free-form per the contract, so all of this is
|
|
72
|
+
contract-compatible; a truncation marker or a nonzero count means the
|
|
73
|
+
re-derivation disagreed with harbor and the affected attributions were
|
|
74
|
+
REFUSED, not guessed.
|
|
75
|
+
"""
|
|
76
|
+
|
|
77
|
+
from __future__ import annotations
|
|
78
|
+
|
|
79
|
+
from dataclasses import dataclass, field
|
|
80
|
+
from typing import Any
|
|
81
|
+
|
|
82
|
+
from loguru import logger
|
|
83
|
+
|
|
84
|
+
from atif_converter.domain.codex_edges import build_codex_edges
|
|
85
|
+
|
|
86
|
+
_SOURCE_UUIDS = "source_uuids"
|
|
87
|
+
_TRUNCATED_AT = "enrichment_truncated_at_step"
|
|
88
|
+
#: Message records that reached no step: the silent end of the positional walk.
|
|
89
|
+
_LEFTOVER_MESSAGES = "enrichment_leftover_messages"
|
|
90
|
+
|
|
91
|
+
#: ATIF ``source`` values keyed by the raw Codex message role. ``developer``
|
|
92
|
+
#: and anything unknown land on ``system`` — the flattening named by
|
|
93
|
+
#: ``DEVELOPER_ROLE_FLATTENED_TO_SYSTEM``.
|
|
94
|
+
_ROLE_TO_SOURCE: dict[str, str] = {"user": "user", "assistant": "agent"}
|
|
95
|
+
|
|
96
|
+
#: Payload types that CONSUME a pending reasoning record, mirroring harbor's
|
|
97
|
+
#: normalization: each one clears ``pending_reasoning`` after reading it.
|
|
98
|
+
_REASONING_CONSUMERS: frozenset[str] = frozenset(
|
|
99
|
+
{
|
|
100
|
+
"message",
|
|
101
|
+
"web_search_call",
|
|
102
|
+
"function_call",
|
|
103
|
+
"custom_tool_call",
|
|
104
|
+
"function_call_output",
|
|
105
|
+
"custom_tool_call_output",
|
|
106
|
+
}
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
def _payload(record: dict[str, Any]) -> dict[str, Any]:
|
|
111
|
+
payload = record.get("payload")
|
|
112
|
+
return payload if isinstance(payload, dict) else {}
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
def message_text(payload: dict[str, Any]) -> str:
|
|
116
|
+
"""Replicate harbor's ``Codex._extract_message_text``.
|
|
117
|
+
|
|
118
|
+
Every content block's ``text`` field, concatenated with NO separator —
|
|
119
|
+
which is what harbor does, and the reason a two-block developer message
|
|
120
|
+
reproduces as one run-together string rather than two lines.
|
|
121
|
+
"""
|
|
122
|
+
content = payload.get("content")
|
|
123
|
+
if not isinstance(content, list):
|
|
124
|
+
return ""
|
|
125
|
+
parts: list[str] = []
|
|
126
|
+
for block in content:
|
|
127
|
+
if isinstance(block, dict):
|
|
128
|
+
text = block.get("text")
|
|
129
|
+
if isinstance(text, str):
|
|
130
|
+
parts.append(text)
|
|
131
|
+
return "".join(parts)
|
|
132
|
+
|
|
133
|
+
|
|
134
|
+
def _source_for_role(role: Any) -> str:
|
|
135
|
+
return _ROLE_TO_SOURCE.get(role, "system") if isinstance(role, str) else "system"
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
@dataclass(slots=True)
|
|
139
|
+
class _RolloutIndex:
|
|
140
|
+
"""Everything the three passes need, derived from the raw records once."""
|
|
141
|
+
|
|
142
|
+
#: call_id -> record ids naming it (the call and its output).
|
|
143
|
+
by_call_id: dict[str, list[str]] = field(default_factory=dict)
|
|
144
|
+
#: api_call_id -> assistant message / web-search record ids in that request.
|
|
145
|
+
by_api_call: dict[str, list[str]] = field(default_factory=dict)
|
|
146
|
+
#: Non-assistant message records in file order: (record id, source, text).
|
|
147
|
+
non_agent_messages: list[tuple[str, str, str]] = field(default_factory=list)
|
|
148
|
+
#: record id or call_id -> the reasoning record it consumed.
|
|
149
|
+
reasoning_for: dict[str, str] = field(default_factory=dict)
|
|
150
|
+
#: ``compacted`` records seen.
|
|
151
|
+
compaction_count: int = 0
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
def _index_rollout(records: list[tuple[dict[str, Any], str]]) -> _RolloutIndex:
|
|
155
|
+
"""Re-derive harbor's grouping over the raw records, ids attached.
|
|
156
|
+
|
|
157
|
+
The ``api_call_id`` counter replicates harbor's exactly: it starts at
|
|
158
|
+
``api_call_1``, and a ``token_count`` event advances it only when the
|
|
159
|
+
request it closes produced model output. Anything that diverges here
|
|
160
|
+
diverges from the step ids in ``step.extra.api_call_id``, which is what
|
|
161
|
+
pass 2 joins on — so a mismatch shows up as an unattributed step rather
|
|
162
|
+
than as a wrong attribution.
|
|
163
|
+
"""
|
|
164
|
+
index = _RolloutIndex()
|
|
165
|
+
edges = build_codex_edges(records)
|
|
166
|
+
api_call_number = 1
|
|
167
|
+
current_api_call = "api_call_1"
|
|
168
|
+
saw_model_output = False
|
|
169
|
+
pending_reasoning: str | None = None
|
|
170
|
+
#: call_id -> the api_call id current when the CALL was seen, because
|
|
171
|
+
#: harbor's normalized event for a tool call is emitted at its OUTPUT and
|
|
172
|
+
#: carries the call's request id, not the one live at output time.
|
|
173
|
+
pending_call_api: dict[str, str] = {}
|
|
174
|
+
|
|
175
|
+
for (record, _source_file), edge in zip(records, edges, strict=True):
|
|
176
|
+
record_id = str(edge["uuid"])
|
|
177
|
+
record_type = record.get("type")
|
|
178
|
+
payload = _payload(record)
|
|
179
|
+
|
|
180
|
+
if record_type == "compacted":
|
|
181
|
+
index.compaction_count += 1
|
|
182
|
+
continue
|
|
183
|
+
if record_type == "event_msg":
|
|
184
|
+
if payload.get("type") == "token_count" and saw_model_output:
|
|
185
|
+
api_call_number += 1
|
|
186
|
+
current_api_call = f"api_call_{api_call_number}"
|
|
187
|
+
saw_model_output = False
|
|
188
|
+
continue
|
|
189
|
+
if record_type != "response_item":
|
|
190
|
+
continue
|
|
191
|
+
|
|
192
|
+
item_type = payload.get("type")
|
|
193
|
+
if item_type == "reasoning":
|
|
194
|
+
pending_reasoning = record_id
|
|
195
|
+
continue
|
|
196
|
+
|
|
197
|
+
if _str(item_type) in _REASONING_CONSUMERS and pending_reasoning is not None:
|
|
198
|
+
key = payload.get("call_id") if "call" in str(item_type) else record_id
|
|
199
|
+
index.reasoning_for[str(key) if isinstance(key, str) and key else record_id] = (
|
|
200
|
+
pending_reasoning
|
|
201
|
+
)
|
|
202
|
+
pending_reasoning = None
|
|
203
|
+
|
|
204
|
+
if item_type == "message":
|
|
205
|
+
role = payload.get("role")
|
|
206
|
+
source = _source_for_role(role)
|
|
207
|
+
if role == "assistant":
|
|
208
|
+
index.by_api_call.setdefault(current_api_call, []).append(record_id)
|
|
209
|
+
saw_model_output = True
|
|
210
|
+
else:
|
|
211
|
+
index.non_agent_messages.append((record_id, source, message_text(payload)))
|
|
212
|
+
continue
|
|
213
|
+
|
|
214
|
+
if item_type == "web_search_call":
|
|
215
|
+
index.by_api_call.setdefault(current_api_call, []).append(record_id)
|
|
216
|
+
saw_model_output = True
|
|
217
|
+
continue
|
|
218
|
+
|
|
219
|
+
call_id = payload.get("call_id")
|
|
220
|
+
if isinstance(call_id, str) and call_id:
|
|
221
|
+
index.by_call_id.setdefault(call_id, []).append(record_id)
|
|
222
|
+
if _str(item_type) in {"function_call", "custom_tool_call"}:
|
|
223
|
+
pending_call_api[call_id] = current_api_call
|
|
224
|
+
saw_model_output = True
|
|
225
|
+
|
|
226
|
+
return index
|
|
227
|
+
|
|
228
|
+
|
|
229
|
+
def _step_call_ids(step: dict[str, Any]) -> set[str]:
|
|
230
|
+
"""Every ``tool_call_id`` and ``source_call_id`` this step names."""
|
|
231
|
+
ids: set[str] = set()
|
|
232
|
+
for call in step.get("tool_calls") or []:
|
|
233
|
+
if isinstance(call, dict):
|
|
234
|
+
call_id = call.get("tool_call_id")
|
|
235
|
+
if isinstance(call_id, str) and call_id:
|
|
236
|
+
ids.add(call_id)
|
|
237
|
+
observation = step.get("observation")
|
|
238
|
+
if isinstance(observation, dict):
|
|
239
|
+
for result in observation.get("results") or []:
|
|
240
|
+
if isinstance(result, dict):
|
|
241
|
+
call_id = result.get("source_call_id")
|
|
242
|
+
if isinstance(call_id, str) and call_id:
|
|
243
|
+
ids.add(call_id)
|
|
244
|
+
return ids
|
|
245
|
+
|
|
246
|
+
|
|
247
|
+
def _step_extra(step: dict[str, Any]) -> dict[str, Any]:
|
|
248
|
+
extra = step.get("extra")
|
|
249
|
+
return extra if isinstance(extra, dict) else {}
|
|
250
|
+
|
|
251
|
+
|
|
252
|
+
def _step_id(step: dict[str, Any]) -> int | None:
|
|
253
|
+
step_id = step.get("step_id")
|
|
254
|
+
return step_id if isinstance(step_id, int) else None
|
|
255
|
+
|
|
256
|
+
|
|
257
|
+
def _attach(step: dict[str, Any], ids: list[str]) -> None:
|
|
258
|
+
"""Merge ``ids`` into the step's ``extra.source_uuids``, order-preserving."""
|
|
259
|
+
if not ids:
|
|
260
|
+
return
|
|
261
|
+
extra = step.get("extra")
|
|
262
|
+
if not isinstance(extra, dict):
|
|
263
|
+
extra = {}
|
|
264
|
+
step["extra"] = extra
|
|
265
|
+
existing = extra.get(_SOURCE_UUIDS)
|
|
266
|
+
merged: list[str] = [str(value) for value in existing] if isinstance(existing, list) else []
|
|
267
|
+
for value in ids:
|
|
268
|
+
if value not in merged:
|
|
269
|
+
merged.append(value)
|
|
270
|
+
extra[_SOURCE_UUIDS] = merged
|
|
271
|
+
|
|
272
|
+
|
|
273
|
+
def enrich_codex_trajectory(
|
|
274
|
+
trajectory: dict[str, Any],
|
|
275
|
+
records: list[tuple[dict[str, Any], str]],
|
|
276
|
+
*,
|
|
277
|
+
copy_input: bool = True,
|
|
278
|
+
) -> dict[str, Any]:
|
|
279
|
+
"""Return ``trajectory`` with Codex source attribution and cache totals.
|
|
280
|
+
|
|
281
|
+
Parameters
|
|
282
|
+
----------
|
|
283
|
+
trajectory
|
|
284
|
+
The harbor Codex trajectory as a dict (``model_dump(mode="json")``).
|
|
285
|
+
records
|
|
286
|
+
The rollout's ``(record, source_file)`` pairs in FILE order — the same
|
|
287
|
+
pairs :func:`~atif_converter.domain.codex_edges.build_codex_edges`
|
|
288
|
+
reads, so the ids attributed here are the ids in ``edges.jsonl``.
|
|
289
|
+
copy_input
|
|
290
|
+
Deep-copy before mutating. The production caller passes ``False``
|
|
291
|
+
because the harbor trajectory has no other consumer and a rollout's
|
|
292
|
+
step list is large.
|
|
293
|
+
"""
|
|
294
|
+
if copy_input:
|
|
295
|
+
import copy
|
|
296
|
+
|
|
297
|
+
trajectory = copy.deepcopy(trajectory)
|
|
298
|
+
|
|
299
|
+
steps = trajectory.get("steps")
|
|
300
|
+
if not isinstance(steps, list):
|
|
301
|
+
return trajectory
|
|
302
|
+
|
|
303
|
+
index = _index_rollout(records)
|
|
304
|
+
|
|
305
|
+
# Passes 1 and 2 — identity. Neither can misattribute: both are lookups.
|
|
306
|
+
for step in steps:
|
|
307
|
+
if not isinstance(step, dict):
|
|
308
|
+
continue
|
|
309
|
+
attributed: list[str] = []
|
|
310
|
+
call_ids = sorted(_step_call_ids(step))
|
|
311
|
+
for call_id in call_ids:
|
|
312
|
+
attributed.extend(index.by_call_id.get(call_id, ()))
|
|
313
|
+
api_call_id = _step_extra(step).get("api_call_id")
|
|
314
|
+
if isinstance(api_call_id, str):
|
|
315
|
+
attributed.extend(index.by_api_call.get(api_call_id, ()))
|
|
316
|
+
if step.get("reasoning_content"):
|
|
317
|
+
for key in (*call_ids, *attributed):
|
|
318
|
+
reasoning_id = index.reasoning_for.get(key)
|
|
319
|
+
if reasoning_id is not None and reasoning_id not in attributed:
|
|
320
|
+
attributed.append(reasoning_id)
|
|
321
|
+
_attach(step, attributed)
|
|
322
|
+
|
|
323
|
+
# Pass 3 — position, over the non-assistant message steps only.
|
|
324
|
+
truncated_at, leftover_messages = _attribute_non_agent_messages(steps, index.non_agent_messages)
|
|
325
|
+
|
|
326
|
+
extra = trajectory.get("extra")
|
|
327
|
+
if not isinstance(extra, dict):
|
|
328
|
+
extra = {}
|
|
329
|
+
_record_cache_total(trajectory, extra)
|
|
330
|
+
if index.compaction_count:
|
|
331
|
+
extra["codex_compaction_count"] = index.compaction_count
|
|
332
|
+
if truncated_at is not None:
|
|
333
|
+
extra[_TRUNCATED_AT] = truncated_at
|
|
334
|
+
logger.warning(
|
|
335
|
+
"codex enrichment: message attribution refused at step {} — the "
|
|
336
|
+
"positional walk disagreed with harbor's step list; later steps "
|
|
337
|
+
"carry no source_uuids rather than a guessed one",
|
|
338
|
+
truncated_at,
|
|
339
|
+
)
|
|
340
|
+
unattributed = sum(
|
|
341
|
+
1 for step in steps if isinstance(step, dict) and not _step_extra(step).get(_SOURCE_UUIDS)
|
|
342
|
+
)
|
|
343
|
+
if unattributed:
|
|
344
|
+
extra["enrichment_unattributed_steps"] = unattributed
|
|
345
|
+
if leftover_messages:
|
|
346
|
+
extra[_LEFTOVER_MESSAGES] = leftover_messages
|
|
347
|
+
logger.warning(
|
|
348
|
+
"codex enrichment: {} message record(s) attributed to no step — the "
|
|
349
|
+
"loss report counts them convertible, so this is real unreported loss",
|
|
350
|
+
leftover_messages,
|
|
351
|
+
)
|
|
352
|
+
if extra:
|
|
353
|
+
trajectory["extra"] = extra
|
|
354
|
+
return trajectory
|
|
355
|
+
|
|
356
|
+
|
|
357
|
+
def _attribute_non_agent_messages(
|
|
358
|
+
steps: list[Any],
|
|
359
|
+
messages: list[tuple[str, str, str]],
|
|
360
|
+
) -> tuple[int | None, int]:
|
|
361
|
+
"""Pair user/system steps with user/developer message records, 1:1.
|
|
362
|
+
|
|
363
|
+
Returns the step id the walk refused at (``None`` when every such step
|
|
364
|
+
paired cleanly) and the count of message records left unconsumed. A step is
|
|
365
|
+
refused when the next unconsumed record's ATIF source or harbor-extracted
|
|
366
|
+
text disagrees with it.
|
|
367
|
+
|
|
368
|
+
THE TWO ENDS OF THE WALK FAIL DIFFERENTLY, and both have to be reported.
|
|
369
|
+
Running out of RECORDS refuses the step, which is loud. Running out of
|
|
370
|
+
STEPS is silent: the remaining records attribute to nothing, and the loss
|
|
371
|
+
report still counts them convertible, so it would understate the loss with
|
|
372
|
+
no marker anywhere. The leftover count is that marker.
|
|
373
|
+
"""
|
|
374
|
+
cursor = 0
|
|
375
|
+
for step in steps:
|
|
376
|
+
if not isinstance(step, dict):
|
|
377
|
+
continue
|
|
378
|
+
source = step.get("source")
|
|
379
|
+
if _str(source) not in {"user", "system"}:
|
|
380
|
+
continue
|
|
381
|
+
if cursor >= len(messages):
|
|
382
|
+
return _step_id(step), 0
|
|
383
|
+
record_id, record_source, record_text = messages[cursor]
|
|
384
|
+
if record_source != source or record_text != (step.get("message") or ""):
|
|
385
|
+
return _step_id(step), len(messages) - cursor
|
|
386
|
+
_attach(step, [record_id])
|
|
387
|
+
cursor += 1
|
|
388
|
+
return None, len(messages) - cursor
|
|
389
|
+
|
|
390
|
+
|
|
391
|
+
def _record_cache_total(trajectory: dict[str, Any], extra: dict[str, Any]) -> None:
|
|
392
|
+
"""Surface the cache-creation total under the Claude Code key.
|
|
393
|
+
|
|
394
|
+
harbor puts the Codex figure in
|
|
395
|
+
``final_metrics.extra.total_cache_write_input_tokens`` and the Claude Code
|
|
396
|
+
figure in ``final_metrics.extra.cache_creation_input_tokens``. One name
|
|
397
|
+
reaches SQL, so the Codex value is republished under
|
|
398
|
+
``trajectory.extra["cache_creation_total"]`` — the key the enrichment pass
|
|
399
|
+
for Claude Code already writes and the corpus already carries.
|
|
400
|
+
"""
|
|
401
|
+
final_metrics = trajectory.get("final_metrics")
|
|
402
|
+
if not isinstance(final_metrics, dict):
|
|
403
|
+
return
|
|
404
|
+
metrics_extra = final_metrics.get("extra")
|
|
405
|
+
if not isinstance(metrics_extra, dict):
|
|
406
|
+
return
|
|
407
|
+
total = metrics_extra.get("total_cache_write_input_tokens")
|
|
408
|
+
if isinstance(total, int):
|
|
409
|
+
extra["cache_creation_total"] = total
|
|
410
|
+
|
|
411
|
+
|
|
412
|
+
def _str(value: object) -> str | None:
|
|
413
|
+
"""``value`` when it is a string, else ``None``: a set-membership key that cannot raise.
|
|
414
|
+
|
|
415
|
+
A transcript field can hold any JSON value, and an object or a list is unhashable, so
|
|
416
|
+
testing one against a set raised ``TypeError`` past the converter (found by the fuzzer).
|
|
417
|
+
A non-string value never matched one of these string sets, so no result changes.
|
|
418
|
+
"""
|
|
419
|
+
return value if isinstance(value, str) else None
|