attune-ai 2.0.0__py3-none-any.whl

attune/models/token_estimator.py

@@ -0,0 +1,420 @@
+"""Token Estimation Service
+
+Pre-flight token estimation for cost prediction using tiktoken.
+Provides accurate token counts before workflow execution.
+
+Copyright 2025 Smart AI Memory, LLC
+Licensed under Fair Source 0.9
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import Any
+
+from attune.config import _validate_file_path
+
+# Try to import tiktoken, fall back to heuristic if not available
+try:
+    import tiktoken
+
+    TIKTOKEN_AVAILABLE = True
+except ImportError:
+    TIKTOKEN_AVAILABLE = False
+
+
+# Heuristic fallback: ~4 tokens per word, ~0.25 tokens per character
+TOKENS_PER_CHAR_HEURISTIC = 0.25
+
+
+@functools.lru_cache(maxsize=4)
+def _get_encoding(model_id: str) -> Any:
+    """Get tiktoken encoding for a model, with caching."""
+    if not TIKTOKEN_AVAILABLE:
+        return None
+
+    # Map model IDs to encoding names
+    if "claude" in model_id.lower() or "anthropic" in model_id.lower():
+        # Claude uses cl100k_base-like encoding
+        return tiktoken.get_encoding("cl100k_base")
+    if "gpt-4" in model_id.lower() or "gpt-3.5" in model_id.lower() or "o1" in model_id.lower():
+        return tiktoken.encoding_for_model("gpt-4")
+    # Default to cl100k_base for unknown models
+    return tiktoken.get_encoding("cl100k_base")
+
+
+def estimate_tokens(text: str, model_id: str = "claude-sonnet-4-5-20250514") -> int:
+    """Estimate token count for text using accurate token counting.
+
+    Uses empathy_llm_toolkit's token counter which leverages tiktoken for fast,
+    accurate local counting (~98% accurate). Falls back to heuristic if unavailable.
+
+    Args:
+        text: The text to count tokens for
+        model_id: The model ID to use for encoding selection
+
+    Returns:
+        Accurate token count
+
+    Raises:
+        ValueError: If model_id is empty
+
+    """
+    # Pattern 1: String ID validation
+    if not model_id or not model_id.strip():
+        raise ValueError("model_id cannot be empty")
+
+    if not text:
+        return 0
+
+    # Use new accurate token counting from attune_llm
+    try:
+        from attune_llm.utils.tokens import count_tokens
+
+        return count_tokens(text, model=model_id, use_api=False)
+    except ImportError:
+        # Fallback to tiktoken if toolkit not available
+        if TIKTOKEN_AVAILABLE:
+            try:
+                encoding = _get_encoding(model_id)
+                if encoding:
+                    return len(encoding.encode(text))
+            except Exception:
+                pass  # Fall through to heuristic
+
+        # Last resort: heuristic fallback
+        return max(1, int(len(text) * TOKENS_PER_CHAR_HEURISTIC))
+
+
+def estimate_workflow_cost(
+    workflow_name: str,
+    input_text: str,
+    provider: str = "anthropic",
+    target_path: str | None = None,
+) -> dict[str, Any]:
+    """Estimate total workflow cost before execution.
+
+    Analyzes workflow stages and estimates token usage and cost for each,
+    providing a cost range for the full workflow.
+
+    Args:
+        workflow_name: Name of the workflow (e.g., "security-audit", "test-gen")
+        input_text: The input text/code to be processed
+        provider: LLM provider (anthropic, openai, ollama, hybrid)
+        target_path: Optional path for file-based workflows
+
+    Returns:
+        Dictionary with cost estimates:
+        {
+            "workflow": str,
+            "provider": str,
+            "input_tokens": int,
+            "stages": [...],
+            "total_min": float,
+            "total_max": float,
+            "display": str,
+            "risk": "low" | "medium" | "high"
+        }
+
+    Raises:
+        ValueError: If workflow_name or provider is empty
+
+    """
+    from .registry import get_model, get_supported_providers
+
+    # Pattern 1: String ID validation
+    if not workflow_name or not workflow_name.strip():
+        raise ValueError("workflow_name cannot be empty")
+    if not provider or not provider.strip():
+        raise ValueError("provider cannot be empty")
+
+    # Validate provider
+    if provider not in get_supported_providers():
+        provider = "anthropic"  # Default fallback
+
+    # Workflow stage configurations by workflow name
+    # Based on actual workflow implementations
+    WORKFLOW_STAGES = {
+        "security-audit": [
+            {"name": "identify_vulnerabilities", "tier": "capable"},
+            {"name": "analyze_risk", "tier": "capable"},
+            {"name": "generate_report", "tier": "cheap"},
+        ],
+        "code-review": [
+            {"name": "analyze_code", "tier": "capable"},
+            {"name": "generate_feedback", "tier": "capable"},
+            {"name": "summarize", "tier": "cheap"},
+        ],
+        "test-gen": [
+            {"name": "analyze_code", "tier": "capable"},
+            {"name": "generate_tests", "tier": "capable"},
+            {"name": "review_tests", "tier": "cheap"},
+        ],
+        "doc-gen": [
+            {"name": "analyze_structure", "tier": "cheap"},
+            {"name": "generate_documentation", "tier": "capable"},
+        ],
+        "bug-predict": [
+            {"name": "analyze_patterns", "tier": "capable"},
+            {"name": "predict_risks", "tier": "capable"},
+        ],
+        "refactor-plan": [
+            {"name": "identify_issues", "tier": "capable"},
+            {"name": "plan_refactoring", "tier": "capable"},
+            {"name": "review_plan", "tier": "cheap"},
+        ],
+        "perf-audit": [
+            {"name": "analyze_performance", "tier": "capable"},
+            {"name": "identify_bottlenecks", "tier": "capable"},
+            {"name": "generate_recommendations", "tier": "cheap"},
+        ],
+        "health-check": [
+            {"name": "scan_codebase", "tier": "cheap"},
+            {"name": "analyze_health", "tier": "capable"},
+            {"name": "generate_report", "tier": "cheap"},
+        ],
+        "pr-review": [
+            {"name": "analyze_diff", "tier": "capable"},
+            {"name": "check_quality", "tier": "capable"},
+            {"name": "generate_review", "tier": "capable"},
+        ],
+        "pro-review": [
+            {"name": "deep_analysis", "tier": "premium"},
+            {"name": "generate_insights", "tier": "capable"},
+        ],
+    }
+
+    # Get stage configuration for this workflow
+    stages_config = WORKFLOW_STAGES.get(
+        workflow_name,
+        [
+            {"name": "analyze", "tier": "capable"},
+            {"name": "generate", "tier": "capable"},
+            {"name": "review", "tier": "cheap"},
+        ],
+    )
+
+    # Estimate input tokens
+    input_tokens = estimate_tokens(input_text)
+
+    # If we have a target path, estimate additional content
+    if target_path:
+        try:
+            import os
+
+            # Validate path to prevent path traversal attacks
+            validated_target = _validate_file_path(target_path)
+
+            if os.path.isfile(validated_target):
+                with open(validated_target, encoding="utf-8", errors="ignore") as f:
+                    file_content = f.read()
+                input_tokens += estimate_tokens(file_content)
+            elif os.path.isdir(validated_target):
+                # Estimate based on directory size (rough heuristic)
+                total_chars = 0
+                for root, _, files in os.walk(validated_target):
+                    for file in files[:50]:  # Limit to first 50 files
+                        if file.endswith((".py", ".js", ".ts", ".tsx", ".jsx")):
+                            try:
+                                filepath = os.path.join(root, file)
+                                validated_filepath = _validate_file_path(filepath)
+                                with open(
+                                    validated_filepath, encoding="utf-8", errors="ignore"
+                                ) as f:
+                                    total_chars += len(f.read())
+                            except (ValueError, OSError):
+                                pass
+                input_tokens += int(total_chars * TOKENS_PER_CHAR_HEURISTIC)
+        except (ValueError, OSError):
+            pass  # Keep original estimate
+
+    # Output multipliers by stage type
+    output_multipliers = {
+        "identify": 0.3,
+        "analyze": 0.8,
+        "generate": 2.0,
+        "review": 0.5,
+        "summarize": 0.3,
+        "fix": 1.5,
+        "test": 1.5,
+        "document": 1.0,
+    }
+
+    estimates = []
+    total_min = 0.0
+    total_max = 0.0
+
+    for stage in stages_config:
+        stage_name = stage.get("name", "unknown")
+        tier = stage.get("tier", "capable")
+
+        # Get model for this tier
+        try:
+            model_info = get_model(provider, tier)
+        except Exception:
+            # Fallback to capable tier
+            model_info = get_model(provider, "capable")
+
+        if model_info is None:
+            # Skip stage if no model available
+            continue
+
+        # Estimate output tokens based on stage type
+        multiplier = 1.0
+        for stage_type, mult in output_multipliers.items():
+            if stage_type in stage_name.lower():
+                multiplier = mult
+                break
+
+        est_output = int(input_tokens * multiplier)
+
+        # Calculate cost
+        cost = (input_tokens / 1_000_000) * model_info.input_cost_per_million + (
+            est_output / 1_000_000
+        ) * model_info.output_cost_per_million
+
+        estimates.append(
+            {
+                "stage": stage_name,
+                "tier": tier,
+                "model": model_info.id,
+                "estimated_input_tokens": input_tokens,
+                "estimated_output_tokens": est_output,
+                "estimated_cost": round(cost, 6),
+            },
+        )
+
+        # Accumulate with variance (80% - 120%)
+        total_min += cost * 0.8
+        total_max += cost * 1.2
+
+    # Determine risk level
+    if total_max > 1.0:
+        risk = "high"
+    elif total_max > 0.10:
+        risk = "medium"
+    else:
+        risk = "low"
+
+    return {
+        "workflow": workflow_name,
+        "provider": provider,
+        "input_tokens": input_tokens,
+        "stages": estimates,
+        "total_min": round(total_min, 4),
+        "total_max": round(total_max, 4),
+        "display": f"${total_min:.3f} - ${total_max:.3f}",
+        "risk": risk,
+        "tiktoken_available": TIKTOKEN_AVAILABLE,
+    }
+
+
+def estimate_single_call_cost(
+    text: str,
+    task_type: str,
+    provider: str = "anthropic",
+) -> dict[str, Any]:
+    """Estimate cost for a single LLM call.
+
+    Args:
+        text: Input text
+        task_type: Type of task (e.g., "summarize", "generate_code")
+        provider: LLM provider
+
+    Returns:
+        Cost estimate dictionary
+
+    Raises:
+        ValueError: If task_type or provider is empty
+
+    """
+    from .registry import get_model
+    from .tasks import get_tier_for_task
+
+    # Pattern 1: String ID validation
+    if not task_type or not task_type.strip():
+        raise ValueError("task_type cannot be empty")
+    if not provider or not provider.strip():
+        raise ValueError("provider cannot be empty")
+
+    input_tokens = estimate_tokens(text)
+
+    # Get tier for task
+    tier = get_tier_for_task(task_type)
+    model_info = get_model(provider, tier.value)
+
+    if model_info is None:
+        # Return a fallback estimate if no model available
+        return {
+            "task_type": task_type,
+            "tier": tier.value,
+            "model": "unknown",
+            "provider": provider,
+            "input_tokens": input_tokens,
+            "estimated_output_tokens": input_tokens,
+            "estimated_cost": 0.0,
+            "display": "$0.0000",
+        }
+
+    # Estimate output based on task type
+    output_multipliers = {
+        "summarize": 0.3,
+        "classify": 0.1,
+        "generate_code": 1.5,
+        "fix_bug": 1.2,
+        "review": 0.5,
+        "document": 0.8,
+    }
+
+    multiplier = output_multipliers.get(task_type, 1.0)
+    est_output = int(input_tokens * multiplier)
+
+    cost = (input_tokens / 1_000_000) * model_info.input_cost_per_million + (
+        est_output / 1_000_000
+    ) * model_info.output_cost_per_million
+
+    return {
+        "task_type": task_type,
+        "tier": tier.value,
+        "model": model_info.id,
+        "provider": provider,
+        "input_tokens": input_tokens,
+        "estimated_output_tokens": est_output,
+        "estimated_cost": round(cost, 6),
+        "display": f"${cost:.4f}",
+    }
+
+
+# CLI support
+if __name__ == "__main__":
+    import argparse
+    import json
+    import sys
+
+    parser = argparse.ArgumentParser(description="Estimate workflow costs")
+    parser.add_argument("workflow", help="Workflow name (e.g., security-audit)")
+    parser.add_argument("--input", "-i", help="Input text or file path")
+    parser.add_argument("--provider", "-p", default="anthropic", help="LLM provider")
+    parser.add_argument("--target", "-t", help="Target path for file-based workflows")
+
+    args = parser.parse_args()
+
+    # Read input
+    input_text = ""
+    if args.input:
+        try:
+            validated_input = _validate_file_path(args.input)
+            with open(validated_input) as f:
+                input_text = f.read()
+        except (FileNotFoundError, ValueError):
+            input_text = args.input
+
+    result = estimate_workflow_cost(
+        workflow_name=args.workflow,
+        input_text=input_text,
+        provider=args.provider,
+        target_path=args.target,
+    )
+
+    print(json.dumps(result, indent=2))
+    sys.exit(0)

attune/models/validation.py

@@ -0,0 +1,280 @@
+"""Configuration Validation for Multi-Model Workflows
+
+Provides schema validation for workflow configurations:
+- Required field validation
+- Type checking
+- Value range validation
+- Provider/tier existence checks
+
+Copyright 2025 Smart-AI-Memory
+Licensed under Fair Source License 0.9
+"""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from attune.config import _validate_file_path
+
+from .registry import MODEL_REGISTRY, ModelTier
+
+
+@dataclass
+class ValidationError:
+    """A single validation error."""
+
+    path: str
+    message: str
+    severity: str = "error"  # "error" | "warning"
+
+    def __str__(self) -> str:
+        return f"[{self.severity.upper()}] {self.path}: {self.message}"
+
+
+@dataclass
+class ValidationResult:
+    """Result of configuration validation."""
+
+    valid: bool
+    errors: list[ValidationError] = field(default_factory=list)
+    warnings: list[ValidationError] = field(default_factory=list)
+
+    def add_error(self, path: str, message: str) -> None:
+        """Add an error."""
+        self.errors.append(ValidationError(path, message, "error"))
+        self.valid = False
+
+    def add_warning(self, path: str, message: str) -> None:
+        """Add a warning."""
+        self.warnings.append(ValidationError(path, message, "warning"))
+
+    def __str__(self) -> str:
+        lines = []
+        if self.valid:
+            lines.append("Configuration is valid")
+        else:
+            lines.append("Configuration has errors")
+
+        for error in self.errors:
+            lines.append(f"  {error}")
+        for warning in self.warnings:
+            lines.append(f"  {warning}")
+
+        return "\n".join(lines)
+
+
+class ConfigValidator:
+    """Validator for multi-model workflow configurations.
+
+    Validates:
+    - Provider names exist in registry
+    - Tier names are valid
+    - Required fields are present
+    - Numeric values are in valid ranges
+    """
+
+    # Valid provider names from registry
+    VALID_PROVIDERS = set(MODEL_REGISTRY.keys())
+
+    # Valid tier names
+    VALID_TIERS = {tier.value for tier in ModelTier}
+
+    # Schema for workflow config
+    WORKFLOW_SCHEMA = {
+        "name": {"type": str, "required": True},
+        "description": {"type": str, "required": False},
+        "default_provider": {"type": str, "required": False},
+        "stages": {"type": list, "required": False},
+    }
+
+    # Schema for stage config
+    STAGE_SCHEMA = {
+        "name": {"type": str, "required": True},
+        "tier": {"type": str, "required": True},
+        "provider": {"type": str, "required": False},
+        "timeout_ms": {"type": int, "required": False, "min": 0, "max": 600000},
+        "max_retries": {"type": int, "required": False, "min": 0, "max": 10},
+    }
+
+    def validate_workflow_config(self, config: dict[str, Any]) -> ValidationResult:
+        """Validate a workflow configuration dictionary.
+
+        Args:
+            config: Workflow configuration dict
+
+        Returns:
+            ValidationResult with any errors or warnings
+
+        """
+        result = ValidationResult(valid=True)
+
+        # Check required fields
+        for field_name, spec in self.WORKFLOW_SCHEMA.items():
+            if spec.get("required") and field_name not in config:
+                result.add_error(field_name, f"Required field '{field_name}' is missing")
+
+        # Validate types
+        for field_name, value in config.items():
+            if field_name in self.WORKFLOW_SCHEMA:
+                spec = self.WORKFLOW_SCHEMA[field_name]
+                expected_type = spec.get("type")
+                if expected_type is not None:
+                    # Cast to type for isinstance check
+                    type_cls = (
+                        expected_type if isinstance(expected_type, type) else type(expected_type)
+                    )
+                    if not isinstance(value, type_cls):
+                        type_name = getattr(type_cls, "__name__", str(type_cls))
+                        result.add_error(
+                            field_name,
+                            f"Expected {type_name}, got {type(value).__name__}",
+                        )
+
+        # Validate default_provider
+        if "default_provider" in config:
+            provider = config["default_provider"]
+            if provider not in self.VALID_PROVIDERS:
+                result.add_error(
+                    "default_provider",
+                    f"Unknown provider '{provider}'. Valid: {sorted(self.VALID_PROVIDERS)}",
+                )
+
+        # Validate stages
+        if "stages" in config and isinstance(config["stages"], list):
+            for i, stage in enumerate(config["stages"]):
+                stage_path = f"stages[{i}]"
+                self._validate_stage(stage, stage_path, result)
+
+        return result
+
+    def _validate_stage(self, stage: dict[str, Any], path: str, result: ValidationResult) -> None:
+        """Validate a single stage configuration."""
+        if not isinstance(stage, dict):
+            result.add_error(path, "Stage must be a dictionary")
+            return
+
+        # Check required fields
+        for field_name, spec in self.STAGE_SCHEMA.items():
+            if spec.get("required") and field_name not in stage:
+                result.add_error(f"{path}.{field_name}", "Required field is missing")
+
+        # Validate tier
+        if "tier" in stage:
+            tier = stage["tier"]
+            if tier not in self.VALID_TIERS:
+                result.add_error(
+                    f"{path}.tier",
+                    f"Unknown tier '{tier}'. Valid: {sorted(self.VALID_TIERS)}",
+                )
+
+        # Validate provider
+        if "provider" in stage:
+            provider = stage["provider"]
+            if provider not in self.VALID_PROVIDERS:
+                result.add_error(
+                    f"{path}.provider",
+                    f"Unknown provider '{provider}'. Valid: {sorted(self.VALID_PROVIDERS)}",
+                )
+
+        # Validate numeric ranges
+        for field_name in ["timeout_ms", "max_retries"]:
+            if field_name in stage:
+                value = stage[field_name]
+                spec = self.STAGE_SCHEMA[field_name]
+
+                if not isinstance(value, int):
+                    result.add_error(
+                        f"{path}.{field_name}",
+                        f"Expected integer, got {type(value).__name__}",
+                    )
+                else:
+                    min_val = spec.get("min")
+                    max_val = spec.get("max")
+                    if isinstance(min_val, int | float) and value < min_val:
+                        result.add_error(
+                            f"{path}.{field_name}",
+                            f"Value {value} below minimum {min_val}",
+                        )
+                    if isinstance(max_val, int | float) and value > max_val:
+                        result.add_error(
+                            f"{path}.{field_name}",
+                            f"Value {value} above maximum {max_val}",
+                        )
+
+    def validate_provider_tier(self, provider: str, tier: str) -> ValidationResult:
+        """Validate that a provider/tier combination exists.
+
+        Args:
+            provider: Provider name
+            tier: Tier name
+
+        Returns:
+            ValidationResult
+
+        """
+        result = ValidationResult(valid=True)
+
+        if provider not in self.VALID_PROVIDERS:
+            result.add_error("provider", f"Unknown provider '{provider}'")
+            return result
+
+        if tier not in self.VALID_TIERS:
+            result.add_error("tier", f"Unknown tier '{tier}'")
+            return result
+
+        # Check if combination exists in registry
+        if tier not in MODEL_REGISTRY.get(provider, {}):
+            result.add_warning(
+                "provider_tier",
+                f"Provider '{provider}' may not have tier '{tier}' configured",
+            )
+
+        return result
+
+
+def validate_config(config: dict[str, Any]) -> ValidationResult:
+    """Convenience function to validate a workflow config.
+
+    Args:
+        config: Configuration dictionary
+
+    Returns:
+        ValidationResult
+
+    """
+    validator = ConfigValidator()
+    return validator.validate_workflow_config(config)
+
+
+def validate_yaml_file(file_path: str) -> ValidationResult:
+    """Validate a YAML configuration file.
+
+    Args:
+        file_path: Path to YAML file
+
+    Returns:
+        ValidationResult
+
+    """
+    import yaml
+
+    result = ValidationResult(valid=True)
+
+    try:
+        validated_path = _validate_file_path(str(file_path))
+        with open(validated_path) as f:
+            config = yaml.safe_load(f)
+    except FileNotFoundError:
+        result.add_error("file", f"File not found: {file_path}")
+        return result
+    except ValueError as e:
+        result.add_error("file", f"Invalid file path: {e}")
+        return result
+    except yaml.YAMLError as e:
+        result.add_error("yaml", f"Invalid YAML: {e}")
+        return result
+
+    if config is None:
+        result.add_error("file", "Empty configuration file")
+        return result
+
+    return validate_config(config)