litellm/tests/test_litellm/test_compression.py

"""
Unit tests for litellm.compress().
"""

import os
import importlib

import pytest

import litellm
from litellm.compression.scoring.bm25 import bm25_score_messages
from litellm.compression.scoring.embedding_scorer import embedding_score_messages
from litellm.compression.content_detection import detect_content_type
from litellm.compression.message_stubbing import extract_key, stub_message
from litellm.compression.retrieval_tool import build_retrieval_tool
from litellm.types.utils import CallTypes

CALL_TYPE = CallTypes.completion
ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages


# ---------------------------------------------------------------------------
# BM25 scorer
# ---------------------------------------------------------------------------


def test_bm25_relevance_ranking():
    query = "Fix the authentication bug in the login handler"
    messages = [
        {
            "role": "user",
            "content": "def login_handler(): authentication check bug fix",
        },
        {"role": "user", "content": "def render_template(name): css styling layout"},
        {"role": "user", "content": "def verify(): authentication token bug handler"},
    ]
    scores = bm25_score_messages(query, messages)
    # Messages sharing query terms should score higher than unrelated ones
    assert scores[0] > scores[1]
    assert scores[2] > scores[1]


def test_bm25_empty_query():
    scores = bm25_score_messages("", [{"role": "user", "content": "hello"}])
    assert scores == [0.0]


def test_bm25_empty_messages():
    scores = bm25_score_messages("query", [])
    assert scores == []


def test_bm25_empty_content():
    scores = bm25_score_messages("query", [{"role": "user", "content": ""}])
    assert scores == [0.0]


# ---------------------------------------------------------------------------
# Content detection
# ---------------------------------------------------------------------------


def test_detect_code():
    code = """
import os
from pathlib import Path

def main():
    class Foo:
        pass
    return Foo()
"""
    assert detect_content_type(code) == "code"


def test_detect_json():
    assert detect_content_type('{"key": "value", "num": 42}') == "json"
    assert detect_content_type("[1, 2, 3]") == "json"


def test_detect_text():
    assert detect_content_type("This is a plain text paragraph about dogs.") == "text"


def test_detect_empty():
    assert detect_content_type("") == "text"


# ---------------------------------------------------------------------------
# Message stubbing
# ---------------------------------------------------------------------------


def test_extract_key_with_filename():
    msg = {"role": "user", "content": "# auth.py\ndef authenticate():\n    pass"}
    used: set = set()
    key = extract_key(msg, fallback_index=0, used_keys=used)
    assert key == "auth.py"


def test_extract_key_fallback():
    msg = {"role": "user", "content": "Some random content without a filename"}
    used: set = set()
    key = extract_key(msg, fallback_index=5, used_keys=used)
    assert key == "message_5"


def test_extract_key_duplicates():
    used: set = set()
    msg = {"role": "user", "content": "# auth.py\ncode here"}
    k1 = extract_key(msg, fallback_index=0, used_keys=used)
    k2 = extract_key(msg, fallback_index=1, used_keys=used)
    assert k1 == "auth.py"
    assert k2 == "auth.py_2"


def test_stub_message():
    msg = {"role": "user", "content": "line1\nline2\nline3"}
    stubbed = stub_message(msg, "test_key")
    assert stubbed["role"] == "user"
    assert "test_key" in stubbed["content"]
    assert "litellm_content_retrieve" in stubbed["content"]
    assert "3 lines" in stubbed["content"]


# ---------------------------------------------------------------------------
# Retrieval tool
# ---------------------------------------------------------------------------


def test_retrieval_tool_schema():
    tool = build_retrieval_tool(["auth.py", "utils.py"])
    assert tool["type"] == "function"
    assert tool["function"]["name"] == "litellm_content_retrieve"
    assert "key" in tool["function"]["parameters"]["properties"]
    assert tool["function"]["parameters"]["properties"]["key"]["enum"] == [
        "auth.py",
        "utils.py",
    ]
    assert tool["function"]["parameters"]["required"] == ["key"]


def test_retrieval_tool_description_lists_keys():
    tool = build_retrieval_tool(["foo.py", "bar.js"])
    desc = tool["function"]["description"]
    assert "foo.py" in desc
    assert "bar.js" in desc


# ---------------------------------------------------------------------------
# compress() — end-to-end
# ---------------------------------------------------------------------------


def test_compress_below_trigger_passthrough():
    messages = [{"role": "user", "content": "hello"}]
    result = litellm.compress(messages, model="gpt-4o", call_type=CALL_TYPE)
    assert result["messages"] == messages
    assert result["cache"] == {}
    assert result["tools"] == []
    assert result["compression_ratio"] == 0.0
    assert result["original_tokens"] == result["compressed_tokens"]


