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