mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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:
parent
d4a3a5e530
commit
3bb4d7e5ca
4 changed files with 278 additions and 1 deletions
124
litellm/responses/compaction.py
Normal file
124
litellm/responses/compaction.py
Normal 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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue