From 2408e09f6a9f7faa047f0fa2a2983c8c0d8c4cff Mon Sep 17 00:00:00 2001 From: Alan Ponnachan <85491837+AlanPonnachan@users.noreply.github.com> Date: Sat, 8 Nov 2025 08:57:28 +0530 Subject: [PATCH] feat: Add support for Anthropic Memory Tool (#16115) * add memory tool in anthropic.py * add memory tool test * make format * update transformation * adding memory to hosted tools * add test * make format --- litellm/llms/anthropic/chat/transformation.py | 60 ++++++++++++------- litellm/types/llms/anthropic.py | 39 ++++++++++-- .../test_anthropic_chat_transformation.py | 25 ++++++++ 3 files changed, 98 insertions(+), 26 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 691b46af8d..96e0acfb44 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -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[ diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 582d238e21..1f20d4b0c7 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -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" \ No newline at end of file + + WEB_FETCH_2025_09_10 = "web-fetch-2025-09-10" + CONTEXT_MANAGEMENT_2025_06_27 = "context-management-2025-06-27" diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index b556b0e5be..f69415fb8e 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -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"