From 3e8f398ec90c78999ea719c2ba54fe763b07a894 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Fri, 30 Oct 2026 15:24:31 -0300 Subject: [PATCH 1/2] fix: preserve Anthropic assistant content Fixes #1118 Signed-off-by: Andre Manoel --- .../src/data_designer/engine/mcp/facade.py | 12 +- .../clients/adapters/anthropic_translation.py | 5 + .../engine/models/clients/types.py | 1 + .../src/data_designer/engine/models/facade.py | 20 ++- .../src/data_designer/engine/models/utils.py | 9 +- .../tests/engine/mcp/test_mcp_facade.py | 13 ++ .../models/clients/test_anthropic_mcp.py | 160 ++++++++++++++++++ .../clients/test_anthropic_translation.py | 22 +++ .../tests/engine/models/test_model_utils.py | 18 ++ 9 files changed, 253 insertions(+), 7 deletions(-) create mode 100644 packages/data-designer-engine/tests/engine/models/clients/test_anthropic_mcp.py diff --git a/packages/data-designer-engine/src/data_designer/engine/mcp/facade.py b/packages/data-designer-engine/src/data_designer/engine/mcp/facade.py index 8defeb9c6..8d583cead 100644 --- a/packages/data-designer-engine/src/data_designer/engine/mcp/facade.py +++ b/packages/data-designer-engine/src/data_designer/engine/mcp/facade.py @@ -170,12 +170,15 @@ def process_completion_response( ChatMessage.as_assistant( content=response_content, reasoning_content=reasoning_content or None, + anthropic_content=message.anthropic_content, ) ] # Has tool calls - execute and return all messages tool_call_dicts = _convert_canonical_tool_calls_to_dicts(tool_calls) - assistant_message = self._build_assistant_tool_message(response_content, tool_call_dicts, reasoning_content) + assistant_message = self._build_assistant_tool_message( + response_content, tool_call_dicts, reasoning_content, anthropic_content=message.anthropic_content + ) tool_messages = self._execute_tool_calls_from_canonical(tool_calls) return [assistant_message, *tool_messages] @@ -215,12 +218,15 @@ def refuse_completion_response( ChatMessage.as_assistant( content=response_content, reasoning_content=reasoning_content or None, + anthropic_content=message.anthropic_content, ) ] # Build assistant message with tool calls (same as normal) tool_call_dicts = _convert_canonical_tool_calls_to_dicts(tool_calls) - assistant_message = self._build_assistant_tool_message(response_content, tool_call_dicts, reasoning_content) + assistant_message = self._build_assistant_tool_message( + response_content, tool_call_dicts, reasoning_content, anthropic_content=message.anthropic_content + ) # Build refusal messages instead of executing tools refusal = refusal_message or DEFAULT_TOOL_REFUSAL_MESSAGE @@ -241,6 +247,7 @@ def _build_assistant_tool_message( response: str | None, tool_calls: list[dict[str, Any]], reasoning_content: str | None = None, + anthropic_content: list[dict[str, Any]] | None = None, ) -> ChatMessage: """Build the assistant message containing tool call requests.""" tool_calls_payload = [ @@ -255,6 +262,7 @@ def _build_assistant_tool_message( content=response or "", reasoning_content=reasoning_content or None, tool_calls=tool_calls_payload, + anthropic_content=anthropic_content, ) def _execute_tool_calls_from_canonical( diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py index 3a98b9645..ca984c73d 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py @@ -126,6 +126,7 @@ def parse_anthropic_response(response_json: dict[str, Any]) -> ChatCompletionRes content="\n".join(text_parts) if text_parts else None, reasoning_content="\n".join(thinking_parts) if thinking_parts else None, tool_calls=tool_calls, + anthropic_content=content_blocks, ) raw_usage = response_json.get("usage") @@ -209,6 +210,10 @@ def translate_non_tool_message(msg: dict[str, Any]) -> dict[str, Any]: if role not in {"user", "assistant"}: raise ValueError(f"Anthropic adapter does not support message role {role!r}.") + # Signed thinking blocks must be replayed with the entire original turn in order. + if role == "assistant" and msg.get("anthropic_content") is not None: + return {"role": role, "content": msg["anthropic_content"]} + content = msg.get("content") if role == "assistant" and msg.get("tool_calls"): translated_content = translate_content_blocks(content) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/types.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/types.py index a6946c34f..d10ae3db8 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/types.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/types.py @@ -54,6 +54,7 @@ class AssistantMessage: reasoning_content: str | None = None tool_calls: list[ToolCall] = field(default_factory=list) images: list[ImagePayload] = field(default_factory=list) + anthropic_content: list[dict[str, Any]] | None = None @dataclass diff --git a/packages/data-designer-engine/src/data_designer/engine/models/facade.py b/packages/data-designer-engine/src/data_designer/engine/models/facade.py index f7e27c4f9..7ccbfc56a 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/facade.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/facade.py @@ -235,7 +235,7 @@ def completion( allow_multiple_choices: bool = True, **kwargs: Any, ) -> ChatCompletionResponse: - message_payloads = [message.to_dict() for message in messages] + message_payloads = [message.to_dict(provider_type=self.model_provider.provider_type) for message in messages] logger.debug( f"Prompting model {self.model_name!r}...", extra={"model": self.model_name, "messages": message_payloads}, @@ -272,7 +272,7 @@ async def acompletion( allow_multiple_choices: bool = True, **kwargs: Any, ) -> ChatCompletionResponse: - message_payloads = [message.to_dict() for message in messages] + message_payloads = [message.to_dict(provider_type=self.model_provider.provider_type) for message in messages] logger.debug( f"Prompting model {self.model_name!r}...", extra={"model": self.model_name, "messages": message_payloads}, @@ -430,7 +430,13 @@ def generate( # No tool calls remaining to process response = (completion_response.message.content or "").strip() reasoning_trace = completion_response.message.reasoning_content - messages.append(ChatMessage.as_assistant(content=response, reasoning_content=reasoning_trace or None)) + messages.append( + ChatMessage.as_assistant( + content=response, + reasoning_content=reasoning_trace or None, + anthropic_content=completion_response.message.anthropic_content, + ) + ) parse_attempts += 1 try: @@ -549,7 +555,13 @@ async def agenerate( response = (completion_response.message.content or "").strip() reasoning_trace = completion_response.message.reasoning_content - messages.append(ChatMessage.as_assistant(content=response, reasoning_content=reasoning_trace or None)) + messages.append( + ChatMessage.as_assistant( + content=response, + reasoning_content=reasoning_trace or None, + anthropic_content=completion_response.message.anthropic_content, + ) + ) parse_attempts += 1 try: diff --git a/packages/data-designer-engine/src/data_designer/engine/models/utils.py b/packages/data-designer-engine/src/data_designer/engine/models/utils.py index 99a63adae..71398d996 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/utils.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/utils.py @@ -101,6 +101,7 @@ class ChatMessage: Each tool call contains 'id', 'type', and 'function' keys. tool_call_id: Optional ID linking a tool response to its corresponding tool call. Required for messages with role='tool'. + anthropic_content: Original Anthropic assistant blocks for request replay only. """ role: Literal["user", "assistant", "system", "tool"] @@ -108,12 +109,14 @@ class ChatMessage: reasoning_content: str | None = None tool_calls: list[dict[str, Any]] = field(default_factory=list) tool_call_id: str | None = None + anthropic_content: list[dict[str, Any]] | None = None - def to_dict(self) -> dict[str, Any]: + def to_dict(self, *, provider_type: str | None = None) -> dict[str, Any]: """Convert the message to a dictionary format for API calls. Content is normalized to a list of ChatML-style blocks to keep a consistent schema across traces and API payloads. + Native assistant blocks are included only for the matching provider's requests. Returns: A dictionary containing the message fields. Only includes non-empty @@ -126,6 +129,8 @@ def to_dict(self) -> dict[str, Any]: result["tool_calls"] = self.tool_calls if self.tool_call_id: result["tool_call_id"] = self.tool_call_id + if provider_type == "anthropic" and self.anthropic_content is not None: + result["anthropic_content"] = self.anthropic_content return result @classmethod @@ -139,6 +144,7 @@ def as_assistant( content: str = "", reasoning_content: str | None = None, tool_calls: list[dict[str, Any]] | None = None, + anthropic_content: list[dict[str, Any]] | None = None, ) -> ChatMessage: """Create an assistant message.""" return cls( @@ -146,6 +152,7 @@ def as_assistant( content=content, reasoning_content=reasoning_content, tool_calls=tool_calls or [], + anthropic_content=anthropic_content, ) @classmethod diff --git a/packages/data-designer-engine/tests/engine/mcp/test_mcp_facade.py b/packages/data-designer-engine/tests/engine/mcp/test_mcp_facade.py index 19bbf143f..81c1ba960 100644 --- a/packages/data-designer-engine/tests/engine/mcp/test_mcp_facade.py +++ b/packages/data-designer-engine/tests/engine/mcp/test_mcp_facade.py @@ -166,6 +166,19 @@ def test_process_completion_preserves_content( assert messages[0].content == "Hello, I can help with that." +@pytest.mark.parametrize("method", ["process_completion_response", "refuse_completion_response"]) +def test_completion_without_tools_preserves_native_content(stub_mcp_facade: MCPFacade, method: str) -> None: + content = [ + {"type": "thinking", "thinking": "", "signature": "signature"}, + {"type": "redacted_thinking", "data": "redacted-data"}, + {"type": "text", "text": " Answer\n"}, + ] + response = ChatCompletionResponse(message=AssistantMessage(content=" Answer\n", anthropic_content=content)) + messages = getattr(stub_mcp_facade, method)(response) + assert len(messages) == 1 + assert messages[0].anthropic_content == content + + def test_process_completion_preserves_reasoning_content( stub_mcp_facade: MCPFacade, mock_completion_response_with_reasoning: ChatCompletionResponse, diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_mcp.py b/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_mcp.py new file mode 100644 index 000000000..d3afcd4a5 --- /dev/null +++ b/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_mcp.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from typing import Any +from unittest.mock import MagicMock + +import httpx +import pytest + +from data_designer.config.mcp import ToolConfig +from data_designer.config.models import ModelConfig, ModelProvider +from data_designer.engine.mcp import io as mcp_io +from data_designer.engine.mcp.facade import DEFAULT_TOOL_REFUSAL_MESSAGE, MCPFacade +from data_designer.engine.mcp.registry import MCPRegistry, MCPToolDefinition, MCPToolResult +from data_designer.engine.model_provider import MCPProviderRegistry, ModelProviderRegistry +from data_designer.engine.models.clients.adapters.anthropic import AnthropicClient +from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode +from data_designer.engine.models.facade import ModelFacade +from data_designer.engine.models.parsers.errors import ParserException + + +def _parse_done(text: str) -> str: + if text != "done": + raise ParserException("Expected done") + return text + + +@pytest.fixture +def anthropic_tool_rounds() -> list[list[dict[str, Any]]]: + return [ + [ + {"type": "thinking", "thinking": "", "signature": f"empty-signature-{turn}"}, + {"type": "text", "text": " Looking up. "}, + {"type": "tool_use", "id": f"call-{turn}-a", "name": "lookup", "input": {"query": "a"}}, + {"type": "thinking", "thinking": " Check another source.\n", "signature": f"signature-{turn}"}, + {"type": "redacted_thinking", "data": f"redacted-data-{turn}"}, + {"type": "tool_use", "id": f"call-{turn}-b", "name": "lookup", "input": {"query": "b"}}, + {"type": "text", "text": " Waiting for results.\n"}, + ] + for turn in range(2) + ] + + +@pytest.fixture +def mcp_calls(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + monkeypatch.setattr( + mcp_io, + "list_tools", + MagicMock( + return_value=(MCPToolDefinition(name="lookup", description="Lookup", input_schema={"type": "object"}),) + ), + ) + calls = MagicMock(return_value=[MCPToolResult(content="Result a"), MCPToolResult(content="Result b")]) + monkeypatch.setattr(mcp_io, "call_tools", calls) + return calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("tool_turn_limit", [1, 2], ids=["refusal", "execution"]) +@pytest.mark.parametrize("retry", ["none", "correction", "restart"]) +async def test_generate_replays_anthropic_content( + is_async: bool, + tool_turn_limit: int, + retry: str, + anthropic_tool_rounds: list[list[dict[str, Any]]], + mcp_calls: MagicMock, + stub_tool_config: ToolConfig, + stub_secret_resolver: MagicMock, + stub_mcp_provider_registry: MCPProviderRegistry, +) -> None: + stub_tool_config.max_tool_call_turns = tool_turn_limit + mcp_registry = MCPRegistry( + tool_configs=[stub_tool_config], + secret_resolver=stub_secret_resolver, + mcp_provider_registry=stub_mcp_provider_registry, + mcp_facade_factory=MCPFacade, + ) + invalid_content = [ + {"type": "thinking", "thinking": "", "signature": "final-signature"}, + {"type": "redacted_thinking", "data": "final-redacted"}, + {"type": "text", "text": " invalid "}, + ] + responses = [{"content": blocks, "stop_reason": "tool_use"} for blocks in anthropic_tool_rounds] + if retry != "none": + responses.append({"content": invalid_content, "stop_reason": "end_turn"}) + responses.append({"content": [{"type": "text", "text": "done"}], "stop_reason": "end_turn"}) + payloads: list[dict[str, Any]] = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + payloads.append(json.loads(request.content)) + return httpx.Response(200, json=responses[len(payloads) - 1]) + + provider = ModelProvider(name="anthropic", endpoint="https://anthropic.test/v1", provider_type="anthropic") + transport = httpx.MockTransport(handle_request) + async with httpx.AsyncClient(transport=transport) as async_client: + with httpx.Client(transport=transport) as sync_client: + client = AnthropicClient( + provider_name=provider.name, + endpoint=provider.endpoint, + concurrency_mode=ClientConcurrencyMode.ASYNC if is_async else ClientConcurrencyMode.SYNC, + async_client=async_client if is_async else None, + sync_client=sync_client if not is_async else None, + ) + model = ModelFacade( + model_config=ModelConfig(alias="claude", model="claude-test", provider=provider.name), + model_provider_registry=ModelProviderRegistry(providers=[provider]), + client=client, + mcp_registry=mcp_registry, + ) + kwargs = { + "parser": _parse_done, + "tool_alias": stub_tool_config.tool_alias, + "max_correction_steps": int(retry == "correction"), + "max_conversation_restarts": int(retry == "restart"), + } + if is_async: + result, trace = await model.agenerate("Use the tools", **kwargs) + else: + result, trace = model.generate("Use the tools", **kwargs) + + assert result == "done" + assert len(payloads) == len(responses) + for request_index, payload in enumerate(payloads[1:], start=1): + assistants = [message for message in payload["messages"] if message["role"] == "assistant"] + assert [message["content"] for message in assistants[:2]] == anthropic_tool_rounds[:request_index] + tool_results = [ + block + for message in payload["messages"] + if message["role"] == "user" and isinstance(message["content"], list) + for block in message["content"] + if block["type"] == "tool_result" + ] + expected_ids = [f"call-{turn}-{suffix}" for turn in range(min(request_index, 2)) for suffix in ("a", "b")] + assert [block["tool_use_id"] for block in tool_results] == expected_ids + for block in tool_results: + expected = ( + DEFAULT_TOOL_REFUSAL_MESSAGE + if block["tool_use_id"].startswith("call-1") and tool_turn_limit == 1 + else f"Result {block['tool_use_id'][-1]}" + ) + assert block["content"] == expected + + if retry == "correction": + assert payloads[-1]["messages"][-2] == {"role": "assistant", "content": invalid_content} + elif retry == "restart": + assert payloads[-1]["messages"] == payloads[-2]["messages"] + assert mcp_calls.call_count == tool_turn_limit + for call in mcp_calls.call_args_list: + assert [(name, args) for _, name, args in call.args[0]] == [ + ("lookup", {"query": "a"}), + ("lookup", {"query": "b"}), + ] + trace_json = json.dumps([message.to_dict() for message in trace]) + assert "signature" not in trace_json + assert "redacted" not in trace_json + assert "anthropic_content" not in trace_json diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py b/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py index 6edc5e983..ee951529e 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py @@ -4,6 +4,7 @@ from __future__ import annotations import json +from typing import Any import pytest @@ -270,6 +271,27 @@ def count_reasoning_text(text: str) -> int: assert json.loads(response.message.tool_calls[0].arguments_json) == {"query": "weather"} +@pytest.mark.parametrize( + "content", + [ + [], + [{"type": "thinking", "thinking": "", "signature": "signature"}], + [{"type": "redacted_thinking", "data": "redacted-data"}], + [{"type": "text", "text": ""}, {"type": "text", "text": " Answer\n"}], + ], + ids=["empty", "empty-thinking", "redacted-thinking", "text"], +) +def test_anthropic_response_content_round_trips_unchanged(content: list[dict[str, Any]]) -> None: + response = parse_anthropic_response({"content": content}) + assert response.message.anthropic_content == content + message = ChatMessage.as_assistant( + content=response.message.content or "", + anthropic_content=response.message.anthropic_content, + ) + request = ChatCompletionRequest(model=MODEL, messages=[message.to_dict(provider_type="anthropic")]) + assert build_anthropic_payload(request)["messages"] == [{"role": "assistant", "content": content}] + + @pytest.mark.parametrize( "stop_reason", ["end_turn", "tool_use", "max_tokens", "model_context_window_exceeded"], diff --git a/packages/data-designer-engine/tests/engine/models/test_model_utils.py b/packages/data-designer-engine/tests/engine/models/test_model_utils.py index e8ba6852d..216b8c872 100644 --- a/packages/data-designer-engine/tests/engine/models/test_model_utils.py +++ b/packages/data-designer-engine/tests/engine/models/test_model_utils.py @@ -3,6 +3,8 @@ from __future__ import annotations +from typing import Any + import pytest from data_designer.config.models import Modality @@ -61,6 +63,22 @@ def test_chat_message_as_tool_accepts_multimodal_content() -> None: assert message.to_dict()["content"] == content +@pytest.mark.parametrize("provider_type", [None, "openai", "anthropic"]) +@pytest.mark.parametrize( + "native_content", + [None, [], [{"type": "thinking", "thinking": "", "signature": "signature"}]], +) +def test_chat_message_serializes_native_content_only_for_anthropic_requests( + provider_type: str | None, native_content: list[dict[str, Any]] | None +) -> None: + message = ChatMessage.as_assistant(content="Answer", anthropic_content=native_content) + expected = {"role": "assistant", "content": [{"type": "text", "text": "Answer"}]} + assert message.to_dict() == expected + if provider_type == "anthropic" and native_content is not None: + expected["anthropic_content"] = native_content + assert message.to_dict(provider_type=provider_type) == expected + + def test_prompt_to_messages_preserves_mixed_media_context_order() -> None: context = [ get_media_url_context(Modality.IMAGE.value, "https://example.com/image.png"), From da79496cc2ed23677070c138a9106b0c38929f50 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Fri, 30 Oct 2026 17:06:36 -0300 Subject: [PATCH 2/2] refactor: isolate provider replay handling Carry opaque provider-tagged replay data separately from canonical messages. Keep provider payload interpretation in the matching adapter. Fixes #1118 Signed-off-by: Andre Manoel --- architecture/models.md | 5 +++ .../src/data_designer/engine/mcp/facade.py | 14 +++---- .../clients/adapters/anthropic_translation.py | 20 ++++++---- .../engine/models/clients/types.py | 20 +++++++++- .../src/data_designer/engine/models/facade.py | 29 +++++++++++---- .../src/data_designer/engine/models/utils.py | 15 +++----- .../tests/engine/mcp/test_mcp_facade.py | 32 +++++++++++----- .../models/clients/test_anthropic_mcp.py | 2 + .../clients/test_anthropic_translation.py | 37 ++++++++++++++++--- .../models/clients/test_openai_compatible.py | 20 ++++++++++ .../engine/models/clients/test_parsing.py | 21 +++++++++++ .../tests/engine/models/test_facade.py | 26 +++++++++++++ .../tests/engine/models/test_model_utils.py | 27 ++++++-------- 13 files changed, 205 insertions(+), 63 deletions(-) diff --git a/architecture/models.md b/architecture/models.md index 36c797df6..e28914a0b 100644 --- a/architecture/models.md +++ b/architecture/models.md @@ -27,6 +27,11 @@ Defines the contract: sync/async chat, embeddings, image generation, `supports_* - **`OpenAICompatibleClient`** — native httpx adapter for OpenAI-compatible endpoints (NIM, vLLM, etc.) - **`AnthropicClient`** — native httpx adapter for the Anthropic Messages API +Assistant messages can carry opaque `ProviderReplayData` tagged by provider type. The model facade and MCP +preserve it through tool turns and retry checkpoints. `ChatCompletionRequest.message_replay_data` aligns +these envelopes with canonical message dictionaries; only the matching adapter interprets its payload. +`ChatMessage.to_dict()` excludes replay data from traces. + ### Client Factory `create_model_client` routes by provider type to the appropriate adapter. Optionally wraps with: diff --git a/packages/data-designer-engine/src/data_designer/engine/mcp/facade.py b/packages/data-designer-engine/src/data_designer/engine/mcp/facade.py index 8d583cead..289186750 100644 --- a/packages/data-designer-engine/src/data_designer/engine/mcp/facade.py +++ b/packages/data-designer-engine/src/data_designer/engine/mcp/facade.py @@ -11,7 +11,7 @@ from data_designer.engine.mcp.errors import DuplicateToolNameError, MCPConfigurationError, MCPToolError from data_designer.engine.mcp.registry import MCPToolDefinition from data_designer.engine.model_provider import MCPProviderRegistry -from data_designer.engine.models.clients.types import ChatCompletionResponse, ToolCall +from data_designer.engine.models.clients.types import ChatCompletionResponse, ProviderReplayData, ToolCall from data_designer.engine.models.utils import ChatMessage from data_designer.engine.secret_resolver import SecretResolver @@ -170,14 +170,14 @@ def process_completion_response( ChatMessage.as_assistant( content=response_content, reasoning_content=reasoning_content or None, - anthropic_content=message.anthropic_content, + replay_data=message.replay_data, ) ] # Has tool calls - execute and return all messages tool_call_dicts = _convert_canonical_tool_calls_to_dicts(tool_calls) assistant_message = self._build_assistant_tool_message( - response_content, tool_call_dicts, reasoning_content, anthropic_content=message.anthropic_content + response_content, tool_call_dicts, reasoning_content, replay_data=message.replay_data ) tool_messages = self._execute_tool_calls_from_canonical(tool_calls) @@ -218,14 +218,14 @@ def refuse_completion_response( ChatMessage.as_assistant( content=response_content, reasoning_content=reasoning_content or None, - anthropic_content=message.anthropic_content, + replay_data=message.replay_data, ) ] # Build assistant message with tool calls (same as normal) tool_call_dicts = _convert_canonical_tool_calls_to_dicts(tool_calls) assistant_message = self._build_assistant_tool_message( - response_content, tool_call_dicts, reasoning_content, anthropic_content=message.anthropic_content + response_content, tool_call_dicts, reasoning_content, replay_data=message.replay_data ) # Build refusal messages instead of executing tools @@ -247,7 +247,7 @@ def _build_assistant_tool_message( response: str | None, tool_calls: list[dict[str, Any]], reasoning_content: str | None = None, - anthropic_content: list[dict[str, Any]] | None = None, + replay_data: ProviderReplayData | None = None, ) -> ChatMessage: """Build the assistant message containing tool call requests.""" tool_calls_payload = [ @@ -262,7 +262,7 @@ def _build_assistant_tool_message( content=response or "", reasoning_content=reasoning_content or None, tool_calls=tool_calls_payload, - anthropic_content=anthropic_content, + replay_data=replay_data, ) def _execute_tool_calls_from_canonical( diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py index ca984c73d..b4e1767cd 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py @@ -18,6 +18,7 @@ ChatCompletionChoice, ChatCompletionRequest, ChatCompletionResponse, + ProviderReplayData, ToolCall, Usage, ) @@ -69,7 +70,7 @@ def merge_system_parts(parts: list[str | list[dict[str, Any]]]) -> str | list[di def build_anthropic_payload(request: ChatCompletionRequest) -> dict[str, Any]: """Build an Anthropic Messages API payload from a canonical request.""" - system_parts, messages = translate_request_messages(request.messages) + system_parts, messages = translate_request_messages(request.messages, request.message_replay_data) payload: dict[str, Any] = { "model": request.model, @@ -126,7 +127,7 @@ def parse_anthropic_response(response_json: dict[str, Any]) -> ChatCompletionRes content="\n".join(text_parts) if text_parts else None, reasoning_content="\n".join(thinking_parts) if thinking_parts else None, tool_calls=tool_calls, - anthropic_content=content_blocks, + replay_data=ProviderReplayData(provider_type="anthropic", payload=content_blocks), ) raw_usage = response_json.get("usage") @@ -149,12 +150,13 @@ def parse_anthropic_response(response_json: dict[str, Any]) -> ChatCompletionRes def translate_request_messages( messages: list[dict[str, Any]], + message_replay_data: list[ProviderReplayData | None] | None = None, ) -> tuple[list[str | list[dict[str, Any]]], list[dict[str, Any]]]: system_parts: list[str | list[dict[str, Any]]] = [] translated_messages: list[dict[str, Any]] = [] pending_tool_results: list[dict[str, Any]] = [] - for msg in messages: + for index, msg in enumerate(messages): role = msg.get("role") if role == "system": system_content = extract_system_content(msg.get("content")) @@ -170,7 +172,8 @@ def translate_request_messages( translated_messages.append({"role": "user", "content": pending_tool_results}) pending_tool_results = [] - translated_messages.append(translate_non_tool_message(msg)) + replay_data = message_replay_data[index] if message_replay_data else None + translated_messages.append(translate_non_tool_message(msg, replay_data)) if pending_tool_results: translated_messages.append({"role": "user", "content": pending_tool_results}) @@ -205,14 +208,17 @@ def extract_system_content(content: Any) -> str | list[dict[str, Any]] | None: return "\n".join(text_parts) -def translate_non_tool_message(msg: dict[str, Any]) -> dict[str, Any]: +def translate_non_tool_message(msg: dict[str, Any], replay_data: ProviderReplayData | None = None) -> dict[str, Any]: role = msg.get("role") if role not in {"user", "assistant"}: raise ValueError(f"Anthropic adapter does not support message role {role!r}.") # Signed thinking blocks must be replayed with the entire original turn in order. - if role == "assistant" and msg.get("anthropic_content") is not None: - return {"role": role, "content": msg["anthropic_content"]} + if role == "assistant" and replay_data is not None and replay_data.provider_type == "anthropic": + content = replay_data.payload + if not isinstance(content, list) or not all(isinstance(block, dict) for block in content): + raise ValueError("Anthropic replay data must be a list of content blocks") + return {"role": role, "content": content} content = msg.get("content") if role == "assistant" and msg.get("tool_calls"): diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/types.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/types.py index d10ae3db8..45a049a29 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/types.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/types.py @@ -48,13 +48,21 @@ class ToolCall: arguments_json: str +@dataclass(frozen=True) +class ProviderReplayData: + """Opaque provider-owned data required to replay an assistant message.""" + + provider_type: str + payload: object + + @dataclass class AssistantMessage: content: str | None = None reasoning_content: str | None = None tool_calls: list[ToolCall] = field(default_factory=list) images: list[ImagePayload] = field(default_factory=list) - anthropic_content: list[dict[str, Any]] | None = None + replay_data: ProviderReplayData | None = None @dataclass @@ -75,6 +83,12 @@ class ChatCompletionRequest: timeout: float | None = None extra_body: dict[str, Any] | None = None extra_headers: dict[str, str] | None = None + # Entries align with messages; an empty list means no replay data. + message_replay_data: list[ProviderReplayData | None] = field(default_factory=list) + + def __post_init__(self) -> None: + if self.message_replay_data and len(self.message_replay_data) != len(self.messages): + raise ValueError("message_replay_data must match the number of messages") @dataclass @@ -152,7 +166,9 @@ class TransportKwargs: - ``headers``: Extra HTTP headers to attach to the outgoing request. """ - _META_FIELDS: ClassVar[frozenset[str]] = frozenset({"extra_body", "extra_headers", "timeout"}) + _META_FIELDS: ClassVar[frozenset[str]] = frozenset( + {"extra_body", "extra_headers", "timeout", "message_replay_data"} + ) body: dict[str, Any] headers: dict[str, str] diff --git a/packages/data-designer-engine/src/data_designer/engine/models/facade.py b/packages/data-designer-engine/src/data_designer/engine/models/facade.py index 7ccbfc56a..fb5779db2 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/facade.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/facade.py @@ -25,6 +25,7 @@ EmbeddingResponse, ImageGenerationRequest, ImageGenerationResponse, + ProviderReplayData, Usage, ) from data_designer.engine.models.errors import ( @@ -235,7 +236,7 @@ def completion( allow_multiple_choices: bool = True, **kwargs: Any, ) -> ChatCompletionResponse: - message_payloads = [message.to_dict(provider_type=self.model_provider.provider_type) for message in messages] + message_payloads = [message.to_dict() for message in messages] logger.debug( f"Prompting model {self.model_name!r}...", extra={"model": self.model_name, "messages": message_payloads}, @@ -245,7 +246,9 @@ def completion( if not allow_multiple_choices: kwargs = self._drop_multi_choice_request_fields(kwargs) try: - request = self._build_chat_completion_request(message_payloads, kwargs) + request = self._build_chat_completion_request( + message_payloads, kwargs, message_replay_data=[message.replay_data for message in messages] + ) response = self._client.completion(request) logger.debug( f"Received completion from model {self.model_name!r}", @@ -272,7 +275,7 @@ async def acompletion( allow_multiple_choices: bool = True, **kwargs: Any, ) -> ChatCompletionResponse: - message_payloads = [message.to_dict(provider_type=self.model_provider.provider_type) for message in messages] + message_payloads = [message.to_dict() for message in messages] logger.debug( f"Prompting model {self.model_name!r}...", extra={"model": self.model_name, "messages": message_payloads}, @@ -282,7 +285,9 @@ async def acompletion( if not allow_multiple_choices: kwargs = self._drop_multi_choice_request_fields(kwargs) try: - request = self._build_chat_completion_request(message_payloads, kwargs) + request = self._build_chat_completion_request( + message_payloads, kwargs, message_replay_data=[message.replay_data for message in messages] + ) response = await self._client.acompletion(request) logger.debug( f"Received completion from model {self.model_name!r}", @@ -434,7 +439,7 @@ def generate( ChatMessage.as_assistant( content=response, reasoning_content=reasoning_trace or None, - anthropic_content=completion_response.message.anthropic_content, + replay_data=completion_response.message.replay_data, ) ) parse_attempts += 1 @@ -559,7 +564,7 @@ async def agenerate( ChatMessage.as_assistant( content=response, reasoning_content=reasoning_trace or None, - anthropic_content=completion_response.message.anthropic_content, + replay_data=completion_response.message.replay_data, ) ) parse_attempts += 1 @@ -839,10 +844,18 @@ def _get_mcp_facade(self, tool_alias: str | None) -> MCPFacade | None: raise MCPConfigurationError(f"Tool alias {tool_alias!r} is not registered.") from exc def _build_chat_completion_request( - self, messages: list[dict[str, Any]], kwargs: dict[str, Any] + self, + messages: list[dict[str, Any]], + kwargs: dict[str, Any], + *, + message_replay_data: list[ProviderReplayData | None], ) -> ChatCompletionRequest: """Build a ChatCompletionRequest from message payloads and consolidated kwargs.""" - request_fields: dict[str, Any] = {"model": self.model_name, "messages": messages} + request_fields: dict[str, Any] = { + "model": self.model_name, + "messages": messages, + "message_replay_data": message_replay_data, + } metadata: dict[str, Any] = {} for key, value in kwargs.items(): diff --git a/packages/data-designer-engine/src/data_designer/engine/models/utils.py b/packages/data-designer-engine/src/data_designer/engine/models/utils.py index 71398d996..47ff47ba7 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/utils.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/utils.py @@ -10,7 +10,7 @@ from pydantic import BaseModel if TYPE_CHECKING: - from data_designer.engine.models.clients.types import ChatCompletionResponse + from data_designer.engine.models.clients.types import ChatCompletionResponse, ProviderReplayData class GenerationTruncationReason(str, Enum): @@ -101,7 +101,7 @@ class ChatMessage: Each tool call contains 'id', 'type', and 'function' keys. tool_call_id: Optional ID linking a tool response to its corresponding tool call. Required for messages with role='tool'. - anthropic_content: Original Anthropic assistant blocks for request replay only. + replay_data: Opaque provider-owned data excluded from trace serialization. """ role: Literal["user", "assistant", "system", "tool"] @@ -109,14 +109,13 @@ class ChatMessage: reasoning_content: str | None = None tool_calls: list[dict[str, Any]] = field(default_factory=list) tool_call_id: str | None = None - anthropic_content: list[dict[str, Any]] | None = None + replay_data: ProviderReplayData | None = None - def to_dict(self, *, provider_type: str | None = None) -> dict[str, Any]: + def to_dict(self) -> dict[str, Any]: """Convert the message to a dictionary format for API calls. Content is normalized to a list of ChatML-style blocks to keep a consistent schema across traces and API payloads. - Native assistant blocks are included only for the matching provider's requests. Returns: A dictionary containing the message fields. Only includes non-empty @@ -129,8 +128,6 @@ def to_dict(self, *, provider_type: str | None = None) -> dict[str, Any]: result["tool_calls"] = self.tool_calls if self.tool_call_id: result["tool_call_id"] = self.tool_call_id - if provider_type == "anthropic" and self.anthropic_content is not None: - result["anthropic_content"] = self.anthropic_content return result @classmethod @@ -144,7 +141,7 @@ def as_assistant( content: str = "", reasoning_content: str | None = None, tool_calls: list[dict[str, Any]] | None = None, - anthropic_content: list[dict[str, Any]] | None = None, + replay_data: ProviderReplayData | None = None, ) -> ChatMessage: """Create an assistant message.""" return cls( @@ -152,7 +149,7 @@ def as_assistant( content=content, reasoning_content=reasoning_content, tool_calls=tool_calls or [], - anthropic_content=anthropic_content, + replay_data=replay_data, ) @classmethod diff --git a/packages/data-designer-engine/tests/engine/mcp/test_mcp_facade.py b/packages/data-designer-engine/tests/engine/mcp/test_mcp_facade.py index 81c1ba960..d2b8ad79b 100644 --- a/packages/data-designer-engine/tests/engine/mcp/test_mcp_facade.py +++ b/packages/data-designer-engine/tests/engine/mcp/test_mcp_facade.py @@ -14,7 +14,12 @@ from data_designer.engine.mcp.facade import DEFAULT_TOOL_REFUSAL_MESSAGE, MCPFacade from data_designer.engine.mcp.registry import MCPToolDefinition, MCPToolResult from data_designer.engine.model_provider import MCPProviderRegistry -from data_designer.engine.models.clients.types import AssistantMessage, ChatCompletionResponse, ToolCall +from data_designer.engine.models.clients.types import ( + AssistantMessage, + ChatCompletionResponse, + ProviderReplayData, + ToolCall, +) def _make_response( @@ -167,16 +172,23 @@ def test_process_completion_preserves_content( @pytest.mark.parametrize("method", ["process_completion_response", "refuse_completion_response"]) -def test_completion_without_tools_preserves_native_content(stub_mcp_facade: MCPFacade, method: str) -> None: - content = [ - {"type": "thinking", "thinking": "", "signature": "signature"}, - {"type": "redacted_thinking", "data": "redacted-data"}, - {"type": "text", "text": " Answer\n"}, - ] - response = ChatCompletionResponse(message=AssistantMessage(content=" Answer\n", anthropic_content=content)) +@pytest.mark.parametrize("tool_calls", [[], [ToolCall(id="call-1", name="lookup", arguments_json="{}")]]) +def test_completion_preserves_opaque_replay_data( + stub_mcp_facade: MCPFacade, method: str, tool_calls: list[ToolCall], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + mcp_io, + "list_tools", + MagicMock(return_value=(MCPToolDefinition(name="lookup", description="Lookup", input_schema={}),)), + ) + monkeypatch.setattr(mcp_io, "call_tools", MagicMock(return_value=[MCPToolResult(content="Result")])) + replay_data = ProviderReplayData(provider_type="future-provider", payload={"opaque": object()}) + response = ChatCompletionResponse( + message=AssistantMessage(content="Answer", tool_calls=tool_calls, replay_data=replay_data) + ) messages = getattr(stub_mcp_facade, method)(response) - assert len(messages) == 1 - assert messages[0].anthropic_content == content + assert len(messages) == 1 + len(tool_calls) + assert messages[0].replay_data is replay_data def test_process_completion_preserves_reasoning_content( diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_mcp.py b/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_mcp.py index d3afcd4a5..d0c81b1f0 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_mcp.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_mcp.py @@ -113,6 +113,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: ) kwargs = { "parser": _parse_done, + "system_prompt": "Be helpful", "tool_alias": stub_tool_config.tool_alias, "max_correction_steps": int(retry == "correction"), "max_conversation_restarts": int(retry == "restart"), @@ -158,3 +159,4 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert "signature" not in trace_json assert "redacted" not in trace_json assert "anthropic_content" not in trace_json + assert "replay_data" not in trace_json diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py b/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py index ee951529e..3ca02a214 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py @@ -30,7 +30,7 @@ translate_tool_result_content, translate_tool_result_message, ) -from data_designer.engine.models.clients.types import ChatCompletionRequest +from data_designer.engine.models.clients.types import ChatCompletionRequest, ProviderReplayData from data_designer.engine.models.usage import TokenCountSource from data_designer.engine.models.utils import ChatMessage @@ -278,20 +278,47 @@ def count_reasoning_text(text: str) -> int: [{"type": "thinking", "thinking": "", "signature": "signature"}], [{"type": "redacted_thinking", "data": "redacted-data"}], [{"type": "text", "text": ""}, {"type": "text", "text": " Answer\n"}], + [{"type": "future-block", "unknown": {"data": [1, None, "unchanged"]}}], ], - ids=["empty", "empty-thinking", "redacted-thinking", "text"], + ids=["empty", "empty-thinking", "redacted-thinking", "text", "future-block"], ) def test_anthropic_response_content_round_trips_unchanged(content: list[dict[str, Any]]) -> None: response = parse_anthropic_response({"content": content}) - assert response.message.anthropic_content == content + assert response.message.replay_data == ProviderReplayData(provider_type="anthropic", payload=content) message = ChatMessage.as_assistant( content=response.message.content or "", - anthropic_content=response.message.anthropic_content, + replay_data=response.message.replay_data, + ) + request = ChatCompletionRequest( + model=MODEL, messages=[message.to_dict()], message_replay_data=[message.replay_data] ) - request = ChatCompletionRequest(model=MODEL, messages=[message.to_dict(provider_type="anthropic")]) assert build_anthropic_payload(request)["messages"] == [{"role": "assistant", "content": content}] +@pytest.mark.parametrize("provider_type", ["openai", "future-provider"]) +def test_anthropic_request_ignores_unmatched_replay_data(provider_type: str) -> None: + message = ChatMessage.as_assistant( + content="Answer", replay_data=ProviderReplayData(provider_type=provider_type, payload=object()) + ) + request = ChatCompletionRequest( + model=MODEL, messages=[message.to_dict()], message_replay_data=[message.replay_data] + ) + assert build_anthropic_payload(request)["messages"] == [ + {"role": "assistant", "content": [{"type": "text", "text": "Answer"}]} + ] + + +@pytest.mark.parametrize("payload", [None, "opaque", {}, ["invalid-block"]]) +def test_anthropic_request_rejects_invalid_matching_replay_data(payload: object) -> None: + request = ChatCompletionRequest( + model=MODEL, + messages=[{"role": "assistant", "content": "Answer"}], + message_replay_data=[ProviderReplayData(provider_type="anthropic", payload=payload)], + ) + with pytest.raises(ValueError, match="Anthropic replay data must be a list of content blocks"): + build_anthropic_payload(request) + + @pytest.mark.parametrize( "stop_reason", ["end_turn", "tool_use", "max_tokens", "model_context_window_exceeded"], diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_openai_compatible.py b/packages/data-designer-engine/tests/engine/models/clients/test_openai_compatible.py index 124171c7a..3f9e2ec8c 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_openai_compatible.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_openai_compatible.py @@ -17,6 +17,7 @@ ChatCompletionRequest, EmbeddingRequest, ImageGenerationRequest, + ProviderReplayData, ) from data_designer.engine.models.usage import TokenCountSource from tests.engine.models.clients.conftest import make_mock_async_client, make_mock_sync_client @@ -82,6 +83,25 @@ def _image_response() -> dict[str, Any]: # --- Chat completion --- +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("provider_type", ["anthropic", "future-provider"]) +async def test_completion_ignores_provider_replay_data(is_async: bool, provider_type: str) -> None: + http_client = make_mock_async_client(_chat_response()) if is_async else make_mock_sync_client(_chat_response()) + client = _make_client(async_client=http_client) if is_async else _make_client(sync_client=http_client) + messages = [{"role": "assistant", "content": "Answer"}] + request = ChatCompletionRequest( + model=MODEL, + messages=messages, + message_replay_data=[ProviderReplayData(provider_type=provider_type, payload=object())], + ) + if is_async: + await client.acompletion(request) + else: + client.completion(request) + assert http_client.post.call_args.kwargs["json"] == {"model": MODEL, "messages": messages} + + def test_completion_maps_canonical_fields() -> None: response_json = _chat_response(content="Hello!", reasoning="step-by-step", reasoning_tokens=3) client = _make_client(sync_client=make_mock_sync_client(response_json)) diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_parsing.py b/packages/data-designer-engine/tests/engine/models/clients/test_parsing.py index bc8ce6a46..e0ea41c6a 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_parsing.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_parsing.py @@ -19,6 +19,7 @@ ChatCompletionResponse, EmbeddingRequest, ImageGenerationRequest, + ProviderReplayData, TransportKwargs, Usage, ) @@ -90,6 +91,26 @@ async def test_aparse_chat_completion_response_preserves_all_choices() -> None: # --- TransportKwargs.from_request: extra_body flattening (default) --- +@pytest.mark.parametrize(("message_count", "replay_count"), [(0, 1), (1, 2), (2, 1)]) +def test_chat_request_rejects_misaligned_replay_data(message_count: int, replay_count: int) -> None: + with pytest.raises(ValueError, match="message_replay_data must match the number of messages"): + ChatCompletionRequest( + model="m", + messages=[{"role": "user", "content": "Hi"}] * message_count, + message_replay_data=[None] * replay_count, + ) + + +def test_transport_kwargs_excludes_opaque_message_replay_data() -> None: + request = ChatCompletionRequest( + model="m", + messages=[{"role": "assistant", "content": "Answer"}], + message_replay_data=[ProviderReplayData(provider_type="future-provider", payload=object())], + temperature=0.2, + ) + assert TransportKwargs.from_request(request).body == {"temperature": 0.2} + + def test_extra_body_keys_are_flattened_into_body() -> None: request = ChatCompletionRequest( model="m", diff --git a/packages/data-designer-engine/tests/engine/models/test_facade.py b/packages/data-designer-engine/tests/engine/models/test_facade.py index a9a95e0f4..f356a6603 100644 --- a/packages/data-designer-engine/tests/engine/models/test_facade.py +++ b/packages/data-designer-engine/tests/engine/models/test_facade.py @@ -19,6 +19,7 @@ EmbeddingResponse, ImageGenerationResponse, ImagePayload, + ProviderReplayData, ToolCall, Usage, ) @@ -876,6 +877,31 @@ def test_completion_forwards_n_to_request( assert request.n == 4 +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_completion_preserves_opaque_replay_data( + stub_model_facade: ModelFacade, stub_model_client: MagicMock, is_async: bool +) -> None: + replay_data = ProviderReplayData(provider_type="future-provider", payload={"opaque": object()}) + messages = [ + ChatMessage.as_user("Question"), + ChatMessage.as_assistant(content="Answer", replay_data=replay_data), + ChatMessage.as_user("Correct it"), + ] + response = _make_response("Corrected") + stub_model_client.completion.return_value = response + stub_model_client.acompletion = AsyncMock(return_value=response) + if is_async: + await stub_model_facade.acompletion(messages) + request = stub_model_client.acompletion.call_args.args[0] + else: + stub_model_facade.completion(messages) + request = stub_model_client.completion.call_args.args[0] + assert request.messages == [message.to_dict() for message in messages] + assert request.message_replay_data == [None, replay_data, None] + assert request.message_replay_data[1] is replay_data + + def test_completion_forwards_tool_choice_and_keeps_configured_extra_body( stub_completion_messages: list[ChatMessage], stub_model_configs: list[Any], diff --git a/packages/data-designer-engine/tests/engine/models/test_model_utils.py b/packages/data-designer-engine/tests/engine/models/test_model_utils.py index 216b8c872..054c8535f 100644 --- a/packages/data-designer-engine/tests/engine/models/test_model_utils.py +++ b/packages/data-designer-engine/tests/engine/models/test_model_utils.py @@ -3,13 +3,16 @@ from __future__ import annotations -from typing import Any - import pytest from data_designer.config.models import Modality from data_designer.config.utils.media_helpers import get_media_base64_context, get_media_url_context -from data_designer.engine.models.clients.types import AssistantMessage, ChatCompletionChoice, ChatCompletionResponse +from data_designer.engine.models.clients.types import ( + AssistantMessage, + ChatCompletionChoice, + ChatCompletionResponse, + ProviderReplayData, +) from data_designer.engine.models.utils import ( ChatMessage, GenerationTruncationReason, @@ -63,20 +66,14 @@ def test_chat_message_as_tool_accepts_multimodal_content() -> None: assert message.to_dict()["content"] == content -@pytest.mark.parametrize("provider_type", [None, "openai", "anthropic"]) -@pytest.mark.parametrize( - "native_content", - [None, [], [{"type": "thinking", "thinking": "", "signature": "signature"}]], -) -def test_chat_message_serializes_native_content_only_for_anthropic_requests( - provider_type: str | None, native_content: list[dict[str, Any]] | None -) -> None: - message = ChatMessage.as_assistant(content="Answer", anthropic_content=native_content) +@pytest.mark.parametrize("provider_type", ["anthropic", "openai", "future-provider"]) +@pytest.mark.parametrize("payload", [None, [], {"opaque": object()}]) +def test_chat_message_excludes_opaque_replay_data(provider_type: str, payload: object) -> None: + replay_data = ProviderReplayData(provider_type=provider_type, payload=payload) + message = ChatMessage.as_assistant(content="Answer", replay_data=replay_data) expected = {"role": "assistant", "content": [{"type": "text", "text": "Answer"}]} assert message.to_dict() == expected - if provider_type == "anthropic" and native_content is not None: - expected["anthropic_content"] = native_content - assert message.to_dict(provider_type=provider_type) == expected + assert message.replay_data is replay_data def test_prompt_to_messages_preserves_mixed_media_context_order() -> None: