codemate-cli 1.0.0__py3-none-any.whl

codemate/ui/streaming.py

@@ -0,0 +1,436 @@
+import asyncio
+from typing import AsyncGenerator, Optional, Dict, Tuple, List
+from rich.console import Console, Group
+from rich.markdown import Markdown
+from rich.live import Live
+from rich.panel import Panel
+from rich.syntax import Syntax
+from rich.text import Text
+import re
+from codemate.utils.error_handler import ErrorHandler
+
+class StreamingHandler:
+    """Handler for streaming responses with rich formatting"""
+    
+    def __init__(self, console: Console):
+        self.console = console
+        self.error_handler = ErrorHandler(console)
+        # Standard markdown code blocks
+        self.code_block_pattern = re.compile(r'```(\w+)?\n(.*?)```', re.DOTALL)
+    
+    def _parse_code_blocks(self, text: str) -> list:
+        """Parse code blocks and identify their type"""
+        blocks = []
+        
+        for match in self.code_block_pattern.finditer(text):
+            language = match.group(1) or 'text'
+            content = match.group(2).strip()
+            
+            # Check for special markers
+            is_file = content.startswith('# FILE:')
+            is_command = content.startswith('# COMMAND')
+            
+            if is_file:
+                # Extract filename from first line
+                first_line = content.split('\n')[0]
+                filename = first_line.replace('# FILE:', '').strip()
+                # Remove the FILE marker line from content
+                code_content = '\n'.join(content.split('\n')[1:]).strip()
+                
+                blocks.append({
+                    'type': 'file_code',
+                    'start': match.start(),
+                    'end': match.end(),
+                    'language': language,
+                    'filename': filename,
+                    'content': code_content
+                })
+            elif is_command:
+                # Remove the COMMAND marker line from content
+                command_content = '\n'.join(content.split('\n')[1:]).strip()
+                
+                blocks.append({
+                    'type': 'command',
+                    'start': match.start(),
+                    'end': match.end(),
+                    'content': command_content
+                })
+            else:
+                # Regular code block
+                blocks.append({
+                    'type': 'regular_code',
+                    'start': match.start(),
+                    'end': match.end(),
+                    'language': language,
+                    'content': content
+                })
+        
+        blocks.sort(key=lambda x: x['start'])
+        return blocks
+    
+    def _render_block(self, block: dict) -> Panel:
+        """Render a single code block with appropriate styling"""
+        
+        if block['type'] == 'file_code':
+            syntax = Syntax(
+                block['content'],
+                block['language'],
+                theme="monokai",
+                line_numbers=False,
+                word_wrap=False
+            )
+            return Panel(
+                syntax,
+                border_style="#48AEF3",
+                padding=(1, 2),
+                title=f"[bold color(#48AEF3)]📄 {block['filename']}[/bold color(#48AEF3)]",
+                title_align="left"
+            )
+        
+        elif block['type'] == 'command':
+            syntax = Syntax(
+                block['content'],
+                "bash",
+                theme="monokai",
+                line_numbers=False,
+                word_wrap=True
+            )
+            return Panel(
+                syntax,
+                border_style="yellow",
+                padding=(1, 2),
+                title="[bold yellow]⚡ COMMAND[/bold yellow]",
+                title_align="left"
+            )
+        
+        elif block['type'] == 'regular_code':
+            syntax = Syntax(
+                block['content'],
+                block['language'],
+                theme="monokai",
+                line_numbers=False,
+                word_wrap=False
+            )
+            return Panel(
+                syntax,
+                border_style="#2FCACE",
+                padding=(1, 2),
+                title=f"[bold color(#2FCACE)]{block['language'].upper()}[/bold color(#2FCACE)]",
+                title_align="left"
+            )
+    
+
+    async def handle_streaming_response_async(
+        self, 
+        generator: AsyncGenerator[Dict[str, str], None]
+    ) -> Tuple[str, Optional[Dict], List[Dict]]:
+        """
+        Handle streaming response - accumulate all content and render once at the end
+        
+        Flow:
+        1. Thinking... (status only, not printed)
+        2. 🔧 Searching knowledge base... (PRINTED to UI)
+        3. ✓ Completed knowledge base search (PRINTED to UI)
+        4. Thinking... (status only, not printed)
+        5. Final formatted message (PRINTED)
+        
+        Returns:
+            Tuple of (message_content, assistant_entry, tool_results)
+        """
+        accumulated_text = ""
+        reasoning_text = ""
+        tool_calls = []
+        tool_results = []
+        active_tools = {}
+        should_stop = False
+        tool_search_shown = False
+        status_context = None
+
+        try:
+            # Show initial "Thinking..." as status only (not printed to UI)
+            status_context = self.console.status("[dim]Thinking...[/dim]", spinner="dots")
+            status_context.start()
+            
+            async for chunk in generator:
+                if should_stop:
+                    break
+                
+                chunk_type = chunk.get("type")
+                
+                # Handle error chunks from API
+                if chunk_type == "error":
+                    if status_context:
+                        status_context.stop()
+                    self.error_handler.handle_streaming_error(chunk)
+                    should_stop = True
+                    break
+                
+                elif chunk_type == "reasoning":
+                    content = chunk.get("content", "")
+                    reasoning_text += content
+                
+                elif chunk_type == "message":
+                    content = chunk.get("content", "")
+                    accumulated_text += content
+                
+                elif chunk_type == "tool_calls":
+                    # Store tool calls for history
+                    tool_calls = chunk.get("tool_calls", [])
+                    
+                    # PRINT to UI: Start of knowledge base search
+                    if tool_calls and not tool_search_shown:
+                        if status_context:
+                            status_context.stop()
+                        self.console.print()
+                        self.console.print("[color(#2FCACE)]🔧 Searching knowledge base...[/color(#2FCACE)]")
+                        tool_search_shown = True
+                
+                elif chunk_type == "tool_call_start":
+                    tool_call = chunk.get("tool_call", {})
+                    tool_id = tool_call.get("id")
+                    tool_name = tool_call.get("name", "unknown")
+                    
+                    if tool_id:
+                        active_tools[tool_id] = tool_name
+                
+                elif chunk_type == "tool_call_complete":
+                    tool_call = chunk.get("tool_call", {})
+                    tool_id = tool_call.get("id")
+                    tool_name = tool_call.get("name")
+                    tool_result = tool_call.get("result")
+                    
+                    # Remove from active tools
+                    if tool_id in active_tools:
+                        del active_tools[tool_id]
+                    
+                    # Store tool result for history
+                    if tool_id and tool_result:
+                        if isinstance(tool_result, dict):
+                            import json
+                            result_content = json.dumps(tool_result)
+                        else:
+                            result_content = str(tool_result)
+                        
+                        tool_results.append({
+                            "role": "tool",
+                            "tool_call_id": tool_id,
+                            "content": result_content,
+                            "name": tool_name
+                        })
+                    
+                    # PRINT to UI: When all tools are complete
+                    if not active_tools and tool_search_shown:
+                        self.console.print("[green]✓ Completed knowledge base search[/green]")
+                        self.console.print()
+                        # Restart "Thinking..." status (not printed, just spinner)
+                        status_context = self.console.status("[dim]Thinking...[/dim]", spinner="dots")
+                        status_context.start()
+                        tool_search_shown = False
+                
+                elif chunk_type == "usage":
+                    # Optionally capture usage stats
+                    pass
+            
+            # Stop any active status before rendering
+            if status_context:
+                status_context.stop()
+            
+            # After streaming completes (only if no error occurred)
+            if not should_stop:
+                # Now PRINT the complete formatted response to UI
+                if accumulated_text.strip() or reasoning_text.strip():
+                    self._render_final_content(accumulated_text, reasoning_text)
+                    self.console.print()  # Add final newline
+
+                # Build assistant history entry
+                assistant_entry = {
+                    "role": "assistant",
+                    "content": accumulated_text.strip()
+                }
+                
+                # Add tool calls if present
+                if tool_calls:
+                    assistant_entry["tool_calls"] = tool_calls
+                
+                # Add reasoning content if present
+                if reasoning_text.strip():
+                    assistant_entry["reasoning_content"] = reasoning_text.strip()
+                
+                return accumulated_text.strip(), assistant_entry, tool_results
+            else:
+                # Error occurred, return empty
+                return "", None, []
+
+        except Exception as e:
+            # Stop the thinking indicator if it's still active
+            if status_context:
+                status_context.stop()
+            self.error_handler.handle_error(e)
+            return "", None, []
+        
+        finally:
+            # Always ensure the generator is properly closed
+            try:
+                await generator.aclose()
+            except Exception:
+                pass
+
+
+    def handle_streaming_response(
+        self,
+        generator: AsyncGenerator[Dict[str, str], None]
+    ) -> Tuple[str, Optional[Dict], List[Dict]]:
+        """Handle streaming response (SYNC wrapper for CLI usage)"""
+        try:
+            # Check if there's already a running event loop
+            try:
+                loop = asyncio.get_running_loop()
+                # If we're in an async context, we need to create a new task and wait for it
+                import concurrent.futures
+                import threading
+                
+                def run_in_thread():
+                    """Run the async method in a separate thread with its own event loop"""
+                    new_loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(new_loop)
+                    try:
+                        return new_loop.run_until_complete(self.handle_streaming_response_async(generator))
+                    finally:
+                        new_loop.close()
+                
+                with concurrent.futures.ThreadPoolExecutor() as executor:
+                    future = executor.submit(run_in_thread)
+                    return future.result()
+                    
+            except RuntimeError:
+                # No running event loop, safe to use asyncio.run()
+                return asyncio.run(self.handle_streaming_response_async(generator))
+                        
+        except Exception as e:
+            # Log for debugging
+            import logging
+            logging.debug(f"Exception in sync streaming wrapper: {e}")
+            self.console.print(f"[red]Error during streaming: {str(e)}[/red]")
+            return "", None, []
+        
+    def _render_final_content(self, text: str, reasoning_text: str = ""):
+        """
+        Render final formatted content after streaming completes
+        """
+        # Render reasoning if present
+        if reasoning_text.strip():
+            reasoning_panel = Panel(
+                Text(reasoning_text, style="dim italic"),
+                border_style="dim",
+                padding=(0, 1),
+                title="[dim]Reasoning[/dim]",
+                title_align="left"
+            )
+            self.console.print(reasoning_panel)
+            self.console.print()
+
+        # Parse all code blocks
+        blocks = self._parse_code_blocks(text)
+        
+        if blocks:
+            # Render with code blocks
+            last_pos = 0
+            
+            for block in blocks:
+                # Add text before this block
+                if block['start'] > last_pos:
+                    pre_text = text[last_pos:block['start']].strip()
+                    if pre_text:
+                        # Render as markdown for better formatting
+                        self.console.print(Markdown(pre_text, code_theme="monokai"))
+                        self.console.print()
+                
+                # Add the rendered block (with panel)
+                self.console.print(self._render_block(block))
+                self.console.print()
+                
+                last_pos = block['end']
+            
+            # Add any remaining text
+            if last_pos < len(text):
+                post_text = text[last_pos:].strip()
+                if post_text:
+                    self.console.print(Markdown(post_text, code_theme="monokai"))
+        else:
+            # No code blocks, render as markdown
+            if text.strip():
+                self.console.print(Markdown(text, code_theme="monokai"))
+    
+    def display_response(self, response_data: dict):
+        """Display a non-streaming response"""
+        try:
+            if "choices" in response_data and len(response_data["choices"]) > 0:
+                message = response_data["choices"][0].get("message", {})
+                content = message.get("content", "")
+                
+                if content:
+                    self._render_final_content(content)
+                else:
+                    self.console.print("[yellow]No response content received.[/yellow]")
+            else:
+                self.console.print("[yellow]Invalid response format.[/yellow]")
+                
+        except Exception as e:
+            self.console.print(f"[red]Error displaying response: {str(e)}[/red]")
+    
+    def display_error(self, error_message: str):
+        """Display an error message"""
+        self.console.print(Panel(
+            f"[red]{error_message}[/red]",
+            border_style="red",
+            title="[bold red]❌ Error[/bold red]",
+            padding=(1, 2)
+        ))
+    
+    def display_code_block(self, code: str, language: str = "python"):
+        """Display a formatted code block"""
+        syntax = Syntax(
+            code,
+            language,
+            theme="monokai",
+            line_numbers=False,
+            word_wrap=False
+        )
+        self.console.print(Panel(
+            syntax,
+            border_style="#48AEF3",
+            padding=(1, 2),
+            title=f"[bold color(#48AEF3)]{language.upper()}[/bold color(#48AEF3)]"
+        ))
+    
+    def display_thinking(self, message: str = "Thinking..."):
+        """Display a thinking/processing message"""
+        return self.console.status(
+            f"[bold color(#2FCACE)]{message}",
+            spinner="dots"
+        )
+
+
+class MarkdownRenderer:
+    """Enhanced markdown renderer with custom styling"""
+    
+    def __init__(self, console: Console):
+        self.console = console
+    
+    def render(self, markdown_text: str, title: Optional[str] = None):
+        """Render markdown with enhanced styling"""
+        markdown = Markdown(
+            markdown_text,
+            code_theme="monokai",
+            inline_code_theme="monokai"
+        )
+        
+        if title:
+            self.console.print(Panel(
+                markdown,
+                title=f"[bold color(#2FCACE)]{title}[/bold color(#2FCACE)]",
+                border_style="#2FCACE",
+                padding=(1, 2)
+            ))
+        else:
+            self.console.print(markdown)

