diff --git a/docs/examples/provider_fields.py b/docs/examples/provider_fields.py new file mode 100644 index 000000000..f96d7c7d1 --- /dev/null +++ b/docs/examples/provider_fields.py @@ -0,0 +1,72 @@ +# pytest: unit +"""Emit a provider-specific wire field without a core-type change. + +A component author can attach `provider_fields` to a `Message` to declare extra +keys that Mellea does not model, targeted at a specific provider. The fields are +merged into the wire message at serialization time. Mellea's known fields always +win on a collision, and a declaration that names a provider the request never +reaches raises a `ValueError` (add `"*"` to opt out of that check). + +See issue #1565 for the design. +""" + +import pytest + +from mellea.helpers.openai_compatible_helpers import message_to_openai_message +from mellea.stdlib.components import Message + + +def targeted_provider_field() -> dict: + """Attach an OpenAI-only `prediction` field and serialize for OpenAI. + + The `"openai"` key matches the OpenAI wire family (openai, litellm, watsonx, + huggingface), so the field lands on the wire message for any of them. + """ + msg = Message( + "user", + "Refactor this function.", + provider_fields={"openai": {"prediction": {"type": "content"}}}, + ) + wire = message_to_openai_message(msg, provider="openai") + assert wire["prediction"] == {"type": "content"} + return wire + + +def portable_field_with_wildcard() -> dict: + """Use `"*"` to declare a field valid on every backend. + + A `"*"` target never raises on a provider mismatch — it is the author's + portability contract that the field is safe to send everywhere. + """ + msg = Message("user", "Hello", provider_fields={"*": {"metadata_tag": "demo"}}) + wire = message_to_openai_message(msg, provider="openai") + assert wire["metadata_tag"] == "demo" + return wire + + +def known_fields_always_win() -> dict: + """An author key that collides with a Mellea-known field is dropped.""" + msg = Message( + "user", + "real content", + provider_fields={"openai": {"content": "hijacked", "extra": "kept"}}, + ) + wire = message_to_openai_message(msg, provider="openai") + assert wire["content"] == "real content" # Mellea's field wins + assert wire["extra"] == "kept" # non-colliding author field lands + return wire + + +def provider_mismatch_raises() -> None: + """Targeting a provider the request never reaches is a hard error.""" + msg = Message("user", "Hi", provider_fields={"ollama": {"keep_alive": "5m"}}) + with pytest.raises(ValueError): + message_to_openai_message(msg, provider="openai") + + +if __name__ == "__main__": + print("Targeted field:", targeted_provider_field()) + print("Wildcard field:", portable_field_with_wildcard()) + print("Known fields win:", known_fields_always_win()) + provider_mismatch_raises() + print("Provider mismatch raised as expected.") diff --git a/mellea/backends/huggingface.py b/mellea/backends/huggingface.py index 1d33742be..cf87f492c 100644 --- a/mellea/backends/huggingface.py +++ b/mellea/backends/huggingface.py @@ -655,7 +655,12 @@ async def _generate_from_intrinsic( # conversation, not multi-turn generation, so reasoning is never replayed # here (no `replay_reasoning=`). The HF chat path serializes via # `to_chat`/`apply_chat_template`, not this helper. - conversation.extend([message_to_openai_message(m) for m in ctx_as_message_list]) + conversation.extend( + [ + message_to_openai_message(m, provider=self._provider) + for m in ctx_as_message_list + ] + ) docs = messages_to_docs(ctx_as_message_list) diff --git a/mellea/backends/litellm.py b/mellea/backends/litellm.py index 061aa516a..b2433bb38 100644 --- a/mellea/backends/litellm.py +++ b/mellea/backends/litellm.py @@ -362,7 +362,9 @@ async def _generate_from_chat_context_standard( replay_flags = should_replay_reasoning(messages, self._provider) conversation.extend( [ - message_to_openai_message(m, self.formatter, replay_reasoning=replay) + message_to_openai_message( + m, self.formatter, replay_reasoning=replay, provider=self._provider + ) for m, replay in zip(messages, replay_flags) ] ) diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index 2babc396c..c5bacc3a0 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -35,6 +35,7 @@ DEFAULT_CHUNK_TIMEOUT, ClientCache, get_current_event_loop, + merge_provider_fields, send_to_queue, should_replay_reasoning, ) @@ -548,6 +549,11 @@ async def generate_from_chat_context( tool_name = m.tool_name or getattr(m, "name", None) if tool_name is not None: message_dict["tool_name"] = tool_name + # Merge any author-declared provider fields (Mellea's known fields win; + # a mismatched target raises). Must run after the known fields are set. + message_dict = merge_provider_fields( + message_dict, m.provider_fields, self._provider + ) conversation.append(message_dict) # Append tool call information if applicable. diff --git a/mellea/backends/openai.py b/mellea/backends/openai.py index 739f5581b..99f7fa123 100644 --- a/mellea/backends/openai.py +++ b/mellea/backends/openai.py @@ -760,7 +760,9 @@ async def _generate_from_intrinsic( # conversation, not multi-turn generation, so reasoning is never replayed # here (no `replay_reasoning=`) — unlike the chat path in # `_generate_from_context`, which applies `should_replay_reasoning`. - conversation.extend([message_to_openai_message(m) for m in messages]) + conversation.extend( + [message_to_openai_message(m, provider=self._provider) for m in messages] + ) docs = messages_to_docs(messages) @@ -992,7 +994,9 @@ async def _generate_from_chat_context_standard( replay_flags = should_replay_reasoning(messages, self._provider) conversation.extend( [ - message_to_openai_message(m, self.formatter, replay_reasoning=replay) + message_to_openai_message( + m, self.formatter, replay_reasoning=replay, provider=self._provider + ) for m, replay in zip(messages, replay_flags) ] ) diff --git a/mellea/backends/utils.py b/mellea/backends/utils.py index b09c5ed3e..401e9ed79 100644 --- a/mellea/backends/utils.py +++ b/mellea/backends/utils.py @@ -19,6 +19,7 @@ from ..core import Context, MelleaLogger, ModelToolCall, Span from ..core.base import AbstractMelleaTool, ModelOutputThunk from ..formatters import ChatFormatter +from ..helpers import merge_provider_fields from ..stdlib.components import Message from .tools import parse_tools, validate_tool_arguments @@ -112,6 +113,9 @@ def to_chat( msg_dict["tool_calls"] = m.tool_calls if m.tool_call_id: msg_dict["tool_call_id"] = m.tool_call_id + # Merge any author-declared provider fields (Mellea's known fields win; + # a mismatched target raises). Must run after the known fields are set. + msg_dict = merge_provider_fields(msg_dict, m.provider_fields, "huggingface") ctx_as_conversation.append(msg_dict) # Check that we ddin't accidentally end up with CBlocks. Only string values can diff --git a/mellea/backends/watsonx.py b/mellea/backends/watsonx.py index 0a7894dd5..38bf0b0ef 100644 --- a/mellea/backends/watsonx.py +++ b/mellea/backends/watsonx.py @@ -45,6 +45,7 @@ chat_completion_delta_merge, extract_model_tool_requests, get_current_event_loop, + merge_provider_fields, send_to_queue, should_replay_reasoning, ) @@ -432,6 +433,11 @@ async def generate_from_chat_context( message_dict["tool_call_id"] = m.tool_call_id if replay and m.thinking: message_dict["reasoning_content"] = m.thinking + # Merge any author-declared provider fields (Mellea's known fields win; + # a mismatched target raises). Must run after the known fields are set. + message_dict = merge_provider_fields( + message_dict, m.provider_fields, self._provider + ) conversation.append(message_dict) if _format is not None: diff --git a/mellea/core/base.py b/mellea/core/base.py index 1080ef94b..f95464705 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -1970,6 +1970,13 @@ class TemplateRepresentation: tool_name (str | None): For a `role="tool"` component, the name of the tool whose result this message carries (e.g. Ollama's tool-result turn keys on it). Defaults to `None`. + provider_fields (dict[str, dict[str, Any]] | None): Optional author-declared + extra wire fields, keyed by provider target. Each outer key is a provider + (matched against the backend's provider string, with the `"openai"` + wire-family alias and a `"*"` wildcard); each inner dict holds fields + merged into the wire message for that target. Mellea's known fields always + win on a key collision. A named target that the request does not hit (and + no `"*"`) is a hard error at serialization. Defaults to `None`. """ @@ -1991,6 +1998,7 @@ class TemplateRepresentation: tool_calls: list[dict[str, Any]] | None = None tool_call_id: str | None = None tool_name: str | None = None + provider_fields: dict[str, dict[str, Any]] | None = None @dataclass diff --git a/mellea/helpers/__init__.py b/mellea/helpers/__init__.py index 67a6930ff..8dfd7c0ae 100644 --- a/mellea/helpers/__init__.py +++ b/mellea/helpers/__init__.py @@ -21,8 +21,10 @@ ) from .event_loop_helper import _run_async_in_thread from .openai_compatible_helpers import ( + OPENAI_COMPATIBLE_WIRE_PROVIDERS, chat_completion_delta_merge, extract_model_tool_requests, + merge_provider_fields, message_to_openai_message, messages_to_docs, should_replay_reasoning, @@ -35,6 +37,7 @@ __all__ = [ "DEFAULT_CHUNK_TIMEOUT", + "OPENAI_COMPATIBLE_WIRE_PROVIDERS", "ClientCache", "_ServerType", "_run_async_in_thread", @@ -43,6 +46,7 @@ "extract_model_tool_requests", "get_current_event_loop", "is_vllm_server_with_structured_output", + "merge_provider_fields", "message_to_openai_message", "messages_to_docs", "send_to_queue", diff --git a/mellea/helpers/openai_compatible_helpers.py b/mellea/helpers/openai_compatible_helpers.py index c366f2003..d6a245c5f 100644 --- a/mellea/helpers/openai_compatible_helpers.py +++ b/mellea/helpers/openai_compatible_helpers.py @@ -12,6 +12,7 @@ from pydantic import BaseModel from ..core.base import AudioBlock, AudioUrlBlock, ImageUrlBlock +from ..core.utils import MelleaLogger if TYPE_CHECKING: from ..core import Formatter, ModelToolCall @@ -19,6 +20,83 @@ from ..stdlib.components import Document, Message +# The providers whose request/wire serialization is OpenAI-compatible — the set a +# `provider_fields` `{"openai": ...}` declaration reaches. This is the *serialization* +# family: `message_to_openai_message` is the shared serializer for `openai`, `litellm`, +# and `huggingface` (huggingface.py), and `watsonx` emits the same wire shape via its +# own inline loop. It is deliberately NOT `Message._parse`'s response tuple +# `("openai", "watsonx", "litellm")`, which groups by *response* shape and excludes HF +# (HF returns token tensors, not a `choices[0].message` dict). Keeping these separate +# prevents a `{"openai": ...}` declaration from wrongly raising on the HuggingFace path. +OPENAI_COMPATIBLE_WIRE_PROVIDERS = frozenset( + {"openai", "litellm", "watsonx", "huggingface"} +) + + +def merge_provider_fields( + base: dict[str, Any], + provider_fields: dict[str, dict[str, Any]] | None, + provider: str, +) -> dict[str, Any]: + """Merge author-declared `provider_fields` into a wire message dict. + + A key matches `provider` iff it is `"*"`, equals `provider` exactly, or is + `"openai"` and `provider` is in `OPENAI_COMPATIBLE_WIRE_PROVIDERS`. Fields from + every matching key are merged into `base`, but only for keys `base` did not + already set — Mellea's known fields always win, and a dropped colliding key is + debug-logged. `"*"` always matches, so its presence never triggers the mismatch + error. + + Args: + base: The wire message dict built from Mellea's known fields; mutated and + returned. + provider_fields: The author's provider-keyed extra fields, or `None`. + provider: The provider string of the backend performing serialization. + + Returns: + `base`, with fields from every matching provider key merged in. + + Raises: + ValueError: If `provider_fields` is non-empty, contains no `"*"`, and no key + matches `provider` — the component targeted a backend it did not hit. + """ + if not provider_fields: + return base + + def _matches(key: str) -> bool: + return ( + key == "*" + or key == provider + or (key == "openai" and provider in OPENAI_COMPATIBLE_WIRE_PROVIDERS) + ) + + matched_any = False + for key, fields in provider_fields.items(): + if not _matches(key): + continue + matched_any = True + for field, value in fields.items(): + if field in base: + MelleaLogger.get_logger().debug( + "provider_fields[%r][%r] collides with a Mellea-known field; " + "dropping the author value (known field wins).", + key, + field, + ) + continue + base[field] = value + + if not matched_any: + raise ValueError( + f"provider_fields declares target(s) {sorted(provider_fields)} but the " + f"request is running on provider {provider!r}, which none of them match. " + 'Add a "*" key to declare the field valid on every backend, or target the ' + "provider this component actually runs on." + ) + + return base + + class ToolCallFunction(TypedDict): """Function details in a tool call.""" @@ -244,7 +322,11 @@ def should_replay_reasoning( def message_to_openai_message( - msg: Message, formatter: Formatter | None = None, *, replay_reasoning: bool = False + msg: Message, + formatter: Formatter | None = None, + *, + replay_reasoning: bool = False, + provider: str = "openai", ) -> dict: """Serialise a Mellea `Message` to the format required by OpenAI-compatible API providers. @@ -258,6 +340,10 @@ def message_to_openai_message( the provider receives the model's prior reasoning. Defaults to `False` (reasoning is stripped), preserving the historical behaviour; callers decide per-turn via their replay policy (see `should_replay_reasoning`). + provider: The calling backend's provider string, used to match the message's + `provider_fields` declaration. Defaults to `"openai"`; OpenAI-compatible + callers pass their own provider (`"litellm"`, `"watsonx"`, `"huggingface"`) + so an `"openai"`-family declaration still reaches the wire. Returns: A dict with `"role"` and `"content"` fields. When the message carries @@ -277,6 +363,8 @@ def message_to_openai_message( ValueError: If the message contains an `AudioUrlBlock`. The OpenAI Chat Completions audio schema does not support audio by URL; fetch the audio and pass it as an `AudioBlock` with base64 data instead. + ValueError: If the message's `provider_fields` names a target that does not + match `provider` and includes no `"*"` (see `merge_provider_fields`). """ # NOTE: `self.formatter.to_chat_messages` explicitly skips `Message` objects. However, we need # to print `Message`s to correctly serialize any documents with the message. Do the printing here. @@ -346,7 +434,8 @@ def message_to_openai_message( if replay_reasoning and msg.thinking: result["reasoning_content"] = msg.thinking - return result + + return merge_provider_fields(result, msg.provider_fields, provider) def messages_to_docs(msgs: list[Message]) -> list[dict[str, str]]: diff --git a/mellea/stdlib/components/chat.py b/mellea/stdlib/components/chat.py index 7b46dba49..ba657884d 100644 --- a/mellea/stdlib/components/chat.py +++ b/mellea/stdlib/components/chat.py @@ -70,6 +70,13 @@ class Message(Component["Message"]): policy. `None` or empty for messages that carry no reasoning (e.g. user turns, or assistant turns from non-thinking models); the replay policy and serializers treat both falsy cases identically. + provider_fields (dict[str, dict[str, Any]] | None): Optional author-declared + extra wire fields Mellea does not model, keyed by provider target. Each + outer key is a provider (matched against the backend's provider string, + with the `"openai"` wire-family alias and a `"*"` wildcard); each inner + dict holds fields merged into the wire message for that target. Mellea's + known fields always win on a key collision, and a named target the request + does not hit (with no `"*"`) raises at serialization. Defaults to `None`. Attributes: Role (type): Type alias for the allowed role literals: `"system"`, @@ -90,8 +97,9 @@ def __init__( tool_call_id: str | None = None, tool_name: str | None = None, thinking: str | None = None, + provider_fields: dict[str, dict[str, Any]] | None = None, ): - """Initialize a Message with a role, text content, optional images, audio, documents, tool calls, an optional tool-call id, an optional tool name, and an optional reasoning trace.""" + """Initialize a Message with a role, text content, optional images, audio, documents, tool calls, an optional tool-call id, an optional tool name, an optional reasoning trace, and optional author-declared provider fields.""" if role not in get_args(Message.Role): raise ValueError( f"Invalid role {role!r}. Must be one of: {list(get_args(Message.Role))}" @@ -106,6 +114,7 @@ def __init__( self._tool_calls = tool_calls self._tool_call_id = tool_call_id self._tool_name = tool_name + self._provider_fields = provider_fields @property def images(self) -> None | list[ImageBlock | ImageUrlBlock]: @@ -132,6 +141,11 @@ def tool_name(self) -> str | None: """Returns the name of the tool whose result this `role="tool"` message carries, if any.""" return self._tool_name + @property + def provider_fields(self) -> dict[str, dict[str, Any]] | None: + """Returns the author-declared provider-keyed extra wire fields, if any.""" + return self._provider_fields + def parts(self) -> list[Span]: """Return the constituent parts of this message, including content, documents, images, and audio. @@ -169,6 +183,7 @@ def format_for_llm(self) -> TemplateRepresentation: tool_calls=self._tool_calls, tool_call_id=self._tool_call_id, tool_name=self._tool_name, + provider_fields=self._provider_fields, ) def __repr__(self) -> str: @@ -224,6 +239,12 @@ def _parse(self, computed: ModelOutputThunk) -> "Message": tool_calls=tool_calls, thinking=thinking, ) + # NOTE: this tuple groups providers by *response* shape (a + # `choices[0].message` dict) and deliberately excludes HuggingFace, which + # returns token tensors here, not a dict. Do NOT unify it with + # `OPENAI_COMPATIBLE_WIRE_PROVIDERS` (openai_compatible_helpers.py), which + # groups by *request/wire* shape and includes HF: they are separate on + # purpose (see #1565). Merging them would break parsing here. if provider in ("openai", "watsonx", "litellm") and isinstance( response, dict ): @@ -318,8 +339,8 @@ def message_from_template_representation( component's declared role and tool metadata are honored consistently across both conversion paths. The representation's `role` overrides `default_role` when set; role validation is deferred to `Message`, which raises `ValueError` for anything - outside `Message.Role`. `thinking`, `tool_calls`, and (for `role="tool"`) - `tool_call_id`/`tool_name` are carried onto the resulting message. + outside `Message.Role`. `thinking`, `tool_calls`, `provider_fields`, and (for + `role="tool"`) `tool_call_id`/`tool_name` are carried onto the resulting message. Args: tr: The template representation returned by the component's `format_for_llm`. @@ -339,6 +360,7 @@ def message_from_template_representation( tool_call_id=tr.tool_call_id, tool_name=tr.tool_name, thinking=tr.thinking, + provider_fields=tr.provider_fields, ) diff --git a/test/backends/test_provider_fields.py b/test/backends/test_provider_fields.py new file mode 100644 index 000000000..e6abbfe3d --- /dev/null +++ b/test/backends/test_provider_fields.py @@ -0,0 +1,151 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for author-pluggable message serialization via `provider_fields` (#1565). + +Covers the round-trip through the IR path, the shared `merge_provider_fields` +helper's match rule / precedence / mismatch error, and that an author-declared +field reaches the wire message on each affected serialization path. +""" + +import logging + +import pytest + +from mellea.core.base import TemplateRepresentation +from mellea.helpers.openai_compatible_helpers import ( + OPENAI_COMPATIBLE_WIRE_PROVIDERS, + merge_provider_fields, + message_to_openai_message, +) +from mellea.stdlib.components import Message +from mellea.stdlib.components.chat import message_from_template_representation + +# --- Round-trip: provider_fields survives IR -> Message -> format_for_llm --- + + +def test_provider_fields_round_trips_through_ir(): + """An author's provider_fields survives Message -> format_for_llm and back.""" + pf = {"openai": {"prediction": {"type": "content"}}} + msg = Message("user", "hi", provider_fields=pf) + + tr = msg.format_for_llm() + assert tr.provider_fields == pf + + rebuilt = message_from_template_representation( + tr, default_role="user", content="hi" + ) + assert rebuilt.provider_fields == pf + + +def test_provider_fields_defaults_to_none(): + """Message and TemplateRepresentation default provider_fields to None.""" + assert Message("user", "hi").provider_fields is None + assert TemplateRepresentation(obj=None, args={}).provider_fields is None + + +# --- merge helper: match rule --- + + +def test_merge_none_and_empty_are_noops(): + """None or empty provider_fields leaves base untouched and never raises.""" + base = {"role": "user", "content": "hi"} + assert merge_provider_fields(dict(base), None, "ollama") == base + assert merge_provider_fields(dict(base), {}, "ollama") == base + + +def test_merge_exact_provider_match(): + """An exact provider key merges its fields into the wire dict.""" + out = merge_provider_fields( + {"role": "user", "content": "hi"}, {"ollama": {"keep_alive": "5m"}}, "ollama" + ) + assert out["keep_alive"] == "5m" + + +def test_merge_wildcard_matches_every_provider(): + """The "*" key merges on any provider.""" + for provider in ("openai", "ollama", "watsonx", "huggingface", "litellm"): + out = merge_provider_fields( + {"role": "user", "content": "hi"}, {"*": {"x": 1}}, provider + ) + assert out["x"] == 1 + + +@pytest.mark.parametrize("provider", sorted(OPENAI_COMPATIBLE_WIRE_PROVIDERS)) +def test_merge_openai_family_alias(provider): + """The "openai" key merges on every OpenAI-compatible wire provider.""" + out = merge_provider_fields( + {"role": "user", "content": "hi"}, {"openai": {"prediction": {}}}, provider + ) + assert "prediction" in out + + +def test_openai_family_set_is_the_serialization_family(): + """The wire family is the serialization family (includes huggingface).""" + assert OPENAI_COMPATIBLE_WIRE_PROVIDERS == frozenset( + {"openai", "litellm", "watsonx", "huggingface"} + ) + + +# --- merge helper: mismatch raises --- + + +def test_merge_mismatch_raises(): + """A named provider key that matches nothing (and no "*") raises ValueError.""" + with pytest.raises(ValueError): + merge_provider_fields( + {"role": "user", "content": "hi"}, {"ollama": {"x": 1}}, "openai" + ) + + +def test_merge_wildcard_suppresses_mismatch(): + """Adding "*" makes an otherwise-mismatched set valid; only "*" fields land.""" + out = merge_provider_fields( + {"role": "user", "content": "hi"}, {"ollama": {"x": 1}, "*": {"y": 2}}, "openai" + ) + assert out["y"] == 2 + assert "x" not in out + + +def test_merge_multiple_named_some_match(): + """When several named keys are present, only the matching one's fields land.""" + out = merge_provider_fields( + {"role": "user", "content": "hi"}, + {"ollama": {"x": 1}, "openai": {"y": 2}}, + "openai", + ) + assert out["y"] == 2 + assert "x" not in out + + +# --- merge helper: precedence (known fields win) --- + + +def test_merge_known_fields_win_and_collision_debug_logged(caplog): + """A colliding author key is dropped (Mellea's field wins) and debug-logged.""" + with caplog.at_level(logging.DEBUG): + out = merge_provider_fields( + {"role": "user", "content": "real"}, + {"openai": {"role": "x", "content": "y", "extra": "kept"}}, + "openai", + ) + assert out["role"] == "user" + assert out["content"] == "real" + assert out["extra"] == "kept" + + +# --- per-backend reach: field lands on the wire dict --- + + +def test_reaches_openai_wire(): + """An author field reaches the OpenAI wire dict via message_to_openai_message.""" + msg = Message("user", "hi", provider_fields={"openai": {"prediction": {"t": 1}}}) + wire = message_to_openai_message(msg, provider="openai") + assert wire["prediction"] == {"t": 1} + + +def test_openai_wire_mismatch_raises(): + """message_to_openai_message raises on a provider-mismatched declaration.""" + msg = Message("user", "hi", provider_fields={"ollama": {"x": 1}}) + with pytest.raises(ValueError): + message_to_openai_message(msg, provider="openai")