litellm/tests/test_litellm/test_compression.py

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

import os

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


# ---------------------------------------------------------------------------
# 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")
    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",
        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_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", 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", 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", 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", 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", compression_trigger=2000)
    # Should have compressed — target = 1000
    assert result["compressed_tokens"] <= result["original_tokens"]


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",
        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",
        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", compression_trigger=1000)
    print(result["messages"])
    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}")