client.py

"""Unified OptimizedClient that integrates all optimization modules."""

import logging
from typing import Any, Dict, List, Optional, Union, Generator
from dataclasses import dataclass, field

from .config import (
    OptimizedClientConfig, load_config, ProviderConfig,
    CacheConfig, RouterConfig, MonitorConfig, OptimizerConfig
)
from .cache import ResponseCache, SimpleEmbeddingProvider, OpenAIEmbeddingProvider
from .router import ModelRouter, RoutingDecision, TaskComplexity, Provider
from .monitor import CostMonitor, BudgetConfig, UsageRecord, Alert
from .optimizer import PromptOptimizer, OptimizationResult
from .utils import count_tokens, estimate_cost, format_messages


logger = logging.getLogger(__name__)


@dataclass
class CompletionResult:
    """Result from an optimized completion request."""
    content: str
    model: str
    provider: str
    input_tokens: int
    output_tokens: int
    total_tokens: int
    cost: float
    cached: bool
    cache_type: Optional[str] = None
    routing_decision: Optional[RoutingDecision] = None
    optimization_result: Optional[OptimizationResult] = None
    raw_response: Optional[Any] = None
    metadata: Dict[str, Any] = field(default_factory=dict)


class OptimizedClient:
    """
    Unified client that wraps API calls with optimization features.

    Features:
    - Response caching (exact and semantic)
    - Smart model routing based on task complexity
    - Cost monitoring and budget management
    - Prompt optimization
    """

    def __init__(
        self,
        config: Optional[OptimizedClientConfig] = None,
        config_path: Optional[str] = None,
        # Direct configuration options
        openai_api_key: Optional[str] = None,
        anthropic_api_key: Optional[str] = None,
        openrouter_api_key: Optional[str] = None,
        enable_cache: bool = True,
        enable_router: bool = True,
        enable_monitor: bool = True,
        enable_optimizer: bool = False,
        default_provider: str = "openrouter",
        default_model: Optional[str] = None,
    ):
        """
        Initialize the OptimizedClient.

        Args:
            config: Full configuration object
            config_path: Path to configuration file
            openai_api_key: OpenAI API key (direct access)
            anthropic_api_key: Anthropic API key (direct access)
            openrouter_api_key: OpenRouter API key (unified access to all models)
            enable_cache: Enable response caching
            enable_router: Enable smart routing
            enable_monitor: Enable cost monitoring
            enable_optimizer: Enable prompt optimization
            default_provider: Default API provider
            default_model: Default model to use
        """
        # Load or create configuration
        if config:
            self.config = config
        elif config_path:
            self.config = load_config(config_path)
        else:
            self.config = load_config()

        # Apply direct overrides
        if openai_api_key:
            if "openai" not in self.config.providers:
                self.config.providers["openai"] = ProviderConfig(name="openai")
            self.config.providers["openai"].api_key = openai_api_key

        if anthropic_api_key:
            if "anthropic" not in self.config.providers:
                self.config.providers["anthropic"] = ProviderConfig(name="anthropic")
            self.config.providers["anthropic"].api_key = anthropic_api_key

        if openrouter_api_key:
            if "openrouter" not in self.config.providers:
                self.config.providers["openrouter"] = ProviderConfig(
                    name="openrouter",
                    api_base="https://openrouter.ai/api/v1",
                    extra_headers={
                        "HTTP-Referer": "https://agent-zero.ai/",
                        "X-Title": "Agent Zero"
                    }
                )
            self.config.providers["openrouter"].api_key = openrouter_api_key

        self.config.cache.enabled = enable_cache
        self.config.router.enabled = enable_router
        self.config.monitor.enabled = enable_monitor
        self.config.optimizer.enabled = enable_optimizer
        self.config.default_provider = default_provider

        if default_model:
            self.config.router.default_model = default_model

        # Initialize components
        self._init_components()

        # API clients (lazy loaded)
        self._openai_client = None
        self._anthropic_client = None

    def _init_components(self):
        """Initialize optimization components."""
        # Cache
        if self.config.cache.enabled:
            embedding_provider = None
            if self.config.cache.embedding_provider == "openai":
                openai_key = self.config.providers.get("openai", ProviderConfig(name="openai")).get_api_key()
                if openai_key:
                    embedding_provider = OpenAIEmbeddingProvider(
                        api_key=openai_key,
                        model=self.config.cache.embedding_model
                    )
            if not embedding_provider:
                embedding_provider = SimpleEmbeddingProvider()

            self.cache = ResponseCache(
                db_path=self.config.cache.db_path,
                default_ttl=self.config.cache.default_ttl,
                similarity_threshold=self.config.cache.similarity_threshold,
                embedding_provider=embedding_provider,
                max_entries=self.config.cache.max_entries
            )
        else:
            self.cache = None

        # Router
        if self.config.router.enabled:
            allowed_providers = None
            if self.config.router.allowed_providers:
                allowed_providers = [
                    Provider(p) for p in self.config.router.allowed_providers
                    if p in ["openai", "anthropic"]
                ]

            self.router = ModelRouter(
                default_model=self.config.router.default_model,
                allowed_providers=allowed_providers,
                cost_optimization=self.config.router.cost_optimization,
                max_cost_per_request=self.config.router.max_cost_per_request
            )
        else:
            self.router = None

        # Monitor
        if self.config.monitor.enabled:
            budget = BudgetConfig(
                daily_limit=self.config.monitor.daily_limit,
                weekly_limit=self.config.monitor.weekly_limit,
                monthly_limit=self.config.monitor.monthly_limit,
                per_request_limit=self.config.monitor.per_request_limit,
                alert_threshold=self.config.monitor.alert_threshold,
                hard_cap=self.config.monitor.hard_cap
            )
            self.monitor = CostMonitor(
                db_path=self.config.monitor.db_path,
                budget=budget,
                alert_callback=self._handle_alert
            )
        else:
            self.monitor = None

        # Optimizer
        if self.config.optimizer.enabled:
            self.optimizer = PromptOptimizer(
                aggressive=self.config.optimizer.aggressive,
                preserve_formatting=self.config.optimizer.preserve_formatting,
                min_savings_threshold=self.config.optimizer.min_savings_threshold
            )
        else:
            self.optimizer = None

    def _handle_alert(self, alert: Alert):
        """Handle budget alerts."""
        logger.warning(f"Budget alert: {alert.message}")

    @property
    def openai(self):
        """Get OpenAI client."""
        if self._openai_client is None:
            try:
                from openai import OpenAI
                provider_config = self.config.providers.get("openai", ProviderConfig(name="openai"))
                self._openai_client = OpenAI(
                    api_key=provider_config.get_api_key(),
                    base_url=provider_config.base_url,
                    timeout=provider_config.timeout,
                    max_retries=provider_config.max_retries
                )
            except ImportError:
                raise ImportError("openai package required: pip install openai")
        return self._openai_client

    @property
    def anthropic(self):
        """Get Anthropic client."""
        if self._anthropic_client is None:
            try:
                from anthropic import Anthropic
                provider_config = self.config.providers.get("anthropic", ProviderConfig(name="anthropic"))
                self._anthropic_client = Anthropic(
                    api_key=provider_config.get_api_key(),
                    base_url=provider_config.base_url,
                    timeout=provider_config.timeout,
                    max_retries=provider_config.max_retries
                )
            except ImportError:
                raise ImportError("anthropic package required: pip install anthropic")
        return self._anthropic_client

    def chat(
        self,
        messages: List[Dict[str, str]],
        model: Optional[str] = None,
        provider: Optional[str] = None,
        use_cache: Optional[bool] = None,
        use_router: Optional[bool] = None,
        optimize_prompt: Optional[bool] = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> CompletionResult:
        """
        Send a chat completion request with optimizations.

        Args:
            messages: List of message dicts with role and content
            model: Model to use (auto-selected if None and router enabled)
            provider: Provider to use (auto-selected based on model)
            use_cache: Override cache setting
            use_router: Override router setting
            optimize_prompt: Override optimizer setting
            max_tokens: Maximum tokens in response
            temperature: Sampling temperature
            metadata: Additional metadata for routing/monitoring
            **kwargs: Additional provider-specific arguments

        Returns:
            CompletionResult with response and metadata
        """
        metadata = metadata or {}
        optimization_result = None
        routing_decision = None

        # Format messages for caching/optimization
        prompt_text = format_messages(messages)

        # 1. Optimize prompt if enabled
        if (optimize_prompt is True) or (optimize_prompt is None and self.config.optimizer.auto_optimize and self.optimizer):
            optimization_result = self.optimizer.optimize(prompt_text)
            if optimization_result.savings_percent >= self.config.optimizer.min_savings_threshold * 100:
                # Apply optimization to last user message
                for msg in reversed(messages):
                    if msg.get("role") == "user":
                        msg["content"] = optimization_result.optimized_text
                        break
                prompt_text = optimization_result.optimized_text

        # 2. Check cache
        use_cache = use_cache if use_cache is not None else self.config.cache.enabled
        if use_cache and self.cache:
            cached = self.cache.get(
                prompt_text,
                model=model or "",
                semantic_search=self.config.cache.semantic_search
            )
            if cached:
                response_text, cache_type = cached

                # Estimate tokens for cached response
                input_tokens = count_tokens(prompt_text)
                output_tokens = count_tokens(response_text)
                cost = estimate_cost(input_tokens, output_tokens, model or self.config.router.default_model)

                # Record cache hit savings
                if self.monitor:
                    self.monitor.record_usage(
                        model=model or "cached",
                        provider="cache",
                        input_tokens=input_tokens,
                        output_tokens=output_tokens,
                        cost=0.0,  # No cost for cached
                        cached=True,
                        metadata={"cache_type": cache_type}
                    )
                    self.cache.record_savings(input_tokens + output_tokens, cost)

                return CompletionResult(
                    content=response_text,
                    model=model or "cached",
                    provider="cache",
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    total_tokens=input_tokens + output_tokens,
                    cost=0.0,
                    cached=True,
                    cache_type=cache_type,
                    optimization_result=optimization_result,
                    metadata={"cost_saved": cost}
                )

        # 3. Route to model if enabled
        use_router = use_router if use_router is not None else self.config.router.enabled
        if use_router and self.router and not model:
            force_provider = Provider(provider) if provider else None
            routing_decision = self.router.route(
                prompt_text,
                metadata=metadata,
                expected_output_tokens=max_tokens // 2,
                force_provider=force_provider
            )
            model = routing_decision.model
            provider = routing_decision.provider.value

        # Set defaults
        model = model or self.config.router.default_model
        provider = provider or self._infer_provider(model)

        # 4. Check budget
        if self.monitor:
            input_tokens = count_tokens(prompt_text)
            estimated_cost = estimate_cost(input_tokens, max_tokens // 2, model)
            allowed, reason = self.monitor.check_budget(estimated_cost)
            if not allowed:
                raise BudgetExceededError(reason)

        # 5. Make API call
        if provider == "openai":
            result = self._call_openai(messages, model, max_tokens, temperature, **kwargs)
        elif provider == "anthropic":
            result = self._call_anthropic(messages, model, max_tokens, temperature, **kwargs)
        else:
            raise ValueError(f"Unknown provider: {provider}")

        # 6. Record usage
        if self.monitor:
            self.monitor.record_usage(
                model=model,
                provider=provider,
                input_tokens=result.input_tokens,
                output_tokens=result.output_tokens,
                cost=result.cost,
                cached=False,
                metadata=metadata
            )

        # 7. Cache response
        if use_cache and self.cache:
            self.cache.set(
                prompt_text,
                result.content,
                model=model,
                tokens_saved=result.total_tokens,
                cost_saved=result.cost
            )

        # Add routing and optimization info
        result.routing_decision = routing_decision
        result.optimization_result = optimization_result

        return result

    def _infer_provider(self, model: str) -> str:
        """Infer provider from model name."""
        model_lower = model.lower()
        if "gpt" in model_lower or "o1" in model_lower:
            return "openai"
        elif "claude" in model_lower:
            return "anthropic"
        return self.config.default_provider

    def _call_openai(
        self,
        messages: List[Dict[str, str]],
        model: str,
        max_tokens: int,
        temperature: float,
        **kwargs
    ) -> CompletionResult:
        """Make OpenAI API call."""
        response = self.openai.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            **kwargs
        )

        content = response.choices[0].message.content or ""
        usage = response.usage

        input_tokens = usage.prompt_tokens if usage else count_tokens(format_messages(messages))
        output_tokens = usage.completion_tokens if usage else count_tokens(content)
        cost = estimate_cost(input_tokens, output_tokens, model)

        return CompletionResult(
            content=content,
            model=model,
            provider="openai",
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            total_tokens=input_tokens + output_tokens,
            cost=cost,
            cached=False,
            raw_response=response
        )

    def _call_anthropic(
        self,
        messages: List[Dict[str, str]],
        model: str,
        max_tokens: int,
        temperature: float,
        **kwargs
    ) -> CompletionResult:
        """Make Anthropic API call."""
        # Extract system message if present
        system = None
        chat_messages = []
        for msg in messages:
            if msg.get("role") == "system":
                system = msg.get("content", "")
            else:
                chat_messages.append(msg)

        create_kwargs = {
            "model": model,
            "messages": chat_messages,
            "max_tokens": max_tokens,
            **kwargs
        }

        if system:
            create_kwargs["system"] = system

        # Only add temperature if not using extended thinking
        if temperature != 1.0 and "thinking" not in kwargs:
            create_kwargs["temperature"] = temperature

        response = self.anthropic.messages.create(**create_kwargs)

        content = ""
        for block in response.content:
            if hasattr(block, "text"):
                content += block.text

        input_tokens = response.usage.input_tokens
        output_tokens = response.usage.output_tokens
        cost = estimate_cost(input_tokens, output_tokens, model)

        return CompletionResult(
            content=content,
            model=model,
            provider="anthropic",
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            total_tokens=input_tokens + output_tokens,
            cost=cost,
            cached=False,
            raw_response=response
        )

    def complete(
        self,
        prompt: str,
        **kwargs
    ) -> CompletionResult:
        """
        Simple completion interface.

        Args:
            prompt: The prompt text
            **kwargs: Arguments passed to chat()

        Returns:
            CompletionResult
        """
        messages = [{"role": "user", "content": prompt}]
        return self.chat(messages, **kwargs)

    # Convenience methods for direct access to components

    def get_cache_stats(self):
        """Get cache statistics."""
        if self.cache:
            return self.cache.get_stats()
        return None

    def get_usage_summary(self, period: str = "daily"):
        """Get usage summary."""
        if self.monitor:
            if period == "daily":
                return self.monitor.get_daily_summary()
            elif period == "weekly":
                return self.monitor.get_weekly_summary()
            elif period == "monthly":
                return self.monitor.get_monthly_summary()
        return None

    def estimate_cost(self, prompt: str, model: Optional[str] = None) -> Dict[str, float]:
        """Estimate costs for a prompt across models."""
        if self.router:
            return self.router.estimate_costs(prompt)
        model = model or self.config.router.default_model
        tokens = count_tokens(prompt)
        return {model: estimate_cost(tokens, tokens, model)}

    def optimize_prompt(self, prompt: str) -> OptimizationResult:
        """Optimize a prompt."""
        if not self.optimizer:
            self.optimizer = PromptOptimizer()
        return self.optimizer.optimize(prompt)

    def clear_cache(self):
        """Clear the response cache."""
        if self.cache:
            self.cache.clear()

    def export_usage(self, filepath: str, format: str = "json"):
        """Export usage data."""
        if self.monitor:
            if format == "json":
                return self.monitor.export_json(filepath)
            elif format == "csv":
                return self.monitor.export_csv(filepath)
        return None


class BudgetExceededError(Exception):
    """Raised when a request would exceed budget limits."""
    pass