litellm/tests/router_unit_tests/test_router_prompt_caching.py

import traceback
import asyncio
from dotenv import load_dotenv
from fastapi import Request
from datetime import datetime

from litellm import Router
import pytest
import litellm
from unittest.mock import patch, MagicMock, AsyncMock
from create_mock_standard_logging_payload import create_standard_logging_payload
from litellm.types.utils import StandardLoggingPayload
import unittest
from pydantic import BaseModel
from litellm.router_utils.prompt_caching_cache import PromptCachingCache


class ExampleModel(BaseModel):
    field1: str
    field2: int


def test_serialize_pydantic_object():
    model = ExampleModel(field1="value", field2=42)
    serialized = PromptCachingCache.serialize_object(model)
    assert serialized == {"field1": "value", "field2": 42}


def test_serialize_dict():
    obj = {"b": 2, "a": 1}
    serialized = PromptCachingCache.serialize_object(obj)
    assert serialized == '{"a":1,"b":2}'  # JSON string with sorted keys


def test_serialize_nested_dict():
    obj = {"z": {"b": 2, "a": 1}, "x": [1, 2, {"c": 3}]}
    serialized = PromptCachingCache.serialize_object(obj)
    expected = '{"x":[1,2,{"c":3}],"z":{"a":1,"b":2}}'  # JSON string with sorted keys
    assert serialized == expected


def test_serialize_list():
    obj = ["item1", {"a": 1, "b": 2}, 42]
    serialized = PromptCachingCache.serialize_object(obj)
    expected = ["item1", '{"a":1,"b":2}', 42]
    assert serialized == expected


def test_serialize_fallback():
    obj = 12345  # Simple non-serializable object
    serialized = PromptCachingCache.serialize_object(obj)
    assert serialized == 12345


def test_serialize_non_serializable():
    class CustomClass:
        def __str__(self):
            return "custom_object"

    obj = CustomClass()
    serialized = PromptCachingCache.serialize_object(obj)
    assert serialized == "custom_object"  # Fallback to string conversion


@pytest.mark.asyncio
async def test_router_prompt_caching_same_cacheable_prefix_routes_to_same_deployment():
    """
    End-to-end test to validate prompt caching routing through LiteLLM Router.

    Tests that requests with same cacheable content but different user messages
    route to the same deployment (for prompt caching).

    This reproduces the issue where requests with same cacheable prefix but different
    user messages should route to the same deployment, but previously didn't because
    the cache key included the entire messages array instead of just the cacheable prefix.
    """
    from litellm.types.llms.openai import AllMessageValues

    def create_messages(user_content: str) -> list[AllMessageValues]:
        """
        Create messages matching the user's exact scenario.

        Message structure:
        - BLOCK 1: System message, first content block (no cache_control)
                  → INCLUDED (comes before the last cacheable block)
        - BLOCK 2: System message, second content block (WITH cache_control)
                  → INCLUDED (this IS the last cacheable block)
        - USER MESSAGE: User message (no cache_control)
                  → NOT included (comes after last cacheable block)
        """
        return [
            {
                "role": "system",
                "content": [
                    # BLOCK 1: No cache_control → INCLUDED (all blocks up to last cacheable are included)
                    {
                        "type": "text",
                        "text": "You are an AI assistant tasked with analyzing legal documents.",
                    },
                    # BLOCK 2: Has cache_control → INCLUDED (this is the last cacheable block)
                    {
                        "type": "text",
                        "text": "Here 3 is the full text of a complex legal agreement"
                        * 400,
                        "cache_control": {"type": "ephemeral"},
                    },
                ],
            },
            {
                "role": "user",
                # USER MESSAGE: No cache_control → NOT included (comes after last cacheable block)
                "content": user_content,
            },
        ]

    # Create router with multiple deployments
    router = Router(
        model_list=[
            {
                "model_name": "test-model",
                "litellm_params": {
                    "model": "gpt-5-mini",
                    "api_base": "https://exampleopenaiendpoint-production-0ee2.up.railway.app/v1",
                    "api_key": f"test-key-{i}",
                },
                "model_info": {"id": f"deployment-{i}"},
            }
            for i in range(1, 4)
        ],
        routing_strategy="simple-shuffle",
        optional_pre_call_checks=["prompt_caching"],
    )

    # Create test messages matching user's exact scenario
    # Same cacheable prefix (system blocks 1+2) but different user messages
    messages1 = create_messages(
        "what are the key terms and conditions in this agreement?"
    )
    messages2 = create_messages("how many words are there?")
    messages3 = create_messages("how many sentences are there?")

    cache = PromptCachingCache(cache=router.cache)

    # Test 1: Cache keys should be same (same cacheable prefix, different user messages)
    key1 = PromptCachingCache.get_prompt_caching_cache_key(messages1, None)
    key2 = PromptCachingCache.get_prompt_caching_cache_key(messages2, None)
    key3 = PromptCachingCache.get_prompt_caching_cache_key(messages3, None)

    assert key1 is not None, "Cache key should not be None"
    assert (
        key1 == key2 == key3
    ), "Cache keys should be the same for same cacheable prefix"

    # Make first request
    try:
        response1 = await router.acompletion(model="test-model", messages=messages1)
        model_id_1 = response1._hidden_params.get("model_id", "unknown")
    except Exception:
        # If API call fails, we can still test the cache key logic
        model_id_1 = "unknown"

    await asyncio.sleep(1)  # Wait for cache write

    # Test 2: Cache lookup should work for messages2 (same cacheable prefix)
    cached_2 = await cache.async_get_model_id(messages2, None)
    # Cache should be found if first request succeeded
    if model_id_1 != "unknown":
        assert (
            cached_2 is not None
        ), "Cache lookup should work for same cacheable prefix"

    # Make second request
    try:
        response2 = await router.acompletion(model="test-model", messages=messages2)
        model_id_2 = response2._hidden_params.get("model_id", "unknown")
    except Exception:
        model_id_2 = "unknown"

    await asyncio.sleep(1)  # Wait for cache write

    # Make third request
    try:
        response3 = await router.acompletion(model="test-model", messages=messages3)
        model_id_3 = response3._hidden_params.get("model_id", "unknown")
    except Exception:
        model_id_3 = "unknown"

    # Test 3: All requests should route to same deployment (if API calls succeeded)
    if model_id_1 != "unknown" and model_id_2 != "unknown" and model_id_3 != "unknown":
        assert (
            model_id_1 == model_id_2 == model_id_3
        ), f"All requests should route to same deployment, but got: {model_id_1}, {model_id_2}, {model_id_3}"


