Skip to content

Commit 174e685

Browse files
keenborder786casparbmdrxy
authored
feat(anthropic): dynamic mapping of Max Tokens for Anthropic (#31946)
- **Description:** Dynamic mapping of `max_tokens` as per the choosen anthropic model. - **Issue:** Fixes #31605 @ccurme --------- Co-authored-by: Caspar Broekhuizen <[email protected]> Co-authored-by: Mason Daugherty <[email protected]>
1 parent 9721684 commit 174e685

File tree

4 files changed

+80
-3
lines changed

4 files changed

+80
-3
lines changed

libs/partners/anthropic/langchain_anthropic/chat_models.py

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
88
from functools import cached_property
99
from operator import itemgetter
10-
from typing import Any, Callable, Literal, Optional, Union, cast
10+
from typing import Any, Callable, Final, Literal, Optional, Union, cast
1111

1212
import anthropic
1313
from langchain_core._api import beta, deprecated
@@ -61,6 +61,32 @@
6161
}
6262

6363

64+
_MODEL_DEFAULT_MAX_OUTPUT_TOKENS: Final[dict[str, int]] = {
65+
"claude-opus-4-1": 32000,
66+
"claude-opus-4": 32000,
67+
"claude-sonnet-4": 64000,
68+
"claude-3-7-sonnet": 64000,
69+
"claude-3-5-sonnet": 8192,
70+
"claude-3-5-haiku": 8192,
71+
"claude-3-haiku": 4096,
72+
}
73+
_FALLBACK_MAX_OUTPUT_TOKENS: Final[int] = 4096
74+
75+
76+
def _default_max_tokens_for(model: str | None) -> int:
77+
"""Return the default max output tokens for an Anthropic model (with fallback).
78+
79+
Can find the Max Tokens limits here: https://docs.anthropic.com/en/docs/about-claude/models/overview#model-comparison-table
80+
"""
81+
if not model:
82+
return _FALLBACK_MAX_OUTPUT_TOKENS
83+
84+
parts = model.split("-")
85+
family = "-".join(parts[:-1]) if len(parts) > 1 else model
86+
87+
return _MODEL_DEFAULT_MAX_OUTPUT_TOKENS.get(family, _FALLBACK_MAX_OUTPUT_TOKENS)
88+
89+
6490
class AnthropicTool(TypedDict):
6591
"""Anthropic tool definition."""
6692

@@ -1205,7 +1231,7 @@ def get_weather(location: str) -> str:
12051231
model: str = Field(alias="model_name")
12061232
"""Model name to use."""
12071233

1208-
max_tokens: int = Field(default=1024, alias="max_tokens_to_sample")
1234+
max_tokens: Optional[int] = Field(default=None, alias="max_tokens_to_sample")
12091235
"""Denotes the number of tokens to predict per generation."""
12101236

12111237
temperature: Optional[float] = None
@@ -1343,6 +1369,15 @@ def _get_ls_params(
13431369
ls_params["ls_stop"] = ls_stop
13441370
return ls_params
13451371

1372+
@model_validator(mode="before")
1373+
@classmethod
1374+
def set_default_max_tokens(cls, values: dict[str, Any]) -> Any:
1375+
"""Set default max_tokens."""
1376+
if values.get("max_tokens") is None:
1377+
model = values.get("model") or values.get("model_name")
1378+
values["max_tokens"] = _default_max_tokens_for(model)
1379+
return values
1380+
13461381
@model_validator(mode="before")
13471382
@classmethod
13481383
def build_extra(cls, values: dict) -> Any:

libs/partners/anthropic/tests/integration_tests/test_chat_models.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -901,7 +901,10 @@ class color_picker(BaseModel):
901901

902902
@pytest.mark.vcr
903903
def test_web_search() -> None:
904-
llm = ChatAnthropic(model="claude-3-5-sonnet-latest") # type: ignore[call-arg]
904+
llm = ChatAnthropic(
905+
model="claude-3-5-sonnet-latest", # type: ignore[call-arg]
906+
max_tokens=1024,
907+
)
905908

906909
tool = {"type": "web_search_20250305", "name": "web_search", "max_uses": 1}
907910
llm_with_tools = llm.bind_tools([tool])

libs/partners/anthropic/tests/unit_tests/test_chat_models.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,45 @@ def test_anthropic_proxy_from_environment() -> None:
111111
assert llm.anthropic_proxy == explicit_proxy
112112

113113

114+
def test_set_default_max_tokens() -> None:
115+
"""Test the set_default_max_tokens function."""
116+
# Test claude-opus-4 models
117+
llm = ChatAnthropic(model="claude-opus-4-20250514", anthropic_api_key="test")
118+
assert llm.max_tokens == 32000
119+
120+
# Test claude-sonnet-4 models
121+
llm = ChatAnthropic(model="claude-sonnet-4-latest", anthropic_api_key="test")
122+
assert llm.max_tokens == 64000
123+
124+
# Test claude-3-7-sonnet models
125+
llm = ChatAnthropic(model="claude-3-7-sonnet-latest", anthropic_api_key="test")
126+
assert llm.max_tokens == 64000
127+
128+
# Test claude-3-5-sonnet models
129+
llm = ChatAnthropic(model="claude-3-5-sonnet-latest", anthropic_api_key="test")
130+
assert llm.max_tokens == 8192
131+
132+
# Test claude-3-5-haiku models
133+
llm = ChatAnthropic(model="claude-3-5-haiku-latest", anthropic_api_key="test")
134+
assert llm.max_tokens == 8192
135+
136+
# Test claude-3-haiku models (should default to 4096)
137+
llm = ChatAnthropic(model="claude-3-haiku-latest", anthropic_api_key="test")
138+
assert llm.max_tokens == 4096
139+
140+
# Test that existing max_tokens values are preserved
141+
llm = ChatAnthropic(
142+
model="claude-3-5-sonnet-latest", max_tokens=2048, anthropic_api_key="test"
143+
)
144+
assert llm.max_tokens == 2048
145+
146+
# Test that explicitly set max_tokens values are preserved
147+
llm = ChatAnthropic(
148+
model="claude-3-5-sonnet-latest", max_tokens=4096, anthropic_api_key="test"
149+
)
150+
assert llm.max_tokens == 4096
151+
152+
114153
@pytest.mark.requires("anthropic")
115154
def test_anthropic_model_name_param() -> None:
116155
llm = ChatAnthropic(model_name="foo") # type: ignore[call-arg, call-arg]

0 commit comments

Comments
 (0)