code-puppy 0.0.88__tar.gz → 0.0.90__tar.gz

{code_puppy-0.0.88 → code_puppy-0.0.90}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: code-puppy
-Version: 0.0.88
+Version: 0.0.90
 Summary: Code generation agent
 Author: Michael Pfaffenberger
 License: MIT

{code_puppy-0.0.88 → code_puppy-0.0.90}/code_puppy/main.py

@@ -1,6 +1,7 @@
 import argparse
 import asyncio
 import os
+import random
 import sys
 
 from dotenv import load_dotenv
@@ -17,6 +18,7 @@ from code_puppy.command_line.prompt_toolkit_completion import (
 )
 from code_puppy.config import ensure_config_exists
 from code_puppy.state_management import get_message_history, set_message_history
+from code_puppy.status_display import StatusDisplay
 
 # Initialize rich console for pretty output
 from code_puppy.tools.common import console
@@ -194,17 +196,167 @@ async def interactive_mode(history_file_path: str) -> None:
             try:
                 prettier_code_blocks()
                 local_cancelled = False
-
+                
+                # Initialize status display for tokens per second and loading messages
+                status_display = StatusDisplay(console)
+                
+                # Print a message indicating we're about to start processing
+                console.print("\nStarting task processing...")
+                
+                async def track_tokens_from_messages():
+                    """
+                    Track real token counts from message history.
+                    
+                    This async function runs in the background and periodically checks
+                    the message history for new tokens. When new tokens are detected,
+                    it updates the StatusDisplay with the incremental count to calculate
+                    an accurate tokens-per-second rate.
+                    
+                    It also looks for SSE stream time_info data to get precise token rate
+                    calculations using the formula: completion_tokens * 1 / completion_time
+                    
+                    The function continues running until status_display.is_active becomes False.
+                    """
+                    from code_puppy.message_history_processor import estimate_tokens_for_message
+                    import json
+                    import re
+                    
+                    last_token_total = 0
+                    last_sse_data = None
+                    
+                    while status_display.is_active:
+                        # Get real token count from message history
+                        messages = get_message_history()
+                        if messages:
+                            # Calculate total tokens across all messages
+                            current_token_total = sum(estimate_tokens_for_message(msg) for msg in messages)
+                            
+                            # If tokens increased, update the display with the incremental count
+                            if current_token_total > last_token_total:
+                                status_display.update_token_count(current_token_total - last_token_total)
+                                last_token_total = current_token_total
+                            
+                            # Try to find SSE stream data in assistant messages
+                            for msg in messages:
+                                # Handle different message types (dict or ModelMessage objects)
+                                if hasattr(msg, 'role') and msg.role == 'assistant':
+                                    # ModelMessage object with role attribute
+                                    content = msg.content if hasattr(msg, 'content') else ''
+                                elif isinstance(msg, dict) and msg.get('role') == 'assistant':
+                                    # Dictionary with 'role' key
+                                    content = msg.get('content', '')
+                                # Support for ModelRequest/ModelResponse objects
+                                elif hasattr(msg, 'message') and hasattr(msg.message, 'role') and msg.message.role == 'assistant':
+                                    # Access content through the message attribute
+                                    content = msg.message.content if hasattr(msg.message, 'content') else ''
+                                else:
+                                    # Skip if not an assistant message or unrecognized format
+                                    continue
+                                
+                                # Convert content to string if it's not already
+                                if not isinstance(content, str):
+                                    try:
+                                        content = str(content)
+                                    except:
+                                        continue
+                                
+                                # Look for SSE usage data pattern in the message content
+                                sse_matches = re.findall(r'\{\s*"usage".*?"time_info".*?\}', content, re.DOTALL)
+                                for match in sse_matches:
+                                    try:
+                                        # Parse the JSON data
+                                        sse_data = json.loads(match)
+                                        if sse_data != last_sse_data:  # Only process new data
+                                            # Check if we have time_info and completion_tokens
+                                            if 'time_info' in sse_data and 'completion_time' in sse_data['time_info'] and \
+                                               'usage' in sse_data and 'completion_tokens' in sse_data['usage']:
+                                                completion_time = float(sse_data['time_info']['completion_time'])
+                                                completion_tokens = int(sse_data['usage']['completion_tokens'])
+                                                
+                                                # Update rate using the accurate SSE data
+                                                if completion_time > 0 and completion_tokens > 0:
+                                                    status_display.update_rate_from_sse(completion_tokens, completion_time)
+                                                    last_sse_data = sse_data
+                                    except (json.JSONDecodeError, KeyError, ValueError):
+                                        # Ignore parsing errors and continue
+                                        pass
+                        
+                        # Small sleep interval for responsive updates without excessive CPU usage
+                        await asyncio.sleep(0.1)
+                
+                async def wrap_agent_run(original_run, *args, **kwargs):
+                    """
+                    Wraps the agent's run method to enable token tracking.
+                    
+                    This wrapper preserves the original functionality while allowing
+                    us to track tokens as they are generated by the model. No additional
+                    logic is needed here since the token tracking happens in a separate task.
+                    
+                    Args:
+                        original_run: The original agent.run method
+                        *args, **kwargs: Arguments to pass to the original run method
+                        
+                    Returns:
+                        The result from the original run method
+                    """
+                    result = await original_run(*args, **kwargs)
+                    return result
+                
                 async def run_agent_task():
