OpenJarvis/tests/learning/routing/test_complexity.py

"""Tests for query complexity analyzer."""

from __future__ import annotations

from openjarvis.learning.routing.complexity import (
    ComplexityQueryAnalyzer,
    ComplexityResult,
    adjust_tokens_for_model,
    is_thinking_model,
    score_complexity,
)


class TestScoreComplexity:
    def test_trivial_query(self) -> None:
        result = score_complexity("Hi")
        assert result.tier == "trivial"
        assert result.score < 0.15

    def test_simple_query(self) -> None:
        result = score_complexity("What is the capital of France?")
        assert result.tier in ("trivial", "simple")
        assert result.score < 0.30

    def test_code_signal(self) -> None:
        result = score_complexity("def hello():\n    return 'world'")
        assert result.signals["has_code"] is True
        assert result.score > 0.0

    def test_math_signal(self) -> None:
        result = score_complexity("Solve the integral of x^2 dx")
        assert result.signals["has_math"] is True
        assert result.score > 0.0

    def test_code_and_math_gives_high_domain(self) -> None:
        query = "```python\nimport numpy\n```\nSolve the integral of x^2"
        result = score_complexity(query)
        assert result.signals["has_code"] is True
        assert result.signals["has_math"] is True
        assert result.signals["domain"] == 1.0

    def test_reasoning_signal(self) -> None:
        result = score_complexity("Explain why the sky is blue step by step")
        assert result.signals["has_reasoning"] is True

    def test_multi_step_signal(self) -> None:
        result = score_complexity("First do X, then do Y, then do Z")
        assert result.signals["has_multi_step"] is True

    def test_reasoning_and_multi_step_combined(self) -> None:
        result = score_complexity(
            "Explain why X works, then analyze Y, then compare them"
        )
        assert result.signals["has_reasoning"] is True
        assert result.signals["has_multi_step"] is True
        assert result.signals["reasoning"] == 1.0

    def test_multi_part_questions(self) -> None:
        query = "What is X? What is Y? What is Z? What is W?"
        result = score_complexity(query)
        assert result.signals["n_questions"] == 4
        assert result.signals["multi_part"] == 1.0

    def test_creative_signal(self) -> None:
        result = score_complexity("Write an essay about climate change")
        assert result.signals["has_creative"] is True

    def test_very_complex_query(self) -> None:
        query = (
            "Explain step by step how to solve the integral of x^2, "
            "then write Python code to compute it numerically. "
            "1. Derive the analytical solution\n"
            "2. Implement numerical integration\n"
            "3. Compare the results and analyze the error"
        )
        result = score_complexity(query)
        assert result.tier in ("complex", "very_complex")
        assert result.score >= 0.55

    def test_score_clamped_to_unit_interval(self) -> None:
        result = score_complexity("x" * 2000)
        assert 0.0 <= result.score <= 1.0

    def test_result_type(self) -> None:
        result = score_complexity("hello")
        assert isinstance(result, ComplexityResult)
        assert isinstance(result.score, float)
        assert isinstance(result.tier, str)
        assert isinstance(result.suggested_max_tokens, int)
        assert isinstance(result.signals, dict)

    def test_token_tiers_increase_with_complexity(self) -> None:
        trivial = score_complexity("Hi")
        complex_ = score_complexity(
            "Explain step by step how to solve the integral of x^2 "
            "and write code to compute the derivative of the matrix equation"
        )
        assert complex_.suggested_max_tokens >= trivial.suggested_max_tokens

    def test_subtask_counting(self) -> None:
        query = "1. First task\n2. Second task\n- Bullet item"
        result = score_complexity(query)
        assert result.signals["n_subtasks"] == 3


class TestIsThinkingModel:
    def test_deepseek_r1(self) -> None:
        assert is_thinking_model("deepseek-r1-32b") is True

    def test_o1_model(self) -> None:
        assert is_thinking_model("o1-preview") is True

    def test_o3_model(self) -> None:
        assert is_thinking_model("o3-mini") is True

    def test_qwq(self) -> None:
        assert is_thinking_model("qwq-32b") is True

    def test_regular_model(self) -> None:
        assert is_thinking_model("llama-3.1-8b") is False

    def test_gpt4(self) -> None:
        assert is_thinking_model("gpt-4o") is False


class TestAdjustTokensForModel:
    def test_thinking_model_doubles(self) -> None:
        assert adjust_tokens_for_model(1024, "deepseek-r1-32b") == 2048

    def test_regular_model_unchanged(self) -> None:
        assert adjust_tokens_for_model(1024, "llama-3.1-8b") == 1024

    def test_no_model_unchanged(self) -> None:
        assert adjust_tokens_for_model(1024) == 1024

    def test_none_model_unchanged(self) -> None:
        assert adjust_tokens_for_model(1024, None) == 1024


class TestComplexityQueryAnalyzer:
    def test_returns_routing_context(self) -> None:
        analyzer = ComplexityQueryAnalyzer()
        ctx = analyzer.analyze("Hello world")
        assert ctx.query == "Hello world"
        assert ctx.query_length == len("Hello world")

    def test_complexity_score_populated(self) -> None:
        analyzer = ComplexityQueryAnalyzer()
        ctx = analyzer.analyze("Explain step by step how photosynthesis works")
        assert ctx.complexity_score > 0.0
        assert ctx.has_reasoning is True

    def test_code_detection(self) -> None:
        analyzer = ComplexityQueryAnalyzer()
        ctx = analyzer.analyze("def foo(): pass")
        assert ctx.has_code is True

    def test_math_detection(self) -> None:
        analyzer = ComplexityQueryAnalyzer()
        ctx = analyzer.analyze("solve the equation x^2 = 4")
        assert ctx.has_math is True

    def test_urgency_passthrough(self) -> None:
        analyzer = ComplexityQueryAnalyzer()
        ctx = analyzer.analyze("test", urgency=0.9)
        assert ctx.urgency == 0.9

    def test_invalid_urgency_defaults(self) -> None:
        analyzer = ComplexityQueryAnalyzer()
        ctx = analyzer.analyze("test", urgency="invalid")
        assert ctx.urgency == 0.5

    def test_thinking_model_adjusts_tokens(self) -> None:
        analyzer = ComplexityQueryAnalyzer()
        ctx_normal = analyzer.analyze("Hello", model="llama-3.1-8b")
        ctx_thinking = analyzer.analyze("Hello", model="deepseek-r1-32b")
        assert ctx_thinking.suggested_max_tokens == ctx_normal.suggested_max_tokens * 2

    def test_metadata_contains_tier(self) -> None:
        analyzer = ComplexityQueryAnalyzer()
        ctx = analyzer.analyze("Hi")
        assert "complexity_tier" in ctx.metadata
        assert "signals" in ctx.metadata