diff --git a/litellm/responses/compaction.py b/litellm/responses/compaction.py new file mode 100644 index 00000000000..26c84197467 --- /dev/null +++ b/litellm/responses/compaction.py @@ -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 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 ... tags, falling back to the full text.""" + match = re.search(r"(.*?)", 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 diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 5faa8b587c9..208ae26b880 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -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) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 9075373f1cf..9b85e8eb63a 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -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 ) diff --git a/tests/llm_responses_api_testing/test_google_ai_studio_responses_api.py b/tests/llm_responses_api_testing/test_google_ai_studio_responses_api.py index 13801fce7d1..bbebd44b8ba 100644 --- a/tests/llm_responses_api_testing/test_google_ai_studio_responses_api.py +++ b/tests/llm_responses_api_testing/test_google_ai_studio_responses_api.py @@ -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="The user repeated the word cats many times.", + ), + 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()