def test_compress_above_trigger():
    big_messages = [
        {"role": "system", "content": "You are a coding assistant."},
        {
            "role": "user",
            "content": "# auth.py\n" + "def authenticate():\n    pass\n" * 2000,
        },
        {
            "role": "user",
            "content": "# utils.py\n" + "def helper():\n    pass\n" * 2000,
        },
        {
            "role": "user",
            "content": "# readme.md\n" + "This is documentation. " * 2000,
        },
        {"role": "user", "content": "Fix the bug in auth.py"},
    ]

    result = litellm.compress(
        big_messages,
        model="gpt-4o",
        call_type=CALL_TYPE,
        compression_trigger=1000,
        compression_target=500,
    )

    assert result["compressed_tokens"] < result["original_tokens"]
    assert result["compression_ratio"] > 0
    assert len(result["cache"]) > 0
    assert len(result["tools"]) == 1
    assert result["tools"][0]["function"]["name"] == "litellm_content_retrieve"


def test_compress_anthropic_list_content_is_boundary_stable():
    messages = [
        {"role": "system", "content": [{"type": "text", "text": "System prompt"}]},
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "# a.py\n" + "alpha " * 2000},
                {
                    "type": "image_url",
                    "image_url": {"url": "https://example.com/a.png"},
                },
            ],
        },
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "# b.py\n" + "beta " * 2000},
                {
                    "type": "image_url",
                    "image_url": {"url": "https://example.com/b.png"},
                },
            ],
        },
        {
            "role": "user",
            "content": [{"type": "text", "text": "Fix alpha bug in a.py"}],
        },
    ]

    result = litellm.compress(
        messages=messages,
        model="claude-sonnet-4-20250514",
        call_type=ANTHROPIC_CALL_TYPE,
        compression_trigger=1000,
        compression_target=500,
    )

    assert result["compressed_tokens"] < result["original_tokens"]
    assert len(result["messages"]) == len(messages)
    assert [m["role"] for m in result["messages"]] == [m["role"] for m in messages]
    assert len(result["cache"]) > 0
    assert len(result["tools"]) == 1
    assert result["tools"][0]["type"] == "custom"
    assert result["tools"][0]["name"] == "litellm_content_retrieve"
    assert "input_schema" in result["tools"][0]


def test_compress_preserves_system_message():
    messages = [
        {"role": "system", "content": "System prompt. " * 500},
        {"role": "user", "content": "Large file content. " * 5000},
        {"role": "user", "content": "Fix the bug"},
    ]
    result = litellm.compress(
        messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
    )
    assert result["messages"][0]["role"] == "system"
    assert "System prompt" in result["messages"][0]["content"]


def test_compress_preserves_last_user_message():
    messages = [
        {"role": "user", "content": "Big context " * 5000},
        {"role": "user", "content": "Fix the bug in auth.py"},
    ]
    result = litellm.compress(
        messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
    )
    last_user = [m for m in result["messages"] if m["role"] == "user"][-1]
    assert "Fix the bug in auth.py" in last_user["content"]


def test_compress_preserves_last_assistant_message():
    messages = [
        {"role": "user", "content": "Big context " * 5000},
        {"role": "assistant", "content": "I'll help with that. " * 2000},
        {"role": "user", "content": "Now fix the bug"},
    ]
    result = litellm.compress(
        messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
    )
    assistant_msgs = [m for m in result["messages"] if m["role"] == "assistant"]
    assert len(assistant_msgs) >= 1
    # The last assistant message should be preserved (not stubbed)
    last_assistant = assistant_msgs[-1]
    assert "I'll help with that" in last_assistant["content"]


def test_cache_keys_match_stubs():
    messages = [
        {"role": "user", "content": "# auth.py\n" + "code " * 5000},
        {"role": "user", "content": "Fix it"},
    ]
    result = litellm.compress(
        messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
    )
    if result["tools"]:
        tool_desc = result["tools"][0]["function"]["description"]
        for key in result["cache"]:
            assert key in tool_desc


def test_compress_default_target():
    """compression_target defaults to compression_trigger // 2."""
    messages = [
        {"role": "user", "content": "content " * 5000},
        {"role": "user", "content": "query"},
    ]
    result = litellm.compress(
        messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=2000
    )
    # Should have compressed — target = 1000
    assert result["compressed_tokens"] <= result["original_tokens"]


