wafer-core 0.1.43__py3-none-any.whl → 0.1.44__py3-none-any.whl

wafer_core/rollouts/progress_app.py

@@ -25,7 +25,7 @@ import sys
 import time
 from collections.abc import Generator
 from contextlib import contextmanager
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from pathlib import Path
 from typing import Any
 
@@ -38,6 +38,8 @@ RED = "\x1b[31m"
 YELLOW = "\x1b[33m"
 CYAN = "\x1b[36m"
 BOLD = "\x1b[1m"
+BG_SELECTED = "\x1b[48;5;236m"  # Dark gray background for selected row
+BLUE = "\x1b[34m"
 
 # Event types the progress display understands.
 # The "message" field in the JSONL record is the event type discriminator.
@@ -55,6 +57,19 @@ KNOWN_EVENT_TYPES = frozenset({
     "sample_retry",
 })
 
+# Event types in per-sample JSONL files
+SAMPLE_EVENT_TYPES = frozenset({
+    "text_delta",
+    "thinking_delta",
+    "assistant_message",
+    "tool_execution",
+    "llm_call",
+    "modal_progress",
+    "turn",
+    "sample_start",
+    "sample_end",
+})
+
 
 # ── State ────────────────────────────────────────────────────────────────────
 
@@ -181,7 +196,7 @@ def derive_state(events: list[dict[str, Any]]) -> RenderState:
     return state
 
 
-def _format_sample_row(sample: SampleState, width: int) -> str:
+def _format_sample_row(sample: SampleState, width: int, selected: bool = False) -> str:
     """Format a single sample row."""
     name = sample.name[:25].ljust(25)
     turn_info = f"T:{sample.turn}"
@@ -189,24 +204,89 @@ def _format_sample_row(sample: SampleState, width: int) -> str:
     if sample.status == "complete":
         if sample.score is not None:
             icon = f"{GREEN}✓{RESET}" if sample.score > 0.5 else f"{RED}✗{RESET}"
-            return f"  {name} {turn_info:>4}  {icon} score={sample.score:.2f}"
+            row = f"  {name} {turn_info:>4}  {icon} score={sample.score:.2f}"
         else:
-            return f"  {name} {turn_info:>4}  {GREEN}✓{RESET} done"
+            row = f"  {name} {turn_info:>4}  {GREEN}✓{RESET} done"
     elif sample.status == "retry":
-        return f"  {name} {turn_info:>4}  {YELLOW}⟳{RESET} retry (attempt {sample.retry_attempt})"
+        row = f"  {name} {turn_info:>4}  {YELLOW}⟳{RESET} retry (attempt {sample.retry_attempt})"
     else:
         phase = sample.phase or "running"
-        return f"  {name} {turn_info:>4}  {CYAN}{phase}...{RESET}"
+        row = f"  {name} {turn_info:>4}  {CYAN}{phase}...{RESET}"
 
+    if selected:
+        # Add selection highlight
+        row = f"{BG_SELECTED}{row}{RESET}"
 
-# ── Render function ──────────────────────────────────────────────────────────
+    return row
 
 
-def render(state: RenderState, width: int, height: int) -> list[str]:
-    """Render state to list of lines. Pure function."""
-    lines: list[str] = []
+# ── Viewport formatting ──────────────────────────────────────────────────────
 
