tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-16 00:24:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: Add support for Anthropic Memory Tool (#16115)

Browse Source 
* add memory tool in anthropic.py

* add memory tool test

* make format

* update transformation

* adding memory to hosted tools

* add test

* make format
This commit is contained in:
Alan Ponnachan
2025-11-07 19:27:28 -08:00
committed by GitHub
parent 6eb74bd62a
commit 2408e09f6a
3 changed files with 98 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/anthropic/chat/transformation.py
+39 -21
View File
@@ -33,6 +33,7 @@ from litellm.types.llms.anthropic import (
AnthropicThinkingParam,
AnthropicWebSearchTool,
AnthropicWebSearchUserLocation,
AnthropicMemoryTool,
)
from litellm.types.llms.openai import (
REASONING_EFFORT,
@@ -82,9 +83,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
to pass metadata to anthropic, it's {"user_id": "any-relevant-information"}
"""
max_tokens: Optional[int] = (
DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS # anthropic requires a default value (Opus, Sonnet, and Haiku have the same default)
)
max_tokens: Optional[
int
] = DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS # anthropic requires a default value (Opus, Sonnet, and Haiku have the same default)
stop_sequences: Optional[list] = None
temperature: Optional[int] = None
top_p: Optional[int] = None
@@ -118,7 +119,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
return super().get_config()
def get_supported_openai_params(self, model: str):
params = [
"stream",
"stop",
@@ -465,11 +465,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
if mcp_servers:
optional_params["mcp_servers"] = mcp_servers
if param == "tool_choice" or param == "parallel_tool_calls":
_tool_choice: Optional[AnthropicMessagesToolChoice] = (
self._map_tool_choice(
tool_choice=non_default_params.get("tool_choice"),
parallel_tool_use=non_default_params.get("parallel_tool_calls"),
)
_tool_choice: Optional[
AnthropicMessagesToolChoice
] = self._map_tool_choice(
tool_choice=non_default_params.get("tool_choice"),
parallel_tool_use=non_default_params.get("parallel_tool_calls"),
)
if _tool_choice is not None:
@@ -517,6 +517,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
self._add_tools_to_optional_params(
optional_params=optional_params, tools=[hosted_web_search_tool]
)
elif param == "extra_headers":
optional_params["extra_headers"] = value
## handle thinking tokens
self.update_optional_params_with_thinking_tokens(
@@ -575,9 +577,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
text=system_message_block["content"],
)
if "cache_control" in system_message_block:
anthropic_system_message_content["cache_control"] = (
system_message_block["cache_control"]
)
anthropic_system_message_content[
"cache_control"
] = system_message_block["cache_control"]
anthropic_system_message_list.append(
anthropic_system_message_content
)
@@ -591,9 +593,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
)
if "cache_control" in _content:
anthropic_system_message_content["cache_control"] = (
_content["cache_control"]
)
anthropic_system_message_content[
"cache_control"
] = _content["cache_control"]
anthropic_system_message_list.append(
anthropic_system_message_content
@@ -641,13 +643,25 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
)
return tools
def update_headers_with_optional_anthropic_beta(self, headers: dict, optional_params: dict) -> dict:
def update_headers_with_optional_anthropic_beta(
self, headers: dict, optional_params: dict
) -> dict:
"""Update headers with optional anthropic beta."""
_tools = optional_params.get("tools", [])
for tool in _tools:
if tool.get("type", None) and tool.get("type").startswith(ANTHROPIC_HOSTED_TOOLS.WEB_FETCH.value):
headers["anthropic-beta"] = ANTHROPIC_BETA_HEADER_VALUES.WEB_FETCH_2025_09_10.value
if tool.get("type", None) and tool.get("type").startswith(
ANTHROPIC_HOSTED_TOOLS.WEB_FETCH.value
):
headers[
"anthropic-beta"
] = ANTHROPIC_BETA_HEADER_VALUES.WEB_FETCH_2025_09_10.value
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
return headers
def transform_request(
@@ -685,7 +699,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
llm_provider="anthropic",
)
headers = self.update_headers_with_optional_anthropic_beta(headers=headers, optional_params=optional_params)
headers = self.update_headers_with_optional_anthropic_beta(
headers=headers, optional_params=optional_params
)
# Separate system prompt from rest of message
anthropic_system_message_list = self.translate_system_message(messages=messages)
@@ -764,7 +780,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
return _message
def extract_response_content(self, completion_response: dict) -> Tuple[
def extract_response_content(
self, completion_response: dict
) -> Tuple[
str,
Optional[List[Any]],
Optional[
litellm/types/llms/anthropic.py
+34 -5
View File
@@ -4,7 +4,11 @@ from typing import Dict, Iterable, List, Optional, Union
from pydantic import BaseModel
from typing_extensions import Literal, Required, TypedDict
from .openai import ChatCompletionCachedContent, ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock
from .openai import (
ChatCompletionCachedContent,
ChatCompletionThinkingBlock,
ChatCompletionRedactedThinkingBlock,
)
class AnthropicMessagesToolChoice(TypedDict, total=False):
@@ -65,12 +69,19 @@ class AnthropicCodeExecutionTool(TypedDict, total=False):
cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
class AnthropicMemoryTool(TypedDict, total=False):
type: Required[str]
name: Required[Literal["memory"]]
cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
AllAnthropicToolsValues = Union[
AnthropicComputerTool,
AnthropicHostedTools,
AnthropicMessagesTool,
AnthropicWebSearchTool,
AnthropicCodeExecutionTool,
AnthropicMemoryTool,
]
@@ -162,6 +173,7 @@ class AnthropicCitationPageLocation(TypedDict, total=False):
Anthropic citation for page-based references.
Used when citing from documents with page numbers.
"""
type: Literal["page_location"]
cited_text: str # The exact text being cited (not counted towards output tokens)
document_index: int # Index referencing the cited document
@@ -175,6 +187,7 @@ class AnthropicCitationCharLocation(TypedDict, total=False):
Anthropic citation for character-based references.
Used when citing from text with character positions.
"""
type: Literal["char_location"]
cited_text: str # The exact text being cited (not counted towards output tokens)
document_index: int # Index referencing the cited document
@@ -317,7 +330,11 @@ class ContentBlockDelta(TypedDict):
type: Literal["content_block_delta"]
index: int
delta: Union[
ContentTextBlockDelta, ContentJsonBlockDelta, ContentCitationsBlockDelta, ContentThinkingBlockDelta, ContentThinkingSignatureBlockDelta
ContentTextBlockDelta,
ContentJsonBlockDelta,
ContentCitationsBlockDelta,
ContentThinkingBlockDelta,
ContentThinkingSignatureBlockDelta,
]
@@ -360,7 +377,9 @@ class ContentBlockStartText(TypedDict):
content_block: TextBlock
ContentBlockContentBlockDict = Union[ToolUseBlock, TextBlock, ChatCompletionThinkingBlock]
ContentBlockContentBlockDict = Union[
ToolUseBlock, TextBlock, ChatCompletionThinkingBlock
]
ContentBlockStart = Union[ContentBlockStartToolUse, ContentBlockStartText]
@@ -463,7 +482,12 @@ class AnthropicResponse(BaseModel):
"""Conversational role of the generated message. This will always be "assistant"."""
content: List[
Union[AnthropicResponseContentBlockText, AnthropicResponseContentBlockToolUse, AnthropicResponseContentBlockThinking, AnthropicResponseContentBlockRedactedThinking]
Union[
AnthropicResponseContentBlockText,
AnthropicResponseContentBlockToolUse,
AnthropicResponseContentBlockThinking,
AnthropicResponseContentBlockRedactedThinking,
]
]
"""Content generated by the model."""
@@ -502,15 +526,20 @@ class AnthropicThinkingParam(TypedDict, total=False):
type: Literal["enabled"]
budget_tokens: int
class ANTHROPIC_HOSTED_TOOLS(str, Enum):
WEB_SEARCH = "web_search"
BASH = "bash"
TEXT_EDITOR = "text_editor"
CODE_EXECUTION = "code_execution"
WEB_FETCH = "web_fetch"
MEMORY = "memory"
class ANTHROPIC_BETA_HEADER_VALUES(str, Enum):
"""
Known beta header values for Anthropic.
"""
WEB_FETCH_2025_09_10 = "web-fetch-2025-09-10"
WEB_FETCH_2025_09_10 = "web-fetch-2025-09-10"
CONTEXT_MANAGEMENT_2025_06_27 = "context-management-2025-06-27"
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py
+25
View File
@@ -364,3 +364,28 @@ def test_get_supported_params_thinking():
config = AnthropicConfig()
params = config.get_supported_openai_params(model="claude-sonnet-4-20250514")
assert "thinking" in params
def test_anthropic_memory_tool_auto_adds_beta_header():
"""
Tests that LiteLLM automatically adds the required 'anthropic-beta' header
when the memory tool is present, and the user has NOT provided a beta header.
"""
config = AnthropicConfig()
memory_tool = [{"type": "memory_20250818", "name": "memory"}]
messages = [{"role": "user", "content": "Remember this."}]
headers = {}
optional_params = {"tools": memory_tool}
config.transform_request(
model="claude-3-5-sonnet-20240620",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers=headers,
)
assert "anthropic-beta" in headers
assert headers["anthropic-beta"] == "context-management-2025-06-27"