code-puppy 0.0.89__py3-none-any.whl → 0.0.91__py3-none-any.whl

code_puppy/agent.py

@@ -80,6 +80,7 @@ def reload_code_generation_agent():
         output_type=str,
         retries=3,
         history_processors=[message_history_accumulator],
+        toolsets=_load_mcp_servers()
     )
     register_all_tools(agent)
     _code_generation_agent = agent

code_puppy/main.py

@@ -17,15 +17,16 @@ 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
 from code_puppy.version_checker import fetch_latest_version
-from code_puppy.message_history_processor import message_history_processor, prune_interrupted_tool_calls
+from code_puppy.message_history_processor import message_history_processor
 
 
 # from code_puppy.tools import *  # noqa: F403
-
+import logfire
 
 # Define a function to get the secret file path
 def get_secret_file_path():
@@ -37,7 +38,10 @@ def get_secret_file_path():
 
 async def main():
     # Ensure the config directory and puppy.cfg with name info exist (prompt user if needed)
+    logfire.configure(token="pylf_v1_us_8G5nLznQtHMRsL4hsNG5v3fPWKjyXbysrMgrQ1bV1wRP")
+    logfire.instrument_pydantic_ai()
     ensure_config_exists()
+
     current_version = __version__
     latest_version = fetch_latest_version("code-puppy")
     console.print(f"Current version: {current_version}")
@@ -194,17 +198,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 +405,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/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/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.89.dist-info → code_puppy-0.0.91.dist-info}/METADATA

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

{code_puppy-0.0.89.dist-info → code_puppy-0.0.91.dist-info}/RECORD

@@ -1,12 +1,13 @@
 code_puppy/__init__.py,sha256=CWH46ZAmJRmHAbOiAhG07OrWYEcEt4yvDTkZU341Wag,169
-code_puppy/agent.py,sha256=7_1FpGPnw8U632OXP0hLmFIozfVvllF491q8gCpaa8c,3284
+code_puppy/agent.py,sha256=Hwlblm9Om4qZPOur28auxBkdJRAA9Kz-36YPQg7i3SI,3321
 code_puppy/agent_prompts.py,sha256=t3-lqDKrDxCKxFa_va4Suze9BT-JOu1dh9iGiAVNFO4,6828
 code_puppy/config.py,sha256=r5nw5ChOP8xd_K5yo8U5OtO2gy2bFhARiyNtDp1JrwQ,5013
-code_puppy/main.py,sha256=NLA6TI2is3htFRc5BGTwF4Xp68EtnN-6nOExD9D8i_g,11513
-code_puppy/message_history_processor.py,sha256=2k9y4wXkyZ9v4XLpfH6E3TKLzf303i-k1hXBE-aIf9U,9249
+code_puppy/main.py,sha256=_9PWsk1GLVvfH9w8KbcoMXvUCwhzxu6WRqguTIYVccE,21183
+code_puppy/message_history_processor.py,sha256=cL9DgSqpwWk0dSamc1I7NR1n1CNefvKlIMSCSA4YkM8,10034
 code_puppy/model_factory.py,sha256=HXuFHNkVjkCcorAd3ScFmSvBILO932UTq6OmNAqisT8,10898
 code_puppy/models.json,sha256=jr0-LW87aJS79GosVwoZdHeeq5eflPzgdPoMbcqpVA8,2728
 code_puppy/state_management.py,sha256=JkTkmq6f9rl_RHPDoBqJvbAzgaMsIkJf-k38ragItIo,1692
+code_puppy/status_display.py,sha256=tzllkO0ssmRZSn20mFXnEt1BZXH4oP9eHSSN-7Jyrok,7946
 code_puppy/summarization_agent.py,sha256=jHUQe6iYJsMT0ywEwO7CrhUIKEamO5imhAsDwvNuvow,2684
 code_puppy/token_utils.py,sha256=g7Jj6NAy_a2ab7BXpwyhktruR-QlUV670H_mCPZV1N4,2110
 code_puppy/version_checker.py,sha256=aRGulzuY4C4CdFvU1rITduyL-1xTFsn4GiD1uSfOl_Y,396
@@ -23,9 +24,9 @@ code_puppy/tools/common.py,sha256=UkhnfLG1bmd4f9nZCcmno088AtKtAnEES1tydxUN-Fk,32
 code_puppy/tools/file_modifications.py,sha256=BzQrGEacS2NZr2ru9N30x_Qd70JDudBKOAPO1XjBohg,13861
 code_puppy/tools/file_operations.py,sha256=5oJSJCKPLSgJ9YcD1CV9D-SJ0hnzVzmuAs-eiriLPLo,13615
 code_puppy/tools/token_check.py,sha256=F3eygdI8fgb6dfCrSkGw_OLI7cb_Kpa5ILft4BQ7hvY,525
-code_puppy-0.0.89.data/data/code_puppy/models.json,sha256=jr0-LW87aJS79GosVwoZdHeeq5eflPzgdPoMbcqpVA8,2728
-code_puppy-0.0.89.dist-info/METADATA,sha256=Dyf2vyOvwhvrBVfNd1xIvmIgcvCEOKOYK5Hyy6xScT4,6351
-code_puppy-0.0.89.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-code_puppy-0.0.89.dist-info/entry_points.txt,sha256=d8YkBvIUxF-dHNJAj-x4fPEqizbY5d_TwvYpc01U5kw,58
-code_puppy-0.0.89.dist-info/licenses/LICENSE,sha256=31u8x0SPgdOq3izJX41kgFazWsM43zPEF9eskzqbJMY,1075
-code_puppy-0.0.89.dist-info/RECORD,,
+code_puppy-0.0.91.data/data/code_puppy/models.json,sha256=jr0-LW87aJS79GosVwoZdHeeq5eflPzgdPoMbcqpVA8,2728
+code_puppy-0.0.91.dist-info/METADATA,sha256=PaIPBOhLU92HD2QkF1lZVCgFBi8KpTV46pzzAIf620M,6351
+code_puppy-0.0.91.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+code_puppy-0.0.91.dist-info/entry_points.txt,sha256=d8YkBvIUxF-dHNJAj-x4fPEqizbY5d_TwvYpc01U5kw,58
+code_puppy-0.0.91.dist-info/licenses/LICENSE,sha256=31u8x0SPgdOq3izJX41kgFazWsM43zPEF9eskzqbJMY,1075
+code_puppy-0.0.91.dist-info/RECORD,,