Skip to content

Commit 8257513

Browse files
authored
fix: None value handling of flattened generation kwargs for AnthropicChatGenerator (#2753)
* fix: None value handling of flattened generation kwargs for AnthropicChatGenerator * Apply suggestion from @tstadel * Apply suggestion from @tstadel * apply feedback
1 parent 697f686 commit 8257513

2 files changed

Lines changed: 71 additions & 22 deletions

File tree

integrations/anthropic/src/haystack_integrations/components/generators/anthropic/chat/chat_generator.py

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -279,27 +279,31 @@ def _prepare_request_params(
279279

280280
def _resolve_flattened_generation_kwargs(self, generation_kwargs: dict[str, Any]) -> dict[str, Any]:
281281
generation_kwargs = generation_kwargs.copy()
282-
if "disable_parallel_tool_use" in generation_kwargs:
283-
disable_parallel_tool_use = generation_kwargs.pop("disable_parallel_tool_use")
284-
tool_choice = generation_kwargs.setdefault("tool_choice", {})
285-
tool_choice["disable_parallel_tool_use"] = disable_parallel_tool_use
286282

287-
if "parallel_tool_use" in generation_kwargs:
288-
parallel_tool_use = generation_kwargs.pop("parallel_tool_use")
283+
disable_parallel_tool_use = generation_kwargs.pop("disable_parallel_tool_use", None)
284+
parallel_tool_use = generation_kwargs.pop("parallel_tool_use", None)
285+
286+
if disable_parallel_tool_use is not None and parallel_tool_use is not None:
287+
msg = "Cannot set both disable_parallel_tool_use and parallel_tool_use"
288+
raise ValueError(msg)
289+
elif parallel_tool_use is not None:
289290
disable_parallel_tool_use = not parallel_tool_use
291+
292+
if disable_parallel_tool_use is not None:
290293
tool_choice = generation_kwargs.setdefault("tool_choice", {})
291294
tool_choice["disable_parallel_tool_use"] = disable_parallel_tool_use
295+
tool_choice.setdefault("type", "auto") # default value
292296

293-
if "tool_choice_type" in generation_kwargs:
294-
tool_choice_type = generation_kwargs.pop("tool_choice_type")
297+
tool_choice_type = generation_kwargs.pop("tool_choice_type", None)
298+
if tool_choice_type is not None:
295299
tool_choice = generation_kwargs.setdefault("tool_choice", {})
296300
tool_choice["type"] = tool_choice_type
297301

298-
if "thinking_budget_tokens" in generation_kwargs:
299-
thinking_budget_tokens = generation_kwargs.pop("thinking_budget_tokens")
302+
thinking_budget_tokens = generation_kwargs.pop("thinking_budget_tokens", None)
303+
if thinking_budget_tokens is not None:
300304
thinking = generation_kwargs.setdefault("thinking", {})
301305
thinking["budget_tokens"] = thinking_budget_tokens
302-
thinking["type"] = "enabled"
306+
thinking.setdefault("type", "enabled")
303307

304308
return generation_kwargs
305309

integrations/anthropic/tests/test_chat_generator.py

Lines changed: 56 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -344,26 +344,71 @@ def test_run_with_params(self, chat_messages, mock_anthropic_completion):
344344
assert response["replies"][0].meta["model"] == "claude-sonnet-4-5"
345345
assert response["replies"][0].meta["finish_reason"] == "stop"
346346

347-
def test_run_with_flattened_generation_kwargs(self, chat_messages, mock_anthropic_completion):
347+
@pytest.mark.parametrize(
348+
"generation_kwargs,expected_kwargs",
349+
[
350+
(
351+
{
352+
"parallel_tool_use": False,
353+
"tool_choice_type": "any",
354+
"thinking_budget_tokens": 1024,
355+
},
356+
{
357+
"tool_choice": {"disable_parallel_tool_use": True, "type": "any"},
358+
"thinking": {"budget_tokens": 1024, "type": "enabled"},
359+
},
360+
),
361+
(
362+
{
363+
"parallel_tool_use": True,
364+
"tool_choice_type": "all",
365+
},
366+
{
367+
"tool_choice": {"disable_parallel_tool_use": False, "type": "all"},
368+
},
369+
),
370+
(
371+
{
372+
"parallel_tool_use": True,
373+
},
374+
{
375+
"tool_choice": {"disable_parallel_tool_use": False, "type": "auto"},
376+
},
377+
),
378+
(
379+
{
380+
"disable_parallel_tool_use": True,
381+
},
382+
{
383+
"tool_choice": {"disable_parallel_tool_use": True, "type": "auto"},
384+
},
385+
),
386+
(
387+
{
388+
"thinking_budget_tokens": None,
389+
"parallel_tool_use": None,
390+
"tool_choice_type": None,
391+
},
392+
{},
393+
),
394+
],
395+
)
396+
def test_run_with_flattened_generation_kwargs(
397+
self, chat_messages, mock_anthropic_completion, generation_kwargs, expected_kwargs
398+
):
348399
"""
349400
Test that the AnthropicChatGenerator component can run with parameters.
350401
"""
351402
component = AnthropicChatGenerator(
352403
api_key=Secret.from_token("test-api-key"),
353-
generation_kwargs={
354-
"max_tokens": 10,
355-
"thinking_budget_tokens": 1024,
356-
"parallel_tool_use": False,
357-
"tool_choice_type": "any",
358-
},
404+
generation_kwargs=generation_kwargs,
359405
)
360406
component.run(chat_messages)
361407

362408
# Check that the component calls the Anthropic API with the correct parameters
363-
_, kwargs = mock_anthropic_completion.call_args
364-
assert kwargs["max_tokens"] == 10
365-
assert kwargs["thinking"] == {"budget_tokens": 1024, "type": "enabled"}
366-
assert kwargs["tool_choice"] == {"disable_parallel_tool_use": True, "type": "any"}
409+
actual_kwargs = mock_anthropic_completion.call_args.kwargs
410+
assert actual_kwargs.get("tool_choice") == expected_kwargs.get("tool_choice")
411+
assert actual_kwargs.get("thinking") == expected_kwargs.get("thinking")
367412

368413
def test_check_duplicate_tool_names(self, tools):
369414
"""Test that the AnthropicChatGenerator component fails to initialize with duplicate tool names."""

0 commit comments

Comments
 (0)