diff --git a/python/packages/core/agent_framework/openai/__init__.py b/python/packages/core/agent_framework/openai/__init__.py index 949e9cd5edc..f8cb9e140e8 100644 --- a/python/packages/core/agent_framework/openai/__init__.py +++ b/python/packages/core/agent_framework/openai/__init__.py @@ -21,6 +21,8 @@ "RawOpenAIChatClient": ("agent_framework_openai", "agent-framework-openai"), "OpenAIChatCompletionClient": ("agent_framework_openai", "agent-framework-openai"), "OpenAIChatCompletionOptions": ("agent_framework_openai", "agent-framework-openai"), + "OpenAIChatMessagePreparer": ("agent_framework_openai", "agent-framework-openai"), + "OpenAIChatResponseContentsParser": ("agent_framework_openai", "agent-framework-openai"), "RawOpenAIChatCompletionClient": ("agent_framework_openai", "agent-framework-openai"), "OpenAIEmbeddingClient": ("agent_framework_openai", "agent-framework-openai"), "OpenAIEmbeddingOptions": ("agent_framework_openai", "agent-framework-openai"), diff --git a/python/packages/core/agent_framework/openai/__init__.pyi b/python/packages/core/agent_framework/openai/__init__.pyi index e2c9a29ef83..80032744961 100644 --- a/python/packages/core/agent_framework/openai/__init__.pyi +++ b/python/packages/core/agent_framework/openai/__init__.pyi @@ -8,7 +8,9 @@ from agent_framework_openai import ( OpenAIChatClient, OpenAIChatCompletionClient, OpenAIChatCompletionOptions, + OpenAIChatMessagePreparer, OpenAIChatOptions, + OpenAIChatResponseContentsParser, OpenAIContentFilterException, OpenAIContinuationToken, OpenAIEmbeddingClient, @@ -23,7 +25,9 @@ __all__ = [ "OpenAIChatClient", "OpenAIChatCompletionClient", "OpenAIChatCompletionOptions", + "OpenAIChatMessagePreparer", "OpenAIChatOptions", + "OpenAIChatResponseContentsParser", "OpenAIContentFilterException", "OpenAIContinuationToken", "OpenAIEmbeddingClient", diff --git a/python/packages/openai/AGENTS.md b/python/packages/openai/AGENTS.md index 1dd22ad0c97..d242eece65d 100644 --- a/python/packages/openai/AGENTS.md +++ b/python/packages/openai/AGENTS.md @@ -33,6 +33,31 @@ The generic OpenAI clients support both OpenAI and Azure OpenAI routing. Precede explicit Azure inputs (`credential`, `azure_endpoint`, `api_version`) → OpenAI API key (`OPENAI_API_KEY`) → Azure environment fallback (`AZURE_OPENAI_*`). +## Adapting the Chat Completions client to OpenAI-compatible endpoints + +`OpenAIChatCompletionClient` targets the OpenAI Chat Completions wire format and is intentionally +kept free of provider-specific quirks. Many "OpenAI-compatible" providers (OpenRouter, vLLM, +Mistral, DeepSeek, Ollama, …) diverge on the edges — e.g. returning reasoning under +`reasoning` / `reasoning_content` / `reasoning_details`, or `content` as a list of chunks. Rather +than branching in core, the client exposes two optional callables so callers adapt it themselves: + +- `response_parser: OpenAIChatResponseContentsParser` — `(message_or_delta, default_contents) -> contents`. + Post-processes the `Content` items parsed from each response choice. Receives the already-selected + `ChatCompletionMessage` (non-streaming) or `ChoiceDelta` (streaming) — the client resolves the + dispatch — so a parser reads provider fields directly (e.g. `getattr(msg, "reasoning", None)`) without + branching. Use it to surface non-standard fields for display. Applied per choice in both paths. +- `message_preparer: OpenAIChatMessagePreparer` — `(message, default_dicts) -> dicts`. Post-processes + the outgoing request message dicts built from each framework `Message` (called once per `Message`, for + every role including `system`/`developer`). Use it to echo provider-specific fields (e.g. vLLM + `reasoning`) back on later turns for multi-turn continuity. To correlate a surfaced-reasoning `Content` + with the dict the default serializer emitted for it, tag the `Content` via `additional_properties` in + the parser and match against `message.contents` rather than raw request-string matching. + +Both default to `None` (no-op → byte-identical stock OpenAI behavior) and are constructor args on +`RawOpenAIChatCompletionClient` / `OpenAIChatCompletionClient`. Provider round-trips generally need +**both**: the parser surfaces the field for display, the preparer sends it back. Prefer a dedicated +client (e.g. `agent-framework-mistral`) when an endpoint diverges substantially. + ## Dependencies - `agent-framework-core` — core abstractions diff --git a/python/packages/openai/agent_framework_openai/__init__.py b/python/packages/openai/agent_framework_openai/__init__.py index 68e00c39db2..08f93991a34 100644 --- a/python/packages/openai/agent_framework_openai/__init__.py +++ b/python/packages/openai/agent_framework_openai/__init__.py @@ -17,6 +17,8 @@ from ._chat_completion_client import ( OpenAIChatCompletionClient, OpenAIChatCompletionOptions, + OpenAIChatMessagePreparer, + OpenAIChatResponseContentsParser, RawOpenAIChatCompletionClient, ) from ._embedding_client import OpenAIEmbeddingClient, OpenAIEmbeddingOptions @@ -33,7 +35,9 @@ "OpenAIChatClient", "OpenAIChatCompletionClient", "OpenAIChatCompletionOptions", + "OpenAIChatMessagePreparer", "OpenAIChatOptions", + "OpenAIChatResponseContentsParser", "OpenAIContentFilterException", "OpenAIContinuationToken", "OpenAIEmbeddingClient", diff --git a/python/packages/openai/agent_framework_openai/_chat_completion_client.py b/python/packages/openai/agent_framework_openai/_chat_completion_client.py index a28b34e1ce0..95ff1b4d3d0 100644 --- a/python/packages/openai/agent_framework_openai/_chat_completion_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_completion_client.py @@ -16,7 +16,7 @@ ) from datetime import datetime, timezone from itertools import chain -from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeAlias, cast, overload from agent_framework._clients import BaseChatClient from agent_framework._compaction import CompactionStrategy, TokenizerProtocol @@ -50,8 +50,9 @@ from openai.lib._parsing._completions import type_to_response_format_param from openai.types import CompletionUsage from openai.types.chat.chat_completion import ChatCompletion, Choice -from openai.types.chat.chat_completion_chunk import ChatCompletionChunk +from openai.types.chat.chat_completion_chunk import ChatCompletionChunk, ChoiceDelta from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice +from openai.types.chat.chat_completion_message import ChatCompletionMessage from openai.types.chat.chat_completion_message_custom_tool_call import ( ChatCompletionMessageCustomToolCall, ) @@ -139,6 +140,37 @@ def _sanitize_author_name(name: str | None) -> str | None: ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None) +OpenAIChatResponseContentsParser: TypeAlias = Callable[ + ["ChatCompletionMessage | ChoiceDelta", list[Content]], list[Content] +] +"""Hook to customize how a response message/delta is parsed into ``Content`` items. + +Called once per choice (non-streaming) or per streaming update-choice, after the client +has built its default ``Content`` list. Receives the already-selected OpenAI +``ChatCompletionMessage`` (non-streaming) or ``ChoiceDelta`` (streaming) and the +default-parsed contents, and returns the contents to use instead. The client resolves the +streaming/non-streaming dispatch, so a parser can read provider fields directly (e.g. +``getattr(message, "reasoning", None)``) without branching. + +This is the extension point for OpenAI-compatible endpoints that return non-standard fields +(e.g. OpenRouter/vLLM ``reasoning`` / ``reasoning_details`` or Mistral chunked ``content``). +The stock client stays free of provider-specific branches; supply a parser to surface such +data. Return the input list unchanged to opt out for a given choice. +""" + +OpenAIChatMessagePreparer: TypeAlias = Callable[["Message", list[dict[str, Any]]], list[dict[str, Any]]] +"""Hook to customize the outgoing request messages built from a single framework ``Message``. + +Called once per framework ``Message`` after the client has built its default list of OpenAI +message dicts. Receives the source ``Message`` and the default dicts, and returns the dicts to +send instead. + +This is the send-side counterpart of :data:`OpenAIChatResponseContentsParser`. Providers such as +vLLM require reasoning to be echoed back on later turns under the same key it was received; use a +preparer to inject those fields. Return the input list unchanged to opt out. +""" + + # region OpenAI Chat Options TypedDict @@ -252,6 +284,8 @@ def __init__( default_headers: Mapping[str, str] | None = None, async_client: AsyncOpenAI | None = None, instruction_role: str | None = None, + response_parser: OpenAIChatResponseContentsParser | None = None, + message_preparer: OpenAIChatMessagePreparer | None = None, compaction_strategy: CompactionStrategy | None = None, tokenizer: TokenizerProtocol | None = None, additional_properties: dict[str, Any] | None = None, @@ -272,6 +306,10 @@ def __init__( default_headers: Additional HTTP headers. async_client: Pre-configured OpenAI client. instruction_role: Role for instruction messages (for example ``"system"``). + response_parser: Optional hook to customize response parsing into ``Content`` items. + See ``OpenAIChatResponseContentsParser``. + message_preparer: Optional hook to customize outgoing request messages. + See ``OpenAIChatMessagePreparer``. compaction_strategy: Optional per-client compaction override. tokenizer: Optional tokenizer for compaction strategies. additional_properties: Additional properties stored on the client instance. @@ -294,6 +332,8 @@ def __init__( default_headers: Mapping[str, str] | None = None, async_client: AsyncAzureOpenAI | AsyncOpenAI | None = None, instruction_role: str | None = None, + response_parser: OpenAIChatResponseContentsParser | None = None, + message_preparer: OpenAIChatMessagePreparer | None = None, compaction_strategy: CompactionStrategy | None = None, tokenizer: TokenizerProtocol | None = None, additional_properties: dict[str, Any] | None = None, @@ -321,6 +361,10 @@ def __init__( async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup. instruction_role: Role for instruction messages (for example ``"system"``). + response_parser: Optional hook to customize response parsing into ``Content`` items. + See ``OpenAIChatResponseContentsParser``. + message_preparer: Optional hook to customize outgoing request messages. + See ``OpenAIChatMessagePreparer``. compaction_strategy: Optional per-client compaction override. tokenizer: Optional tokenizer for compaction strategies. additional_properties: Additional properties stored on the client instance. @@ -343,6 +387,8 @@ def __init__( default_headers: Mapping[str, str] | None = None, async_client: AsyncOpenAI | None = None, instruction_role: str | None = None, + response_parser: OpenAIChatResponseContentsParser | None = None, + message_preparer: OpenAIChatMessagePreparer | None = None, compaction_strategy: CompactionStrategy | None = None, tokenizer: TokenizerProtocol | None = None, additional_properties: dict[str, Any] | None = None, @@ -376,6 +422,13 @@ def __init__( async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup. instruction_role: Role for instruction messages (for example ``"system"``). + response_parser: Optional hook to customize how each response choice/delta is parsed + into ``Content`` items. Use it to surface non-standard fields from + OpenAI-compatible endpoints (e.g. OpenRouter/vLLM reasoning or Mistral chunked + content) without subclassing. See ``OpenAIChatResponseContentsParser``. + message_preparer: Optional hook to customize the outgoing request messages built from + each framework ``Message``. Use it to echo provider-specific fields (e.g. vLLM + ``reasoning``) back on later turns. See ``OpenAIChatMessagePreparer``. compaction_strategy: Optional per-client compaction override. tokenizer: Optional tokenizer for compaction strategies. additional_properties: Additional properties stored on the client instance. @@ -429,6 +482,8 @@ def __init__( else: self.default_headers = None self.instruction_role = instruction_role + self.response_parser = response_parser + self.message_preparer = message_preparer self._use_azure_client = use_azure_client if use_azure_client: self.OTEL_PROVIDER_NAME = "azure.ai.openai" # type: ignore[misc] @@ -809,6 +864,8 @@ def _parse_response_from_openai(self, response: ChatCompletion, options: Mapping contents.extend(parsed_tool_calls) if reasoning_details := getattr(choice.message, "reasoning_details", None): contents.append(Content.from_text_reasoning(protected_data=json.dumps(reasoning_details))) + if self.response_parser is not None: + contents = list(self.response_parser(choice.message, contents)) messages.append(Message(role="assistant", contents=contents)) return ChatResponse( response_id=response.id, @@ -848,11 +905,15 @@ def _parse_response_update_from_openai( if choice.delta is None: # pyright: ignore[reportUnnecessaryComparison] continue - contents.extend(self._parse_tool_calls_from_openai(choice)) + choice_contents: list[Content] = [] + choice_contents.extend(self._parse_tool_calls_from_openai(choice)) if text_content := self._parse_text_from_openai(choice): - contents.append(text_content) + choice_contents.append(text_content) if reasoning_details := getattr(choice.delta, "reasoning_details", None): - contents.append(Content.from_text_reasoning(protected_data=json.dumps(reasoning_details))) + choice_contents.append(Content.from_text_reasoning(protected_data=json.dumps(reasoning_details))) + if self.response_parser is not None: + choice_contents = list(self.response_parser(choice.delta, choice_contents)) + contents.extend(choice_contents) return ChatResponseUpdate( created_at=datetime.fromtimestamp(chunk.created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"), contents=contents, @@ -897,6 +958,8 @@ def _parse_text_from_openai(self, choice: Choice | ChunkChoice) -> Content | Non """Parse the choice into a Content object with type='text'.""" message = choice.message if isinstance(choice, Choice) else choice.delta if message.content: + if not isinstance(message.content, str): + return None return Content.from_text(text=message.content, raw_representation=choice) if hasattr(message, "refusal") and message.refusal: return Content.from_text(text=message.refusal, raw_representation=choice) @@ -970,7 +1033,18 @@ def _prepare_messages_for_openai( # region Parsers def _prepare_message_for_openai(self, message: Message) -> list[dict[str, Any]]: - """Prepare a chat message for OpenAI.""" + """Prepare a chat message for OpenAI, applying the ``message_preparer`` hook if set. + + The hook is applied here so it runs exactly once per framework ``Message`` for every + role, including ``system`` / ``developer`` messages that build a different shape. + """ + all_messages = self._build_openai_messages(message) + if self.message_preparer is not None: + all_messages = list(self.message_preparer(message, all_messages)) + return all_messages + + def _build_openai_messages(self, message: Message) -> list[dict[str, Any]]: + """Build the default OpenAI message dicts for a framework message (no hook applied).""" # System/developer messages default to plain string content because some # OpenAI-compatible endpoints reject list content for non-user roles. The # exception is a prompt cache breakpoint on a text part: it can only live on @@ -1008,6 +1082,7 @@ def _prepare_message_for_openai(self, message: Message) -> list[dict[str, Any]]: details := message.additional_properties["reasoning_details"] ): args["reasoning_details"] = details + match content.type: case "function_call": if all_messages and "tool_calls" in all_messages[-1]: @@ -1203,6 +1278,8 @@ def __init__( default_headers: Mapping[str, str] | None = None, async_client: AsyncOpenAI | None = None, instruction_role: str | None = None, + response_parser: OpenAIChatResponseContentsParser | None = None, + message_preparer: OpenAIChatMessagePreparer | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None, @@ -1220,6 +1297,10 @@ def __init__( default_headers: Additional HTTP headers. async_client: Pre-configured OpenAI client. instruction_role: Role for instruction messages (for example ``"system"``). + response_parser: Optional hook to customize response parsing into ``Content`` items. + See ``OpenAIChatResponseContentsParser``. + message_preparer: Optional hook to customize outgoing request messages. + See ``OpenAIChatMessagePreparer``. base_url: Base URL override. When not provided explicitly, the constructor reads ``OPENAI_BASE_URL``. env_file_path: Optional ``.env`` file that is checked before the process environment @@ -1243,6 +1324,8 @@ def __init__( default_headers: Mapping[str, str] | None = None, async_client: AsyncAzureOpenAI | AsyncOpenAI | None = None, instruction_role: str | None = None, + response_parser: OpenAIChatResponseContentsParser | None = None, + message_preparer: OpenAIChatMessagePreparer | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None, @@ -1269,6 +1352,10 @@ def __init__( async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup. instruction_role: Role for instruction messages (for example ``"system"``). + response_parser: Optional hook to customize response parsing into ``Content`` items. + See ``OpenAIChatResponseContentsParser``. + message_preparer: Optional hook to customize outgoing request messages. + See ``OpenAIChatMessagePreparer``. env_file_path: Optional ``.env`` file that is checked before process environment variables for ``AZURE_OPENAI_*`` values. env_file_encoding: Encoding for the ``.env`` file. @@ -1287,6 +1374,8 @@ def __init__( default_headers: Mapping[str, str] | None = None, async_client: AsyncOpenAI | None = None, instruction_role: str | None = None, + response_parser: OpenAIChatResponseContentsParser | None = None, + message_preparer: OpenAIChatMessagePreparer | None = None, base_url: str | None = None, azure_endpoint: str | None = None, api_version: str | None = None, @@ -1315,6 +1404,12 @@ def __init__( async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup. instruction_role: Role to use for instruction messages (for example ``"system"``). + response_parser: Optional hook to customize how each response choice/delta is parsed + into ``Content`` items (e.g. to surface OpenRouter/vLLM reasoning or Mistral + chunked content). See ``OpenAIChatResponseContentsParser``. + message_preparer: Optional hook to customize the outgoing request messages built from + each framework ``Message`` (e.g. to echo vLLM ``reasoning`` back on later turns). + See ``OpenAIChatMessagePreparer``. base_url: Base URL override. For OpenAI routing this maps to ``OPENAI_BASE_URL``. For Azure routing this may be used instead of ``azure_endpoint`` when you want to pass the full ``.../openai/v1`` base URL directly. @@ -1382,6 +1477,8 @@ class MyOptions(OpenAIChatCompletionOptions, total=False): default_headers=default_headers, async_client=async_client, instruction_role=instruction_role, + response_parser=response_parser, + message_preparer=message_preparer, env_file_path=env_file_path, env_file_encoding=env_file_encoding, middleware=middleware, diff --git a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py index d933c280f0b..43026fad42c 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py @@ -2468,3 +2468,244 @@ def test_prepare_options_prompt_cache_options_guarded_on_old_openai(monkeypatch: # endregion + + +# region response_parser / message_preparer hooks + +_VLLM_REASONING_FIELD_KEY = "_source_reasoning_field" + + +def _vllm_reasoning_parser(message: Any, contents: list[Content]) -> list[Content]: + """Example response_parser: surface a top-level ``reasoning`` field as reasoning content. + + Receives the already-selected message/delta (no streaming dispatch needed) and tags the + surfaced content with its originating field name so a message_preparer can echo it back + and correlate it robustly. + """ + reasoning = getattr(message, "reasoning", None) + if isinstance(reasoning, str) and reasoning: + return [ + *contents, + Content.from_text_reasoning( + text=reasoning, + additional_properties={_VLLM_REASONING_FIELD_KEY: "reasoning"}, + ), + ] + return contents + + +def _vllm_reasoning_preparer(message: Message, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Example message_preparer: echo surfaced reasoning back under its originating key. + + Correlates via ``message.contents`` markers instead of raw request-string matching: each + marked reasoning content maps to exactly one auto-emitted assistant text dict, removed + one-to-one in order, then the provider field is attached to the final message. + """ + surfaced = [ + (content.additional_properties[_VLLM_REASONING_FIELD_KEY], content.text) + for content in message.contents + if content.type == "text_reasoning" + and _VLLM_REASONING_FIELD_KEY in content.additional_properties + and content.text + ] + if not surfaced: + return messages + + remaining = list(messages) + fields: dict[str, str] = {} + for field_name, text in surfaced: + # Remove exactly one auto-emitted assistant text dict matching this reasoning text. + for i, msg in enumerate(remaining): + if msg.get("role") == "assistant" and "tool_calls" not in msg and msg.get("content") == text: + remaining.pop(i) + break + fields[field_name] = fields.get(field_name, "") + text + + if remaining: + for field_name, value in fields.items(): + remaining[-1][field_name] = value + return remaining + + +def _make_chat_completion(message: ChatCompletionMessage, model: str = "vllm-model") -> ChatCompletion: + return ChatCompletion( + id="test-response", + object="chat.completion", + created=1234567890, + model=model, + choices=[Choice(index=0, message=message, finish_reason="stop")], + ) + + +def test_response_parser_hook_transforms_contents(openai_unit_test_env: dict[str, str]) -> None: + """A response_parser can surface provider-specific fields (e.g. vLLM `reasoning`).""" + client = OpenAIChatCompletionClient(response_parser=_vllm_reasoning_parser) + message = ChatCompletionMessage.model_construct(role="assistant", content="Answer.", reasoning="Thinking...") + + parsed = client._parse_response_from_openai(_make_chat_completion(message), {}) + + reasoning = [c for c in parsed.messages[0].contents if c.type == "text_reasoning"] + assert len(reasoning) == 1 + assert reasoning[0].text == "Thinking..." + + +def test_response_parser_hook_streaming(openai_unit_test_env: dict[str, str]) -> None: + """The response_parser is also applied on the streaming path.""" + from openai.types.chat.chat_completion_chunk import ChatCompletionChunk, ChoiceDelta + from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice + + client = OpenAIChatCompletionClient(response_parser=_vllm_reasoning_parser) + delta = ChoiceDelta.model_construct(role="assistant", content=None, reasoning="step 1") + chunk = ChatCompletionChunk( + id="test-chunk", + object="chat.completion.chunk", + created=1234567890, + model="vllm-model", + choices=[ChunkChoice(index=0, delta=delta, finish_reason=None)], + ) + + update = client._parse_response_update_from_openai(chunk) + + reasoning = [c for c in update.contents if c.type == "text_reasoning"] + assert len(reasoning) == 1 + assert reasoning[0].text == "step 1" + + +def test_message_preparer_hook_transforms_messages(openai_unit_test_env: dict[str, str]) -> None: + """A message_preparer can rewrite the outgoing request messages.""" + + def preparer(message: Message, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + for msg in messages: + msg["custom_field"] = "injected" + return messages + + client = OpenAIChatCompletionClient(message_preparer=preparer) + prepared = client._prepare_message_for_openai(Message(role="assistant", contents=[Content.from_text("hi")])) + + assert prepared[-1]["custom_field"] == "injected" + assert prepared[-1]["content"] == "hi" + + +@pytest.mark.parametrize("role", ["system", "developer"]) +def test_message_preparer_hook_runs_for_system_and_developer(role: str, openai_unit_test_env: dict[str, str]) -> None: + """The message_preparer runs once per Message, including system/developer roles.""" + seen: list[str] = [] + + def preparer(message: Message, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + seen.append(str(message.role)) + for msg in messages: + msg["gateway_field"] = "required" + return messages + + client = OpenAIChatCompletionClient(message_preparer=preparer) + prepared = client._prepare_message_for_openai(Message(role=cast(Any, role), contents=[Content.from_text("sys")])) + + assert seen == [role] + assert prepared[-1]["gateway_field"] == "required" + assert prepared[-1]["content"] == "sys" + + +def test_message_preparer_correlation_does_not_drop_same_text_answer( + openai_unit_test_env: dict[str, str], +) -> None: + """Marker-based correlation removes only one reasoning dict, even if the answer shares its text.""" + client = OpenAIChatCompletionClient(message_preparer=_vllm_reasoning_preparer) + # Answer text is byte-identical to the reasoning text; only the surfaced reasoning + # (marked) content should be echoed back and its single dict removed. + same = "same text" + reasoning = Content.from_text_reasoning(text=same, additional_properties={_VLLM_REASONING_FIELD_KEY: "reasoning"}) + message = Message(role="assistant", contents=[Content.from_text(same), reasoning]) + + prepared = client._prepare_message_for_openai(message) + + # The answer message survives (one dict removed, not both) and carries the echoed field. + assert len(prepared) == 1 + assert prepared[0]["content"] == same + assert prepared[0]["reasoning"] == same + + +def test_hooks_roundtrip_vllm_reasoning(openai_unit_test_env: dict[str, str]) -> None: + """End-to-end: parser surfaces reasoning for display, preparer echoes it back under `reasoning`.""" + client = OpenAIChatCompletionClient( + response_parser=_vllm_reasoning_parser, + message_preparer=_vllm_reasoning_preparer, + ) + message = ChatCompletionMessage.model_construct(role="assistant", content="42.", reasoning="Because reasons.") + + parsed = client._parse_response_from_openai(_make_chat_completion(message), {}) + # Reasoning is surfaced for display. + assert any(c.type == "text_reasoning" and c.text == "Because reasons." for c in parsed.messages[0].contents) + + prepared = client._prepare_message_for_openai(parsed.messages[0]) + # A single assistant message carries the answer plus the reasoning echoed back under its key, + # and the reasoning is NOT duplicated as visible content. + assert len(prepared) == 1 + assert prepared[0]["content"] == "42." + assert prepared[0]["reasoning"] == "Because reasons." + + +def test_no_hooks_keeps_default_behavior(openai_unit_test_env: dict[str, str]) -> None: + """Without hooks, top-level `reasoning` is ignored and `reasoning_details` stays opaque.""" + client = OpenAIChatCompletionClient() + message = ChatCompletionMessage.model_construct( + role="assistant", + content="Answer.", + reasoning="ignored without a parser", + reasoning_details=[{"type": "reasoning.text", "text": "opaque"}], + ) + + parsed = client._parse_response_from_openai(_make_chat_completion(message, model="some-model"), {}) + + reasoning = [c for c in parsed.messages[0].contents if c.type == "text_reasoning"] + # reasoning_details surfaces as a single opaque reasoning content (baseline behavior); no text. + assert len(reasoning) == 1 + assert reasoning[0].text is None + assert reasoning[0].protected_data is not None + + +def test_default_parsing_skips_non_string_content(openai_unit_test_env: dict[str, str]) -> None: + """Structured list `content` is skipped by default parsing (no malformed text Content).""" + client = OpenAIChatCompletionClient() + chunked = [ + {"type": "thinking", "thinking": [{"type": "text", "text": "reasoning"}]}, + {"type": "text", "text": "answer"}, + ] + message = ChatCompletionMessage.model_construct(role="assistant", content=cast(Any, chunked)) + + parsed = client._parse_response_from_openai(_make_chat_completion(message, model="mistral-medium-latest"), {}) + + # The list is not wrapped as a text Content; nothing is emitted for it. + assert not any(c.type == "text" for c in parsed.messages[0].contents) + + +def test_response_parser_can_expand_chunked_content(openai_unit_test_env: dict[str, str]) -> None: + """A response_parser receives the selected message and can expand structured list content.""" + + def chunk_parser(message: Any, contents: list[Content]) -> list[Content]: + if not isinstance(message.content, list): + return contents + expanded = list(contents) + for chunk in message.content: + if chunk.get("type") == "thinking": + text = "".join(part.get("text", "") for part in chunk.get("thinking", [])) + expanded.append(Content.from_text_reasoning(text=text)) + elif chunk.get("type") == "text": + expanded.append(Content.from_text(text=chunk["text"])) + return expanded + + client = OpenAIChatCompletionClient(response_parser=chunk_parser) + chunked = [ + {"type": "thinking", "thinking": [{"type": "text", "text": "reasoning"}]}, + {"type": "text", "text": "answer"}, + ] + message = ChatCompletionMessage.model_construct(role="assistant", content=cast(Any, chunked)) + + parsed = client._parse_response_from_openai(_make_chat_completion(message, model="mistral-medium-latest"), {}) + contents = parsed.messages[0].contents + + assert [c.type for c in contents] == ["text_reasoning", "text"] + assert contents[0].text == "reasoning" + assert contents[1].text == "answer" + + +# endregion