[content] crowd-pilot/crowd_pilot/serialization_utils.py (python)
class ConversationState:
    """
    Mutable state used while constructing conversations.
    """
    conversations: List[List[Dict[str, str]]]
    max_chars_per_conversation: int
    max_chars_per_turn: int
    min_conversation_turns: int
    current_conversation: List[Dict[str, str]] = field(default_factory=list)
    current_chars: int = 0
    files_opened_in_conversation: set[str] = field(default_factory=set)

    def start_new_conversation(self) -> None:
        if self.current_conversation:
            if len(self.current_conversation) >= self.min_conversation_turns:
                self.conversations.append(self.current_conversation)
        self.current_conversation = []
        self.current_chars = 0
        self.files_opened_in_conversation.clear()

    def append_turn(self, turn: Dict[str, str]) -> None:
        value = turn.get("value", "")
        # Enforce a per-turn character budget to prevent pathological single
        # turns (e.g., massive file dumps) from creating multi-megabyte conversations.
        if len(value) > self.max_chars_per_turn:
            value = value[: self.max_chars_per_turn]
            turn["value"] = value
        turn_len = len(value)
        if (
            self.current_conversation  # only start a new conversation on non-empty current conversation
            and self.current_chars + turn_len > self.max_chars_per_conversation
        ):
            self.start_new_conversation()
        self.current_conversation.append(turn)
        self.current_chars += turn_len

    def maybe_capture_file_contents(
        self,
        file_path: str,
        content: str,
    ) -> None:
        """
        Capture the contents of the given file in the current conversation if it hasn't been opened yet.
        """
        if file_path in self.files_opened_in_conversation:
            return
        cmd = f"cat -n {file_path}"
        self.append_turn({
            "from": "Assistant",
            "value": _fenced_block("bash", _clean_text(cmd)),
        })
        output = _line_numbered_output(content)
        self.append_turn({
            "from": "User",
            "value": f"<stdout>
{output}
</stdout>",
        })
        self.files_opened_in_conversation.add(file_path)