def test_compress_nested_tool_result_extracts_text_only():
    messages = [
        {"role": "system", "content": [{"type": "text", "text": "System rules"}]},
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "prefix"},
                {
                    "type": "tool_result",
                    "tool_use_id": "toolu_1",
                    "content": [
                        {"type": "text", "text": "nested text fragment"},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": "https://example.com/secret-tool.png",
                            },
                        },
                    ],
                },
                {
                    "type": "image_url",
                    "image_url": {"url": "https://example.com/top.png"},
                },
                {"type": "text", "text": " " + ("irrelevant " * 3000)},
            ],
        },
        {
            "role": "user",
            "content": [{"type": "text", "text": "final query that must remain"}],
        },
    ]

    result = litellm.compress(
        messages=messages,
        model="claude-sonnet-4-20250514",
        call_type=ANTHROPIC_CALL_TYPE,
        compression_trigger=500,
        compression_target=100,
    )

    cached_text = " ".join(result["cache"].values())
    assert "nested text fragment" in cached_text
    assert "https://example.com/secret-tool.png" not in cached_text
    assert "https://example.com/top.png" not in cached_text


def test_compress_default_call_type_is_completion():
    result = litellm.compress(
        messages=[
            {"role": "user", "content": "Large context " * 4000},
            {"role": "user", "content": "query"},
        ],
        model="gpt-4o",
        compression_trigger=1000,
        compression_target=500,
    )

    assert result["compressed_tokens"] <= result["original_tokens"]
    assert isinstance(result["tools"], list)


def test_compress_forwards_embedding_model_params(monkeypatch):
    captured = {}

    def fake_embedding_score_messages(
        query, messages, model, cache=None, embedding_model_params=None
    ):
        captured["query"] = query
        captured["model"] = model
        captured["embedding_model_params"] = embedding_model_params
        return [0.0] * len(messages)

    monkeypatch.setattr(
        "litellm.compression.scoring.embedding_scorer.embedding_score_messages",
        fake_embedding_score_messages,
    )

    result = litellm.compress(
        messages=[
            {"role": "user", "content": "Authentication code " * 2000},
            {"role": "user", "content": "Fix auth"},
        ],
        model="gpt-4o",
        call_type=CALL_TYPE,
        compression_trigger=1000,
        embedding_model="text-embedding-3-small",
        embedding_model_params={"api_base": "https://example-embeddings.test"},
    )

    assert result["compressed_tokens"] <= result["original_tokens"]
    assert captured["model"] == "text-embedding-3-small"
    assert captured["embedding_model_params"] == {
        "api_base": "https://example-embeddings.test"
    }


def test_embedding_scorer_forwards_embedding_model_params(monkeypatch):
    captured = {}

    class _MockResponse:
        data = [
            {"embedding": [1.0, 0.0]},
            {"embedding": [1.0, 0.0]},
            {"embedding": [0.0, 1.0]},
        ]

    def fake_embedding(**kwargs):
        captured.update(kwargs)
        return _MockResponse()

    monkeypatch.setattr(litellm, "embedding", fake_embedding)

    scores = embedding_score_messages(
        query="auth",
        messages=[
            {"role": "user", "content": "auth code"},
            {"role": "user", "content": "cooking recipe"},
        ],
        model="text-embedding-3-small",
        embedding_model_params={"api_base": "https://example-embeddings.test"},
    )

    assert len(scores) == 2
    assert captured["model"] == "text-embedding-3-small"
    assert captured["api_base"] == "https://example-embeddings.test"


# ---------------------------------------------------------------------------
# Embedding scorer — integration test (skipped without API key)
# ---------------------------------------------------------------------------


@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="Needs OPENAI_API_KEY")
def test_embedding_scorer():
    result = litellm.compress(
        messages=[
            {"role": "user", "content": "Authentication code " * 2000},
            {"role": "user", "content": "Unrelated cooking recipes " * 2000},
            {"role": "user", "content": "Fix auth"},
        ],
        model="gpt-4o",
        call_type=CALL_TYPE,
        compression_trigger=1000,
        embedding_model="text-embedding-3-small",
    )
    assert result["compression_ratio"] > 0
    assert len(result["cache"]) > 0


@pytest.mark.parametrize(
    "final_user_message, expected_content",
    [
        ("How to cook?", "Unrelated cooking recipes "),
        ("Fix auth", "Authentication code "),
    ],
)
def test_simple_compression(final_user_message, expected_content):
    messages = [
        {"role": "user", "content": "Authentication code " * 2000},
        {"role": "user", "content": "Unrelated cooking recipes " * 2000},
        {"role": "user", "content": final_user_message},
    ]
    result = litellm.compress(
        messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
    )
    if expected_content == "Unrelated cooking recipes ":
        assert "Unrelated cooking recipes " in result["messages"][1]["content"]
        assert "Authentication code " not in result["messages"][0]["content"]
    elif expected_content == "Authentication code ":
        assert "Authentication code " in result["messages"][0]["content"]
        assert "Unrelated cooking recipes " not in result["messages"][1]["content"]
    else:
        raise ValueError(f"Unexpected expected_content: {expected_content}")


