code-muse 0.0.1__py3-none-any.whl

code_muse/status_display.py

@@ -0,0 +1,261 @@
+import asyncio
+import threading
+import time
+
+from rich.console import Console
+from rich.live import Live
+from rich.panel import Panel
+from rich.spinner import Spinner
+from rich.text import Text
+
+# Registers the "radar" spinner in Rich's SPINNERS dictionary
+import code_muse.messaging.spinner  # noqa: F401
+
+# Global variable to track current token per second rate
+CURRENT_TOKEN_RATE = 0.0
+# FREE-THREADED: _TOKEN_RATE_LOCK guards sync-only rate updates from display threads.
+_TOKEN_RATE_LOCK = threading.Lock()
+
+
+class StatusDisplay:
+    """
+    Displays real-time status information during model execution,
+    including token per second rate and rotating loading messages.
+    """
+
+    def __init__(self, console: Console):
+        self.console = console
+        self.token_count = 0
+        self.start_time = None
+        self.last_update_time = None
+        self.last_token_count = 0
+        self.current_rate = 0
+        self.is_active = False
+        self.task = None
+        self.live = None
+        self.loading_messages = [
+            "Fetching...",
+            "Sniffing around...",
+            "Wagging tail...",
+            "Pawsing for a moment...",
+            "Chasing tail...",
+            "Digging up results...",
+            "Barking at the data...",
+            "Rolling over...",
+            "Panting with excitement...",
+            "Chewing on it...",
+            "Prancing along...",
+            "Howling at the code...",
+            "Snuggling up to the task...",
+            "Bounding through data...",
+            "Puppy pondering...",
+        ]
+        self.current_message_index = 0
+        self.spinner = Spinner("radar", text="")
+
+    def _calculate_rate(self) -> float:
+        """Calculate the current token rate"""
+        current_time = time.time()
+        if self.last_update_time:
+            time_diff = current_time - self.last_update_time
+            token_diff = self.token_count - self.last_token_count
+            if time_diff > 0:
+                rate = token_diff / time_diff
+                # Smooth the rate calculation with the current rate
+                if self.current_rate > 0:
+                    self.current_rate = (self.current_rate * 0.7) + (rate * 0.3)
+                else:
+                    self.current_rate = rate
+
+                # Only ensure rate is not negative
+                self.current_rate = max(0, self.current_rate)
+
+                # Update the global rate for other components to access
+                global CURRENT_TOKEN_RATE
+                with _TOKEN_RATE_LOCK:
+                    CURRENT_TOKEN_RATE = self.current_rate
+
+        self.last_update_time = current_time
+        self.last_token_count = self.token_count
+        return self.current_rate
+
+    def update_rate_from_sse(
+        self, completion_tokens: int, completion_time: float
+    ) -> None:
+        """Update the token rate directly using SSE time_info data
+
+        Args:
+            completion_tokens: Number of tokens in the completion (from SSE stream)
+            completion_time: Time taken for completion in seconds (from SSE stream)
+        """
+        if completion_time > 0:
+            # Using the direct t/s formula: tokens / time
+            rate = completion_tokens / completion_time
+
+            # Use a lighter smoothing for this more accurate data
+            if self.current_rate > 0:
+                self.current_rate = (self.current_rate * 0.3) + (
+                    rate * 0.7
+                )  # Weight SSE data more heavily
+            else:
+                self.current_rate = rate
+
+            # Update the global rate
+            global CURRENT_TOKEN_RATE
+            with _TOKEN_RATE_LOCK:
+                CURRENT_TOKEN_RATE = self.current_rate
+
+    @staticmethod
+    def get_current_rate() -> float:
+        """Get the current token rate for use in other components"""
+        global CURRENT_TOKEN_RATE
+        with _TOKEN_RATE_LOCK:
+            return CURRENT_TOKEN_RATE
+
+    def update_token_count(self, tokens: int) -> None:
+        """Update the token count and recalculate the rate"""
+        # Reset timing if this is the first update of a new task
+        if self.start_time is None:
+            self.start_time = time.time()
+            self.last_update_time = self.start_time
+            # Reset token counters for new task
+            self.last_token_count = 0
+            self.current_rate = 0.0
+            # Set initial token count
+            self.token_count = tokens if tokens >= 0 else 0
+            return  # Don't calculate rate on first initialization
+
+        # Allow for incremental updates (common for streaming) or absolute updates
+        if tokens > self.token_count or tokens < 0:
+            # Incremental update or reset
+            self.token_count = tokens if tokens >= 0 else 0
+        else:
+            # If tokens <= current count but > 0, treat as incremental
+            # This handles simulated token streaming
+            self.token_count += tokens
+
+        self._calculate_rate()
+
+    def _get_status_panel(self) -> Panel:
+        """Generate a status panel with current rate and animated message"""
+        rate_text = (
+            f"{self.current_rate:.1f} t/s" if self.current_rate > 0 else "Warming up..."
+        )
+
+        # Update spinner
+        self.spinner.update()
+
+        # Rotate through loading messages every few updates
+        if int(time.time() * 2) % 4 == 0:
+            self.current_message_index = (self.current_message_index + 1) % len(
+                self.loading_messages
+            )
+
+        # Create a highly visible status message
+        status_text = Text.assemble(
+            Text(f"⏳ {rate_text} ", style="bold green"),
+            str(self.spinner),
+            Text(
+                f" {self.loading_messages[self.current_message_index]} ⏳",
+                style="bold yellow",
+            ),
+        )
+
+        # Use expanded panel with more visible formatting
+        return Panel(
+            status_text,
+            title="[bold blue]Muse Status[/bold blue]",
+            border_style="bright_blue",
+            expand=False,
+            padding=(1, 2),
+        )
+
+    def _get_status_text(self) -> Text:
+        """Generate a status text with current rate and animated message"""
+        rate_text = (
+            f"{self.current_rate:.1f} t/s" if self.current_rate > 0 else "Warming up..."
+        )
+
+        # Update spinner
+        self.spinner.update()
+
+        # Rotate through loading messages
+        self.current_message_index = (self.current_message_index + 1) % len(
+            self.loading_messages
+        )
+        message = self.loading_messages[self.current_message_index]
+
+        # Create a highly visible status text
+        return Text.assemble(
+            Text(f"[>>] {rate_text}", style="bold green"),
+            Text(f" {message}", style="yellow"),
+        )
+
+    async def _update_display(self) -> None:
+        """Update the display continuously while active using Rich Live display"""
+        # Lazy import to avoid circular dependency during module initialization
+        from code_muse.messaging import emit_info
+
+        # Add a newline to ensure we're below the blue bar
+        emit_info("")
+
+        # Create a Live display that will update in-place
+        with Live(
+            self._get_status_text(),
+            console=self.console,
+            refresh_per_second=2,  # Update twice per second
+            transient=False,  # Keep the final state visible
+        ) as live:
+            # Keep updating the live display while active
+            while self.is_active:
+                live.update(self._get_status_text())
+                await asyncio.sleep(0.5)
+
+    def start(self) -> None:
+        """Start the status display"""
+        if not self.is_active:
+            self.is_active = True
+            self.start_time = time.time()
+            self.last_update_time = self.start_time
+            self.token_count = 0
+            self.last_token_count = 0
+            self.current_rate = 0
+            self.task = asyncio.create_task(self._update_display())
+
+    def stop(self) -> None:
+        """Stop the status display"""
+        # Lazy import to avoid circular dependency during module initialization
+        from code_muse.messaging import emit_info
+
+        if self.is_active:
+            self.is_active = False
+            if self.task:
+                self.task.cancel()
+            self.task = None
+
+            # Print final stats
+            elapsed = time.time() - self.start_time if self.start_time else 0
+            avg_rate = self.token_count / elapsed if elapsed > 0 else 0
+            emit_info(
+                f"Completed: {self.token_count} tokens in {elapsed:.1f}s ({avg_rate:.1f} t/s avg)"
+            )
+
+            # Reset state
+            self.start_time = None
+            self.token_count = 0
+            self.last_update_time = None
+            self.last_token_count = 0
+            self.current_rate = 0
+
+            # Reset global rate to 0 to avoid affecting subsequent tasks
+            global CURRENT_TOKEN_RATE
+            with _TOKEN_RATE_LOCK:
+                CURRENT_TOKEN_RATE = 0.0
+        else:
+            # Even if not active, ensure we print stats when stop is called
+            # This is for testing purposes
+            elapsed = time.time() - self.start_time if self.start_time else 0
+            avg_rate = self.token_count / elapsed if elapsed > 0 else 0
+            emit_info(
+                f"Completed: {self.token_count} tokens in {elapsed:.1f}s ({avg_rate:.1f} t/s avg)"
+            )

