cortex-llm 1.0.0__py3-none-any.whl

cortex/template_registry/__init__.py

@@ -0,0 +1,15 @@
+"""Template Registry System for Cortex.
+
+This module provides intelligent template management for different model types,
+handling various chat formats and output processing requirements.
+"""
+
+from cortex.template_registry.registry import TemplateRegistry
+from cortex.template_registry.auto_detector import TemplateDetector
+from cortex.template_registry.template_profiles.base import BaseTemplateProfile
+
+__all__ = [
+    'TemplateRegistry',
+    'TemplateDetector',
+    'BaseTemplateProfile'
+]

cortex/template_registry/auto_detector.py

@@ -0,0 +1,144 @@
+"""Automatic template detection for models."""
+
+import logging
+from typing import List, Dict, Any, Optional, Tuple
+from pathlib import Path
+
+from cortex.template_registry.template_profiles.base import BaseTemplateProfile, TemplateType
+from cortex.template_registry.template_profiles.standard import (
+    ChatMLProfile, LlamaProfile, AlpacaProfile, SimpleProfile, GemmaProfile
+)
+from cortex.template_registry.template_profiles.complex import ReasoningProfile
+
+logger = logging.getLogger(__name__)
+
+
+class TemplateDetector:
+    """Automatically detect the best template for a model."""
+    
+    def __init__(self):
+        """Initialize the detector with all available profiles."""
+        self.profiles = [
+            ReasoningProfile(),
+            GemmaProfile(),  # Add Gemma profile with high priority
+            ChatMLProfile(),
+            LlamaProfile(),
+            AlpacaProfile(),
+            SimpleProfile()  # Fallback
+        ]
+    
+    def detect_template(
+        self, 
+        model_name: str, 
+        model_path: Optional[Path] = None,
+        tokenizer: Any = None
+    ) -> Tuple[BaseTemplateProfile, float]:
+        """Detect the best template for a model.
+        
+        Args:
+            model_name: Name of the model
+            model_path: Optional path to model files
+            tokenizer: Optional tokenizer object
+            
+        Returns:
+            Tuple of (best_profile, confidence_score)
+        """
+        candidates = []
+        
+        # Check each profile
+        for profile in self.profiles:
+            can_handle, confidence = profile.can_handle(model_name, tokenizer)
+            if can_handle:
+                candidates.append((profile, confidence))
+                logger.debug(f"Profile {profile.config.name} can handle {model_name} with confidence {confidence}")
+        
+        # Sort by confidence
+        candidates.sort(key=lambda x: x[1], reverse=True)
+        
+        if candidates:
+            best_profile, confidence = candidates[0]
+            logger.info(f"Selected {best_profile.config.name} template for {model_name} (confidence: {confidence:.2f})")
+            return best_profile, confidence
+        
+        # Fallback to simple profile
+        logger.warning(f"No specific template found for {model_name}, using simple fallback")
+        return SimpleProfile(), 0.1
+    
+    def detect_from_tokenizer_config(self, tokenizer: Any) -> Optional[TemplateType]:
+        """Try to detect template type from tokenizer configuration."""
+        if not tokenizer:
+            return None
+        
+        try:
+            # Check for chat_template attribute
+            if hasattr(tokenizer, 'chat_template'):
+                template_str = str(tokenizer.chat_template)
+                
+                # Check for known patterns
+                if '<|im_start|>' in template_str:
+                    return TemplateType.CHATML
+                elif '[INST]' in template_str:
+                    return TemplateType.LLAMA
+                elif '<|channel|>' in template_str:
+                    return TemplateType.REASONING
+                elif '### Instruction' in template_str:
+                    return TemplateType.ALPACA
+        except (AttributeError, TypeError) as e:
+            logger.debug(f"Error accessing tokenizer chat_template: {e}")
+        
+        try:
+            # Check tokenizer config
+            if hasattr(tokenizer, 'tokenizer_config'):
+                config = tokenizer.tokenizer_config
+                if isinstance(config, dict):
+                    # Check for model type hints
+                    model_type = config.get('model_type', '').lower()
+                    if 'llama' in model_type:
+                        return TemplateType.LLAMA
+                    elif 'gpt' in model_type:
+                        return TemplateType.OPENAI
+        except (AttributeError, TypeError, KeyError) as e:
+            logger.debug(f"Error accessing tokenizer config: {e}")
+        
+        return None
+    
+    def test_template(
+        self, 
+        profile: BaseTemplateProfile,
+        test_prompt: str = "Hello, how are you?"
+    ) -> Dict[str, Any]:
+        """Test a template profile with a sample prompt.
+        
+        Args:
+            profile: Template profile to test
+            test_prompt: Test prompt to use
+            
+        Returns:
+            Dictionary with test results
+        """
+        messages = [
+            {"role": "user", "content": test_prompt}
+        ]
+        
+        formatted = profile.format_messages(messages)
+        
+        # Simulate a response
+        sample_response = "I'm doing well, thank you for asking! How can I help you today?"
+        
+        # Add reasoning tokens for reasoning profile
+        if profile.config.template_type == TemplateType.REASONING:
+            sample_response = (
+                "<|channel|>analysis<|message|>The user is greeting me and asking about my state. "
+                "I should respond politely.<|end|>"
+                f"<|channel|>final<|message|>{sample_response}<|end|>"
+            )
+        
+        processed = profile.process_response(sample_response)
+        
+        return {
+            'profile_name': profile.config.name,
+            'formatted_prompt': formatted,
+            'raw_response': sample_response,
+            'processed_response': processed,
+            'template_type': profile.config.template_type.value
+        }

