test(compaction): add unit tests (15 pass) and hidden interview test set

This commit is contained in:
Ishaan Jaffer 2026-04-13 16:39:04 -07:00 committed by ryan-crabbe-berri
parent 1ac219c75e
commit e529e3856e
2 changed files with 671 additions and 0 deletions

View file

@ -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']}"

View file

@ -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}"
)