diff --git a/python/packages/autogen-ext/src/autogen_ext/models/_utils/rstrip_last_assistant_message.py b/python/packages/autogen-ext/src/autogen_ext/models/_utils/rstrip_last_assistant_message.py new file mode 100644 index 000000000000..730643cce1fb --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/models/_utils/rstrip_last_assistant_message.py @@ -0,0 +1,21 @@ +from typing import Sequence + +from autogen_core.models import AssistantMessage, LLMMessage + + +def rstrip_last_assistant_message(messages: Sequence[LLMMessage]) -> Sequence[LLMMessage]: + """ + Remove trailing whitespace from the last assistant message, dropping the + message entirely if that leaves it empty. Some providers (e.g. Anthropic) + reject text content blocks that are empty or end in whitespace. + + Only the trailing message is affected; earlier messages in the sequence + are left untouched even if they are also `AssistantMessage`s. + """ + if messages and isinstance(messages[-1], AssistantMessage): + if isinstance(messages[-1].content, str): + messages[-1].content = messages[-1].content.rstrip() + if messages[-1].content == "": + messages = messages[:-1] + + return messages 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..827d4721b2d6 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 @@ -65,6 +65,7 @@ from pydantic import BaseModel, SecretStr from typing_extensions import Self, Unpack +from .._utils.rstrip_last_assistant_message import rstrip_last_assistant_message from . import _model_info from .config import ( AnthropicBedrockClientConfiguration, @@ -538,17 +539,6 @@ def _get_thinking_config(self, extra_create_args: Mapping[str, Any]) -> Dict[str return {"thinking": thinking_config} - def _rstrip_last_assistant_message(self, messages: Sequence[LLMMessage]) -> Sequence[LLMMessage]: - """ - Remove the last assistant message if it is empty. - """ - # 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() - - return messages - async def create( self, messages: Sequence[LLMMessage], @@ -586,7 +576,7 @@ async def create( # Merge continuous system messages into a single message messages = self._merge_system_messages(messages) - messages = self._rstrip_last_assistant_message(messages) + messages = rstrip_last_assistant_message(messages) for message in messages: if isinstance(message, SystemMessage): @@ -804,7 +794,7 @@ async def create_stream( # Merge continuous system messages into a single message messages = self._merge_system_messages(messages) - messages = self._rstrip_last_assistant_message(messages) + messages = rstrip_last_assistant_message(messages) for message in messages: if isinstance(message, SystemMessage): diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py index a80e912534ab..8dbad5a5924d 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py @@ -75,6 +75,7 @@ from .._utils.normalize_stop_reason import normalize_stop_reason from .._utils.parse_r1_content import parse_r1_content +from .._utils.rstrip_last_assistant_message import rstrip_last_assistant_message from . import _model_info from ._transformation import ( get_transformer, @@ -487,17 +488,6 @@ def __init__( def create_from_config(cls, config: Dict[str, Any]) -> ChatCompletionClient: return OpenAIChatCompletionClient(**config) - def _rstrip_last_assistant_message(self, messages: Sequence[LLMMessage]) -> Sequence[LLMMessage]: - """ - Remove the last assistant message if it is empty. - """ - # 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() - - return messages - def _process_create_args( self, messages: Sequence[LLMMessage], @@ -610,7 +600,7 @@ def _process_create_args( # in that case, for ad-hoc, we using startswith instead of model_family for code consistency if create_args.get("model", "unknown").startswith("claude-"): # When Claude models last message is AssistantMessage, It could not end with whitespace - messages = self._rstrip_last_assistant_message(messages) + messages = rstrip_last_assistant_message(messages) oai_messages_nested = [ to_oai_type( 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..8da79b2951f1 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 @@ -18,6 +18,7 @@ ) from autogen_core.models._types import LLMMessage from autogen_core.tools import FunctionTool +from autogen_ext.models._utils.rstrip_last_assistant_message import rstrip_last_assistant_message from autogen_ext.models.anthropic import ( AnthropicBedrockChatCompletionClient, AnthropicChatCompletionClient, @@ -828,14 +829,48 @@ def test_mock_rstrip_trailing_whitespace_at_last_assistant_content() -> None: AssistantMessage(content="foobar ", source="assistant"), ] - # 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] + # This will crash if rstrip_last_assistant_message is not applied to "content" + result = rstrip_last_assistant_message(messages) assert isinstance(result[-1].content, str) assert result[-1].content == "foobar" +def test_mock_rstrip_removes_whitespace_only_last_assistant_message() -> None: + """If the last assistant message is whitespace-only, rstrip leaves an empty + string, which Anthropic's API rejects (text content blocks must be non-empty). + The message should be dropped entirely instead.""" + messages: list[LLMMessage] = [ + UserMessage(content="foo", source="user"), + UserMessage(content="bar", source="user"), + AssistantMessage(content=" ", source="assistant"), + ] + + result = rstrip_last_assistant_message(messages) + + assert len(result) == 2 + assert isinstance(result[-1], UserMessage) + + +def test_mock_rstrip_only_affects_trailing_assistant_message() -> None: + """An interleaved (non-trailing) assistant message with trailing whitespace + must be left untouched, even if it would otherwise be empty after rstrip.""" + messages: list[LLMMessage] = [ + UserMessage(content="foo", source="user"), + AssistantMessage(content=" ", source="assistant"), + UserMessage(content="bar", source="user"), + AssistantMessage(content="baz ", source="assistant"), + ] + + result = rstrip_last_assistant_message(messages) + + assert len(result) == 4 + assert isinstance(result[1], AssistantMessage) + assert result[1].content == " " + assert isinstance(result[-1], AssistantMessage) + assert result[-1].content == "baz" + + @pytest.mark.asyncio async def test_anthropic_tool_choice_with_actual_api() -> None: """Test tool_choice parameter with actual Anthropic API endpoints.""" diff --git a/python/packages/autogen-ext/tests/models/test_openai_model_client.py b/python/packages/autogen-ext/tests/models/test_openai_model_client.py index ba79795d1ed7..c78d51b38da4 100644 --- a/python/packages/autogen-ext/tests/models/test_openai_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_openai_model_client.py @@ -23,6 +23,7 @@ ) from autogen_core.models._model_client import ModelFamily from autogen_core.tools import BaseTool, FunctionTool +from autogen_ext.models._utils.rstrip_last_assistant_message import rstrip_last_assistant_message from autogen_ext.models.openai import AzureOpenAIChatCompletionClient, OpenAIChatCompletionClient from autogen_ext.models.openai._model_info import resolve_model from autogen_ext.models.openai._openai_client import ( @@ -2635,14 +2636,29 @@ def test_rstrip_railing_whitespace_at_last_assistant_content() -> None: AssistantMessage(content="foobar ", source="assistant"), ] - # This will crash if _rstrip_railing_whitespace_at_last_assistant_content is not applied to "content" - dummy_client = OpenAIChatCompletionClient(model="claude-3-5-haiku-20241022", api_key="dummy-key") - result = dummy_client._rstrip_last_assistant_message(messages) # pyright: ignore[reportPrivateUsage] + # This will crash if rstrip_last_assistant_message is not applied to "content" + result = rstrip_last_assistant_message(messages) assert isinstance(result[-1].content, str) assert result[-1].content == "foobar" +def test_rstrip_removes_whitespace_only_last_assistant_message() -> None: + """If the last assistant message is whitespace-only, rstrip leaves an empty + string, which Anthropic's API rejects (text content blocks must be non-empty). + The message should be dropped entirely instead.""" + messages: list[LLMMessage] = [ + UserMessage(content="foo", source="user"), + UserMessage(content="bar", source="user"), + AssistantMessage(content=" ", source="assistant"), + ] + + result = rstrip_last_assistant_message(messages) + + assert len(result) == 2 + assert isinstance(result[-1], UserMessage) + + def test_find_model_family() -> None: assert _find_model_family("openai", "gpt-4") == ModelFamily.GPT_4 assert _find_model_family("openai", "gpt-4-latest") == ModelFamily.GPT_4