Model Routing

class ComplexityRouter:
    def __init__(self):
        self.model_tiers = {
            "simple": {"model": "gpt-3.5-turbo", "max_tokens": 500},
            "moderate": {"model": "gpt-4-turbo", "max_tokens": 2000},
            "complex": {"model": "gpt-4", "max_tokens": 4000},
        }

    def classify_complexity(self, query, task_type):
        """Classify the complexity of a query."""
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            temperature=0.0,
            messages=[
                {
                    "role": "system",
                    "content": """Classify this query as simple, moderate, or complex.

Simple: Single factual question, classification, simple lookup
Moderate: Multi-step reasoning, comparison, analysis
Complex: Deep analysis, multi-source correlation, creative generation

Return only: simple, moderate, or complex"""
                },
                {"role": "user", "content": query}
            ]
        )
        complexity = response["choices"][0]["message"]["content"].strip()
        return self.model_tiers.get(complexity, self.model_tiers["moderate"])

    def route(self, query, task_type):
        """Route a query to the appropriate model."""
        tier = self.classify_complexity(query, task_type)
        return tier

class LatencyRouter:
    def __init__(self):
        self.model_latencies = {
            "gpt-3.5-turbo": {"p50": 500, "p99": 2000},
            "gpt-4-turbo": {"p50": 1500, "p99": 5000},
            "gpt-4": {"p50": 3000, "p99": 10000},
        }

    def route_by_latency(self, query, latency_slo_ms):
        """Route to a model that meets the latency SLO."""
        for model_name, latency in sorted(
            self.model_latencies.items(),
            key=lambda x: x[1]["p99"]
        ):
            if latency["p99"] <= latency_slo_ms:
                return {
                    "model": model_name,
                    "expected_latency": latency["p50"],
                    "meets_slo": True
                }

        # No model meets the SLO, use the fastest
        fastest = min(self.model_latencies.items(), key=lambda x: x[1]["p99"])
        return {
            "model": fastest[0],
            "expected_latency": fastest[1]["p50"],
            "meets_slo": False
        }

class ModelRouter:
    def __init__(self):
        self.routes = {
            "alert_classification": {
                "primary": "gpt-3.5-turbo",
                "fallback": "gpt-4-turbo",
                "local": "llama-3-8b",
            },
            "runbook_lookup": {
                "primary": "gpt-4-turbo",
                "fallback": "gpt-4",
                "local": "llama-3-70b",
            },
            "rca_analysis": {
                "primary": "gpt-4",
                "fallback": "gpt-4-turbo",
                "local": None,
            },
        }

    def route_with_fallback(self, task_type, query):
        """Route with automatic fallback on failure."""
        config = self.routes.get(task_type, self.routes["alert_classification"])

        for model_key in ["primary", "fallback", "local"]:
            model = config[model_key]
            if model is None:
                continue

            try:
                result = call_model(model, query)
                return {
                    "model": model,
                    "result": result,
                    "attempt": model_key
                }
            except Exception as e:
                log_fallback(task_type, model, str(e))
                continue

        return {"error": "All models failed"}