def test_extract_cacheable_prefix_with_string_content_and_message_level_cache_control():
    """
    Test that extract_cacheable_prefix correctly handles messages where:
    - content is a string (not a list of content blocks)
    - cache_control is a sibling key at the message level

    This is a valid message format per LiteLLM's ChatCompletionUserMessage type:
    {"role": "user", "content": "...", "cache_control": {"type": "ephemeral"}}

    Regression test for issue #19228.
    """
    # Test case 1: Single message with string content and message-level cache_control
    messages_string_content = [
        {"role": "system", "content": "You are a helpful assistant"},
        {
            "role": "user",
            "content": "This is a large message that should be cached",
            "cache_control": {"type": "ephemeral", "ttl": "5m"},
        },
    ]

    result = PromptCachingCache.extract_cacheable_prefix(messages_string_content)

    # Should return both messages (system + user with cache_control)
    assert len(result) == 2, f"Expected 2 messages, got {len(result)}"
    assert result[0]["role"] == "system"
    assert result[1]["role"] == "user"
    assert result[1]["content"] == "This is a large message that should be cached"
    assert result[1].get("cache_control") == {"type": "ephemeral", "ttl": "5m"}


def test_extract_cacheable_prefix_with_string_content_no_cache_control():
    """
    Test that extract_cacheable_prefix returns empty list when:
    - content is a string
    - no cache_control is present
    """
    messages_no_cache = [
        {"role": "system", "content": "You are a helpful assistant"},
        {"role": "user", "content": "Hello"},
    ]

    result = PromptCachingCache.extract_cacheable_prefix(messages_no_cache)

    # Should return empty list (no cacheable content)
    assert len(result) == 0, f"Expected 0 messages, got {len(result)}"


def test_extract_cacheable_prefix_mixed_string_and_list_content():
    """
    Test that extract_cacheable_prefix handles messages with a mix of:
    - String content with message-level cache_control
    - List content with block-level cache_control

    The last cache_control (regardless of format) should determine the cacheable prefix.
    """
    # Message with string content + cache_control, followed by message with list content + cache_control
    messages_mixed = [
        {"role": "system", "content": "You are a helpful assistant"},
        {
            "role": "user",
            "content": "First cached message",
            "cache_control": {"type": "ephemeral"},
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "Second cached message in list format",
                    "cache_control": {"type": "ephemeral"},
                }
            ],
        },
        {"role": "user", "content": "This should not be in the prefix"},
    ]

    result = PromptCachingCache.extract_cacheable_prefix(messages_mixed)

    # Should include first 3 messages (up to and including the last cache_control)
    assert len(result) == 3, f"Expected 3 messages, got {len(result)}"
    assert result[0]["role"] == "system"
    assert result[1]["content"] == "First cached message"
    assert isinstance(result[2]["content"], list)