def _clean_text(text: str) -> str:
    # Normalize line endings and strip trailing spaces; preserve tabs/newlines.
    return text.replace("
", "
").replace("
", "
").rstrip()


def _fenced_block(language: Optional[str], content: str) -> str:
    lang = (language or "").lower()
    return f"```{lang}
{content}
```
"


def _apply_change(content: str, offset: int, length: int, new_text: str) -> str:
    # Mirrors crowd_code_player.replay_file.apply_change
    base = str(content)
    text = str(new_text) if pd.notna(new_text) else ""
    text = text.replace("\
", "
").replace("\
", "
")
    if offset > len(base):
        base = base + (" " * (offset - len(base)))
    return base[:offset] + text + base[offset + length:]


def _apply_backspaces(text: str) -> str:
    out: List[str] = []
    for ch in text:
        if ch == "\b":  # \x08
            if out:
                out.pop()
        else:
            out.append(ch)
    return "".join(out)


def _normalize_terminal_output(raw: str) -> str:
    """
    Normalize PTY/terminal output for training:
      - Apply backspaces (\x08)
      - Strip OSC (window title/shell integration) first, keeping BEL/ST terminators intact
      - Resolve carriage returns (
) by keeping the last rewrite per line
      - Strip CSI (coloring etc.)
      - Finally drop any remaining BEL (\x07)
    """
    if not raw:
        return raw
    s = _apply_backspaces(raw)
    # Remove OSC sequences that are properly terminated (BEL or ST)
    s = _ANSI_OSC_TERMINATED_RE.sub("", s)
    # Fallback: drop any unterminated OSC up to end-of-line only
    s = "
".join(_ANSI_OSC_LINE_FALLBACK_RE.sub("", line) for line in s.split("
"))
    # Resolve carriage returns per line:
    # - If there are multiple rewrites, keep the last non-empty conversation
    # - If it's CRLF (ending with '
' before '
'), keep the content before '
'
    resolved_lines: List[str] = []
    for seg in s.split("
"):
        parts = seg.split("
")
        chosen = ""
        # pick last non-empty part if available; else last part
        for p in reversed(parts):
            if p != "":
                chosen = p
                break
        if chosen == "" and parts:
            chosen = parts[-1]
        resolved_lines.append(chosen)
    s = "
".join(resolved_lines)
    # Strip ANSI escape sequences
    s = _ANSI_CSI_RE.sub("", s)
    # Remove any remaining BEL beeps
    s = s.replace("\x07", "")
    return s


def _line_numbered_output(content: str, start_line: Optional[int] = None, end_line: Optional[int] = None) -> str:
    # FIXME (f.srambical): check whether this corresponds **exactly** to the output of cat -n {file_path} | sed -n '{vstart},{vend}p'
    lines = content.splitlines()
    total = len(lines)
    if total == 0:
        return ""
    s = 1 if start_line is None else max(1, min(start_line, total))
    e = total if end_line is None else max(1, min(end_line, total))
    assert e >= s, "End line number cannot be less than start line number! Likely a bug in the line numbering computation."
    buf: List[str] = []
    for idx in range(s, e + 1):
        buf.append(f"{idx:6}	{lines[idx - 1]}")
    return "
".join(buf)


def _compute_viewport(total_lines: int, center_line: int, radius: int) -> Tuple[int, int]:
    if total_lines <= 0:
        return (1, 0)
    start = max(1, center_line - radius)
    end = min(total_lines, center_line + radius)
    assert end >= start, "Viewport cannot have negative width! Likely a bug in the viewport computation."
    return (start, end)


def _escape_single_quotes_for_sed(text: str) -> str:
    # Close quote, add an escaped single quote, reopen quote: '"'"'
    return text.replace("'", "'\"'\"'")


def _compute_changed_block_lines(
    before: str, after: str
) -> Tuple[int, int, int, int, List[str]]:
    """
    Return 1-based start and end line numbers in 'before' that should be
    replaced, 1-based start and end line numbers in 'after' that contain
    the replacement, and the replacement lines from 'after'.

    For pure deletions, the replacement list may be empty.
    """
    before_lines = before.splitlines()
    after_lines = after.splitlines()
    sm = difflib.SequenceMatcher(a=before_lines, b=after_lines, autojunk=False)
    opcodes = [op for op in sm.get_opcodes() if op[0] != "equal"]
    assert opcodes, "Opcode list cannot be empty! Likely a bug in the diff computation."

    first = opcodes[0]
    last = opcodes[-1]
    # i1/i2 refer to 'before' indices, j1/j2 to 'after'
    start_before = max(1, first[1] + 1)
    end_before = last[2]  # no increment since we go from 'exclusive' to 'inclusive' indexing
    start_after = max(1, first[3] + 1)
    end_after = last[4]
    replacement_lines = after_lines[first[3] : last[4]]
    return (start_before, end_before, start_after, end_after, replacement_lines)


def session_to_nemo_conversations(
    df: pd.DataFrame,
    max_chars_per_conversation: int,
    max_chars_per_turn: int,
    min_conversation_turns: int,
    viewport_radius: int = 10,
    normalize_terminal_output: bool = True,
    coalesce_radius: int = 5,
) -> List[List[Dict[str, str]]]:
    """
    Convert a session DataFrame to one or more NeMo conversations.

    - Conversations are created by approximately limiting the total characters
      across all `value` fields to `max_chars_per_conversation`.
    - When a new conversation starts (after the first), the first time a file is
      referenced in that conversation we re-log the full file contents with
      `cat -n <file>` and numbered output so that each conversation is self-contained.
    """
    file_states: Dict[str, str] = {}
    per_file_viewport: Dict[str, Optional[Tuple[int, int]]] = {}

    conversations: List[List[Dict[str, str]]] = []
    conversation_state = ConversationState(
        conversations=conversations,
        max_chars_per_conversation=max_chars_per_conversation,
        max_chars_per_turn=max_chars_per_turn,
        min_conversation_turns=min_conversation_turns,
    )

    terminal_output_buffer: List[str] = []
    pending_edits_before: Dict[str, Optional[str]] = {}
    pending_edit_regions: Dict[str, Optional[Tuple[int, int]]] = {}

    def _flush_terminal_output_buffer() -> None:
        if not terminal_output_buffer:
            return
        aggregated = "".join(terminal_output_buffer)
        out = aggregated
        if normalize_terminal_output:
            out = _normalize_terminal_output(out)
        cleaned = _clean_text(out)
        if cleaned.strip():
            conversation_state.append_turn({
                "from": "User",
                "value": f"<stdout>
{cleaned}
</stdout>",
            })
        terminal_output_buffer.clear()

    def _flush_pending_edit_for_file(target_file: str) -> None:
        before_snapshot = pending_edits_before.get(target_file)
        if before_snapshot is None:
            return
        after_state = file_states.get(target_file, "")
        if before_snapshot.rstrip("
") == after_state.rstrip("
"):
            pending_edits_before[target_file] = None
            pending_edit_regions[target_file] = None
            return
        (
            start_before,
            end_before,
            start_after,
            end_after,
            repl_lines,
        ) = _compute_changed_block_lines(before_snapshot, after_state)
        before_total_lines = len(before_snapshot.splitlines())
        if end_before < start_before:
            escaped_lines = [_escape_single_quotes_for_sed(line) for line in repl_lines]
            sed_payload = "
".join(escaped_lines)
            if start_before <= max(1, before_total_lines):
                sed_cmd = f"sed -i '{start_before}i\\
{sed_payload}' {target_file}"
            else:
                sed_cmd = f"sed -i '$a\\
{sed_payload}' {target_file}"
        elif not repl_lines:
            sed_cmd = f"sed -i '{start_before},{end_before}d' {target_file}"
        else:
            escaped_lines = [_escape_single_quotes_for_sed(line) for line in repl_lines]
            sed_payload = "
".join(escaped_lines)
            sed_cmd = f"sed -i '{start_before},{end_before}c\\
{sed_payload}' {target_file}"
        total_lines = len(after_state.splitlines())
        center = (start_after + end_after) // 2
        vp = _compute_viewport(total_lines, center, viewport_radius)
        per_file_viewport[target_file] = vp
        vstart, vend = vp
        conversation_state.maybe_capture_file_contents(target_file, before_snapshot)
        chained_cmd = f"{sed_cmd} && cat -n {target_file} | sed -n '{vstart},{vend}p'"
        conversation_state.append_turn({
            "from": "Assistant",
            "value": _fenced_block("bash", _clean_text(chained_cmd)),
        })
        viewport_output = _line_numbered_output(after_state, vstart, vend)
        conversation_state.append_turn({
            "from": "User",
            "value": f"<stdout>
{viewport_output}
</stdout>",
        })
        pending_edits_before[target_file] = None
        pending_edit_regions[target_file] = None

    def _flush_all_pending_edits() -> None:
        for fname in list(pending_edits_before.keys()):
            _flush_pending_edit_for_file(fname)

    for i in range(len(df)):
        row = df.iloc[i]
        file_path: str = row["File"]
        event_type = row["Type"]

        match event_type:
            case "tab":
                _flush_all_pending_edits()
                _flush_terminal_output_buffer()
                text = row["Text"]
                if pd.notna(text):
                    content = str(text).replace("\
", "
").replace("\
", "
")
                    file_states[file_path] = content
                    cmd = f"cat -n {file_path}"
                    conversation_state.append_turn({
                        "from": "Assistant",
                        "value": _fenced_block("bash", _clean_text(cmd)),
                    })
                    output = _line_numbered_output(content)
                    conversation_state.append_turn({
                        "from": "User",
                        "value": f"<stdout>
{output}
</stdout>",
                    })
                    conversation_state.files_opened_in_conversation.add(file_path)
                else:
                    # File switch without content snapshot: show current viewport only
                    content = file_states.get(file_path, "")
                    total_lines = len(content.splitlines())
                    vp = per_file_viewport.get(file_path)
                    if not vp or vp[1] == 0:
                        vp = _compute_viewport(total_lines, 1, viewport_radius)
                        per_file_viewport[file_path] = vp
                    if vp:
                        vstart, vend = vp
                        conversation_state.maybe_capture_file_contents(file_path, content)
                        cmd = f"cat -n {file_path} | sed -n '{vstart},{vend}p'"
                        conversation_state.append_turn({
                            "from": "Assistant",
                            "value": _fenced_block("bash", _clean_text(cmd)),
                        })
                        viewport_output = _line_numbered_output(content, vstart, vend)
                        conversation_state.append_turn({
                            "from": "User",
                            "value": f"<stdout>
{viewport_output}
</stdout>",
                        })

            case "content":
                _flush_terminal_output_buffer()
                offset = int(row["RangeOffset"])
                length = int(row["RangeLength"])
                new_text = row["Text"]
                before = file_states.get(file_path, "")
                # Approximate current edit region in line space
                new_text_str = str(new_text) if pd.notna(new_text) else ""
                start_line_current = before[:offset].count("
") + 1
                deleted_conversation = before[offset:offset + length]
                lines_added = new_text_str.count("
")
                lines_deleted = deleted_conversation.count("
")
                region_start = start_line_current
                region_end = start_line_current + max(lines_added, lines_deleted, 0)
                # Flush pending edits if this edit is far from the pending region
                current_region = pending_edit_regions.get(file_path)
                if current_region is not None:
                    rstart, rend = current_region
                    if region_start < (rstart - coalesce_radius) or region_start > (rend + coalesce_radius):
                        _flush_pending_edit_for_file(file_path)
                        current_region = None
                after = _apply_change(before, offset, length, new_text)
                if pending_edits_before.get(file_path) is None:
                    pending_edits_before[file_path] = before
                # Update/initialize region union
                if current_region is None:
                    pending_edit_regions[file_path] = (region_start, max(region_start, region_end))
                else:
                    rstart, rend = current_region
                    pending_edit_regions[file_path] = (min(rstart, region_start), max(rend, region_end))
                file_states[file_path] = after

            case "selection_command" | "selection_mouse" | "selection_keyboard":
                # During an edit burst (pending edits), suppress flush and viewport emissions
                if pending_edits_before.get(file_path) is None:
                    _flush_terminal_output_buffer()
                else:
                    # Skip emitting viewport while edits are pending to avoid per-keystroke sed/cat spam
                    continue
                offset = int(row["RangeOffset"])
                content = file_states.get(file_path, "")
                total_lines = len(content.splitlines())
                target_line = content[:offset].count("
") + 1
                vp = per_file_viewport.get(file_path)
                should_emit = False
                if not vp or vp[1] == 0:
                    vp = _compute_viewport(total_lines, target_line, viewport_radius)
                    per_file_viewport[file_path] = vp
                    should_emit = True
                else:
                    vstart, vend = vp
                    if target_line < vstart or target_line > vend:
                        vp = _compute_viewport(total_lines, target_line, viewport_radius)
                        per_file_viewport[file_path] = vp
                        should_emit = True
                if should_emit and vp:
                    vstart, vend = vp
                    conversation_state.maybe_capture_file_contents(file_path, content)
                    cmd = f"cat -n {file_path} | sed -n '{vstart},{vend}p'"
                    conversation_state.append_turn({
                        "from": "Assistant",
                        "value": _fenced_block("bash", _clean_text(cmd)),
                    })
                    viewport_output = _line_numbered_output(content, vstart, vend)
                    conversation_state.append_turn({
                        "from": "User",
                        "value": f"<stdout>
{viewport_output}
</stdout>",
                    })

            case "terminal_command":
                _flush_all_pending_edits()
                _flush_terminal_output_buffer()
                command = row["Text"]
                command_str = str(command).replace("\
", "
").replace("\
", "
")
                conversation_state.append_turn({
                    "from": "Assistant",
                    "value": _fenced_block("bash", _clean_text(command_str)),
                })

            case "terminal_output":
                output = row["Text"]
                raw_output = str(output).replace("\
", "
").replace("\
", "
")
                terminal_output_buffer.append(raw_output)

            case "terminal_focus":
                _flush_all_pending_edits()
                _flush_terminal_output_buffer()
                # No-op for bash transcript; focus changes don't emit commands/output
                pass

            case "git_branch_checkout":
                _flush_all_pending_edits()
                _flush_terminal_output_buffer()
                branch_info = row["Text"]
                branch_str = str(branch_info).replace("\
", "
").replace("\
", "
")
                cleaned = _clean_text(branch_str)
                m = re.search(r"to '([^']+)'", cleaned)
                if not m:
                    raise ValueError(f"Could not extract branch name from git checkout message: {cleaned}")
                branch_name = m.group(1).strip()
                # Safe-quote branch if it contains special characters
                if re.search(r"[^A-Za-z0-9._/\\-]", branch_name):
                    branch_name = "'" + branch_name.replace("'", "'\"'\"'") + "'"
                cmd = f"git checkout {branch_name}"
                conversation_state.append_turn({
                    "from": "Assistant",
                    "value": _fenced_block("bash", _clean_text(cmd)),
                })

            case _:
                raise ValueError(f"Unknown event type: {event_type}")

    _flush_all_pending_edits()
    _flush_terminal_output_buffer()
    if conversation_state.current_conversation:
        if len(conversation_state.current_conversation) >= min_conversation_turns:
            conversations.append(conversation_state.current_conversation)
    return conversations