+                    """
+                    Main task runner for the agent with token tracking.
+                    
+                    This function:  
+                    1. Sets up the agent with token tracking
+                    2. Starts the status display showing token rate
+                    3. Runs the agent with the user's task
+                    4. Ensures proper cleanup of all resources
+                    
+                    Returns the agent's result or raises any exceptions that occurred.
+                    """
+                    # Token tracking task reference for cleanup
+                    token_tracking_task = None
+                    
                     try:
+                        # Initialize the agent
                         agent = get_code_generation_agent()
+                        
+                        # Start status display
+                        status_display.start()
+                        
+                        # Start token tracking
+                        token_tracking_task = asyncio.create_task(track_tokens_from_messages())
+                        
+                        # Create a wrapper for the agent's run method
+                        original_run = agent.run
+                        
+                        async def wrapped_run(*args, **kwargs):
+                            return await wrap_agent_run(original_run, *args, **kwargs)
+                        
+                        agent.run = wrapped_run
+                        
+                        # Run the agent with MCP servers
                         async with agent.run_mcp_servers():
-                            return await agent.run(
-                                task, message_history=get_message_history()
+                            result = await agent.run(
+                                task,
+                                message_history=get_message_history()
                             )
+                            return result
                     except Exception as e:
                         console.log("Task failed", e)
-
+                        raise
+                    finally:
+                        # Clean up resources
+                        if status_display.is_active:
+                            status_display.stop()
+                        if token_tracking_task and not token_tracking_task.done():
+                            token_tracking_task.cancel()
+                    if not agent_task.done():
+                        set_message_history(
+                            message_history_processor(
+                                get_message_history()
+                            )
+                        )
                 agent_task = asyncio.create_task(run_agent_task())
 
                 import signal
@@ -251,11 +403,20 @@ async def interactive_mode(history_file_path: str) -> None:
 
                 if local_cancelled:
                     console.print("Task canceled by user")
+                    # Ensure status display is stopped if canceled
+                    if status_display.is_active:
+                        status_display.stop()
                 else:
-                    agent_response = result.output
-                    console.print(agent_response)
-                    filtered = message_history_processor(get_message_history())
-                    set_message_history(filtered)
+                    if result is not None and hasattr(result, 'output'):
+                        agent_response = result.output
+                        console.print(agent_response)
+                        filtered = message_history_processor(get_message_history())
+                        set_message_history(filtered)
+                    else:
+                        console.print("[yellow]No result received from the agent[/yellow]")
+                        # Still process history if possible
+                        filtered = message_history_processor(get_message_history())
+                        set_message_history(filtered)
 
                 # Show context status
                 console.print(

{code_puppy-0.0.88 → code_puppy-0.0.90}/code_puppy/message_history_processor.py

@@ -17,6 +17,13 @@ from code_puppy.tools.common import console
 from code_puppy.model_factory import ModelFactory
 from code_puppy.config import get_model_name
 
+# Import the status display to get token rate info
+try:
+    from code_puppy.status_display import StatusDisplay
+    STATUS_DISPLAY_AVAILABLE = True
+except ImportError:
+    STATUS_DISPLAY_AVAILABLE = False
+
 # Import summarization agent
 try:
     from code_puppy.summarization_agent import (
@@ -246,9 +253,25 @@ def message_history_processor(messages: List[ModelMessage]) -> List[ModelMessage
     model_max = get_model_context_length()
 
     proportion_used = total_current_tokens / model_max
+    
+    # Include token per second rate if available
+    token_rate_info = ""
+    if STATUS_DISPLAY_AVAILABLE:
+        current_rate = StatusDisplay.get_current_rate()
+        if current_rate > 0:
+            # Format with improved precision when using SSE data
+            if current_rate > 1000:
+                token_rate_info = f", {current_rate:.0f} t/s"
+            else:
+                token_rate_info = f", {current_rate:.1f} t/s"
+    
+    # Print blue status bar - ALWAYS at top
     console.print(f"""
-[bold white on blue] Tokens in context: {total_current_tokens}, total model capacity: {model_max}, proportion used: {proportion_used:.2f}
+[bold white on blue] Tokens in context: {total_current_tokens}, total model capacity: {model_max}, proportion used: {proportion_used:.2f}{token_rate_info}
 """)
+    
+    # Print extra line to ensure separation
+    console.print("\n")
 
     if proportion_used > 0.85:
         summary = summarize_messages(messages)

code_puppy-0.0.90/code_puppy/status_display.py

@@ -0,0 +1,209 @@
+import asyncio
+import random
+import time
+from datetime import datetime
+from typing import List, Optional
+
+from rich.console import Console, RenderableType
+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
+
+
+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
+                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
+            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
+        return CURRENT_TOKEN_RATE
+
+    def update_token_count(self, tokens: int) -> None:
+        """Update the token count and recalculate the rate"""
+        if self.start_time is None:
+            self.start_time = time.time()
+            self.last_update_time = self.start_time
+        
+        # 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"),
+            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]Code Puppy 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"""
+        # Add a newline to ensure we're below the blue bar
+        self.console.print("\n")
+        
+        # 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"""
+        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
+            self.console.print(f"[dim]Completed: {self.token_count} tokens in {elapsed:.1f}s ({avg_rate:.1f} t/s avg)[/dim]")
+            
+            # Reset
+            self.start_time = None
+            self.token_count = 0

{code_puppy-0.0.88 → code_puppy-0.0.90}/code_puppy/tools/file_operations.py

@@ -25,6 +25,7 @@ class ListedFile(BaseModel):
 
 class ListFileOutput(BaseModel):
     files: List[ListedFile]
+    error: str | None = None
 
 
 def _list_files(
@@ -270,7 +271,7 @@ def _grep(context: RunContext, search_string: str, directory: str = ".") -> Grep
                                 **{
                                     "file_path": file_path,
                                     "line_number": line_number,
-                                    "line_content": line_content.rstrip("\n\r"),
+                                    "line_content": line_content.rstrip("\n\r")[2048:],
                                 }
                             )
                             matches.append(match_info)
@@ -311,7 +312,15 @@ def _grep(context: RunContext, search_string: str, directory: str = ".") -> Grep
 def list_files(
     context: RunContext, directory: str = ".", recursive: bool = True
 ) -> ListFileOutput:
-    return _list_files(context, directory, recursive)
+    list_files_output = _list_files(context, directory, recursive)
+    tokenizer = get_tokenizer()
+    num_tokens = len(tokenizer.encode(list_files_output.model_dump_json()))
+    if num_tokens > 10000:
+        return ListFileOutput(
+            files=[],
+            error="Too many files - tokens exceeded. Try listing non-recursively"
+        )
+    return list_files_output
 
 
 def read_file(context: RunContext, file_path: str = "", start_line: int | None = None, num_lines: int | None = None) -> ReadFileOutput:

{code_puppy-0.0.88 → code_puppy-0.0.90}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "code-puppy"
-version = "0.0.88"
+version = "0.0.90"
 description = "Code generation agent"
 readme = "README.md"
 requires-python = ">=3.10"