core-python/tests/test_orchestrator.py

"""Tests for master-model routing and topology metadata."""

from unittest.mock import patch

from maris_core.orchestrator.routing import (
    build_system_prompt,
    detect_route,
    get_specialist_branches,
)


def test_detect_route_prefers_code_branch_for_programming_requests() -> None:
    route = detect_route("Lūdzu uzraksti Python API klientu ar retry loģiku")

    assert route.capability == "code_generation"
    assert route.branch == "coder"
    assert route.profile == "coder"
    assert route.target_studio == "/code"


def test_detect_route_prefers_code_branch_for_calculator_app_requests() -> None:
    route = detect_route("Uzprogrammē kalkulatoru ar pogām un rezultāta lauku")

    assert route.capability == "code_generation"
    assert route.branch == "coder"


def test_detect_route_falls_back_to_master_chat() -> None:
    route = detect_route("Pastāsti man par Latvijas ekonomiku")

    assert route.capability == "text_chat"
    assert route.branch == "master"
    assert route.profile == "general"


def test_detect_route_keeps_programming_explanations_in_chat() -> None:
    route = detect_route("Paskaidro, kā strādā REST API autentifikācija bez koda piemēra")

    assert route.capability == "text_chat"
    assert route.branch == "master"


def test_detect_route_does_not_send_generic_improvement_requests_to_autonomous() -> None:
    route = detect_route("Uzlabo šo tekstu, lai tas būtu skaidrāks klientam")

    assert route.capability == "text_chat"
    assert route.branch == "master"


def test_detect_route_keeps_workflow_explanations_in_chat() -> None:
    route = detect_route("Paskaidro, kā darbojas workflow orchestration Maris AI sistēmā")

    assert route.capability == "text_chat"
    assert route.branch == "master"


def test_detect_route_keeps_media_topic_conversations_in_chat() -> None:
    route = detect_route("Es gribu vienkārši sarunāties ar Maris AI par mūziku un video idejām")

    assert route.capability == "text_chat"
    assert route.branch == "master"


def test_detect_route_keeps_logo_design_explanations_in_chat() -> None:
    route = detect_route("Pastāsti man par logo dizaina principiem pirms ģenerēšanas")

    assert route.capability == "text_chat"
    assert route.branch == "master"


def test_detect_route_carries_session_context() -> None:
    route = detect_route(
        "Turpinām šo pašu sarunu",
        session_id="session-123",
        persona_id="strategist",
    )

    assert route.session_context.session_id == "session-123"
    assert route.session_context.persona_id == "strategist"
    assert route.session_context.memory_enabled is True


def test_detect_route_prefers_vision_branch_when_visual_input_exists() -> None:
    route = detect_route(
        "Kas ir redzams šajā attēlā?",
        session_id="session-vision",
        has_vision_input=True,
    )

    assert route.capability == "vision_analysis"
    assert route.branch == "vision"
    assert route.profile == "vision"


def test_detect_route_prefers_image_branch_for_logo_design_requests() -> None:
    with patch.dict("os.environ", {"IMAGE_MODEL": "MarisUK/maris-ai-image"}, clear=False):
        route = detect_route("Uzlabo dizainu un uztaisi jaunu logo manam produktam")

    assert route.capability == "image_generation"
    assert route.branch == "vision"


def test_detect_route_prefers_autonomous_branch_for_execution_style_improvement_requests() -> None:
    route = detect_route("Veic labojumus un uzlabojumus manā darba plūsmā")

    assert route.capability == "autonomous_tasks"
    assert route.branch == "planner"


def test_build_system_prompt_changes_with_profile() -> None:
    prompt = build_system_prompt("planner")

    assert "plānotājs" in prompt.lower() or "prioritizē" in prompt.lower()


def test_build_system_prompt_uses_persona_specific_profile_defaults() -> None:
    prompt = build_system_prompt("general", persona_id="strategist")

    assert "prioritizē" in prompt.lower() or "plānotājs" in prompt.lower()


def test_specialist_branches_include_master_and_multimodal_paths() -> None:
    branches = get_specialist_branches()
    capabilities = {branch.capability for branch in branches}

    assert "text_chat" in capabilities
    assert "image_generation" in capabilities
    assert "video_generation" in capabilities


def test_detect_route_uses_llm_result_when_available() -> None:
    fake_pipeline = lambda messages, max_new_tokens, temperature: [  # noqa: E731
        {
            "generated_text": messages
            + [
                {
                    "role": "assistant",
                    "content": (
                        '{"capability":"image_generation","reasoning":"Lietotājs prasa attēlu.",'
                        '"confidence":0.94}'
                    ),
                }
            ]
        }
    ]

    with (
        patch("maris_core.text.generate.get_pipeline", return_value=fake_pipeline),
        patch.dict("os.environ", {"IMAGE_MODEL": "MarisUK/maris-ai-image"}, clear=False),
    ):
        route = detect_route("Uzzīmē man kinoplakātu ar neona gaismām")

    assert route.capability == "image_generation"
    assert route.branch == "vision"
    assert route.profile == "image"