feat(responses): add universal server-side context compaction

When input tokens exceed a configurable threshold, the Responses API
now automatically summarizes conversation history before sending it to
the LLM. This reduces token usage for long conversations while
preserving context continuity across all providers.

- Add compaction module with summarization logic and threshold checks
- Integrate compaction into the async response handler, prepending a
  compaction output item when summarization occurs
- Support compaction input items in the transformation layer so
  follow-up turns can consume prior summaries
- Forward custom_llm_provider through compaction to fix provider
  resolution for models using provider-prefixed names
- Add mock unit test validating the two-call flow (summarize then
  complete) and response structure

Made-with: Cursor
This commit is contained in:
SwiftWinds 2026-04-03 12:23:44 -07:00
parent d4a3a5e530
commit 3bb4d7e5ca
4 changed files with 278 additions and 1 deletions

View file

@ -0,0 +1,124 @@
"""
Universal server-side context compaction for the Responses API.
Summarizes conversation history when input tokens exceed a configured threshold,
using the same model the caller is already using. Works across all providers.
"""
import re
from typing import Any, Dict, List, Optional, Tuple
import litellm
from litellm._logging import verbose_logger
SUMMARIZATION_SYSTEM_PROMPT = (
"You have written a partial transcript for the initial task above. "
"Please write a summary of the transcript. The purpose of this summary is "
"to provide continuity so you can continue to make progress towards solving "
"the task in a future context, where the raw history above may not be "
"accessible and will be replaced with this summary. Write down anything "
"that would be helpful, including the state, next steps, learnings etc. "
"You must wrap your summary in a <summary></summary> block."
)
MIN_COMPACT_THRESHOLD = 1000
def _get_compact_threshold(context_management: List[Dict[str, Any]]) -> Optional[int]:
"""Extract the compact_threshold from a context_management list, if present."""
for entry in context_management:
if not isinstance(entry, dict):
continue
if entry.get("type") == "compaction":
threshold = entry.get("compact_threshold")
if threshold is not None:
return max(int(threshold), MIN_COMPACT_THRESHOLD)
return None
return None
def _extract_summary(text: str) -> str:
"""Pull content out of <summary>...</summary> tags, falling back to the full text."""
match = re.search(r"<summary>(.*?)</summary>", text, re.DOTALL)
if match:
return match.group(1).strip()
return text.strip()
def _serialize_messages_for_summary(messages: List[Dict[str, Any]]) -> str:
"""Render a message list into a readable transcript for the summarizer."""
parts: List[str] = []
for msg in messages:
role = msg.get("role", "unknown")
content = msg.get("content", "")
if isinstance(content, list):
text_parts = []
for block in content:
if isinstance(block, dict):
text_parts.append(block.get("text", str(block)))
else:
text_parts.append(str(block))
content = "\n".join(text_parts)
parts.append(f"[{role}]: {content}")
return "\n\n".join(parts)
async def maybe_compact_context(
messages: List[Dict[str, Any]],
model: str,
context_management: List[Dict[str, Any]],
custom_llm_provider: Optional[str] = None,
) -> Tuple[List[Dict[str, Any]], Optional[str]]:
"""
Check whether compaction should trigger and, if so, summarize ALL messages.
Returns:
(messages, summary_text):
- If compaction triggered: messages is [{role: "user", content: summary}],
summary_text is the raw summary string.
- If not triggered: messages is unchanged, summary_text is None.
"""
threshold = _get_compact_threshold(context_management)
if threshold is None:
return messages, None
token_count = litellm.token_counter(model=model, messages=messages)
verbose_logger.debug(
"compaction: token_count=%d, threshold=%d", token_count, threshold
)
if token_count <= threshold:
return messages, None
verbose_logger.info(
"compaction: triggering summarization (tokens=%d > threshold=%d)",
token_count,
threshold,
)
transcript = _serialize_messages_for_summary(messages)
summarization_messages: List[Dict[str, Any]] = [
{"role": "user", "content": transcript},
{"role": "user", "content": SUMMARIZATION_SYSTEM_PROMPT},
]
acompletion_kwargs: Dict[str, Any] = {
"model": model,
"messages": summarization_messages,
}
if custom_llm_provider is not None:
acompletion_kwargs["custom_llm_provider"] = custom_llm_provider
summary_response = await litellm.acompletion(**acompletion_kwargs)
summary_text_raw = summary_response.choices[0].message.content or ""
summary_text = _extract_summary(summary_text_raw)
verbose_logger.debug("compaction: summary length=%d chars", len(summary_text))
compacted_messages: List[Dict[str, Any]] = [
{"role": "user", "content": summary_text},
]
return compacted_messages, summary_text

View file