-    # Header
+
+def format_sample_event(
+    event: dict[str, Any], last_turn: int | None = None
+) -> tuple[str | None, int | None]:
+    """Format a per-sample JSONL event into a readable line.
+
+    Returns (formatted_line, new_turn_number).
+    formatted_line is None if the event should not be displayed.
+    new_turn_number is the turn number if this event changes it, else None.
+    """
+    msg = event.get("message", "")
+
+    if msg == "text_delta":
+        # Streaming text — return raw text for accumulation
+        return event.get("text", ""), None
+
+    elif msg == "thinking_delta":
+        # Extended thinking — show dimmed
+        text = event.get("text", "")
+        return (f"{DIM}{text}{RESET}" if text else None), None
+
+    elif msg == "assistant_message":
+        # Final assistant message — skip if we've been streaming
+        # The caller handles dedup
+        return None, None
+
+    elif msg == "tool_execution":
+        tool = event.get("tool_name", "tool")
+        duration = event.get("duration_ms", 0)
+        status = event.get("result_summary", "")
+        is_error = event.get("is_error", False)
+        color = RED if is_error else CYAN
+        return f"{color}→ {tool}{RESET} ({duration / 1000:.1f}s) {status}", None
+
+    elif msg == "llm_call":
+        provider = event.get("provider", "")
+        model = event.get("model", "")
+        tokens_in = event.get("tokens_in", 0)
+        tokens_out = event.get("tokens_out", 0)
+        duration = event.get("duration_ms", 0)
+        model_short = model.split("/")[-1] if model else provider
+        return (
+            f"{DIM}◆ {model_short} {tokens_in}→{tokens_out} tok ({duration / 1000:.1f}s){RESET}",
+            None,
+        )
+
+    elif msg == "modal_progress":
+        phase = event.get("phase", "")
+        return (f"{YELLOW}⟳ {phase}...{RESET}" if phase else None), None
+
+    elif msg == "turn":
+        turn = event.get("turn")
+        # Only show turn separator when turn number changes
+        if turn is not None and turn != last_turn:
+            return f"{DIM}── turn {turn} ──{RESET}", turn
+        return None, turn if turn is not None else last_turn
+
+    return None, None
+
+
+# ── Render function ──────────────────────────────────────────────────────────
+
+
+def render_header(state: RenderState, width: int) -> str:
+    """Render the progress header line."""
     if state.gepa_iter is not None:
         header = (
             f"GEPA iter {state.gepa_iter}/{state.gepa_total or '?'} │ best: {state.gepa_best:.0%}"
@@ -231,7 +311,14 @@ def render(state: RenderState, width: int, height: int) -> list[str]:
         else:
             header = f"{state.eval_name}: {state.completed} samples [{_format_time(elapsed)}]"
 
-    lines.append(header[:width])
+    return header[:width]
+
+
+def render(state: RenderState, width: int, height: int) -> list[str]:
+    """Render state to list of lines. Pure function."""
+    lines: list[str] = []
+
+    lines.append(render_header(state, width))
 
     # Active samples, sorted by most recently updated
     active = [s for s in state.samples.values() if s.status != "complete" and s.phase]
@@ -263,9 +350,32 @@ class Model:
     events: tuple[dict[str, Any], ...] = ()
     done: bool = False
 
+    # Viewport state
+    focus: str = "list"  # "list" or "viewport"
+    selected_idx: int = 0  # Index into active samples list
+    viewport_lines: tuple[str, ...] = ()  # Formatted lines for viewport
+    viewport_scroll: int = 0  # Scroll position in viewport
+    viewport_auto_follow: bool = True  # Auto-scroll to bottom on new content
+    tailing_sample_id: str | None = None  # Which sample file we're tailing
+
+    # For streaming text accumulation
+    streaming_text: str = ""
+
+    # For deduping turn separators
+    last_turn: int | None = None
+
 
 @dataclass(frozen=True)
 class NewEvent:
+    """New line from events.jsonl."""
+
+    line: str
+
+
+@dataclass(frozen=True)
+class SampleEvent:
+    """New line from per-sample JSONL file."""
+
     line: str
 
 
@@ -276,32 +386,249 @@ def _parse_event(line: str) -> dict[str, Any] | None:
         return None
 
 
+def _get_active_samples(events: tuple[dict[str, Any], ...]) -> list[SampleState]:
+    """Get list of active (non-complete) samples from events."""
+    state = derive_state(list(events))
+    active = [s for s in state.samples.values() if s.status != "complete" and s.phase]
+    active.sort(key=lambda s: s.last_update, reverse=True)
+    return active
+
+
+def _get_selected_sample_id(events: tuple[dict[str, Any], ...], selected_idx: int) -> str | None:
+    """Get the sample_id of the currently selected sample."""
+    active = _get_active_samples(events)
+    if 0 <= selected_idx < len(active):
+        return active[selected_idx].id
+    return None
+
+
 def update(model: Model, msg: object) -> tuple[Model, Any]:
     from ._pytui import Cmd, KeyPress
 
     match msg:
         case KeyPress(key="q" | "\x03"):
             return model, Cmd.quit()
+
+        # Focus switching
+        case KeyPress(key="\t"):  # Tab
+            new_focus = "viewport" if model.focus == "list" else "list"
+            return replace(model, focus=new_focus), Cmd.none()
+
+        # Navigation in list focus
+        case KeyPress(key="j" | "\x1b[B") if model.focus == "list":  # Down
+            active = _get_active_samples(model.events)
+            new_idx = min(model.selected_idx + 1, len(active) - 1) if active else 0
+            new_sample_id = _get_selected_sample_id(model.events, new_idx)
+            # Reset viewport when selection changes
+            if new_sample_id != model.tailing_sample_id:
+                return replace(
+                    model,
+                    selected_idx=new_idx,
+                    tailing_sample_id=new_sample_id,
+                    viewport_lines=(),
+                    viewport_scroll=0,
+                    viewport_auto_follow=True,
+                    streaming_text="",
+                    last_turn=None,
+                ), Cmd.none()
+            return replace(model, selected_idx=new_idx), Cmd.none()
+
+        case KeyPress(key="k" | "\x1b[A") if model.focus == "list":  # Up
+            new_idx = max(model.selected_idx - 1, 0)
+            new_sample_id = _get_selected_sample_id(model.events, new_idx)
+            if new_sample_id != model.tailing_sample_id:
+                return replace(
+                    model,
+                    selected_idx=new_idx,
+                    tailing_sample_id=new_sample_id,
+                    viewport_lines=(),
+                    viewport_scroll=0,
+                    viewport_auto_follow=True,
+                    streaming_text="",
+                    last_turn=None,
+                ), Cmd.none()
+            return replace(model, selected_idx=new_idx), Cmd.none()
+
+        # Navigation in viewport focus
+        case KeyPress(key="j" | "\x1b[B") if model.focus == "viewport":  # Down
+            new_scroll = model.viewport_scroll + 1
+            return replace(
+                model, viewport_scroll=new_scroll, viewport_auto_follow=False
+            ), Cmd.none()
+
+        case KeyPress(key="k" | "\x1b[A") if model.focus == "viewport":  # Up
+            new_scroll = max(0, model.viewport_scroll - 1)
+            return replace(
+                model, viewport_scroll=new_scroll, viewport_auto_follow=False
+            ), Cmd.none()
+
+        case KeyPress(key="G") if model.focus == "viewport":  # Jump to bottom
+            new_scroll = max(0, len(model.viewport_lines) - 1)
+            return replace(model, viewport_scroll=new_scroll, viewport_auto_follow=True), Cmd.none()
+
+        case KeyPress(key="g"):  # gg = top (simplified, just g for now)
+            if model.focus == "viewport":
+                return replace(model, viewport_scroll=0, viewport_auto_follow=False), Cmd.none()
+
+        # Events from events.jsonl
         case NewEvent(line=line):
             event = _parse_event(line)
             if event is None:
                 return model, Cmd.none()
             new_events = model.events + (event,)
             if event.get("message") == "eval_end":
-                return Model(events=new_events, done=True), Cmd.quit()
-            return Model(events=new_events, done=model.done), Cmd.none()
+                return replace(model, events=new_events, done=True), Cmd.quit()
+
+            # Auto-select first sample if none selected
+            new_model = replace(model, events=new_events)
+            if model.tailing_sample_id is None:
+                sample_id = _get_selected_sample_id(new_events, model.selected_idx)
+                if sample_id:
+                    new_model = replace(new_model, tailing_sample_id=sample_id)
+
+            return new_model, Cmd.none()
+
+        # Events from per-sample JSONL file
+        case SampleEvent(line=line):
+            event = _parse_event(line)
+            if event is None:
+                return model, Cmd.none()
+
+            formatted, new_turn = format_sample_event(event, model.last_turn)
+            msg_type = event.get("message", "")
+
+            # Handle streaming text accumulation
+            if msg_type == "text_delta":
+                new_streaming = model.streaming_text + (formatted or "")
+                return replace(model, streaming_text=new_streaming), Cmd.none()
+
+            # When we get a non-delta event, flush accumulated streaming text
+            new_lines = list(model.viewport_lines)
+            if model.streaming_text:
+                # Wrap streaming text to reasonable width and add
+                for line in model.streaming_text.split("\n"):
+                    if line:
+                        new_lines.append(line)
+                new_streaming = ""
+            else:
+                new_streaming = model.streaming_text
+
+            # Add the formatted event line
+            if formatted and msg_type != "text_delta":
+                new_lines.append(formatted)
+
+            # Auto-scroll if enabled
+            new_scroll = model.viewport_scroll
+            if model.viewport_auto_follow:
+                new_scroll = max(0, len(new_lines) - 1)
+
+            return replace(
+                model,
+                viewport_lines=tuple(new_lines),
+                viewport_scroll=new_scroll,
+                streaming_text=new_streaming,
+                last_turn=new_turn if new_turn is not None else model.last_turn,
+            ), Cmd.none()
+
     return model, Cmd.none()
 
 
 def view(model: Model, width: int, height: int) -> list[str]:
+    """Render the split view: sample list on top, viewport on bottom."""
     state = derive_state(list(model.events))
-    lines = render(state, width, height)
+    lines: list[str] = []
 
-    # Pad to fill screen, add footer
-    footer = f"{DIM}q: quit{RESET}"
+    # Header
+    lines.append(render_header(state, width))
+
+    # Calculate layout
+    # Sample list gets ~1/3 of space (min 4 lines), viewport gets rest
+    list_height = max(4, min(height // 3, 10))
+    viewport_height = height - list_height - 3  # -3 for header, separator, footer
+
+    # Active samples with selection highlight
+    active = _get_active_samples(model.events)
+
+    # Render sample list
+    for i, sample in enumerate(active[: list_height - 1]):
+        selected = i == model.selected_idx
+        row = _format_sample_row(sample, width, selected=selected)
+        lines.append(row)
+
+    # Pad list area
+    while len(lines) < list_height:
+        lines.append("")
+
+    # Show hidden count
+    if len(active) > list_height - 1:
+        hidden = len(active) - (list_height - 1)
+        lines[list_height - 1] = f"{DIM}  ... and {hidden} more{RESET}"
+
+    # Separator with focus indicator and selected sample name
+    selected_name = ""
+    if model.tailing_sample_id:
+        active = _get_active_samples(model.events)
+        if 0 <= model.selected_idx < len(active):
+            selected_name = f" [{active[model.selected_idx].name[:20]}]"
+
+    if model.focus == "list":
+        label = "samples"
+    else:
+        label = "viewport"
+    # Build separator to fill width
+    prefix = f"─── {label}{selected_name} "
+    fill = "─" * max(0, width - len(prefix))
+    lines.append(f"{DIM}{prefix}{fill}{RESET}")
+
+    # Viewport
+    if model.tailing_sample_id:
+        # Include any streaming text as a "virtual" line at the end
+        display_lines = list(model.viewport_lines)
+        if model.streaming_text:
+            # Show streaming text (last 200 chars to keep it manageable)
+            stream_preview = model.streaming_text[-200:]
+            if len(model.streaming_text) > 200:
+                stream_preview = "..." + stream_preview
+            display_lines.append(f"{CYAN}{stream_preview}{RESET}")
+
+        if display_lines:
+            # Clamp scroll
+            max_scroll = max(0, len(display_lines) - viewport_height)
+            scroll = min(model.viewport_scroll, max_scroll)
+
+            # Render visible lines
+            visible = display_lines[scroll : scroll + viewport_height]
+            for line in visible:
+                # Truncate long lines (but preserve ANSI - crude truncation for now)
+                # TODO: proper ANSI-aware truncation
+                if len(line) > width + 20:  # Allow some slack for ANSI codes
+                    lines.append(line[: width + 10] + "...")
+                else:
+                    lines.append(line)
+        else:
+            lines.append(f"{DIM}  (waiting for events...){RESET}")
+
+        # Pad viewport
+        while len(lines) < list_height + 1 + viewport_height:
+            lines.append("")
+    else:
+        # No sample selected
+        lines.append(f"{DIM}  (select a sample to view its output){RESET}")
+        while len(lines) < list_height + 1 + viewport_height:
+            lines.append("")
+
+    # Footer
+    if model.focus == "list":
+        footer = f"{DIM}j/k: select  Tab: focus viewport  q: quit{RESET}"
+    else:
+        follow_status = "follow" if model.viewport_auto_follow else "manual"
+        footer = f"{DIM}j/k: scroll  G: bottom  Tab: focus list  [{follow_status}]  q: quit{RESET}"
+
+    # Pad to fill screen
     while len(lines) < height - 1:
         lines.append("")
     lines.append(footer)
+
     return lines[:height]
 
 
@@ -375,16 +702,22 @@ if __name__ == "__main__":
     from ._pytui import App, Cmd, Sub
 
     events_file = sys.argv[1]
+    output_dir = Path(events_file).parent
+    samples_dir = output_dir / "samples"
 
-    # Model needs the path for subscriptions. We use a non-frozen wrapper
-    # to thread the path through — subscriptions() needs it but it's static.
-    _events_path = events_file
-
-    # Patch subscriptions to close over the path
     def _subscriptions(model: Model) -> Sub:
         if model.done:
             return Sub.none()
-        return Sub.file_tail(_events_path, lambda line: NewEvent(line=line))
+
+        subs = [Sub.file_tail(events_file, lambda line: NewEvent(line=line))]
+
+        # Tail the selected sample's file if we have one
+        if model.tailing_sample_id:
+            sample_file = samples_dir / f"{model.tailing_sample_id}.jsonl"
+            if sample_file.exists():
+                subs.append(Sub.file_tail(str(sample_file), lambda line: SampleEvent(line=line)))
+
+        return Sub.batch(*subs)
 
     App(
         init=(Model(), Cmd.none()),

{wafer_core-0.1.43.dist-info → wafer_core-0.1.44.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-core
-Version: 0.1.43
+Version: 0.1.44
 Summary: Core utilities and environments for Wafer GPU kernel optimization
 Requires-Python: >=3.10
 Requires-Dist: aiohttp>=3.9.0

{wafer_core-0.1.43.dist-info → wafer_core-0.1.44.dist-info}/RECORD

@@ -362,7 +362,7 @@ wafer_core/rollouts/models.py,sha256=BrMnUpTA9_HnOghpedzdLUm9Di3FyJouzOkK4PPBg_k
 wafer_core/rollouts/paths.py,sha256=9XtrA9ylhb5LttMFe2DE7X0IHeUMjuGUerII9OscYec,3436
 wafer_core/rollouts/pipeline.py,sha256=vlJTYE3ZX2XScpF9pmtv91K8Q0g8uLmcbI5jn6b5Hzg,15319
 wafer_core/rollouts/progress.py,sha256=szA9cvWT2xUxGVhF9BaAqJMmKDqMAUlxImxcOpcnqbY,29228
-wafer_core/rollouts/progress_app.py,sha256=ir3fk8OZd09Zs7e7EC2WjQs_H-6W-zCbZAjSCDIcKEo,13517
+wafer_core/rollouts/progress_app.py,sha256=JSx4FKn_H8Thi2IK19mr2RvpH15BZuUi3nC0h6qMcD8,26218
 wafer_core/rollouts/prompt.py,sha256=EDmGb0rhWwke7tokIcO8dukc3q5c8x0n5Omi5CpAQmA,11022
 wafer_core/rollouts/providers.py,sha256=dcGJh1p30hstVbCDDtJ902lyafkg81DKjcOzb0uuKS0,1400
 wafer_core/rollouts/remote.py,sha256=cAYpRCONlsTeRxzLiegAUfjZWGtqBNwZTHehMhk5ldA,8816
@@ -723,6 +723,6 @@ wafer_core/utils/modal_execution/modal_app.py,sha256=VfS2cX8gHtnlPXemmMcEwDPeQdh
 wafer_core/utils/modal_execution/modal_config.py,sha256=7cGX9TGqilQ3qxI3OFGXV5orjtyRU-PEDOJ4vP2oxno,4421
 wafer_core/utils/modal_execution/modal_execution.py,sha256=gChjnV6jqA3A7IRP3DfvV5cSfm_MN0X4f7JZufXgdZE,24594
 wafer_core/utils/modal_execution/test_modal.py,sha256=_jqou_hrLs1Daf1590Pnb0a_lXMMa2rczAPpW9HpoNQ,8153
-wafer_core-0.1.43.dist-info/METADATA,sha256=mvNNv5AucqwdnYMP2OO2NaUmPe1xAk5re_YSwjIu4s8,1477
-wafer_core-0.1.43.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-wafer_core-0.1.43.dist-info/RECORD,,
+wafer_core-0.1.44.dist-info/METADATA,sha256=onunSBewljqBQ2pzzBa7iv3KYK5sAVUHDLByQORNqLc,1477
+wafer_core-0.1.44.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+wafer_core-0.1.44.dist-info/RECORD,,