Litellm anthropic contex management param support (#16528)

* Add support for anthropic context_management param

* Add support for anthropic context_management param

* Add context management in response

* fix review changes

---------

Co-authored-by: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
This commit is contained in:
Sameer Kankute 2025-11-16 02:30:09 +05:30 • committed by GitHub
parent 4eaa95a55f
commit 59d17cc468
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 214 additions and 10 deletions

View file

@ -953,6 +953,30 @@ except Exception as e:
s/o @[Shekhar Patnaik](https://www.linkedin.com/in/patnaikshekhar) for requesting this!
### Context Management (Beta)
Anthropic’s [context editing](https://docs.claude.com/en/docs/build-with-claude/context-editing) API lets you automatically clear older tool results or thinking blocks. LiteLLM now forwards the native `context_management` payload when you call Anthropic models, and automatically attaches the required `context-management-2025-06-27` beta header.
```python
from litellm import completion
response = completion(
model="anthropic/claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Summarize the latest tool results"}],
context_management={
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 30000},
"keep": {"type": "tool_uses", "value": 3},
"clear_at_least": {"type": "input_tokens", "value": 5000},
"exclude_tools": ["web_search"],
}
]
},
)
```
### Anthropic Hosted Tools (Computer, Text Editor, Web Search, Memory)

View file

@ -129,7 +129,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
"parallel_tool_calls",
"response_format",
"user",
"web_search_options",
"web_search_options"
]
if "claude-3-7-sonnet" in model or supports_reasoning(
@ -646,6 +646,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
return tools
def _ensure_context_management_beta_header(self, headers: dict) -> None:
beta_value = ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value
existing_beta = headers.get("anthropic-beta")
if existing_beta is None:
headers["anthropic-beta"] = beta_value
return
existing_values = [beta.strip() for beta in existing_beta.split(",")]
if beta_value not in existing_values:
headers["anthropic-beta"] = f"{existing_beta}, {beta_value}"
def update_headers_with_optional_anthropic_beta(
self, headers: dict, optional_params: dict
) -> dict:
@ -661,9 +671,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
elif tool.get("type", None) and tool.get("type").startswith(
ANTHROPIC_HOSTED_TOOLS.MEMORY.value
):
headers["anthropic-beta"] = (
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value
)
headers[
"anthropic-beta"
] = ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value
if optional_params.get("context_management") is not None:
self._ensure_context_management_beta_header(headers)
return headers
def transform_request(
@ -973,13 +985,21 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
):
text_content = prefix_prompt + text_content
context_management: Optional[Dict] = completion_response.get(
"context_management"
)
provider_specific_fields: Dict[str, Any] = {
"citations": citations,
"thinking_blocks": thinking_blocks,
}
if context_management is not None:
provider_specific_fields["context_management"] = context_management
_message = litellm.Message(
tool_calls=tool_calls,
content=text_content or None,
provider_specific_fields={
"citations": citations,
"thinking_blocks": thinking_blocks,
},
provider_specific_fields=provider_specific_fields,
thinking_blocks=thinking_blocks,
reasoning_content=reasoning_content,
)
@ -1012,6 +1032,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
model_response.created = int(time.time())
model_response.model = completion_response["model"]
context_management_response = completion_response.get("context_management")
if context_management_response is not None:
_hidden_params["context_management"] = context_management_response
try:
model_response.__dict__["context_management"] = (
context_management_response
)
except Exception:
pass
model_response._hidden_params = _hidden_params
return model_response

View file

@ -6,7 +6,10 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
from litellm.types.llms.anthropic import AnthropicMessagesRequest
from litellm.types.llms.anthropic import (
ANTHROPIC_BETA_HEADER_VALUES,
AnthropicMessagesRequest,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -32,6 +35,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
"tools",
"tool_choice",
"thinking",
"context_management",
# TODO: Add Anthropic `metadata` support
# "metadata",
]
@ -71,6 +75,11 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
if "content-type" not in headers:
headers["content-type"] = "application/json"
headers = self._update_headers_with_optional_anthropic_beta(
headers=headers,
context_management=optional_params.get("context_management"),
)
return headers, api_base
def transform_anthropic_messages_request(
@ -142,3 +151,18 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
request_body=request_body,
litellm_logging_obj=litellm_logging_obj,
)
@staticmethod
def _update_headers_with_optional_anthropic_beta(
headers: dict, context_management: Optional[Dict]
) -> dict:
if context_management is None:
return headers
existing_beta = headers.get("anthropic-beta")
beta_value = ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value
if existing_beta is None:
headers["anthropic-beta"] = beta_value
elif beta_value not in [beta.strip() for beta in existing_beta.split(",")]:
headers["anthropic-beta"] = f"{existing_beta}, {beta_value}"
return headers

View file

@ -1,5 +1,5 @@
from enum import Enum
from typing import Dict, Iterable, List, Optional, Union
from typing import Any, Dict, Iterable, List, Optional, Union
from pydantic import BaseModel
from typing_extensions import Literal, Required, TypedDict
@ -282,6 +282,7 @@ class AnthropicMessagesRequestOptionalParams(TypedDict, total=False):
top_k: Optional[int]
top_p: Optional[float]
mcp_servers: Optional[List[AnthropicMcpServerTool]]
context_management: Optional[Dict[str, Any]]
class AnthropicMessagesRequest(AnthropicMessagesRequestOptionalParams, total=False):

View file

@ -11,6 +11,9 @@ sys.path.insert(
from unittest.mock import MagicMock, patch
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
AnthropicMessagesConfig,
)
from litellm.types.utils import PromptTokensDetailsWrapper, ServerToolUse
@ -389,3 +392,125 @@ def test_anthropic_memory_tool_auto_adds_beta_header():
assert "anthropic-beta" in headers
assert headers["anthropic-beta"] == "context-management-2025-06-27"
def _sample_context_management_payload():
return {
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 30000},
"keep": {"type": "tool_uses", "value": 3},
"clear_at_least": {"type": "input_tokens", "value": 5000},
"exclude_tools": ["web_search"],
"clear_tool_inputs": False,
}
]
}
def test_anthropic_messages_validate_adds_beta_header():
config = AnthropicMessagesConfig()
headers, _ = config.validate_anthropic_messages_environment(
headers={},
model="claude-sonnet-4-20250514",
messages=[{"role": "user", "content": [{"type": "text", "text": "Hi"}]}],
optional_params={"context_management": _sample_context_management_payload()},
litellm_params={},
)
assert headers["anthropic-beta"] == "context-management-2025-06-27"
def test_anthropic_messages_transform_includes_context_management():
config = AnthropicMessagesConfig()
payload = _sample_context_management_payload()
headers = {
"x-api-key": "test",
"anthropic-version": "2023-06-01",
"content-type": "application/json",
}
result = config.transform_anthropic_messages_request(
model="claude-sonnet-4-20250514",
messages=[{"role": "user", "content": [{"type": "text", "text": "Hi"}]}],
anthropic_messages_optional_request_params={
"max_tokens": 512,
"context_management": payload,
},
litellm_params={},
headers=headers,
)
assert result["context_management"] == payload
def test_anthropic_chat_headers_add_context_management_beta():
config = AnthropicConfig()
headers = config.update_headers_with_optional_anthropic_beta(
headers={},
optional_params={"context_management": _sample_context_management_payload()},
)
assert headers["anthropic-beta"] == "context-management-2025-06-27"
def test_anthropic_chat_transform_request_includes_context_management():
config = AnthropicConfig()
headers = {}
result = config.transform_request(
model="claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Hello"}],
optional_params={
"context_management": _sample_context_management_payload(),
"max_tokens": 256,
},
litellm_params={},
headers=headers,
)
assert result["context_management"] == _sample_context_management_payload()
def test_transform_parsed_response_includes_context_management_metadata():
import httpx
from litellm.types.utils import ModelResponse
config = AnthropicConfig()
context_management_payload = {
"applied_edits": [
{
"type": "clear_tool_uses_20250919",
"cleared_tool_uses": 2,
"cleared_input_tokens": 5000,
}
]
}
completion_response = {
"id": "msg_context_management_test",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-20250514",
"content": [{"type": "text", "text": "Done."}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {
"input_tokens": 10,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"output_tokens": 5,
},
"context_management": context_management_payload,
}
raw_response = httpx.Response(
status_code=200,
headers={},
)
model_response = ModelResponse()
result = config.transform_parsed_response(
completion_response=completion_response,
raw_response=raw_response,
model_response=model_response,
json_mode=False,
prefix_prompt=None,
)
assert result.__dict__.get("context_management") == context_management_payload
provider_fields = result.choices[0].message.provider_specific_fields
assert provider_fields and provider_fields["context_management"] == context_management_payload