@ -2,9 +2,10 @@
Handler for transforming responses api requests to litellm.completion requests
"""
from typing import Any, Coroutine, Dict, Optional, Union
from typing import Any, Coroutine, Dict, List, Optional, Union
import litellm
from litellm.responses.compaction import maybe_compact_context
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
LiteLLMCompletionStreamingIterator,
)
@ -109,6 +110,17 @@ class LiteLLMCompletionTransformationHandler:
acompletion_args.update(kwargs)
acompletion_args.update(litellm_completion_request)
context_management = acompletion_args.pop("context_management", None)
summary_text: Optional[str] = None
if context_management:
compacted_messages, summary_text = await maybe_compact_context(
messages=acompletion_args["messages"],
model=acompletion_args["model"],
context_management=context_management,
custom_llm_provider=acompletion_args.get("custom_llm_provider"),
)
acompletion_args["messages"] = compacted_messages
litellm_completion_response: Union[
ModelResponse, litellm.CustomStreamWrapper
] = await litellm.acompletion(
@ -122,6 +134,11 @@ class LiteLLMCompletionTransformationHandler:
responses_api_request=responses_api_request,
)
if summary_text is not None:
responses_api_response.output = _prepend_compaction_output(
summary_text, responses_api_response.output
)
return responses_api_response
elif isinstance(litellm_completion_response, litellm.CustomStreamWrapper):
@ -138,3 +155,14 @@ class LiteLLMCompletionTransformationHandler:
raise ValueError(
f"Unexpected response type: {type(litellm_completion_response)}"
)
def _prepend_compaction_output(
summary_text: str, existing_output: List[Any]
) -> List[Any]:
"""Prepend a compaction output item before the existing output items."""
compaction_item = {
"type": "compaction",
"content": summary_text,
}
return [compaction_item] + list(existing_output)

View file

@ -376,6 +376,16 @@ class LiteLLMCompletionResponsesConfig:
elif isinstance(input, list):
existing_tool_call_ids: Set[str] = set()
for _input in input:
if isinstance(_input, dict) and _input.get("type") == "compaction":
messages.clear()
compaction_content = _input.get("content", "")
messages.append(
ChatCompletionSystemMessage(
role="system", content=compaction_content
)
)
continue
chat_completion_messages = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message(
input_item=_input
)

View file

@ -249,6 +249,121 @@ async def test_gemini_3_responses_api_streaming_with_thought_signatures():
print(f"✅ Collected {len(chunks)} streaming chunks")
@pytest.mark.asyncio
async def test_mock_google_ai_studio_compaction():
"""
Test that universal compaction works for Google AI Studio via the Responses API.
Sends a very large input ("cats " * 100_000) with a low compact_threshold (50_000).
Mocks litellm.acompletion so no real API call is made:
- 1st call: summarization (triggered by compaction)
- 2nd call: the actual completion using ONLY the summary
Validates the response output has a compaction item followed by a text item.
"""
request_model = "gemini/gemini-2.5-flash"
large_input = "cats " * 100_000
summary_response = litellm.ModelResponse(
id="summary-id",
created=1000000000,
model=request_model,
object="chat.completion",
choices=[
litellm.utils.Choices(
index=0,
message=litellm.utils.Message(
role="assistant",
content="<summary>The user repeated the word cats many times.</summary>",
),
finish_reason="stop",
)
],
)
final_response = litellm.ModelResponse(
id="final-id",
created=1000000001,
model=request_model,
object="chat.completion",
choices=[
litellm.utils.Choices(
index=0,
message=litellm.utils.Message(
role="assistant",
content="Based on the summary, you were talking about cats.",
),
finish_reason="stop",
)
],
)
call_count = 0
async def mock_acompletion(*args, **kwargs):
nonlocal call_count
call_count += 1
if call_count == 1:
return summary_response
return final_response
with patch("litellm.acompletion", new_callable=AsyncMock) as mock_ac:
mock_ac.side_effect = mock_acompletion
response = await litellm.aresponses(
model=request_model,
input=large_input,
context_management=[
{"type": "compaction", "compact_threshold": 50000}
],
)
assert call_count == 2, f"Expected 2 acompletion calls, got {call_count}"
# 1st call: summarization — messages should contain the large input
first_call_kwargs = mock_ac.call_args_list[0][1]
first_msgs = first_call_kwargs.get("messages", [])
assert any(
"cats" in str(m.get("content", "")) for m in first_msgs
), "Summarization call should contain the original input"
# 2nd call: actual completion — messages should contain ONLY the summary
second_call_kwargs = mock_ac.call_args_list[1][1]
second_msgs = second_call_kwargs.get("messages", [])
assert len(second_msgs) == 1, (
f"After compaction, completion should receive exactly 1 message (the summary), "
f"got {len(second_msgs)}"
)
assert "cats many times" in str(second_msgs[0].get("content", "")), (
"Completion call should see the extracted summary, not the original input"
)
# Validate response structure
from litellm.types.llms.openai import ResponsesAPIResponse
assert isinstance(response, ResponsesAPIResponse)
assert len(response.output) >= 2, (
f"Response output should have at least 2 items (compaction + text), "
f"got {len(response.output)}"
)
# First output item should be the compaction block
compaction_item = response.output[0]
if isinstance(compaction_item, dict):
assert compaction_item["type"] == "compaction"
assert "cats many times" in compaction_item["content"]
else:
assert getattr(compaction_item, "type", None) == "compaction"
# Second output item should be the text response
text_item = response.output[1]
if isinstance(text_item, dict):
assert text_item.get("type") == "message"
else:
assert getattr(text_item, "type", None) == "message"
print("compaction test passed: response output =", json.dumps(response.output, indent=2, default=str))
class TestGoogleAIStudioResponsesAPITest(BaseResponsesAPITest):
def get_base_completion_call_args(self):
#litellm._turn_on_debug()