feat: Allow max tokens to be customizable in Claude

PiperOrigin-RevId: 789901925

Author: shukladivyansh

src/google/adk/models/anthropic_llm.py

@@ -46,8 +46,6 @@
 
 logger = logging.getLogger("google_adk." + __name__)
 
-MAX_TOKEN = 8192
-
 
 class ClaudeRequest(BaseModel):
  system_instruction: str
@@ -245,9 +243,11 @@ class Claude(BaseLlm):
 
  Attributes:
  model: The name of the Claude model.
+ max_tokens: The maximum number of tokens to generate.
  """
 
  model: str = "claude-3-5-sonnet-v2@20241022"
+ max_tokens: int = 8192
 
  @staticmethod
  @override
@@ -284,7 +284,7 @@ async def generate_content_async(
  messages=messages,
  tools=tools,
  tool_choice=tool_choice,
- max_tokens=MAX_TOKEN,
+ max_tokens=self.max_tokens,
  )
  yield message_to_generate_content_response(message)
 

tests/unittests/models/test_anthropic_llm.py

@@ -122,3 +122,32 @@ async def mock_coro():
  assert len(responses) == 1
  assert isinstance(responses[0], LlmResponse)
  assert responses[0].content.parts[0].text == "Hello, how can I help you?"
+
+
+@pytest.mark.asyncio
+async def test_generate_content_async_with_max_tokens(
+ llm_request, generate_content_response, generate_llm_response
+):
+ claude_llm = Claude(model="claude-3-5-sonnet-v2@20241022", max_tokens=4096)
+ with mock.patch.object(claude_llm, "_anthropic_client") as mock_client:
+ with mock.patch.object(
+ anthropic_llm,
+ "message_to_generate_content_response",
+ return_value=generate_llm_response,
+ ):
+ # Create a mock coroutine that returns the generate_content_response.
+ async def mock_coro():
+ return generate_content_response
+
+ # Assign the coroutine to the mocked method
+ mock_client.messages.create.return_value = mock_coro()
+
+ _ = [
+ resp
+ async for resp in claude_llm.generate_content_async(
+ llm_request, stream=False
+ )
+ ]
+ mock_client.messages.create.assert_called_once()
+ _, kwargs = mock_client.messages.create.call_args
+ assert kwargs["max_tokens"] == 4096