cortex/template_registry/config_manager.py

@@ -0,0 +1,234 @@
+"""Configuration manager for template registry."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Dict, Any, Optional, List
+from datetime import datetime
+from dataclasses import dataclass, asdict
+
+logger = logging.getLogger(__name__)
+
+
+# Configuration schema for validation
+CONFIG_SCHEMA = {
+    "version": str,
+    "models": dict,
+    "global_settings": dict
+}
+
+MODEL_CONFIG_SCHEMA = {
+    "detected_type": str,
+    "user_preference": str,
+    "custom_filters": list,
+    "show_reasoning": bool,
+    "last_updated": str,
+    "confidence": float
+}
+
+
+@dataclass
+class ModelTemplateConfig:
+    """Configuration for a specific model's template."""
+    detected_type: str
+    user_preference: str
+    custom_filters: List[str]
+    show_reasoning: bool = False
+    last_updated: str = ""
+    confidence: float = 0.0
+    
+    def __post_init__(self):
+        if not self.last_updated:
+            self.last_updated = datetime.now().isoformat()
+
+
+@dataclass
+class GlobalSettings:
+    """Global template settings."""
+    auto_detect: bool = True
+    prompt_on_unknown: bool = True
+    verbose_mode: bool = False
+    cache_templates: bool = True
+    default_fallback: str = "simple"
+
+
+class TemplateConfigManager:
+    """Manage template configurations and user preferences."""
+    
+    def __init__(self, config_path: Optional[Path] = None):
+        """Initialize the configuration manager."""
+        self.config_path = config_path or (Path.home() / ".cortex" / "template_config.json")
+        self.config_path.parent.mkdir(parents=True, exist_ok=True)
+        self.config = self._load_config()
+    
+    def _validate_config_structure(self, config: Dict[str, Any]) -> bool:
+        """Validate configuration structure matches expected schema.
+        
+        Args:
+            config: Configuration dictionary to validate
+            
+        Returns:
+            True if valid, False otherwise
+        """
+        # Check top-level keys
+        for key, expected_type in CONFIG_SCHEMA.items():
+            if key not in config:
+                logger.warning(f"Missing required config key: {key}")
+                return False
+            if not isinstance(config[key], expected_type):
+                logger.warning(f"Invalid type for {key}: expected {expected_type}, got {type(config[key])}")
+                return False
+        
+        # Validate model configs
+        if "models" in config:
+            for model_name, model_config in config["models"].items():
+                if not isinstance(model_config, dict):
+                    logger.warning(f"Invalid model config for {model_name}")
+                    return False
+                # Validate model config structure
+                for key, expected_type in MODEL_CONFIG_SCHEMA.items():
+                    if key in model_config and not isinstance(model_config[key], expected_type):
+                        logger.warning(f"Invalid type for {model_name}.{key}")
+                        return False
+        
+        return True
+    
+    def _load_config(self) -> Dict[str, Any]:
+        """Load configuration from file."""
+        if self.config_path.exists():
+            try:
+                with open(self.config_path, 'r') as f:
+                    config = json.load(f)
+                    
+                # Validate loaded configuration
+                if self._validate_config_structure(config):
+                    return config
+                else:
+                    logger.warning("Invalid config structure, using defaults")
+                    return self._get_default_config()
+                    
+            except json.JSONDecodeError as e:
+                logger.error(f"Invalid JSON in template config: {e}")
+            except IOError as e:
+                logger.error(f"Error reading template config file: {e}")
+            except PermissionError as e:
+                logger.error(f"Permission denied reading template config: {e}")
+        
+        # Return default config
+        return self._get_default_config()
+    
+    def _get_default_config(self) -> Dict[str, Any]:
+        """Get default configuration."""
+        return {
+            "version": "1.0",
+            "models": {},
+            "global_settings": asdict(GlobalSettings())
+        }
+    
+    def _save_config(self) -> None:
+        """Save configuration to file."""
+        try:
+            # Ensure parent directory exists
+            self.config_path.parent.mkdir(parents=True, exist_ok=True)
+            
+            # Write to temporary file first for atomicity
+            temp_path = self.config_path.with_suffix('.tmp')
+            with open(temp_path, 'w') as f:
+                json.dump(self.config, f, indent=2)
+            
+            # Atomic rename
+            temp_path.replace(self.config_path)
+            logger.debug(f"Saved template config to {self.config_path}")
+            
+        except IOError as e:
+            logger.error(f"IO error saving template config: {e}")
+        except PermissionError as e:
+            logger.error(f"Permission denied saving template config: {e}")
+        except OSError as e:
+            logger.error(f"OS error saving template config: {e}")
+    
+    def get_model_config(self, model_name: str) -> Optional[ModelTemplateConfig]:
+        """Get configuration for a specific model."""
+        if model_name in self.config.get("models", {}):
+            data = self.config["models"][model_name]
+            return ModelTemplateConfig(
+                detected_type=data.get("detected_type", "unknown"),
+                user_preference=data.get("user_preference", "auto"),
+                custom_filters=data.get("custom_filters", []),
+                show_reasoning=data.get("show_reasoning", False),
+                last_updated=data.get("last_updated", ""),
+                confidence=data.get("confidence", 0.0)
+            )
+        return None
+    
+    def save_model_config(self, model_name: str, config: ModelTemplateConfig) -> None:
+        """Save configuration for a specific model."""
+        if "models" not in self.config:
+            self.config["models"] = {}
+        
+        self.config["models"][model_name] = {
+            "detected_type": config.detected_type,
+            "user_preference": config.user_preference,
+            "custom_filters": config.custom_filters,
+            "show_reasoning": config.show_reasoning,
+            "last_updated": datetime.now().isoformat(),
+            "confidence": config.confidence
+        }
+        
+        self._save_config()
+        logger.info(f"Saved template configuration for {model_name}")
+    
+    def get_global_settings(self) -> GlobalSettings:
+        """Get global settings."""
+        data = self.config.get("global_settings", {})
+        return GlobalSettings(
+            auto_detect=data.get("auto_detect", True),
+            prompt_on_unknown=data.get("prompt_on_unknown", True),
+            verbose_mode=data.get("verbose_mode", False),
+            cache_templates=data.get("cache_templates", True),
+            default_fallback=data.get("default_fallback", "simple")
+        )
+    
+    def update_global_settings(self, **kwargs) -> None:
+        """Update global settings."""
+        settings = self.get_global_settings()
+        
+        for key, value in kwargs.items():
+            if hasattr(settings, key):
+                setattr(settings, key, value)
+        
+        self.config["global_settings"] = asdict(settings)
+        self._save_config()
+    
+    def list_configured_models(self) -> List[str]:
+        """List all models with saved configurations."""
+        return list(self.config.get("models", {}).keys())
+    
+    def remove_model_config(self, model_name: str) -> bool:
+        """Remove configuration for a model."""
+        if model_name in self.config.get("models", {}):
+            del self.config["models"][model_name]
+            self._save_config()
+            logger.info(f"Removed template configuration for {model_name}")
+            return True
+        return False
+    
+    def export_config(self) -> Dict[str, Any]:
+        """Export the entire configuration."""
+        return self.config.copy()
+    
+    def import_config(self, config: Dict[str, Any]) -> None:
+        """Import a configuration.
+        
+        Args:
+            config: Configuration dictionary to import
+            
+        Raises:
+            ValueError: If configuration is invalid
+        """
+        if not self._validate_config_structure(config):
+            raise ValueError("Invalid configuration structure")
+        
+        self.config = config
+        self._save_config()
+        logger.info("Imported template configuration")

