feat: Allow max tokens to be customizable in Claude

PiperOrigin-RevId: 789901925
This commit is contained in:
Divyansh Shukla
2025-08-01 14:43:57 -07:00
committed by Copybara-Service
parent 2bb20411f4
commit 7556ebc76a
2 changed files with 32 additions and 3 deletions
+3 -3
View File
@@ -46,8 +46,6 @@ __all__ = ["Claude"]
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 @@ class Claude(BaseLlm):
messages=messages,
tools=tools,
tool_choice=tool_choice,
max_tokens=MAX_TOKEN,
max_tokens=self.max_tokens,
)
yield message_to_generate_content_response(message)