Merge nucleic/plucky-north-vole-sdna into dev

export_swe_chat.py

@@ -1,11 +1,11 @@
 #!/usr/bin/env python3
-"""Export three-turn SWE-chat candidates without retaining later-turn text.
+"""Export first-prompt/first-response SWE-chat candidates without later text.
 
 The source snapshot is gated and deliberately stays below ``.artifacts/``. This
 importer does not download it: callers supply an already accepted, revision-pinned
 Parquet snapshot. It reads Parquet in record batches, joins the small sessions table
 only for repository/user grouping, and writes an unlabeled JSONL that contains the
-first prompt plus hashes (never text) for the two teacher-context prompts.
+first prompt plus a teacher-response hash for later labeling.
 """
 
 from __future__ import annotations
@@ -13,7 +13,7 @@ from __future__ import annotations
 import argparse
 import json
 import sys
-from collections import Counter, defaultdict
+from collections import Counter
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
@@ -45,23 +45,23 @@ class Candidate:
     session_id: str
     repo_id: str | None
     user_id: str | None
-    turns: tuple[Turn, Turn, Turn]
+    prompt_turn: Turn
+    response_turn: Turn
 
     def json(self, revision: str) -> dict[str, Any]:
-        first, second, third = self.turns
+        first, response = self.prompt_turn, self.response_turn
         return {
             "schemaVersion": SCHEMA_VERSION,
             "repoID": self.repo_id,
             "userID": self.user_id,
             "sessionID": self.session_id,
-            "sourceTurnIDs": [turn.turn_id for turn in self.turns],
+            "sourceTurnIDs": [first.turn_id, response.turn_id],
             "sourceRevision": revision,
             "promptHash": prompt_hash(first.prompt),
-            "contextPromptHashes": [prompt_hash(second.prompt), prompt_hash(third.prompt)],
+            "teacherResponseHash": prompt_hash(response.prompt),
             "prompt": first.prompt,
-            # This ignored pre-labeling file is the only artifact allowed to carry
-            # later text. Canonical labeled JSONL contains only the seven data fields.
-            "teacherContext": [second.prompt, third.prompt],
+            # Response text lives only in this ignored pre-labeling artifact.
+            "teacherResponse": response.prompt,
         }
 
 
@@ -136,11 +136,36 @@ def select_candidates(
     max_per_repo: int,
     max_per_user: int,
 ) -> tuple[list[Candidate], dict[str, int]]:
-    """Apply documented row filters and deterministic session-level selection."""
+    """Test-friendly pair selection. Production uses two streaming passes below."""
+
+    rows = list(conversation_rows)
+    prompts, funnel = first_user_prompts(rows)
+    preliminary = attach_first_responses(rows, prompts, sessions=sessions, funnel=funnel)
+    return cap_and_dedupe(preliminary, funnel, max_per_repo=max_per_repo, max_per_user=max_per_user)
+
+
+def _turn(row: dict[str, Any], *, funnel: Counter[str], prefix: str) -> Turn | None:
+    session_id = _as_text(row.get("session_id"))
+    turn_id = _as_text(row.get("turn_id"))
+    content = row.get("content")
+    conversation_turn_number = _as_int(row.get("conversation_turn_number"))
+    turn_number = _as_int(row.get("turn_number"))
+    if (
+        session_id is None or turn_id is None or not isinstance(content, str)
+        or not content.strip() or "\x00" in content
+        or conversation_turn_number is None or turn_number is None
+    ):
+        funnel[f"rejected{prefix}MalformedOrEmpty"] += 1
+        return None
+    return Turn(session_id, turn_id, conversation_turn_number, turn_number, content)
+
+
+def first_user_prompts(rows: Iterable[dict[str, Any]]) -> tuple[dict[str, Turn], Counter[str]]:
+    """First streaming pass: retain the earliest eligible user prompt per session."""
 
     funnel: Counter[str] = Counter()
-    by_session: dict[str, list[Turn]] = defaultdict(list)
-    for row in conversation_rows:
+    prompts: dict[str, Turn] = {}
+    for row in rows:
         funnel["conversationRows"] += 1
         if row.get("turn_type") != "user_prompt":
             funnel["rejectedTurnType"] += 1
@@ -149,49 +174,77 @@ def select_candidates(
             funnel["rejectedRole"] += 1
             continue
         if row.get("is_conversational") is not True:
-            funnel["rejectedNonConversational"] += 1
+            funnel["rejectedUserNonConversational"] += 1
             continue
         if row.get("is_continuation") is True:
             funnel["rejectedContinuation"] += 1
             continue
-        session_id = _as_text(row.get("session_id"))
-        turn_id = _as_text(row.get("turn_id"))
-        prompt = row.get("content")
-        conversation_turn_number = _as_int(row.get("conversation_turn_number"))
-        turn_number = _as_int(row.get("turn_number"))
-        if (
-            session_id is None or turn_id is None or not isinstance(prompt, str)
-            or not prompt.strip() or "\x00" in prompt
-            or conversation_turn_number is None or turn_number is None
-        ):
-            funnel["rejectedMalformedOrEmpty"] += 1
+        turn = _turn(row, funnel=funnel, prefix="User")
+        if turn is None:
             continue
-        # Retain only the first three checked ordinals while streaming.  This bounds
-        # memory by sessions × 3, not by every eligible prompt in the large config.
-        turns = by_session[session_id]
-        turns.append(Turn(session_id, turn_id, conversation_turn_number, turn_number, prompt))
-        turns.sort(key=lambda turn: (turn.conversation_turn_number, turn.turn_number, turn.turn_id))
-        del turns[3:]
-        funnel["eligibleRows"] += 1
+        previous = prompts.get(turn.session_id)
+        if previous is None or (turn.conversation_turn_number, turn.turn_number, turn.turn_id) < (
+            previous.conversation_turn_number, previous.turn_number, previous.turn_id
+        ):
+            prompts[turn.session_id] = turn
+        funnel["eligibleUserPrompts"] += 1
+    funnel["sessionsWithEligibleFirstPrompt"] = len(prompts)
+    return prompts, funnel
+
+
+def attach_first_responses(
+    rows: Iterable[dict[str, Any]],
+    prompts: dict[str, Turn],
+    *,
+    sessions: dict[str, tuple[str | None, str | None]],
+    funnel: Counter[str],
+) -> list[Candidate]:
+    """Second pass: find the immediately following conversational assistant response."""
+
+    responses: dict[str, Turn] = {}
+    for row in rows:
+        if row.get("turn_type") != "assistant_response" or row.get("role") != "assistant":
+            continue
+        if row.get("is_conversational") is not True:
+            funnel["rejectedAssistantNonConversational"] += 1
+            continue
+        session_id = _as_text(row.get("session_id"))
+        prompt = prompts.get(session_id or "")
+        if prompt is None:
+            continue
+        turn = _turn(row, funnel=funnel, prefix="Assistant")
+        if turn is None or turn.conversation_turn_number != prompt.conversation_turn_number + 1:
+            continue
+        previous = responses.get(turn.session_id)
+        if previous is not None:
+            raise DataError(f"ambiguous assistant response after first prompt in session {turn.session_id!r}")
+        responses[turn.session_id] = turn
 
     preliminary: list[Candidate] = []
-    for session_id, turns in sorted(by_session.items()):
-        ordered = sorted(turns, key=lambda turn: (turn.conversation_turn_number, turn.turn_number, turn.turn_id))
-        if len(ordered) < 3:
-            funnel["sessionsFewerThanThreeEligiblePrompts"] += 1
-            continue
-        ordinals = [(turn.conversation_turn_number, turn.turn_number) for turn in ordered[:3]]
-        if len(set(ordinals)) != len(ordinals):
-            funnel["sessionsAmbiguousTurnOrder"] += 1
+    for session_id, prompt in sorted(prompts.items()):
+        response = responses.get(session_id)
+        if response is None:
+            funnel["sessionsWithoutFirstAssistantResponse"] += 1
             continue
         repo_id, user_id = sessions.get(session_id, (None, None))
-        preliminary.append(Candidate(session_id, repo_id, user_id, tuple(ordered[:3])))
-    funnel["sessionsWithThreeEligiblePrompts"] = len(preliminary)
+        preliminary.append(Candidate(session_id, repo_id, user_id, prompt, response))
+    funnel["sessionsWithPromptAndFirstAssistantResponse"] = len(preliminary)
+    return preliminary
+
+
+def cap_and_dedupe(
+    preliminary: Sequence[Candidate],
+    funnel: Counter[str],
+    *,
+    max_per_repo: int,
+    max_per_user: int,
+) -> tuple[list[Candidate], dict[str, int]]:
+    """Dedupe normalized first prompts and cap repository/user concentration."""
 
     first_by_hash: dict[str, str] = {}
     deduped: list[Candidate] = []
     for candidate in preliminary:
-        digest = prompt_hash(candidate.turns[0].prompt)
+        digest = prompt_hash(candidate.prompt_turn.prompt)
         if digest in first_by_hash:
             funnel["rejectedDuplicateFirstPrompt"] += 1
             continue
@@ -232,14 +285,16 @@ def export(
     if max_per_repo <= 0 or max_per_user <= 0:
         raise DataError("source concentration caps must be positive")
     sessions = session_metadata(parquet_rows(sessions_path, ("session_id", "repo_id", "user_id")))
-    candidates, funnel = select_candidates(
-        parquet_rows(
-            conversations,
-            ("session_id", "turn_id", "conversation_turn_number", "turn_number", "turn_type", "role", "is_conversational", "is_continuation", "content"),
-        ),
-        sessions=sessions,
-        max_per_repo=max_per_repo,
-        max_per_user=max_per_user,
+    columns = (
+        "session_id", "turn_id", "conversation_turn_number", "turn_number",
+        "turn_type", "role", "is_conversational", "is_continuation", "content",
+    )
+    prompts, funnel = first_user_prompts(parquet_rows(conversations, columns))
+    preliminary = attach_first_responses(
+        parquet_rows(conversations, columns), prompts, sessions=sessions, funnel=funnel
+    )
+    candidates, funnel = cap_and_dedupe(
+        preliminary, funnel, max_per_repo=max_per_repo, max_per_user=max_per_user
     )
     output.parent.mkdir(parents=True, exist_ok=True)
     output.write_text("".join(f"{canonical_json(candidate.json(revision))}\n" for candidate in candidates), encoding="utf-8")
@@ -254,10 +309,10 @@ def export(
             "rawFiles": [{"path": str(path), "sha256": file_sha256(path)} for path in sorted([*conversations, *sessions_path])],
         },
         "selection": {
-            "rowFilter": "turn_type=user_prompt, role=user, is_conversational=true, not is_continuation",
-            "perSession": "first three eligible non-empty prompts ordered by conversation_turn_number then turn_number",
+            "rowFilter": "first user_prompt with role=user, is_conversational=true, not is_continuation",
+            "perSession": "first eligible user prompt plus the immediately following conversational assistant_response",
             "studentText": "first prompt only",
-            "teacherContext": "second and third prompts retained only in ignored candidate JSONL until labeling",
+            "teacherContext": "first assistant response retained only in ignored candidate JSONL until labeling",
             "dedupe": "exact normalized first prompt",
             "maxPerRepo": max_per_repo,
             "maxPerUser": max_per_user,