codemate/utils/__init__.py

@@ -0,0 +1,21 @@
+"""
+Utility functions for CodeMate CLI
+"""
+
+from codemate.utils.errors import (
+    CodeMateError,
+    APIError,
+    ConfigError,
+    AuthenticationError,
+    NetworkError,
+    handle_errors,
+)
+
+__all__ = [
+    "CodeMateError",
+    "APIError",
+    "ConfigError",
+    "AuthenticationError",
+    "NetworkError",
+    "handle_errors",
+]

codemate/utils/auth.py

@@ -0,0 +1,164 @@
+"""
+Authentication utilities
+"""
+
+import os
+import hashlib
+from typing import Optional
+from pathlib import Path
+
+
+class APIKeyManager:
+    """Manage API key storage and validation"""
+    
+    @staticmethod
+    def validate_key_format(api_key: str) -> bool:
+        """
+        Validate API key format
+        
+        Args:
+            api_key: API key to validate
+            
+        Returns:
+            True if format is valid
+        """
+        if not api_key or not isinstance(api_key, str):
+            return False
+        
+        # Basic validation - adjust based on your API key format
+        if len(api_key) < 10:
+            return False
+        
+        # Check for common prefixes
+        valid_prefixes = ['sk-', 'key-', 'api-']
+        has_valid_prefix = any(api_key.startswith(prefix) for prefix in valid_prefixes)
+        
+        return has_valid_prefix or len(api_key) >= 20
+    
+    @staticmethod
+    def mask_key(api_key: str, show_first: int = 8, show_last: int = 4) -> str:
+        """
+        Mask API key for display
+        
+        Args:
+            api_key: API key to mask
+            show_first: Number of characters to show at start
+            show_last: Number of characters to show at end
+            
+        Returns:
+            Masked key string
+        """
+        if not api_key:
+            return "Not set"
+        
+        if len(api_key) <= show_first + show_last:
+            return "*" * len(api_key)
+        
+        return f"{api_key[:show_first]}...{api_key[-show_last:]}"
+    
+    @staticmethod
+    def hash_key(api_key: str) -> str:
+        """
+        Create a hash of the API key for logging/tracking
+        
+        Args:
+            api_key: API key to hash
+            
+        Returns:
+            SHA256 hash of the key
+        """
+        return hashlib.sha256(api_key.encode()).hexdigest()[:16]
+
+
+class EnvironmentAuth:
+    """Handle authentication from environment variables"""
+    
+    @staticmethod
+    def get_api_key() -> Optional[str]:
+        """Get API key from environment"""
+        return os.environ.get('CODEMATE_API_KEY')
+    
+    @staticmethod
+    def get_endpoint() -> Optional[str]:
+        """Get endpoint from environment"""
+        return os.environ.get('CODEMATE_ENDPOINT')
+    
+    @staticmethod
+    def set_env_var(key: str, value: str):
+        """Set environment variable"""
+        os.environ[key] = value
+    
+    @staticmethod
+    def has_env_auth() -> bool:
+        """Check if environment authentication is configured"""
+        return bool(os.environ.get('CODEMATE_API_KEY'))
+
+
+class SecureStorage:
+    """Secure storage for sensitive data"""
+    
+    def __init__(self, storage_dir: Path):
+        self.storage_dir = storage_dir
+        self.storage_dir.mkdir(parents=True, exist_ok=True)
+    
+    def store_secure(self, key: str, value: str):
+        """
+        Store data securely (basic implementation)
+        
+        For production, consider using system keyring:
+        - macOS: Keychain
+        - Windows: Credential Manager
+        - Linux: Secret Service API
+        """
+        file_path = self.storage_dir / f".{key}"
+        
+        # Set restrictive permissions (Unix-like systems)
+        try:
+            file_path.touch(mode=0o600, exist_ok=True)
+        except Exception:
+            pass  # Windows doesn't support mode parameter
+        
+        with open(file_path, 'w') as f:
+            f.write(value)
+    
+    def retrieve_secure(self, key: str) -> Optional[str]:
+        """Retrieve securely stored data"""
+        file_path = self.storage_dir / f".{key}"
+        
+        if not file_path.exists():
+            return None
+        
+        try:
+            with open(file_path, 'r') as f:
+                return f.read().strip()
+        except Exception:
+            return None
+    
+    def delete_secure(self, key: str):
+        """Delete securely stored data"""
+        file_path = self.storage_dir / f".{key}"
+        file_path.unlink(missing_ok=True)
+
+
+def verify_api_connection(api_key: str, endpoint: str) -> bool:
+    """
+    Verify API connection with given credentials
+    
+    Args:
+        api_key: API key to test
+        endpoint: API endpoint to test
+        
+    Returns:
+        True if connection successful
+    """
+    import httpx
+    
+    try:
+        with httpx.Client(timeout=5.0) as client:
+            response = client.get(
+                f"{endpoint}/health",
+                headers={"Authorization": f"Bearer {api_key}"}
+            )
+            return response.status_code == 200
+    except Exception:
+        return False