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..f76fa8a1cad 100644 --- a/python/packages/openai/AGENTS.md +++ b/python/packages/openai/AGENTS.md @@ -33,6 +33,26 @@ 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` — `(choice, default_contents) -> contents`. + Post-processes the `Content` items parsed from each response choice / streaming delta. Use it to + surface non-standard fields for display. Applied per choice in both streaming and non-streaming paths. +- `message_preparer: OpenAIChatMessagePreparer` — `(message, default_dicts) -> dicts`. Post-processes + the outgoing request message dicts built from each framework `Message`. Use it to echo + provider-specific fields (e.g. vLLM `reasoning`) back on later turns for multi-turn continuity. + +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..0ac0e070129 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 @@ -139,6 +139,32 @@ def _sanitize_author_name(name: str | None) -> str | None: ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None) +OpenAIChatResponseContentsParser: TypeAlias = Callable[["Choice | ChunkChoice", list[Content]], list[Content]] +"""Hook to customize how a response choice/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 raw OpenAI ``Choice`` / streaming +``ChunkChoice`` and the default-parsed contents, and returns the contents to use instead. + +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 +278,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 +300,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 +326,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 +355,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 +381,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 +416,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 +476,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 +858,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, contents)) messages.append(Message(role="assistant", contents=contents)) return ChatResponse( response_id=response.id, @@ -848,11 +899,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, 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, @@ -1008,6 +1063,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]: @@ -1088,6 +1144,8 @@ def _prepare_message_for_openai(self, message: Message) -> list[dict[str, Any]]: for text_item in text_items ) + if self.message_preparer is not None: + all_messages = list(self.message_preparer(message, all_messages)) return all_messages def _prepare_content_for_openai(self, content: Content) -> dict[str, Any]: @@ -1203,6 +1261,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 +1280,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 +1307,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 +1335,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 +1357,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 +1387,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 +1460,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..25d47dcfb26 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,140 @@ def test_prepare_options_prompt_cache_options_guarded_on_old_openai(monkeypatch: # endregion + + +# region response_parser / message_preparer hooks + +_VLLM_REASONING_KEY = "vllm_reasoning" + + +def _vllm_reasoning_parser(choice: Any, contents: list[Content]) -> list[Content]: + """Example response_parser: surface a top-level ``reasoning`` field as reasoning content.""" + message = choice.message if hasattr(choice, "message") else choice.delta + reasoning = getattr(message, "reasoning", None) + if isinstance(reasoning, str) and reasoning: + return [ + *contents, + Content.from_text_reasoning(text=reasoning, additional_properties={_VLLM_REASONING_KEY: True}), + ] + 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 vLLM's ``reasoning`` key. + + The default serializer would replay the surfaced reasoning as visible assistant text, so this + drops those auto-emitted messages and attaches the reasoning to the final assistant message. + """ + reasoning_texts = { + content.text + for content in message.contents + if content.type == "text_reasoning" and content.additional_properties.get(_VLLM_REASONING_KEY) and content.text + } + if not reasoning_texts: + return messages + filtered = [m for m in messages if m.get("content") not in reasoning_texts] + if filtered: + filtered[-1]["reasoning"] = "".join(sorted(reasoning_texts)) + return filtered + + +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" + + +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 + + +# endregion