cortex/template_registry/interactive.py

@@ -0,0 +1,260 @@
+"""Interactive template setup for user-friendly configuration."""
+
+from typing import Optional, Dict, Any, List
+from rich.console import Console
+from rich.table import Table
+from rich.prompt import Prompt, Confirm
+
+from cortex.template_registry.template_profiles.base import BaseTemplateProfile, TemplateType
+from cortex.template_registry.auto_detector import TemplateDetector
+from cortex.template_registry.config_manager import ModelTemplateConfig
+
+
+class InteractiveTemplateSetup:
+    """Interactive template configuration wizard."""
+    
+    def __init__(self, console: Optional[Console] = None):
+        """Initialize the interactive setup."""
+        self.console = console or Console()
+        self.detector = TemplateDetector()
+    
+    def setup_model_template(
+        self, 
+        model_name: str,
+        tokenizer: Any = None,
+        current_config: Optional[ModelTemplateConfig] = None
+    ) -> ModelTemplateConfig:
+        """Interactive setup for model template.
+        
+        Args:
+            model_name: Name of the model
+            tokenizer: Optional tokenizer object
+            current_config: Existing configuration if any
+            
+        Returns:
+            Updated model template configuration
+        """
+        self.console.print(f"\n✓ Model loaded: [bold cyan]{model_name}[/bold cyan]")
+        
+        # Detect template
+        profile, confidence = self.detector.detect_template(model_name, tokenizer=tokenizer)
+        
+        if confidence < 0.5:
+            self.console.print("\n⚠️  [yellow]Template configuration needed for optimal performance[/yellow]")
+        
+        self.console.print(f"\nDetecting template format...")
+        self.console.print(f"✓ Found: [green]{profile.config.description}[/green] (confidence: {confidence:.0%})")
+        
+        # Check if this is a reasoning model
+        is_reasoning = profile.config.template_type == TemplateType.REASONING
+        
+        if is_reasoning:
+            self.console.print("\n[yellow]Note: This model includes internal reasoning/analysis in its output.[/yellow]")
+        
+        # Show options
+        self.console.print("\nHow would you like to handle this model's output?\n")
+        
+        options = []
+        if is_reasoning:
+            options = [
+                ("simple", "Simple mode - Hide internal reasoning (recommended)", True),
+                ("full", "Full mode - Show all model outputs", False),
+                ("custom", "Custom - Configure manually", False),
+                ("test", "Test - See examples of each mode", False)
+            ]
+        else:
+            options = [
+                ("auto", "Automatic - Use detected template", True),
+                ("custom", "Custom - Configure manually", False),
+                ("test", "Test - See examples with different templates", False)
+            ]
+        
+        # Display options
+        for i, (key, desc, recommended) in enumerate(options, 1):
+            marker = " [green](recommended)[/green]" if recommended else ""
+            self.console.print(f"[{i}] {desc}{marker}")
+        
+        # Get user choice
+        choice = Prompt.ask(
+            f"\nSelect option (1-{len(options)})",
+            default="1",
+            choices=[str(i) for i in range(1, len(options) + 1)]
+        )
+        
+        selected_key = options[int(choice) - 1][0]
+        
+        # Handle selection
+        if selected_key == "test":
+            self._show_template_tests(model_name, profile)
+            # Recurse to get actual selection
+            return self.setup_model_template(model_name, tokenizer, current_config)
+        
+        elif selected_key == "custom":
+            return self._custom_setup(model_name, profile)
+        
+        else:
+            # Create configuration
+            config = ModelTemplateConfig(
+                detected_type=profile.config.template_type.value,
+                user_preference=selected_key,
+                custom_filters=profile.config.custom_filters,
+                show_reasoning=(selected_key == "full") if is_reasoning else False,
+                confidence=confidence
+            )
+            
+            self.console.print(f"\n✓ Template configured: [green]{selected_key} mode[/green]")
+            self.console.print("✓ Configuration saved for future use")
+            self.console.print("\n[dim]Tip: Use /template to adjust settings anytime[/dim]")
+            
+            return config
+    
+    def _show_template_tests(self, model_name: str, detected_profile: BaseTemplateProfile) -> None:
+        """Show examples of different template modes."""
+        self.console.print("\n[bold]Testing different template modes:[/bold]\n")
+        
+        test_prompt = "What is 2+2?"
+        
+        # Test different profiles
+        profiles_to_test = []
+        
+        if detected_profile.config.template_type == TemplateType.REASONING:
+            # Test with and without reasoning
+            simple_profile = detected_profile.__class__()
+            simple_profile.config.show_reasoning = False
+            
+            full_profile = detected_profile.__class__()
+            full_profile.config.show_reasoning = True
+            
+            profiles_to_test = [
+                ("Simple Mode", simple_profile),
+                ("Full Mode", full_profile)
+            ]
+        else:
+            # Test different template types
+            from cortex.template_registry.template_profiles.standard import (
+                ChatMLProfile, LlamaProfile, SimpleProfile
+            )
+            
+            profiles_to_test = [
+                ("Detected", detected_profile),
+                ("ChatML", ChatMLProfile()),
+                ("Llama", LlamaProfile()),
+                ("Simple", SimpleProfile())
+            ]
+        
+        for name, profile in profiles_to_test:
+            result = self.detector.test_template(profile, test_prompt)
+            
+            self.console.print(f"[bold cyan]{name}:[/bold cyan]")
+            self.console.print("─" * 40)
+            
+            # Show formatted prompt
+            self.console.print("[dim]Formatted prompt:[/dim]")
+            self.console.print(f"  {result['formatted_prompt'][:100]}..." if len(result['formatted_prompt']) > 100 else f"  {result['formatted_prompt']}")
+            
+            # Show processed response
+            self.console.print("[dim]Output:[/dim]")
+            self.console.print(f"  {result['processed_response']}")
+            self.console.print()
+    
+    def _custom_setup(self, model_name: str, detected_profile: BaseTemplateProfile) -> ModelTemplateConfig:
+        """Custom template configuration."""
+        self.console.print("\n[bold]Custom Template Configuration[/bold]\n")
+        
+        # Select template type
+        template_types = [
+            ("chatml", "ChatML format"),
+            ("llama", "Llama format"),
+            ("alpaca", "Alpaca format"),
+            ("reasoning", "Reasoning/CoT format"),
+            ("simple", "Simple format")
+        ]
+        
+        self.console.print("Available template types:")
+        for i, (key, desc) in enumerate(template_types, 1):
+            self.console.print(f"[{i}] {desc}")
+        
+        choice = Prompt.ask(
+            "Select template type",
+            default="1",
+            choices=[str(i) for i in range(1, len(template_types) + 1)]
+        )
+        
+        selected_type = template_types[int(choice) - 1][0]
+        
+        # Configure filters
+        custom_filters = []
+        if Confirm.ask("Configure custom output filters?", default=False):
+            filters_input = Prompt.ask("Enter tokens to filter (comma-separated)")
+            custom_filters = [f.strip() for f in filters_input.split(",") if f.strip()]
+        
+        # Show reasoning option
+        show_reasoning = False
+        if selected_type == "reasoning":
+            show_reasoning = Confirm.ask("Show internal reasoning/analysis?", default=False)
+        
+        config = ModelTemplateConfig(
+            detected_type=selected_type,
+            user_preference="custom",
+            custom_filters=custom_filters,
+            show_reasoning=show_reasoning,
+            confidence=1.0  # User manually configured
+        )
+        
+        self.console.print("\n✓ Custom template configured")
+        return config
+    
+    def show_current_config(self, model_name: str, config: ModelTemplateConfig) -> None:
+        """Display current configuration for a model."""
+        table = Table(title=f"Template Configuration for {model_name}")
+        table.add_column("Setting", style="cyan")
+        table.add_column("Value", style="green")
+        
+        table.add_row("Template Type", config.detected_type)
+        table.add_row("User Preference", config.user_preference)
+        table.add_row("Show Reasoning", str(config.show_reasoning))
+        table.add_row("Custom Filters", ", ".join(config.custom_filters) if config.custom_filters else "None")
+        table.add_row("Confidence", f"{config.confidence:.0%}")
+        table.add_row("Last Updated", config.last_updated)
+        
+        self.console.print(table)
+    
+    def quick_adjust_template(self, model_name: str, config: ModelTemplateConfig) -> ModelTemplateConfig:
+        """Quick adjustment interface for template settings."""
+        self.console.print(f"\n[bold]Adjust template for {model_name}[/bold]\n")
+        
+        self.show_current_config(model_name, config)
+        
+        self.console.print("\n[1] Toggle reasoning display")
+        self.console.print("[2] Change template type")
+        self.console.print("[3] Edit filters")
+        self.console.print("[4] Reset to defaults")
+        self.console.print("[0] Cancel")
+        
+        choice = Prompt.ask("Select option", choices=["0", "1", "2", "3", "4"])
+        
+        if choice == "1":
+            config.show_reasoning = not config.show_reasoning
+            self.console.print(f"✓ Reasoning display: [green]{'enabled' if config.show_reasoning else 'disabled'}[/green]")
+        
+        elif choice == "2":
+            return self._custom_setup(model_name, None)
+        
+        elif choice == "3":
+            filters_input = Prompt.ask("Enter tokens to filter (comma-separated)", default=",".join(config.custom_filters))
+            config.custom_filters = [f.strip() for f in filters_input.split(",") if f.strip()]
+            self.console.print("✓ Filters updated")
+        
+        elif choice == "4":
+            # Reset to detected defaults
+            profile, confidence = self.detector.detect_template(model_name)
+            config = ModelTemplateConfig(
+                detected_type=profile.config.template_type.value,
+                user_preference="auto",
+                custom_filters=profile.config.custom_filters,
+                show_reasoning=False,
+                confidence=confidence
+            )
+            self.console.print("✓ Reset to defaults")
+        
+        return config