From e529e3856e5340ea0a3d0c72e70384f97afcf27a Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 13 Apr 2026 16:39:04 -0700 Subject: [PATCH] test(compaction): add unit tests (15 pass) and hidden interview test set --- .../responses/test_compaction_unit.py | 330 +++++++++++++++++ .../test_hidden_compaction_interview.py | 341 ++++++++++++++++++ 2 files changed, 671 insertions(+) create mode 100644 tests/test_litellm/responses/test_compaction_unit.py create mode 100644 tests/test_litellm/responses/test_hidden_compaction_interview.py diff --git a/tests/test_litellm/responses/test_compaction_unit.py b/tests/test_litellm/responses/test_compaction_unit.py new file mode 100644 index 00000000000..cc79e66751e --- /dev/null +++ b/tests/test_litellm/responses/test_compaction_unit.py @@ -0,0 +1,330 @@ +""" +Unit tests for the compaction implementation. +These tests mock LLM calls and run without API keys. +""" +import base64 +import json +from typing import Any, Dict, List +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) + + +# --------------------------------------------------------------------------- +# Token counter tests +# --------------------------------------------------------------------------- + + +def test_cheap_token_counter_string(): + result = LiteLLMCompletionResponsesConfig._cheap_token_counter("hello world") + assert result == len("hello world") // 4 + + +def test_cheap_token_counter_list(): + input_items = [ + {"type": "message", "role": "user", "content": "word " * 100}, + {"type": "message", "role": "assistant", "content": "reply " * 50}, + ] + count = LiteLLMCompletionResponsesConfig._cheap_token_counter(input_items) + expected = (len("word " * 100) + len("reply " * 50)) // 4 + assert count == expected + + +def test_cheap_token_counter_compaction_item(): + summary = "This is a summary." + encrypted = base64.b64encode(summary.encode()).decode() + input_items = [ + {"type": "compaction", "encrypted_content": encrypted}, + {"type": "message", "role": "user", "content": "hello"}, + ] + count = LiteLLMCompletionResponsesConfig._cheap_token_counter(input_items) + # Should count decoded length of encrypted content, not raw length + assert count > 0 + + +def test_cheap_token_counter_for_messages(): + messages = [ + {"role": "user", "content": "word " * 100}, + {"role": "assistant", "content": "answer " * 50}, + ] + count = LiteLLMCompletionResponsesConfig._cheap_token_counter_for_messages(messages) + expected = (len("word " * 100) + len("answer " * 50)) // 4 + assert count == expected + + +# --------------------------------------------------------------------------- +# should_execute_compaction tests +# --------------------------------------------------------------------------- + + +def test_should_execute_compaction_over_threshold(): + result = LiteLLMCompletionResponsesConfig.should_execute_compaction( + input_token_size=5000, + context_management=[{"type": "compaction", "compact_threshold": 1000}], + ) + assert result is True + + +def test_should_execute_compaction_under_threshold(): + result = LiteLLMCompletionResponsesConfig.should_execute_compaction( + input_token_size=500, + context_management=[{"type": "compaction", "compact_threshold": 1000}], + ) + assert result is False + + +def test_should_execute_compaction_no_context_management(): + result = LiteLLMCompletionResponsesConfig.should_execute_compaction( + input_token_size=999999, + context_management=None, + ) + assert result is False + + +def test_should_execute_compaction_empty_list(): + result = LiteLLMCompletionResponsesConfig.should_execute_compaction( + input_token_size=999999, + context_management=[], + ) + assert result is False + + +def test_should_execute_compaction_exact_threshold(): + # At exactly the threshold: should compact + result = LiteLLMCompletionResponsesConfig.should_execute_compaction( + input_token_size=1000, + context_management=[{"type": "compaction", "compact_threshold": 1000}], + ) + assert result is True + + +# --------------------------------------------------------------------------- +# Compaction item decrypt in input transformation +# --------------------------------------------------------------------------- + + +def test_compaction_item_in_input_decrypts_to_message(): + """When input contains type:compaction, it must be converted to a context message.""" + summary = "We were discussing the payment service incident." + encrypted = base64.b64encode(summary.encode()).decode() + compaction_item = { + "type": "compaction", + "id": "cmp_123", + "encrypted_content": encrypted, + } + messages = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=compaction_item + ) + assert len(messages) == 1 + msg = messages[0] + content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", "") + assert summary in str(content), ( + f"Decoded summary not found in message content. Got: {content}" + ) + + +def test_compaction_item_empty_encrypted_content_returns_empty(): + """Compaction item with no encrypted_content → empty list, not an error.""" + compaction_item = {"type": "compaction", "id": "cmp_123", "encrypted_content": ""} + messages = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=compaction_item + ) + assert messages == [] + + +# --------------------------------------------------------------------------- +# _apply_compaction_to_messages (with mocked LLM) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_apply_compaction_returns_compaction_item_and_new_messages(): + """_apply_compaction_to_messages returns (new_msgs, compaction_item_dict).""" + messages = [ + {"role": "user", "content": "Start of conversation."}, + {"role": "assistant", "content": "Sure, let's talk."}, + {"role": "user", "content": "What is the answer?"}, # last user msg kept + ] + + fake_model_response = MagicMock() + fake_model_response.choices = [MagicMock()] + fake_model_response.choices[0].message.content = "The conversation was about debugging." + + from litellm.types.utils import ModelResponse + + fake_model_response.__class__ = ModelResponse + + with patch("litellm.acompletion", new=AsyncMock(return_value=fake_model_response)): + new_messages, compaction_item = await LiteLLMCompletionResponsesConfig._apply_compaction_to_messages( + model="claude-3-haiku-20240307", + messages=messages, + ) + + # The compaction item must have required fields + assert compaction_item.get("type") == "compaction" + assert compaction_item.get("id", "").startswith("cmp_") + assert compaction_item.get("encrypted_content"), "encrypted_content must be non-empty" + + # The encrypted content must be base64-decodable + decoded = base64.b64decode(compaction_item["encrypted_content"] + "==").decode("utf-8") + assert len(decoded) > 0 + + # The new messages must contain the last user message + assert new_messages[-1]["role"] == "user" + assert new_messages[-1]["content"] == "What is the answer?" + + # The first new message should be the summary + assert "[Previous conversation summary:" in new_messages[0]["content"] + + +@pytest.mark.asyncio +async def test_apply_compaction_keeps_last_user_message(): + """The last user message must be in the tail, not summarized away.""" + CURRENT_QUESTION = "What is the answer to the secret question?" + messages = [ + {"role": "user", "content": "word " * 500}, + {"role": "assistant", "content": "OK."}, + {"role": "user", "content": CURRENT_QUESTION}, + ] + + fake_model_response = MagicMock() + fake_model_response.choices = [MagicMock()] + fake_model_response.choices[0].message.content = "Summary of prior conversation." + + from litellm.types.utils import ModelResponse + + fake_model_response.__class__ = ModelResponse + + with patch("litellm.acompletion", new=AsyncMock(return_value=fake_model_response)): + new_messages, _ = await LiteLLMCompletionResponsesConfig._apply_compaction_to_messages( + model="claude-3-haiku-20240307", + messages=messages, + ) + + last_user = [m for m in new_messages if m.get("role") == "user"][-1] + assert last_user["content"] == CURRENT_QUESTION, ( + "Current question must be preserved verbatim in new messages" + ) + + +# --------------------------------------------------------------------------- +# Case 3: no recompaction logic (threshold check) +# --------------------------------------------------------------------------- + + +def test_no_recompaction_when_summary_msg_plus_new_under_threshold(): + """After decryption, the summary msg + new user msg should be small → no recompact.""" + summary = "Short summary." # ~3 tokens + summary_msg = {"role": "user", "content": f"[Previous conversation summary: {summary}]"} + new_user_msg = {"role": "user", "content": "What is 3+3?"} + messages = [summary_msg, new_user_msg] + + token_count = LiteLLMCompletionResponsesConfig._cheap_token_counter_for_messages(messages) + should_compact = LiteLLMCompletionResponsesConfig.should_execute_compaction( + input_token_size=token_count, + context_management=[{"type": "compaction", "compact_threshold": 200000}], + ) + assert not should_compact, ( + f"Should NOT compact small messages ({token_count} tokens) under high threshold" + ) + + +# --------------------------------------------------------------------------- +# Case 2: full end-to-end with mocked LLM (async_response_api_handler) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_handler_injects_compaction_item_in_response(): + """Full handler path: compaction triggers, item injected at output[0].""" + from litellm.responses.litellm_completion_transformation.handler import ( + LiteLLMCompletionTransformationHandler, + ) + from litellm.types.llms.openai import ResponsesAPIResponse + + fat_content = "word " * 5000 + request_input = [ + {"type": "message", "role": "user", "content": "Start."}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": fat_content}], + }, + {"type": "message", "role": "user", "content": "Continue."}, + ] + responses_api_request = { + "context_management": [{"type": "compaction", "compact_threshold": 100}], + "store": False, + } + + # Fake the summarization call + fake_summary_response = MagicMock() + fake_summary_response.choices = [MagicMock()] + fake_summary_response.choices[0].message.content = "The user discussed a start and wanted to continue." + + from litellm.types.utils import ModelResponse as MR + + fake_summary_response.__class__ = MR + + # Fake the actual LLM call — set reasoning_content=None so Pydantic doesn't choke + choice_mock = MagicMock() + choice_mock.message.content = "OK, continuing." + choice_mock.message.role = "assistant" + choice_mock.message.reasoning_content = None # must be falsy to skip reasoning path + choice_mock.message.tool_calls = None + choice_mock.finish_reason = "stop" + choice_mock.index = 0 + + fake_llm_response = MagicMock(spec=MR) + fake_llm_response.choices = [choice_mock] + fake_llm_response.model = "claude-3-haiku-20240307" + fake_llm_response.id = "resp_test_123" + fake_llm_response.created = 1234567890 + fake_llm_response.object = "chat.completion" + fake_llm_response.usage = MagicMock() + fake_llm_response.usage.prompt_tokens = 10 + fake_llm_response.usage.completion_tokens = 5 + fake_llm_response.usage.total_tokens = 15 + + call_count = {"n": 0} + + async def fake_acompletion(*args, **kwargs): + call_count["n"] += 1 + if call_count["n"] == 1: + # First call is summarization + return fake_summary_response + # Second call is actual LLM + return fake_llm_response + + from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams + + litellm_completion_request = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="claude-3-haiku-20240307", + input=request_input, + responses_api_request=ResponsesAPIOptionalRequestParams(**responses_api_request), + ) + + handler = LiteLLMCompletionTransformationHandler() + with patch("litellm.acompletion", new=AsyncMock(side_effect=fake_acompletion)): + response = await handler.async_response_api_handler( + litellm_completion_request=litellm_completion_request, + request_input=request_input, + responses_api_request=ResponsesAPIOptionalRequestParams(**responses_api_request), + ) + + assert isinstance(response, ResponsesAPIResponse) + assert response.output is not None + assert len(response.output) > 0 + + first_item = response.output[0] + first_type = first_item.get("type") if isinstance(first_item, dict) else getattr(first_item, "type", None) + assert first_type == "compaction", ( + f"Expected compaction at output[0], got: {first_type}" + ) + + # The summarization call should have been made (2 total calls: summarize + actual) + assert call_count["n"] == 2, f"Expected 2 LLM calls, got {call_count['n']}" diff --git a/tests/test_litellm/responses/test_hidden_compaction_interview.py b/tests/test_litellm/responses/test_hidden_compaction_interview.py new file mode 100644 index 00000000000..82d119b1723 --- /dev/null +++ b/tests/test_litellm/responses/test_hidden_compaction_interview.py @@ -0,0 +1,341 @@ +""" +Hidden test set for the /v1/responses compaction V0 interview exercise. +NOT included in the candidate's starter branch. +Run against the candidate's submitted implementation. +""" +import json +import os +from typing import Any, Dict + +import pytest + +import litellm +from litellm.types.llms.openai import ResponsesAPIResponse + +# --------------------------------------------------------------------------- +# T0 — Baseline: no context_management param → normal response shape +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_baseline_no_compaction_param(): + """Baseline: no context_management → normal response, no compaction fields, output[0] is a message""" + response = await litellm.aresponses( + model="groq/llama-3.3-70b-versatile", + input=[ + {"type": "message", "role": "user", "content": "What is the capital of France?"}, + ], + store=False, + max_output_tokens=50, + ) + assert response is not None + assert response.output is not None + assert len(response.output) > 0 + first_item = response.output[0] + first_type = first_item.get("type") if isinstance(first_item, dict) else getattr(first_item, "type", None) + assert first_type == "message", f"Expected first output item to be 'message', got '{first_type}'" + assert not any( + (item.get("type") if isinstance(item, dict) else getattr(item, "type", None)) == "compaction" + for item in response.output + ) + + +# --------------------------------------------------------------------------- +# T1 — No-op when conversation is under the compact_threshold (Case 1) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_compaction_noop_under_threshold(): + """Case 1: small conversation + high threshold → no compaction item in output""" + response = await litellm.aresponses( + model="groq/llama-3.3-70b-versatile", + input=[ + {"type": "message", "role": "user", "content": "What is 2+2?"}, + ], + context_management=[{"type": "compaction", "compact_threshold": 200000}], + store=False, + max_output_tokens=50, + ) + assert response is not None + output_types = [ + item.get("type") if isinstance(item, dict) else getattr(item, "type", None) + for item in response.output + ] + assert "compaction" not in output_types, ( + "Compaction item should NOT appear when conversation is under the threshold" + ) + assert "message" in output_types + + +# --------------------------------------------------------------------------- +# T2 — Compaction triggers when over threshold (Case 2) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_compaction_triggers_over_threshold(): + """Case 2: large conversation, low threshold → compaction item with encrypted_content in output""" + fat_message = "word " * 5000 + response = await litellm.aresponses( + model="groq/llama-3.3-70b-versatile", + input=[ + {"type": "message", "role": "user", "content": "Start of conversation."}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": fat_message}], + }, + {"type": "message", "role": "user", "content": "Now answer: what is 1+1?"}, + ], + context_management=[{"type": "compaction", "compact_threshold": 100}], + store=False, + max_output_tokens=50, + ) + assert response is not None + output_by_type: Dict[str, Any] = {} + for item in response.output: + t = item.get("type") if isinstance(item, dict) else getattr(item, "type", None) + if t and t not in output_by_type: + output_by_type[t] = item + + assert "compaction" in output_by_type, "Compaction item missing from output" + cmp_item = output_by_type["compaction"] + enc = cmp_item.get("encrypted_content") if isinstance(cmp_item, dict) else getattr(cmp_item, "encrypted_content", None) + assert enc, "Compaction item must have non-empty encrypted_content" + cmp_id = cmp_item.get("id") if isinstance(cmp_item, dict) else getattr(cmp_item, "id", None) + assert cmp_id, "Compaction item must have an id" + assert "message" in output_by_type, "Model response missing from output" + + +# --------------------------------------------------------------------------- +# T2b — Compaction item is FIRST in output array (ordering) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_compaction_item_is_first_in_output(): + """Compaction item must be output[0]; model message must follow""" + fat_message = "word " * 5000 + response = await litellm.aresponses( + model="groq/llama-3.3-70b-versatile", + input=[ + {"type": "message", "role": "user", "content": "Start."}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": fat_message}], + }, + {"type": "message", "role": "user", "content": "Continue."}, + ], + context_management=[{"type": "compaction", "compact_threshold": 100}], + store=False, + max_output_tokens=50, + ) + first_item = response.output[0] + first_type = first_item.get("type") if isinstance(first_item, dict) else getattr(first_item, "type", None) + assert first_type == "compaction", ( + f"Expected compaction at output[0], got: {first_type}" + ) + remaining_types = [ + item.get("type") if isinstance(item, dict) else getattr(item, "type", None) + for item in response.output[1:] + ] + assert "message" in remaining_types, "Model response missing after compaction item" + + +# --------------------------------------------------------------------------- +# T3 — Encrypted content survives a round-trip (Case 2 → 3) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_encrypted_content_round_trip(): + """encrypted_content from turn 1 can be sent back as input in turn 2 without error""" + fat_message = "word " * 5000 + + resp1 = await litellm.aresponses( + model="groq/llama-3.3-70b-versatile", + input=[ + {"type": "message", "role": "user", "content": "Let's talk about debugging."}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": fat_message}], + }, + {"type": "message", "role": "user", "content": "Continue."}, + ], + context_management=[{"type": "compaction", "compact_threshold": 100}], + store=False, + max_output_tokens=50, + ) + cmp_items = [ + item for item in resp1.output + if (item.get("type") if isinstance(item, dict) else getattr(item, "type", None)) == "compaction" + ] + assert cmp_items, "Expected compaction in turn 1" + cmp_item = cmp_items[0] + if not isinstance(cmp_item, dict): + cmp_item = dict(cmp_item) + + resp2 = await litellm.aresponses( + model="groq/llama-3.3-70b-versatile", + input=[ + cmp_item, + {"type": "message", "role": "user", "content": "What were we talking about?"}, + ], + context_management=[{"type": "compaction", "compact_threshold": 200000}], + store=False, + max_output_tokens=100, + ) + assert resp2 is not None + assert resp2.output is not None + assert len(resp2.output) > 0 + + +# --------------------------------------------------------------------------- +# T3b — LLM receives DECRYPTED content on second turn (decrypt-before-LLM) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_llm_receives_decrypted_content_on_second_turn(): + """ + The LLM must receive the decoded conversation, not the raw encrypted blob. + We embed a known fact in the compacted history, then ask about it in turn 2. + If the LLM can answer, decryption happened correctly. + """ + fat_prefix = "word " * 3000 + known_fact = "The answer to the secret question is BANANA." + + resp1 = await litellm.aresponses( + model="groq/llama-3.3-70b-versatile", + input=[ + {"type": "message", "role": "user", "content": f"Remember this: {known_fact}"}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": fat_prefix + " I have noted that."}], + }, + {"type": "message", "role": "user", "content": "Good, keep that in mind."}, + ], + context_management=[{"type": "compaction", "compact_threshold": 100}], + store=False, + max_output_tokens=100, + ) + cmp_items = [ + item for item in resp1.output + if (item.get("type") if isinstance(item, dict) else getattr(item, "type", None)) == "compaction" + ] + assert cmp_items, "Expected compaction to trigger in turn 1" + cmp_item = cmp_items[0] + if not isinstance(cmp_item, dict): + cmp_item = dict(cmp_item) + + resp2 = await litellm.aresponses( + model="groq/llama-3.3-70b-versatile", + input=[ + cmp_item, + {"type": "message", "role": "user", "content": "What is the answer to the secret question?"}, + ], + context_management=[{"type": "compaction", "compact_threshold": 200000}], + store=False, + max_output_tokens=100, + ) + assert resp2 is not None + text_output = "" + for item in resp2.output: + item_type = item.get("type") if isinstance(item, dict) else getattr(item, "type", None) + if item_type == "message": + content_list = item.get("content") if isinstance(item, dict) else getattr(item, "content", []) + if isinstance(content_list, list): + for block in content_list: + if isinstance(block, dict) and block.get("type") == "output_text": + text_output += block.get("text", "") + elif isinstance(content_list, str): + text_output += content_list + + assert "BANANA" in text_output.upper(), ( + f"LLM didn't receive decrypted context. Got: '{text_output[:200]}'" + ) + + +# --------------------------------------------------------------------------- +# T4 — No recompaction when compacted + new messages stay under threshold (Case 3) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_no_recompaction_when_under_threshold(): + """Case 3: compaction item + small new message, total under threshold → no new compaction""" + import base64 + + fake_summary = "This is a summary of the previous conversation about debugging." + fake_encrypted = base64.b64encode(fake_summary.encode()).decode() + prior_cmp = { + "type": "compaction", + "id": "cmp_prior_001", + "encrypted_content": fake_encrypted, + } + + response = await litellm.aresponses( + model="groq/llama-3.3-70b-versatile", + input=[ + prior_cmp, + {"type": "message", "role": "user", "content": "What is 3+3?"}, + ], + context_management=[{"type": "compaction", "compact_threshold": 200000}], + store=False, + max_output_tokens=50, + ) + assert response is not None + new_cmp_items = [ + item for item in response.output + if (item.get("type") if isinstance(item, dict) else getattr(item, "type", None)) == "compaction" + and (item.get("id") if isinstance(item, dict) else getattr(item, "id", None)) != "cmp_prior_001" + ] + assert not new_cmp_items, ( + "Should NOT re-compact when compacted context + new messages are under threshold" + ) + output_types = [ + item.get("type") if isinstance(item, dict) else getattr(item, "type", None) + for item in response.output + ] + assert "message" in output_types + + +# --------------------------------------------------------------------------- +# T7 — Absolute no-op when context_management is absent (regression guard) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_no_context_management_is_noop(): + """No context_management param → response has no compaction fields anywhere""" + response = await litellm.aresponses( + model="groq/llama-3.3-70b-versatile", + input=[{"type": "message", "role": "user", "content": "What is the capital of France?"}], + store=False, + max_output_tokens=50, + ) + assert response is not None + output_types = [ + item.get("type") if isinstance(item, dict) else getattr(item, "type", None) + for item in response.output + ] + assert "compaction" not in output_types + + +# --------------------------------------------------------------------------- +# T8 — compact_threshold = 0 handled gracefully (no unhandled crash) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_compact_threshold_zero_handled_gracefully(): + """compact_threshold=0 should not raise an unhandled exception""" + try: + response = await litellm.aresponses( + model="groq/llama-3.3-70b-versatile", + input=[{"type": "message", "role": "user", "content": "Hello."}], + context_management=[{"type": "compaction", "compact_threshold": 0}], + store=False, + max_output_tokens=50, + ) + assert response is not None + except Exception as e: + err = str(e).lower() + assert "compact_threshold" in err or "threshold" in err or "compaction" in err, ( + f"Expected a clear validation error about compact_threshold, got: {e}" + )