diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py index 6f68cf7b8cbc..2e248a8652e4 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py @@ -540,14 +540,16 @@ def _get_thinking_config(self, extra_create_args: Mapping[str, Any]) -> Dict[str def _rstrip_last_assistant_message(self, messages: Sequence[LLMMessage]) -> Sequence[LLMMessage]: """ - Remove the last assistant message if it is empty. + Remove trailing assistant messages before sending Anthropic requests. """ - # When Claude models last message is AssistantMessage, It could not end with whitespace - if isinstance(messages[-1], AssistantMessage): - if isinstance(messages[-1].content, str): - messages[-1].content = messages[-1].content.rstrip() + trimmed_messages = list(messages) + while trimmed_messages and isinstance(trimmed_messages[-1], AssistantMessage): + trimmed_messages.pop() - return messages + if len(trimmed_messages) == len(messages): + return messages + + return trimmed_messages async def create( self, diff --git a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py index 6238eee3bbe5..99b50fd3023d 100644 --- a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py @@ -800,8 +800,8 @@ async def test_anthropic_empty_assistant_content_string() -> None: @pytest.mark.asyncio -async def test_anthropic_trailing_whitespace_at_last_assistant_content() -> None: - """Test that an empty assistant content string is handled correctly.""" +async def test_anthropic_trailing_assistant_message() -> None: + """Test that a trailing assistant message is handled correctly.""" api_key = os.getenv("ANTHROPIC_API_KEY") if not api_key: pytest.skip("ANTHROPIC_API_KEY not found in environment variables") @@ -821,19 +821,35 @@ async def test_anthropic_trailing_whitespace_at_last_assistant_content() -> None assert isinstance(result.content, str) -def test_mock_rstrip_trailing_whitespace_at_last_assistant_content() -> None: +def test_mock_rstrip_removes_trailing_assistant_messages() -> None: messages: list[LLMMessage] = [ UserMessage(content="foo", source="user"), UserMessage(content="bar", source="user"), AssistantMessage(content="foobar ", source="assistant"), + AssistantMessage(content="baz", source="assistant"), + ] + + dummy_client = AnthropicChatCompletionClient(model="claude-3-5-haiku-20241022", api_key="dummy-key") + result = dummy_client._rstrip_last_assistant_message(messages) # pyright: ignore[reportPrivateUsage] + + assert result == messages[:2] + assert len(messages) == 4 + assert messages[-1].content == "baz" + + +def test_mock_rstrip_preserves_non_trailing_assistant_messages() -> None: + messages: list[LLMMessage] = [ + UserMessage(content="foo", source="user"), + AssistantMessage(content="bar", source="assistant"), + UserMessage(content="baz", source="user"), ] - # This will crash if _rstrip_railing_whitespace_at_last_assistant_content is not applied to "content" dummy_client = AnthropicChatCompletionClient(model="claude-3-5-haiku-20241022", api_key="dummy-key") result = dummy_client._rstrip_last_assistant_message(messages) # pyright: ignore[reportPrivateUsage] - assert isinstance(result[-1].content, str) - assert result[-1].content == "foobar" + assert list(result) == messages + assert result[-1].content == "baz" + assert result[-2].content == "bar" @pytest.mark.asyncio