OpenJarvis/tests/intelligence/test_routing_models.py

"""Tests for router behavior with the extended model catalog."""

from __future__ import annotations

import pytest

from openjarvis.intelligence.model_catalog import register_builtin_models
from openjarvis.intelligence.router import HeuristicRouter, build_routing_context
from openjarvis.learning._stubs import RoutingContext

# New local model keys for testing
NEW_LOCAL_MODELS = [
    "gpt-oss:120b",   # 117B total, 5.1B active, MoE
    "qwen3:8b",       # 8.2B, dense
    "glm-4.7-flash",  # 30B total, 3.0B active, MoE
    "trinity-mini",   # 26B total, 3.0B active, MoE
]

# Cloud model keys
CLOUD_MODELS = [
    "gpt-5-mini",
    "claude-opus-4-6",
    "claude-sonnet-4-6",
    "claude-haiku-4-5",
    "gemini-2.5-pro",
    "gemini-3-flash",
]


def _setup_models() -> None:
    """Register builtin models needed for the tests."""
    register_builtin_models()


class TestRouterWithNewModels:
    """Router behavior when using the new local models."""

    def test_short_query_routes_to_smallest(self) -> None:
        _setup_models()
        router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
        ctx = RoutingContext(query="hi", query_length=2)
        selected = router.select_model(ctx)
        # qwen3:8b is 8.2B -- the smallest by parameter_count_b
        assert selected == "qwen3:8b"

    def test_code_query_routes_to_largest(self) -> None:
        _setup_models()
        router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
        ctx = RoutingContext(
            query="def merge_sort(arr):", query_length=22, has_code=True
        )
        selected = router.select_model(ctx)
        # No model with "code"/"coder" in name, falls to largest → gpt-oss:120b (117B)
        assert selected == "gpt-oss:120b"

    def test_code_query_with_coder_available(self) -> None:
        _setup_models()
        models = NEW_LOCAL_MODELS + ["deepseek-coder-v2:16b"]
        router = HeuristicRouter(available_models=models)
        ctx = RoutingContext(
            query="import numpy as np", query_length=18, has_code=True
        )
        selected = router.select_model(ctx)
        assert selected == "deepseek-coder-v2:16b"

    def test_math_query_routes_to_largest(self) -> None:
        _setup_models()
        router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
        ctx = RoutingContext(
            query="solve the integral of x^2 dx", query_length=29, has_math=True
        )
        selected = router.select_model(ctx)
        assert selected == "gpt-oss:120b"

    def test_long_context_routes_to_largest(self) -> None:
        _setup_models()
        router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
        ctx = RoutingContext(query="x" * 501, query_length=501)
        selected = router.select_model(ctx)
        assert selected == "gpt-oss:120b"

    def test_high_urgency_routes_to_smallest(self) -> None:
        _setup_models()
        router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
        ctx = RoutingContext(
            query="solve the integral of x^2", query_length=25,
            has_math=True, urgency=0.9,
        )
        selected = router.select_model(ctx)
        # High urgency overrides everything → smallest params → qwen3:8b (8.2B)
        assert selected == "qwen3:8b"

    def test_reasoning_query_routes_to_largest(self) -> None:
        _setup_models()
        router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
        ctx = build_routing_context(
            "Please explain step by step how neural networks learn"
        )
        selected = router.select_model(ctx)
        assert selected == "gpt-oss:120b"

    def test_medium_query_uses_default(self) -> None:
        _setup_models()
        router = HeuristicRouter(
            available_models=NEW_LOCAL_MODELS,
            default_model="glm-4.7-flash",
        )
        # Medium length query, no code/math/reasoning → rule 6 default
        ctx = RoutingContext(query="Tell me about the weather today", query_length=60)
        selected = router.select_model(ctx)
        assert selected == "glm-4.7-flash"


class TestRouterCloudFallback:
    """Router behavior when only cloud models are available."""

    def test_no_local_falls_to_cloud(self) -> None:
        _setup_models()
        router = HeuristicRouter(available_models=CLOUD_MODELS)
        ctx = RoutingContext(query="hi", query_length=2)
        selected = router.select_model(ctx)
        # Should return one of the cloud models (smallest params)
        assert selected in CLOUD_MODELS

    def test_cloud_model_selection_with_math(self) -> None:
        _setup_models()
        router = HeuristicRouter(available_models=CLOUD_MODELS)
        ctx = RoutingContext(query="solve x", query_length=7, has_math=True)
        selected = router.select_model(ctx)
        # All cloud models have parameter_count_b=0 so falls to first
        assert selected in CLOUD_MODELS

    def test_empty_models_returns_fallback(self) -> None:
        router = HeuristicRouter(
            available_models=[],
            fallback_model="gpt-5-mini",
        )
        ctx = RoutingContext(query="hello", query_length=5)
        assert router.select_model(ctx) == "gpt-5-mini"


class TestRouterParameterized:
    """Parametrized cross-product tests for model/query combinations."""

    @pytest.mark.parametrize("query,expected_is_largest", [
        ("hi", False),
        ("solve the integral of sin(x)", True),
        ("def foo(): pass", True),  # code → largest when no coder
        ("x" * 501, True),
    ])
    def test_query_type_selects_expected_size(
        self, query: str, expected_is_largest: bool,
    ) -> None:
        _setup_models()
        router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
        ctx = build_routing_context(query)
        selected = router.select_model(ctx)
        if expected_is_largest:
            assert selected == "gpt-oss:120b"
        else:
            assert selected == "qwen3:8b"

    @pytest.mark.parametrize("model_id", NEW_LOCAL_MODELS)
    def test_single_model_always_returns_it(self, model_id: str) -> None:
        _setup_models()
        router = HeuristicRouter(available_models=[model_id])
        ctx = RoutingContext(query="hello world", query_length=11)
        assert router.select_model(ctx) == model_id

    @pytest.mark.parametrize("urgency", [0.85, 0.9, 1.0])
    def test_high_urgency_always_smallest(self, urgency: float) -> None:
        _setup_models()
        router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
        ctx = RoutingContext(
            query="complex reasoning task", query_length=23,
            has_math=True, urgency=urgency,
        )
        assert router.select_model(ctx) == "qwen3:8b"