mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
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:
parent
4eaa95a55f
commit
59d17cc468
5 changed files with 214 additions and 10 deletions
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue