newcode 0.1.1__py3-none-any.whl

code_puppy/status_display.py

@@ -0,0 +1,257 @@
+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
+
+# Global variable to track current token per second rate
+CURRENT_TOKEN_RATE = 0.0
+_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("dots", 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 cyan"),
+            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]Agent 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 cyan"),
+            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_puppy.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_puppy.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_puppy/summarization_agent.py

@@ -0,0 +1,176 @@
+import asyncio
+import atexit
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from typing import List
+
+from pydantic_ai import Agent
+
+from code_puppy.config import (
+    get_global_model_name,
+)
+from code_puppy.model_factory import ModelFactory, make_model_settings
+
+# Keep a module-level agent reference to avoid rebuilding per call
+_summarization_agent = None
+_agent_lock = threading.Lock()
+
+# Safe sync runner for async agent.run calls
+# Avoids "event loop is already running" by offloading to a separate thread loop when needed
+_thread_pool: ThreadPoolExecutor | None = None
+
+# Reload counter
+_reload_count = 0
+
+
+def _ensure_thread_pool():
+    global _thread_pool
+    # Check if pool is None OR if it's been shutdown
+    if _thread_pool is None or _thread_pool._shutdown:
+        _thread_pool = ThreadPoolExecutor(
+            max_workers=1, thread_name_prefix="summarizer-loop"
+        )
+    return _thread_pool
+
+
+def _shutdown_thread_pool():
+    global _thread_pool
+    if _thread_pool is not None:
+        _thread_pool.shutdown(wait=False)
+        _thread_pool = None
+
+
+atexit.register(_shutdown_thread_pool)
+
+
+async def _run_agent_async(agent: Agent, prompt: str, message_history: List):
+    return await agent.run(prompt, message_history=message_history)
+
+
+class SummarizationError(Exception):
+    """Raised when summarization fails with details about the failure."""
+
+    def __init__(self, message: str, original_error: Exception | None = None):
+        self.original_error = original_error
+        super().__init__(message)
+
+
+def run_summarization_sync(prompt: str, message_history: List) -> List:
+    """Run the summarization agent synchronously.
+
+    Raises:
+        SummarizationError: If summarization fails for any reason.
+    """
+    try:
+        agent = get_summarization_agent()
+    except Exception as e:
+        raise SummarizationError(
+            f"Failed to initialize summarization agent: {type(e).__name__}: {e}",
+            original_error=e,
+        ) from e
+
+    # Handle claude-code models: prepend system prompt to user prompt
+    from code_puppy.model_utils import prepare_prompt_for_model
+
+    model_name = get_global_model_name()
+    prepared = prepare_prompt_for_model(
+        model_name, _get_summarization_instructions(), prompt
+    )
+    prompt = prepared.user_prompt
+
+    def _run_in_thread():
+        """
+        Run the async agent in a dedicated thread with its own event loop.
+        Uses run_until_complete instead of asyncio.run to avoid shutting down
+        the default executor (which breaks DBOS in the main thread).
+        Does NOT touch global event loop state.
+        """
+        loop = asyncio.new_event_loop()
+        try:
+            coro = agent.run(prompt, message_history=message_history)
+            return loop.run_until_complete(coro)
+        finally:
+            # Clean up without shutting down the default executor
+            try:
+                # Cancel pending tasks
+                pending = asyncio.all_tasks(loop)
+                for task in pending:
+                    task.cancel()
+                if pending:
+                    loop.run_until_complete(
+                        asyncio.gather(*pending, return_exceptions=True)
+                    )
+                loop.run_until_complete(loop.shutdown_asyncgens())
+            finally:
+                loop.close()
+
+    try:
+        # Always use thread pool since we're likely in an existing event loop
+        pool = _ensure_thread_pool()
+        result = pool.submit(_run_in_thread).result()
+        return result.new_messages()
+    except Exception as e:
+        error_type = type(e).__name__
+        error_msg = str(e) if str(e) else "(no details available)"
+        raise SummarizationError(
+            f"LLM call failed during summarization: [{error_type}] {error_msg}",
+            original_error=e,
+        ) from e
+
+
+def _get_summarization_instructions() -> str:
+    """Get the system instructions for the summarization agent."""
+    return """You are a message summarization expert. Your task is to summarize conversation messages
+while preserving important context and information. The summaries should be concise but capture the essential content
+and intent of the original messages. This is to help manage token usage in a conversation history
+while maintaining context for the AI to continue the conversation effectively.
+
+When summarizing:
+1. Keep summary concise but informative
+2. Preserve important context and key information and decisions
+3. Keep any important technical details
+4. Don't summarize the system message
+5. Make sure all tool calls and responses are summarized, as they are vital
+6. Focus on token usage efficiency and system message preservation"""
+
+
+def reload_summarization_agent():
+    """Create a specialized agent for summarizing messages when context limit is reached."""
+    from code_puppy.model_utils import prepare_prompt_for_model
+
+    models_config = ModelFactory.load_config()
+    model_name = get_global_model_name()
+    model = ModelFactory.get_model(model_name, models_config)
+
+    # Handle claude-code models: swap instructions (prompt prepending happens in run_summarization_sync)
+    instructions = _get_summarization_instructions()
+    prepared = prepare_prompt_for_model(
+        model_name, instructions, "", prepend_system_to_user=False
+    )
+    instructions = prepared.instructions
+
+    model_settings = make_model_settings(model_name)
+
+    agent = Agent(
+        model=model,
+        instructions=instructions,
+        output_type=str,
+        retries=1,  # Fewer retries for summarization
+        model_settings=model_settings,
+    )
+    # NOTE: We intentionally DON'T wrap in DBOSAgent here.
+    # Summarization is a simple one-shot call that doesn't need durable execution,
+    # and DBOSAgent causes async event loop conflicts with run_sync().
+    return agent
+
+
+def get_summarization_agent(force_reload=True):
+    """
+    Retrieve the summarization agent with the currently set MODEL_NAME.
+    Forces a reload if the model has changed, or if force_reload is passed.
+    """
+    global _summarization_agent
+    with _agent_lock:
+        if force_reload or _summarization_agent is None:
+            _summarization_agent = reload_summarization_agent()
+        return _summarization_agent