code_muse/stream_parser/__init__.py

@@ -0,0 +1,76 @@
+"""Streaming text parsers for assistant markup.
+
+This module ports Codex's Rust ``utils/stream-parser`` crate to idiomatic
+Python 3.11+.  Parsers are composable: one parser can wrap another,
+delegating and merging output.
+"""
+
+from code_muse.stream_parser.assistant_text_parser import (
+    AssistantTextChunk,
+    AssistantTextStreamParser,
+)
+from code_muse.stream_parser.citation_parser import (
+    CITATION_CLOSE,
+    CITATION_OPEN,
+    CitationStreamParser,
+    strip_citations,
+)
+from code_muse.stream_parser.inline_hidden_tag_parser import (
+    ExtractedInlineTag,
+    InlineHiddenTagParser,
+    InlineTagSpec,
+)
+from code_muse.stream_parser.proposed_plan_parser import (
+    ProposedPlanParser,
+    ProposedPlanSegment,
+    ProposedPlanSegmentType,
+    extract_proposed_plan_text,
+    strip_proposed_plan_blocks,
+)
+from code_muse.stream_parser.stream_text_chunk import StreamTextChunk
+from code_muse.stream_parser.stream_text_parser import StreamTextParser
+from code_muse.stream_parser.tagged_line_parser import (
+    TaggedLineParser,
+    TaggedLineSegment,
+    TaggedLineSegmentNormal,
+    TaggedLineSegmentTagDelta,
+    TaggedLineSegmentTagEnd,
+    TaggedLineSegmentTagStart,
+    TagSpec,
+)
+from code_muse.stream_parser.utf8_stream_parser import (
+    IncompleteUtf8AtEof,
+    InvalidUtf8,
+    Utf8StreamParser,
+    Utf8StreamParserError,
+)
+
+__all__ = [
+    "AssistantTextChunk",
+    "AssistantTextStreamParser",
+    "CITATION_CLOSE",
+    "CITATION_OPEN",
+    "CitationStreamParser",
+    "ExtractedInlineTag",
+    "IncompleteUtf8AtEof",
+    "InlineHiddenTagParser",
+    "InlineTagSpec",
+    "InvalidUtf8",
+    "ProposedPlanParser",
+    "ProposedPlanSegment",
+    "ProposedPlanSegmentType",
+    "StreamTextChunk",
+    "StreamTextParser",
+    "TagSpec",
+    "TaggedLineParser",
+    "TaggedLineSegment",
+    "TaggedLineSegmentNormal",
+    "TaggedLineSegmentTagDelta",
+    "TaggedLineSegmentTagEnd",
+    "TaggedLineSegmentTagStart",
+    "Utf8StreamParser",
+    "Utf8StreamParserError",
+    "extract_proposed_plan_text",
+    "strip_citations",
+    "strip_proposed_plan_blocks",
+]

