llmopt/engine/optimization_engine.py

"""
Optimization Engine — determines the optimal inference configuration.

Solves:
    x* = argmin J(x) = α·Cost(x) + β·Tokens(x) - γ·Quality(x)

Subject to:
    Capability(model) >= Complexity(query)
    Quality >= threshold
    Latency <= max_latency

V1: deterministic rule engine with fixed per-budget-mode weights.
V2: BayesianWeightOptimizer (Optuna) learns α,β,γ from past outcomes.
    Falls back to V1 fixed weights if optuna is not installed.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from llmopt.registry.model_registry import ModelRegistry, ModelSpec
from llmopt.estimator.complexity_estimator import ComplexityResult

logger = logging.getLogger(__name__)

DATA_DIR = Path(__file__).parent.parent.parent / "data"
BAYES_STUDY_PATH = DATA_DIR / "bayesian_study.json"


# ---------------------------------------------------------------------------
# User constraints schema
# ---------------------------------------------------------------------------

@dataclass
class UserConstraints:
    """
    Caller-supplied constraints that bound the optimization search space.
    """
    budget_mode: str = "balanced"       # cheap | balanced | quality
    max_cost_per_request: Optional[float] = None    # USD hard cap
    max_latency_score: Optional[float] = None       # lower = faster model
    quality_threshold: float = 0.60                 # min acceptable quality proxy
    exclude_providers: Optional[list[str]] = None   # e.g. ["ollama"] for cloud-only
    only_providers: Optional[list[str]] = None      # e.g. ["openai"]
    prefer_local: bool = False                      # prefer Ollama models
    compression_enabled: Optional[bool] = None      # None = auto-decide

    def __post_init__(self):
        if self.exclude_providers is None:
            self.exclude_providers = []
        if self.only_providers is None:
            self.only_providers = []


# ---------------------------------------------------------------------------
# Optimization result schema
# ---------------------------------------------------------------------------

@dataclass
class OptimizationResult:
    selected_model: str
    provider: str
    estimated_cost: float               # USD for this request
    estimated_input_tokens: int
    estimated_output_tokens: int
    max_tokens: int                     # hard cap to pass to the API
    compression_enabled: bool
    system_prompt_style: str            # "verbose" | "concise" | "minimal"
    rationale: list[str]
    fallback_model: Optional[str] = None
    objective_score: float = 0.0        # lower is better

    def to_dict(self) -> dict:
        return self.__dict__.copy()


# ---------------------------------------------------------------------------
# Objective weights per budget mode
# ---------------------------------------------------------------------------

_BUDGET_WEIGHTS = {
    #              α (cost)  β (tokens)  γ (quality)
    "cheap":     (0.60,     0.30,       0.10),
    "balanced":  (0.40,     0.20,       0.40),
    "quality":   (0.10,     0.10,       0.80),
}

# max_tokens caps per output-length bucket
_MAX_TOKENS_MAP = {
    "short":     150,
    "medium":    400,
    "long":      900,
    "very_long": 1800,
}


# ---------------------------------------------------------------------------
# Bayesian Weight Optimizer (V2)
# ---------------------------------------------------------------------------

class BayesianWeightOptimizer:
    """
    Uses Optuna to find optimal α,β,γ weights for J(x) based on
    accumulated feedback (cost vs quality trade-offs from past requests).

    Falls back gracefully to fixed V1 weights if optuna is not installed.
    """

    def __init__(self):
        self._optuna_available = False
        self._studies: dict = {}  # one study per budget_mode
        try:
            import optuna  # type: ignore
            optuna.logging.set_verbosity(optuna.logging.WARNING)
            self._optuna_available = True
            logger.info("BayesianWeightOptimizer: Optuna available. Using Bayesian weight tuning.")
        except ImportError:
            logger.info("BayesianWeightOptimizer: Optuna not installed. Using V1 fixed weights.")

    def get_weights(self, budget_mode: str) -> tuple[float, float, float]:
        """
        Returns (α, β, γ) weights for the given budget mode.
        Uses Bayesian optimization if optuna is available and we have
        enough feedback history, otherwise falls back to V1 fixed weights.
        """
        if not self._optuna_available:
            return _BUDGET_WEIGHTS.get(budget_mode, _BUDGET_WEIGHTS["balanced"])

        # Load saved trials
        history = self._load_history(budget_mode)
        if len(history) < 5:
            # Not enough data yet — use V1 defaults but still warm up
            logger.debug(f"Bayesian: Only {len(history)} trials for '{budget_mode}', using V1 defaults.")
            return _BUDGET_WEIGHTS.get(budget_mode, _BUDGET_WEIGHTS["balanced"])

        try:
            import optuna  # type: ignore

            study_key = budget_mode
            if study_key not in self._studies:
                self._studies[study_key] = optuna.create_study(direction="minimize")
                # Seed with historical trials
                for trial_data in history:
                    self._studies[study_key].add_trial(
                        optuna.trial.create_trial(
                            params={"alpha": trial_data["alpha"], "beta": trial_data["beta"], "gamma": trial_data["gamma"]},
                            distributions={
                                "alpha": optuna.distributions.FloatDistribution(0.05, 0.90),
                                "beta":  optuna.distributions.FloatDistribution(0.05, 0.60),
                                "gamma": optuna.distributions.FloatDistribution(0.05, 0.90),
                            },
                            value=trial_data["outcome"],
                        )
                    )

            study = self._studies[study_key]
            best = study.best_params
            α = best["alpha"]
            β = best["beta"]
            γ = best["gamma"]
            logger.debug(f"Bayesian weights for '{budget_mode}': α={α:.3f} β={β:.3f} γ={γ:.3f}")
            return α, β, γ

        except Exception as e:
            logger.warning(f"Bayesian weight retrieval failed: {e}. Using V1 defaults.")
            return _BUDGET_WEIGHTS.get(budget_mode, _BUDGET_WEIGHTS["balanced"])

    def record_outcome(
        self,
        budget_mode: str,
        alpha: float, beta: float, gamma: float,
        actual_cost: float,
        quality_score: float,
    ) -> None:
        """
        Records the outcome of a request. The 'outcome' score is what
        we want to minimize: actual cost weighted against quality.
        Call this after receiving a response + evaluation score.
        """
        # Composite outcome: high cost = bad, low quality = bad
        # Normalise: assume max_cost ~$0.02, quality in [1,10] → [0,1]
        cost_norm = min(actual_cost / 0.02, 1.0)
        quality_norm = quality_score / 10.0
        outcome = cost_norm - quality_norm  # minimise this

        history = self._load_history(budget_mode)
        history.append({
            "alpha": alpha, "beta": beta, "gamma": gamma,
            "actual_cost": actual_cost,
            "quality_score": quality_score,
            "outcome": outcome,
        })
        self._save_history(budget_mode, history)
        # Invalidate the in-memory study so it reloads next time
        self._studies.pop(budget_mode, None)

    def _load_history(self, budget_mode: str) -> list:
        if not BAYES_STUDY_PATH.exists():
            return []
        try:
            data = json.loads(BAYES_STUDY_PATH.read_text())
            return data.get(budget_mode, [])
        except Exception:
            return []

    def _save_history(self, budget_mode: str, history: list) -> None:
        existing = {}
        if BAYES_STUDY_PATH.exists():
            try:
                existing = json.loads(BAYES_STUDY_PATH.read_text())
            except Exception:
                pass
        existing[budget_mode] = history
        BAYES_STUDY_PATH.parent.mkdir(parents=True, exist_ok=True)
        BAYES_STUDY_PATH.write_text(json.dumps(existing, indent=2))


# ---------------------------------------------------------------------------
# Engine
# ---------------------------------------------------------------------------

class OptimizationEngine:
    """
    Core decision engine.  Selects model + config that minimizes
    J(x) = α·Cost + β·Tokens - γ·Quality under user constraints.

    V2: Uses BayesianWeightOptimizer to learn optimal α,β,γ weights over time.
    """

    def __init__(self, registry: ModelRegistry):
        self.registry = registry
        self.bayes = BayesianWeightOptimizer()

    def optimize(
        self,
        complexity: ComplexityResult,
        output_length_bucket: str,
        constraints: Optional[UserConstraints] = None,
    ) -> OptimizationResult:
        if constraints is None:
            constraints = UserConstraints()

        α, β, γ = self.bayes.get_weights(constraints.budget_mode)
        logger.debug(f"Using weights α={α:.3f} β={β:.3f} γ={γ:.3f} for mode '{constraints.budget_mode}'")

        # --- 1. Build candidate set ---
        candidates = self.registry.capable_of(
            complexity=complexity.score,
            min_reasoning=complexity.required_reasoning * 0.85,   # 15% headroom
            min_coding=complexity.required_coding * 0.85,
            min_math=complexity.required_math * 0.85,
            exclude_providers=constraints.exclude_providers,
            only_providers=constraints.only_providers if constraints.only_providers else None,
        )

        if not candidates:
            # Hard fallback: use the most capable model in registry
            candidates = self.registry.all_models()
            candidates = [max(candidates, key=lambda m: m.capability_score)]

        # --- 2. Apply hard filters ---
        candidates = self._apply_hard_filters(candidates, constraints, complexity)

        if not candidates:
            candidates = self.registry.all_models()

        # --- 3. Score candidates via objective function ---
        scored = []
        for model in candidates:
            obj, rationale = self._objective(
                model,
                α, β, γ,
                complexity.estimated_input_tokens,
                complexity.estimated_output_tokens,
            )
            scored.append((obj, model, rationale))

        scored.sort(key=lambda x: x[0])
        best_score, best_model, best_rationale = scored[0]
        fallback = scored[1][1].model_name if len(scored) > 1 else None

        # --- 4. Determine configuration ---
        max_tokens = _MAX_TOKENS_MAP.get(output_length_bucket, 400)
        compression = self._should_compress(constraints, complexity, constraints.budget_mode)
        system_prompt_style = self._system_prompt_style(constraints.budget_mode, complexity.score)

        estimated_cost = best_model.cost_per_request(
            complexity.estimated_input_tokens,
            complexity.estimated_output_tokens,
        )

        return OptimizationResult(
            selected_model=best_model.model_name,
            provider=best_model.provider,
            estimated_cost=round(estimated_cost, 6),
            estimated_input_tokens=complexity.estimated_input_tokens,
            estimated_output_tokens=complexity.estimated_output_tokens,
            max_tokens=max_tokens,
            compression_enabled=compression,
            system_prompt_style=system_prompt_style,
            rationale=best_rationale,
            fallback_model=fallback,
            objective_score=round(best_score, 4),
        )

    # ------------------------------------------------------------------
    # Objective function  J(x) = α·Cost + β·Tokens - γ·Quality
    # ------------------------------------------------------------------

    def _objective(
        self,
        model: ModelSpec,
        α: float, β: float, γ: float,
        input_tokens: int,
        output_tokens: int,
    ) -> tuple[float, list[str]]:
        """
        Normalized objective score.  Lower is better.
        Costs are normalized to [0,1] against registry max values.
        """
        max_cost = self._max_cost_in_registry(input_tokens, output_tokens)
        max_tokens = input_tokens + output_tokens

        cost_norm = model.cost_per_request(input_tokens, output_tokens) / (max_cost + 1e-9)
        token_norm = (input_tokens + output_tokens) / (max_tokens + 1e-9)  # uniform here; prompt optimizer changes this
        quality_norm = model.capability_score  # higher = better

        J = α * cost_norm + β * token_norm - γ * quality_norm
        rationale = [
            f"model={model.model_name}",
            f"provider={model.provider}",
            f"capability={model.capability_score:.3f}",
            f"cost_norm={cost_norm:.4f}",
            f"J={J:.4f} (α={α},β={β},γ={γ})",
        ]
        return J, rationale

    # ------------------------------------------------------------------
    # Hard filters
    # ------------------------------------------------------------------

    def _apply_hard_filters(
        self,
        candidates: list[ModelSpec],
        constraints: UserConstraints,
        complexity: ComplexityResult,
    ) -> list[ModelSpec]:
        filtered = []
        for m in candidates:
            # Cost cap
            if constraints.max_cost_per_request is not None:
                est = m.cost_per_request(
                    complexity.estimated_input_tokens,
                    complexity.estimated_output_tokens,
                )
                if est > constraints.max_cost_per_request:
                    continue
            # Latency cap (latency_score: higher = faster)
            if constraints.max_latency_score is not None:
                if m.latency_score < constraints.max_latency_score:
                    continue
            # Quality floor: capability_score >= quality_threshold
            if m.capability_score < constraints.quality_threshold:
                continue
            filtered.append(m)
        return filtered

    # ------------------------------------------------------------------
    # Config helpers
    # ------------------------------------------------------------------

    def _should_compress(
        self,
        constraints: UserConstraints,
        complexity: ComplexityResult,
        budget_mode: str,
    ) -> bool:
        if constraints.compression_enabled is not None:
            return constraints.compression_enabled
        # Auto: compress for long contexts or cheap modes
        return budget_mode in ("cheap", "balanced") or complexity.estimated_input_tokens > 800

    def _system_prompt_style(self, budget_mode: str, complexity_score: float) -> str:
        if budget_mode == "cheap":
            return "minimal"
        elif budget_mode == "quality" or complexity_score > 0.75:
            return "verbose"
        return "concise"

    def _max_cost_in_registry(self, input_tokens: int, output_tokens: int) -> float:
        costs = [
            m.cost_per_request(input_tokens, output_tokens)
            for m in self.registry.all_models()
        ]
        return max(costs) if costs else 1.0