6krrt/tests/test_classifier_input.py

"""The classifier gets the instruction, not the document.

It decides a category and a tier. Handing it a pasted document is harmful
rather than merely wasteful — measured on a ~20k-token prompt, both local
models failed and neither failed cleanly:

    qwen3.5       28.7s, finish_reason=length, empty content
    mistral-nemo  41.8s, echoed the input back inside its JSON

Both surface as unparseable output, i.e. 30-40s of local inference spent to
reach the same `source: "fallback"` an instant failure would have produced.
With the clamp, the same prompts classify correctly in ~2.2s.
"""

from __future__ import annotations

from dispatcher import (
    _classifier_user_content,
    _previous_context,
    clamp_for_classifier,
)


def test_short_input_is_untouched():
    assert clamp_for_classifier("classify me", 8000) == "classify me"


def test_input_at_the_limit_is_untouched():
    text = "x" * 8000
    assert clamp_for_classifier(text, 8000) == text


def test_the_instruction_survives_at_either_end():
    # Which end carries the instruction depends on how the caller phrased it:
    # "Translate this: <doc>" puts it first, "<doc> — translate this" last.
    # The middle of a pasted document is the one part that never carries it.
    text = "HEAD-INSTRUCTION " + ("filler " * 20000) + " TAIL-INSTRUCTION"
    out = clamp_for_classifier(text, 8000)
    assert "HEAD-INSTRUCTION" in out
    assert "TAIL-INSTRUCTION" in out


def test_clamped_output_is_bounded():
    text = "y" * 500_000
    out = clamp_for_classifier(text, 8000)
    # Head + tail + the elision marker, not the original.
    assert len(out) < 8200
    assert len(out) < len(text)


def test_the_elision_is_visible_to_the_model():
    # The model should be able to tell it is seeing an excerpt, rather than
    # silently reasoning about a document that appears to jump mid-sentence.
    out = clamp_for_classifier("z" * 50_000, 8000)
    assert "elided" in out


def test_zero_disables_clamping():
    # Escape hatch: a deployment with a large-context classifier and a reason
    # to use it should not have to patch code.
    text = "w" * 100_000
    assert clamp_for_classifier(text, 0) == text


def test_negative_is_treated_as_disabled_not_as_a_crash():
    text = "w" * 100_000
    assert clamp_for_classifier(text, -1) == text


# --- context framing (ported from llmrouter) -------------------------------

def test_framing_without_context_is_just_the_task():
    assert _classifier_user_content("Refactor this", None, True) == "Refactor this"
    assert _classifier_user_content("Refactor this", None, False) == "Refactor this"


def test_framing_puts_prior_turn_first_when_enabled():
    out = _classifier_user_content("Try now?", "Design a distributed system", True)
    assert out.startswith("Context: Design a distributed system")
    assert "\n---\nMessage: Try now?\n\n" in out
    # The framing instruction is appended.
    assert ("A short follow-up message ('Yes', 'Try now?', 'Go ahead') "
            "continues the prior turn") in out


def test_legacy_framing_puts_task_first():
    out = _classifier_user_content("Try now?", "Design a distributed system", False)
    assert out.startswith("Try now?")
    assert "--- context ---" in out
    assert "Design a distributed system" in out


def test_framing_input_still_feeds_the_clamp():
    # The framed input is what clamp_for_classifier trims, so a huge prior turn
    # is still bounded to head + tail and the elision stays visible.
    context = "context blob " * 50_000
    framed = _classifier_user_content("Try now?", context, True)
    assert len(framed) > 8000
    clamped = clamp_for_classifier(framed, 8000)
    assert "Context:" in clamped
    assert "Message: Try now?" in clamped
    assert "elided" in clamped


# --- previous-message extraction ------------------------------------------

def test_previous_context_is_the_message_before_last_user():
    messages = [
        {"role": "user", "content": "first"},
        {"role": "assistant", "content": "partly done"},
        {"role": "user", "content": "Try now?"},
    ]
    assert _previous_context(messages) == "partly done"


def test_previous_context_skips_user_messages():
    # Only the message before the LAST user is returned; earlier user messages
    # are skipped, matching llmrouter ("short follow-up inherits from context").
    messages = [
        {"role": "user", "content": "earlier"},
        {"role": "user", "content": "current"},
    ]
    assert _previous_context(messages) == ""


def test_previous_context_empty_when_no_prior_message():
    assert _previous_context([{"role": "user", "content": "only"}]) == ""
    assert _previous_context([]) == ""


def test_previous_context_handles_content_blocks():
    messages = [
        {
            "role": "assistant",
            "content": [
                {"type": "text", "text": "head"},
                {"type": "tool_result", "content": "output"},
            ],
        },
        {"role": "user", "content": "Go"},
    ]
    ctx = _previous_context(messages)
    assert "head" in ctx
    assert "output" in ctx


def test_previous_context_excludes_tool_role_messages():
    # Standard tool loop: assistant(tool_call) → tool(result) → user("ok fix it")
    # Should not include raw tool output as context.
    messages = [
        {"role": "assistant", "content": "Let me check", "tool_calls": [{"id": "t1", "type": "function", "function": {"name": "read_file", "arguments": "{}"}}]},
        {"role": "tool", "tool_call_id": "t1", "content": "/home/alee/file.py\nline 1\nline 2\nline 3"},
        {"role": "user", "content": "ok fix it"},
    ]
    assert _previous_context(messages) == "Let me check"
    assert "/home/alee/file.py" not in _previous_context(messages)


def test_previous_context_excludes_system_prompt():
    # [system, user] on the first turn — should never return a system-prompt fragment.
    messages = [
        {"role": "system", "content": "You are a helpful coding assistant with tool access. You have access to read_file, write_file, and run_tests."},
        {"role": "user", "content": "hello"},
    ]
    assert _previous_context(messages) == ""


def test_previous_context_is_truncated():
    messages = [
        {"role": "assistant", "content": "z" * 10_000},
        {"role": "user", "content": "Go"},
    ]
    assert len(_previous_context(messages)) <= 200


def test_previous_context_returns_assistant_text_when_it_exists():
    # [system, assistant("the real prior turn"), user("Yes")]
    # Should return the assistant turn text, truncated to 200.
    messages = [
        {"role": "system", "content": "You are helpful."},
        {"role": "assistant", "content": "I think you should refactor the auth module to use middleware instead of decorators."},
        {"role": "user", "content": "Yes"},
    ]
    ctx = _previous_context(messages)
    assert "refactor" in ctx
    assert str(ctx) == "I think you should refactor the auth module to use middleware instead of decorators."[:200]


# --- classify() caller invariant --------------------------------------------

def test_classify_user_content_invariant_when_no_context():
    # Regression for the removed if/else in classify(): _classifier_user_content
    # already handles None context (returns task), so removing the outer guard
    # must preserve user_content == task.
    assert _classifier_user_content("do something", None, True) == "do something"
    assert _classifier_user_content("do something", None, False) == "do something"