def test_compress_anthropic_drops_irrelevant_tool_exchange_span(monkeypatch):
    compress_module = importlib.import_module("litellm.compression.compress")

    def fake_bm25_score_messages(query, messages):
        assert "final query" in query
        assert len(messages) == 5
        # Prefer idx=0 and de-prioritize the tool exchange span (idx=1,2)
        return [0.95, 0.01, 0.02, 0.8, 1.0]

    def fake_token_counter(model, messages=None, text=None):
        if messages is not None:
            return 1000
        if text is None:
            return 0
        if "final query" in text:
            return 50
        if "assistant_tail" in text:
            return 20
        if "other_blob" in text:
            return 220
        if "tool_payload_relevant" in text:
            return 200
        if text == "":
            return 1
        return 10

    monkeypatch.setattr(
        compress_module, "bm25_score_messages", fake_bm25_score_messages
    )
    monkeypatch.setattr(compress_module, "token_counter", fake_token_counter)

    messages = [
        {"role": "user", "content": "other_blob " * 300},
        {
            "role": "assistant",
            "content": [
                {
                    "type": "tool_use",
                    "id": "toolu_drop",
                    "name": "litellm_content_retrieve",
                    "input": {"key": "message_1"},
                }
            ],
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "tool_result",
                    "tool_use_id": "toolu_drop",
                    "content": [{"type": "text", "text": "tool_payload_relevant"}],
                }
            ],
        },
        {"role": "assistant", "content": "assistant_tail"},
        {"role": "user", "content": "final query"},
    ]

    result = litellm.compress(
        messages=messages,
        model="claude-sonnet-4-20250514",
        call_type=ANTHROPIC_CALL_TYPE,
        compression_trigger=100,
        compression_target=280,
    )

    # idx=1,2 should be dropped atomically (no orphan tool blocks left behind)
    assert len(result["messages"]) == 3
    assert result["messages"][0]["role"] == "user"
    assert "other_blob" in result["messages"][0]["content"]
    assert result["messages"][1]["content"] == "assistant_tail"
    assert result["messages"][2]["content"] == "final query"
    assert result["cache"] == {}


def test_compress_anthropic_keeps_relevant_tool_exchange_span(monkeypatch):
    compress_module = importlib.import_module("litellm.compression.compress")

    def fake_bm25_score_messages(query, messages):
        assert "final query" in query
        assert len(messages) == 5
        # Prefer the tool exchange span over idx=0
        return [0.05, 0.01, 0.92, 0.8, 1.0]

    def fake_token_counter(model, messages=None, text=None):
        if messages is not None:
            return 1000
        if text is None:
            return 0
        if "final query" in text:
            return 50
        if "assistant_tail" in text:
            return 20
        if "other_blob" in text:
            return 220
        if "tool_payload_relevant" in text:
            return 200
        if text == "":
            return 1
        return 10

    monkeypatch.setattr(
        compress_module, "bm25_score_messages", fake_bm25_score_messages
    )
    monkeypatch.setattr(compress_module, "token_counter", fake_token_counter)

    messages = [
        {"role": "user", "content": "other_blob " * 300},
        {
            "role": "assistant",
            "content": [
                {
                    "type": "tool_use",
                    "id": "toolu_keep",
                    "name": "litellm_content_retrieve",
                    "input": {"key": "message_1"},
                }
            ],
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "tool_result",
                    "tool_use_id": "toolu_keep",
                    "content": [{"type": "text", "text": "tool_payload_relevant"}],
                }
            ],
        },
        {"role": "assistant", "content": "assistant_tail"},
        {"role": "user", "content": "final query"},
    ]

    result = litellm.compress(
        messages=messages,
        model="claude-sonnet-4-20250514",
        call_type=ANTHROPIC_CALL_TYPE,
        compression_trigger=100,
        compression_target=280,
    )

    assert len(result["messages"]) == 5
    assert result["messages"][1]["role"] == "assistant"
    assert result["messages"][2]["role"] == "user"
    # idx=0 should be compressed instead
    assert "litellm_content_retrieve" in result["messages"][0]["content"]
    assert len(result["cache"]) == 1


def test_compress_anthropic_malformed_tool_sequence_passes_through():
    messages = [
        {"role": "user", "content": "other_blob " * 300},
        {
            "role": "assistant",
            "content": [
                {
                    "type": "tool_use",
                    "id": "toolu_broken",
                    "name": "litellm_content_retrieve",
                    "input": {"key": "message_1"},
                }
            ],
        },
        {"role": "user", "content": [{"type": "text", "text": "missing tool_result"}]},
        {"role": "user", "content": "final query"},
    ]

    result = litellm.compress(
        messages=messages,
        model="claude-sonnet-4-20250514",
        call_type=ANTHROPIC_CALL_TYPE,
        compression_trigger=100,
        compression_target=280,
    )

    assert result["messages"] == messages
    assert result["cache"] == {}
    assert result["tools"] == []
    assert result["compression_skipped_reason"] == "invalid_anthropic_tool_sequence"