code_muse/stream_parser/assistant_text_parser.py

@@ -0,0 +1,90 @@
+"""Combined assistant text parser: citations + plan segments in one pass."""
+
+from dataclasses import dataclass, field
+
+from code_muse.stream_parser.citation_parser import CitationStreamParser
+from code_muse.stream_parser.proposed_plan_parser import (
+    ProposedPlanParser,
+    ProposedPlanSegment,
+)
+
+
+@dataclass
+class AssistantTextChunk:
+    """Result chunk from :class:`AssistantTextStreamParser`.
+
+    Attributes:
+        visible_text: Text safe to render immediately (tags stripped).
+        citations: Citation bodies extracted in this chunk.
+        plan_segments: Plan boundary / delta segments when plan mode is on.
+    """
+
+    visible_text: str = ""
+    citations: list[str] = field(default_factory=list)
+    plan_segments: list[ProposedPlanSegment] = field(default_factory=list)
+
+    def is_empty(self) -> bool:
+        """Return True when nothing was produced in this chunk."""
+        return not self.visible_text and not self.citations and not self.plan_segments
+
+
+class AssistantTextStreamParser:
+    """Parses assistant text streaming markup in one pass.
+
+    Strips ``<oai-mem-citation>`` tags and extracts citation payloads.
+    In plan mode, also strips ``<proposed_plan>`` blocks and emits plan
+    segments so callers can render or hide plan content independently.
+
+    The two sub-parsers are wired in series: incoming text first goes through
+    :class:`CitationStreamParser`; the resulting visible text is then fed into
+    :class:`ProposedPlanParser` when plan mode is enabled.
+    """
+
+    def __init__(self, plan_mode: bool = False) -> None:
+        self.plan_mode = plan_mode
+        self._citations = CitationStreamParser()
+        self._plan = ProposedPlanParser()
+
+    def push_str(self, chunk: str) -> AssistantTextChunk:
+        """Feed a new text chunk.
+
+        Args:
+            chunk: Raw assistant text delta.
+
+        Returns:
+            Chunk containing visible text, any newly-extracted citations, and
+            any plan segments when plan mode is enabled.
+        """
+        citation_chunk = self._citations.push_str(chunk)
+        out = self._parse_visible_text(citation_chunk.visible_text)
+        out.citations = citation_chunk.extracted
+        return out
+
+    def finish(self) -> AssistantTextChunk:
+        """Flush any buffered state.
+
+        Unterminated citations and plan blocks are auto-closed.
+
+        Returns:
+            Final chunk with trailing visible text, citations, and plan
+            segments.
+        """
+        citation_chunk = self._citations.finish()
+        out = self._parse_visible_text(citation_chunk.visible_text)
+        if self.plan_mode:
+            tail = self._plan.finish()
+            if not tail.is_empty():
+                out.visible_text += tail.visible_text
+                out.plan_segments.extend(tail.extracted)
+        out.citations = citation_chunk.extracted
+        return out
+
+    def _parse_visible_text(self, visible_text: str) -> AssistantTextChunk:
+        """Route citation-visible text through the plan parser when needed."""
+        if not self.plan_mode:
+            return AssistantTextChunk(visible_text=visible_text)
+        plan_chunk = self._plan.push_str(visible_text)
+        return AssistantTextChunk(
+            visible_text=plan_chunk.visible_text,
+            plan_segments=plan_chunk.extracted,
+        )

