From e5cb30f59e4a5fc62499b71797cb44d451b8e3ca Mon Sep 17 00:00:00 2001 From: ZakWork Date: Sun, 16 Mar 2025 01:49:01 +0000 Subject: [PATCH 1/6] fixex R1 reasoning parser for openai client --- .../models/openai/_openai_client.py | 36 +++++++++++-------- 1 file changed, 22 insertions(+), 14 deletions(-) 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 2301655b32d1..b02a32dd7aae 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 @@ -617,6 +617,8 @@ async def create( else: finish_reason = choice.finish_reason content = choice.message.content or "" + if choice.message.model_extra['reasoning_content'] is not None: + thought = choice.message.model_extra['reasoning_content'] logprobs: Optional[List[ChatCompletionTokenLogprob]] = None if choice.logprobs and choice.logprobs.content: @@ -630,8 +632,7 @@ async def create( for x in choice.logprobs.content ] - if isinstance(content, str) and self._model_info["family"] == ModelFamily.R1: - thought, content = parse_r1_content(content) + response = CreateResult( finish_reason=normalize_stop_reason(finish_reason), @@ -725,6 +726,7 @@ async def create_stream( stop_reason = None maybe_model = None content_deltas: List[str] = [] + thought_deltas: List[str] = [] full_tool_calls: Dict[int, FunctionCall] = {} completion_tokens = 0 logprobs: Optional[List[ChatCompletionTokenLogprob]] = None @@ -784,7 +786,10 @@ async def create_stream( # NOTE: for OpenAI, tool_calls and content are mutually exclusive it seems, so we can skip the rest of the loop. # However, this may not be the case for other APIs -- we should expect this may need to be updated. continue - + #thought: str | None = None + if choice.delta.model_extra['reasoning_content'] is not None: + thought_deltas.append(choice.delta.model_extra['reasoning_content']) + yield choice.delta.model_extra['reasoning_content'] # Otherwise, get tool calls if choice.delta.tool_calls is not None: for tool_call_chunk in choice.delta.tool_calls: @@ -837,22 +842,25 @@ async def create_stream( # Detect whether it is a function call or just text. content: Union[str, List[FunctionCall]] thought: str | None = None + # Determine the content and thought based on what was collected if full_tool_calls: - # This is a tool call. + # This is a tool call response content = list(full_tool_calls.values()) - if len(content_deltas) > 1: - # Put additional text content in the thought field. + if content_deltas: + # Store any text alongside tool calls as thoughts thought = "".join(content_deltas) - elif len(content_deltas) > 0: - # This is a text-only content. - content = "".join(content_deltas) else: - warnings.warn("No text content or tool calls are available. Model returned empty result.", stacklevel=2) - content = "" + # This is a text response (possibly with thoughts) + if content_deltas: + content = "".join(content_deltas) + else: + warnings.warn("No text content or tool calls are available. Model returned empty result.", stacklevel=2) + content = "" + + # Always set thoughts if we have any, regardless of other content types + if thought_deltas: + thought = "".join(thought_deltas) - # Parse R1 content if needed. - if isinstance(content, str) and self._model_info["family"] == ModelFamily.R1: - thought, content = parse_r1_content(content) # Create the result. result = CreateResult( From c45d19d76c0f54c56759e10b74fc40f0e8d0d631 Mon Sep 17 00:00:00 2001 From: ZakWork Date: Sun, 16 Mar 2025 19:17:46 +0000 Subject: [PATCH 2/6] Fixed issues with dict and restored local R1 features --- .../models/openai/_openai_client.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) 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 b02a32dd7aae..e602726d7265 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 @@ -615,10 +615,12 @@ async def create( ) finish_reason = "tool_calls" else: + # if not tool_calls, then it is a text response and we populate the content and thought fields. finish_reason = choice.finish_reason content = choice.message.content or "" - if choice.message.model_extra['reasoning_content'] is not None: - thought = choice.message.model_extra['reasoning_content'] + # if there is a reasoning_content field, then we populate the thought field. This is for models such as R1. + if choice.message.model_extra.get('reasoning_content') is not None: + thought = choice.message.model_extra.get('reasoning_content') logprobs: Optional[List[ChatCompletionTokenLogprob]] = None if choice.logprobs and choice.logprobs.content: @@ -632,6 +634,9 @@ async def create( for x in choice.logprobs.content ] + # This is for local R1 models. + if isinstance(content, str) and self._model_info["family"] == ModelFamily.R1 and thought is None: + thought, content = parse_r1_content(content) response = CreateResult( @@ -786,10 +791,11 @@ async def create_stream( # NOTE: for OpenAI, tool_calls and content are mutually exclusive it seems, so we can skip the rest of the loop. # However, this may not be the case for other APIs -- we should expect this may need to be updated. continue - #thought: str | None = None - if choice.delta.model_extra['reasoning_content'] is not None: - thought_deltas.append(choice.delta.model_extra['reasoning_content']) - yield choice.delta.model_extra['reasoning_content'] + # if there is a reasoning_content field, then we populate the thought field. This is for models such as R1. + reasoning_content = choice.delta.model_extra.get('reasoning_content') + if reasoning_content is not None: + thought_deltas.append(reasoning_content) + yield reasoning_content # Otherwise, get tool calls if choice.delta.tool_calls is not None: for tool_call_chunk in choice.delta.tool_calls: @@ -860,7 +866,10 @@ async def create_stream( # Always set thoughts if we have any, regardless of other content types if thought_deltas: thought = "".join(thought_deltas) - + + # This is for local R1 models. + if isinstance(content, str) and self._model_info["family"] == ModelFamily.R1 and thought is None: + thought, content = parse_r1_content(content) # Create the result. result = CreateResult( From ad45981ada8e0fd2a6d8cf98985c9292f41192ec Mon Sep 17 00:00:00 2001 From: ZakWork Date: Sun, 16 Mar 2025 22:25:15 +0000 Subject: [PATCH 3/6] fixed issues with formating and none reference --- .../models/openai/_openai_client.py | 65 +++++++++++-------- 1 file changed, 39 insertions(+), 26 deletions(-) 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 e602726d7265..0bef7ecc4df2 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 @@ -382,7 +382,11 @@ def __init__( elif model_capabilities is not None and model_info is not None: raise ValueError("model_capabilities and model_info are mutually exclusive") elif model_capabilities is not None and model_info is None: - warnings.warn("model_capabilities is deprecated, use model_info instead", DeprecationWarning, stacklevel=2) + warnings.warn( + "model_capabilities is deprecated, use model_info instead", + DeprecationWarning, + stacklevel=2, + ) info = cast(ModelInfo, model_capabilities) info["family"] = ModelFamily.UNKNOWN self._model_info = info @@ -528,7 +532,7 @@ async def create( future = asyncio.ensure_future( self._client.beta.chat.completions.parse( messages=create_params.messages, - tools=create_params.tools if len(create_params.tools) > 0 else NOT_GIVEN, + tools=(create_params.tools if len(create_params.tools) > 0 else NOT_GIVEN), response_format=create_params.response_format, **create_params.create_args, ) @@ -539,7 +543,7 @@ async def create( self._client.chat.completions.create( messages=create_params.messages, stream=False, - tools=create_params.tools if len(create_params.tools) > 0 else NOT_GIVEN, + tools=(create_params.tools if len(create_params.tools) > 0 else NOT_GIVEN), **create_params.create_args, ) ) @@ -618,9 +622,11 @@ async def create( # if not tool_calls, then it is a text response and we populate the content and thought fields. finish_reason = choice.finish_reason content = choice.message.content or "" - # if there is a reasoning_content field, then we populate the thought field. This is for models such as R1. - if choice.message.model_extra.get('reasoning_content') is not None: - thought = choice.message.model_extra.get('reasoning_content') + # if there is a reasoning_content field, then we populate the thought field. This is for models such as R1 - direct from deepseek api. + if choice.message.model_extra is not None: + reasoning_content = choice.message.model_extra.get("reasoning_content") + if reasoning_content is not None: + thought = reasoning_content logprobs: Optional[List[ChatCompletionTokenLogprob]] = None if choice.logprobs and choice.logprobs.content: @@ -634,10 +640,9 @@ async def create( for x in choice.logprobs.content ] - # This is for local R1 models. + # This is for local R1 models. if isinstance(content, str) and self._model_info["family"] == ModelFamily.R1 and thought is None: thought, content = parse_r1_content(content) - response = CreateResult( finish_reason=normalize_stop_reason(finish_reason), @@ -774,9 +779,7 @@ async def create_stream( choice = ( chunk.choices[0] if len(chunk.choices) > 0 - else choice - if chunk.usage is not None and stop_reason is not None - else cast(ChunkChoice, None) + else (choice if chunk.usage is not None and stop_reason is not None else cast(ChunkChoice, None)) ) # for liteLLM chunk usage, do the following hack keeping the pervious chunk.stop_reason (if set). @@ -791,11 +794,12 @@ async def create_stream( # NOTE: for OpenAI, tool_calls and content are mutually exclusive it seems, so we can skip the rest of the loop. # However, this may not be the case for other APIs -- we should expect this may need to be updated. continue - # if there is a reasoning_content field, then we populate the thought field. This is for models such as R1. - reasoning_content = choice.delta.model_extra.get('reasoning_content') - if reasoning_content is not None: - thought_deltas.append(reasoning_content) - yield reasoning_content + # if there is a reasoning_content field, then we populate the thought field. This is for models such as R1. + if choice.delta.model_extra is not None: + reasoning_content = choice.delta.model_extra.get("reasoning_content") + if reasoning_content is not None: + thought_deltas.append(reasoning_content) + yield reasoning_content # Otherwise, get tool calls if choice.delta.tool_calls is not None: for tool_call_chunk in choice.delta.tool_calls: @@ -860,14 +864,17 @@ async def create_stream( if content_deltas: content = "".join(content_deltas) else: - warnings.warn("No text content or tool calls are available. Model returned empty result.", stacklevel=2) + warnings.warn( + "No text content or tool calls are available. Model returned empty result.", + stacklevel=2, + ) content = "" - + # Always set thoughts if we have any, regardless of other content types if thought_deltas: thought = "".join(thought_deltas) - - # This is for local R1 models. + + # This is for local R1 models. if isinstance(content, str) and self._model_info["family"] == ModelFamily.R1 and thought is None: thought, content = parse_r1_content(content) @@ -936,7 +943,7 @@ async def _create_stream_chunks_beta_client( async with self._client.beta.chat.completions.stream( messages=oai_messages, tools=tool_params if len(tool_params) > 0 else NOT_GIVEN, - response_format=response_format if response_format is not None else NOT_GIVEN, + response_format=(response_format if response_format is not None else NOT_GIVEN), **create_args_no_response_format, ) as stream: while True: @@ -989,9 +996,9 @@ def count_tokens(self, messages: Sequence[LLMMessage], *, tools: Sequence[Tool | if isinstance(message, UserMessage) and isinstance(value, list): typed_message_value = cast(List[ChatCompletionContentPartParam], value) - assert len(typed_message_value) == len( - message.content - ), "Mismatch in message content and typed message value" + assert len(typed_message_value) == len(message.content), ( + "Mismatch in message content and typed message value" + ) # We need image properties that are only in the original message for part, content_part in zip(typed_message_value, message.content, strict=False): @@ -1061,7 +1068,11 @@ def remaining_tokens(self, messages: Sequence[LLMMessage], *, tools: Sequence[To @property def capabilities(self) -> ModelCapabilities: # type: ignore - warnings.warn("capabilities is deprecated, use model_info instead", DeprecationWarning, stacklevel=2) + warnings.warn( + "capabilities is deprecated, use model_info instead", + DeprecationWarning, + stacklevel=2, + ) return self._model_info @property @@ -1166,7 +1177,9 @@ class OpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component[OpenA async def main() -> None: # Similar for AzureOpenAIChatCompletionClient. - model_client = OpenAIChatCompletionClient(model="gpt-4o") # assuming OPENAI_API_KEY is set in the environment. + model_client = OpenAIChatCompletionClient( + model="gpt-4o" + ) # assuming OPENAI_API_KEY is set in the environment. messages = [UserMessage(content="Write a very short story about a dragon.", source="user")] From 3dd650983b2b96aa962d57095f4f55a59f808de9 Mon Sep 17 00:00:00 2001 From: ZakWork Date: Mon, 17 Mar 2025 02:48:08 +0000 Subject: [PATCH 4/6] Update test_openai_model_client.py Added tests for the handling of reasoning_content for R1 model --- .../tests/models/test_openai_model_client.py | 125 +++++++++++++++++- 1 file changed, 122 insertions(+), 3 deletions(-) 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 db8dbe1da968..a08e172424ad 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 @@ -21,9 +21,6 @@ ) from autogen_core.models._model_client import ModelFamily from autogen_core.tools import BaseTool, FunctionTool -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 calculate_vision_tokens, convert_tools, to_oai_type from openai.resources.beta.chat.completions import ( # type: ignore AsyncChatCompletionStreamManager as BetaAsyncChatCompletionStreamManager, # type: ignore ) @@ -53,6 +50,14 @@ from openai.types.completion_usage import CompletionUsage from pydantic import BaseModel, Field +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 ( + calculate_vision_tokens, + convert_tools, + to_oai_type, +) + ResponseFormatT = TypeVar("ResponseFormatT", bound=BaseModel) @@ -886,6 +891,120 @@ async def _stream() -> AsyncGenerator[MockChunkEvent, None]: assert response.response == "happy" +@pytest.mark.asyncio +async def test_r1_reasoning_content(monkeypatch: pytest.MonkeyPatch) -> None: + """Test handling of reasoning_content in R1 model. Testing create without streaming.""" + + async def _mock_create(*args: Any, **kwargs: Any) -> ChatCompletion: + return ChatCompletion( + id="test_id", + model="r1", + object="chat.completion", + created=1234567890, + choices=[ + Choice( + index=0, + message=ChatCompletionMessage( + role="assistant", + content="This is the main content", + reasoning_content="This is the reasoning content", + ), + finish_reason="stop", + ) + ], + usage=CompletionUsage( + prompt_tokens=10, + completion_tokens=10, + total_tokens=20, + ), + ) + + # Patch the client creation + + monkeypatch.setattr(AsyncCompletions, "create", _mock_create) + + # Create the client + model_client = OpenAIChatCompletionClient( + model="r1", + api_key="", + model_info={ + "family": ModelFamily.R1, + "vision": False, + "function_calling": False, + "json_output": False, + "structured_output": False, + }, + ) + + # Test the create method + result = await model_client.create([UserMessage(content="Test message", source="user")]) + + # Verify that the content and thought are as expected + assert result.content == "This is the main content" + assert result.thought == "This is the reasoning content" + + +@pytest.mark.asyncio +async def test_r1_reasoning_content_streaming(monkeypatch: pytest.MonkeyPatch) -> None: + """Test that reasoning_content in model_extra is correctly extracted and streamed.""" + + async def _mock_create_stream(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatCompletionChunk, None]: + contentChunks = [None, None, "This is the main content"] + reasoningChunks = ["This is the reasoning content 1", "This is the reasoning content 2", None] + for i in range(len(contentChunks)): + await asyncio.sleep(0.1) + yield ChatCompletionChunk( + id="id", + choices=[ + ChunkChoice( + finish_reason="stop" if i == len(contentChunks) - 1 else None, + index=0, + delta=ChoiceDelta( + content=contentChunks[i], + reasoning_content=reasoningChunks[i], + role="assistant", + ), + ), + ], + created=0, + model="r1", + object="chat.completion.chunk", + usage=CompletionUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0), + ) + + async def _mock_create(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatCompletionChunk, None]: + return _mock_create_stream(*args, **kwargs) + + # Patch the client creation + monkeypatch.setattr(AsyncCompletions, "create", _mock_create) + # Create the client + model_client = OpenAIChatCompletionClient( + model="r1", + api_key="", + model_info={ + "family": ModelFamily.R1, + "vision": False, + "function_calling": False, + "json_output": False, + "structured_output": False, + }, + ) + # Test the create_stream method + chunks: List[str | CreateResult] = [] + async for chunk in model_client.create_stream(messages=[UserMessage(content="Hello", source="user")]): + chunks.append(chunk) + + # Verify that the chunks first stream the reasoning content and then the main content + # Then verify that the final result has the correct content and thought + assert len(chunks) == 4 + assert chunks[0] == "This is the reasoning content 1" + assert chunks[1] == "This is the reasoning content 2" + assert chunks[2] == "This is the main content" + assert isinstance(chunks[3], CreateResult) + assert chunks[3].content == "This is the main content" + assert chunks[3].thought == "This is the reasoning content 1This is the reasoning content 2" + + @pytest.mark.asyncio async def test_r1_think_field(monkeypatch: pytest.MonkeyPatch) -> None: async def _mock_create_stream(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatCompletionChunk, None]: From 8330523a986f5238087782fc2af6a05e69b67fdb Mon Sep 17 00:00:00 2001 From: ZakWork Date: Mon, 17 Mar 2025 08:01:28 +0000 Subject: [PATCH 5/6] Fixed format issues using uv and poe format --- .../autogen_ext/models/openai/_openai_client.py | 10 ++++------ .../tests/models/test_openai_model_client.py | 15 +++++++-------- 2 files changed, 11 insertions(+), 14 deletions(-) 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 0bef7ecc4df2..3acedc93e1d8 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 @@ -996,9 +996,9 @@ def count_tokens(self, messages: Sequence[LLMMessage], *, tools: Sequence[Tool | if isinstance(message, UserMessage) and isinstance(value, list): typed_message_value = cast(List[ChatCompletionContentPartParam], value) - assert len(typed_message_value) == len(message.content), ( - "Mismatch in message content and typed message value" - ) + assert len(typed_message_value) == len( + message.content + ), "Mismatch in message content and typed message value" # We need image properties that are only in the original message for part, content_part in zip(typed_message_value, message.content, strict=False): @@ -1177,9 +1177,7 @@ class OpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component[OpenA async def main() -> None: # Similar for AzureOpenAIChatCompletionClient. - model_client = OpenAIChatCompletionClient( - model="gpt-4o" - ) # assuming OPENAI_API_KEY is set in the environment. + model_client = OpenAIChatCompletionClient(model="gpt-4o") # assuming OPENAI_API_KEY is set in the environment. messages = [UserMessage(content="Write a very short story about a dragon.", source="user")] 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 a08e172424ad..9f4765a28944 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 @@ -21,6 +21,13 @@ ) from autogen_core.models._model_client import ModelFamily from autogen_core.tools import BaseTool, FunctionTool +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 ( + calculate_vision_tokens, + convert_tools, + to_oai_type, +) from openai.resources.beta.chat.completions import ( # type: ignore AsyncChatCompletionStreamManager as BetaAsyncChatCompletionStreamManager, # type: ignore ) @@ -50,14 +57,6 @@ from openai.types.completion_usage import CompletionUsage from pydantic import BaseModel, Field -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 ( - calculate_vision_tokens, - convert_tools, - to_oai_type, -) - ResponseFormatT = TypeVar("ResponseFormatT", bound=BaseModel) From f9f48a0cb5afb6593dfb73d66b8dd318f36d3de3 Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Mon, 17 Mar 2025 10:02:40 -0700 Subject: [PATCH 6/6] fix type checking --- .../autogen-ext/tests/models/test_openai_model_client.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 9f4765a28944..885ab36a5567 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 @@ -906,7 +906,8 @@ async def _mock_create(*args: Any, **kwargs: Any) -> ChatCompletion: message=ChatCompletionMessage( role="assistant", content="This is the main content", - reasoning_content="This is the reasoning content", + # The reasoning content is included in model_extra for hosted R1 models. + reasoning_content="This is the reasoning content", # type: ignore ), finish_reason="stop", ) @@ -960,7 +961,8 @@ async def _mock_create_stream(*args: Any, **kwargs: Any) -> AsyncGenerator[ChatC index=0, delta=ChoiceDelta( content=contentChunks[i], - reasoning_content=reasoningChunks[i], + # The reasoning content is included in model_extra for hosted R1 models. + reasoning_content=reasoningChunks[i], # type: ignore role="assistant", ), ),