code_muse/stream_parser/citation_parser.py

@@ -0,0 +1,76 @@
+"""Stream parser for <oai-mem-citation> tags."""
+
+from code_muse.stream_parser.inline_hidden_tag_parser import (
+    InlineHiddenTagParser,
+    InlineTagSpec,
+)
+from code_muse.stream_parser.stream_text_chunk import StreamTextChunk
+from code_muse.stream_parser.stream_text_parser import StreamTextParser
+
+CITATION_OPEN = "<oai-mem-citation>"
+CITATION_CLOSE = "</oai-mem-citation>"
+
+
+class CitationStreamParser(StreamTextParser[str]):
+    """Stream parser for ``<oai-mem-citation>…</oai-mem-citation>`` tags.
+
+    Thin wrapper around :class:`InlineHiddenTagParser` that returns citation
+    bodies as plain strings and omits the citation tags from visible text.
+    Matching is literal and non-nested.  Unterminated tags auto-close at
+    :meth:`finish`.
+    """
+
+    def __init__(self) -> None:
+        self._inner = InlineHiddenTagParser(
+            [InlineTagSpec("citation", CITATION_OPEN, CITATION_CLOSE)]
+        )
+
+    def push_str(self, chunk: str) -> StreamTextChunk[str]:
+        """Feed a new text chunk.
+
+        Args:
+            chunk: Incoming assistant text delta.
+
+        Returns:
+            Visible text with citation tags stripped, plus any citation
+            payloads extracted from the chunk.
+        """
+        inner = self._inner.push_str(chunk)
+        return StreamTextChunk(
+            visible_text=inner.visible_text,
+            extracted=[tag.content for tag in inner.extracted],
+        )
+
+    def finish(self) -> StreamTextChunk[str]:
+        """Flush any buffered state.
+
+        Unterminated citation tags are auto-closed and their accumulated
+        content is returned as a citation payload.
+
+        Returns:
+            Final visible text and any trailing citations.
+        """
+        inner = self._inner.finish()
+        return StreamTextChunk(
+            visible_text=inner.visible_text,
+            extracted=[tag.content for tag in inner.extracted],
+        )
+
+
+def strip_citations(text: str) -> tuple[str, list[str]]:
+    """Strip citation tags from a complete string.
+
+    Args:
+        text: Full assistant response text.
+
+    Returns:
+        ``(visible_text, citations)`` where ``visible_text`` has all
+        ``<oai-mem-citation>…</oai-mem-citation>`` blocks removed and
+        ``citations`` is the list of extracted citation bodies.
+    """
+    parser = CitationStreamParser()
+    out = parser.push_str(text)
+    tail = parser.finish()
+    visible = out.visible_text + tail.visible_text
+    citations = out.extracted + tail.extracted
+    return visible, citations