From dce9e0e12f04b903e7c59da807a4e2e83f7ab858 Mon Sep 17 00:00:00 2001 From: hongxing Date: Tue, 21 Jul 2026 17:52:36 +0200 Subject: [PATCH] fix(llm): preserve OpenAI tool call extra content --- .../llm/model_clients/base_model_client.py | 2 + .../llm/model_clients/openai_model_client.py | 18 +- .../core/foundation/llm/schema/message.py | 3 + .../foundation/llm/schema/message_chunk.py | 65 +++- .../core/foundation/llm/schema/tool_call.py | 5 +- .../test_openai_tool_call_extra_content.py | 345 ++++++++++++++++++ 6 files changed, 423 insertions(+), 15 deletions(-) create mode 100644 tests/unit_tests/core/foundation/llm/test_openai_tool_call_extra_content.py diff --git a/openjiuwen/core/foundation/llm/model_clients/base_model_client.py b/openjiuwen/core/foundation/llm/model_clients/base_model_client.py index c28f17856..b0099721d 100644 --- a/openjiuwen/core/foundation/llm/model_clients/base_model_client.py +++ b/openjiuwen/core/foundation/llm/model_clients/base_model_client.py @@ -258,6 +258,8 @@ def _convert_messages_to_dict(messages: Union[str, List[BaseMessage], List[dict] "arguments": tc.arguments } }) + if tc.extra_content is not None: + tool_calls_list[-1]["extra_content"] = tc.extra_content msg_dict["tool_calls"] = tool_calls_list if msg.reasoning_content: diff --git a/openjiuwen/core/foundation/llm/model_clients/openai_model_client.py b/openjiuwen/core/foundation/llm/model_clients/openai_model_client.py index 969f31b3c..9fa4eee11 100644 --- a/openjiuwen/core/foundation/llm/model_clients/openai_model_client.py +++ b/openjiuwen/core/foundation/llm/model_clients/openai_model_client.py @@ -557,6 +557,18 @@ async def _astream_with_parser( def _extract_reasoning_content(msg_or_delta: Any) -> Optional[str]: return getattr(msg_or_delta, 'reasoning_content', None) + @staticmethod + def _extract_tool_call_extra_content(tool_call: Any) -> Optional[dict[str, Any]]: + """Extract opaque provider metadata attached to an OpenAI-compatible tool call.""" + extra_content = getattr(tool_call, 'extra_content', None) + if extra_content is None: + model_extra = getattr(tool_call, 'model_extra', None) + if isinstance(model_extra, dict): + extra_content = model_extra.get('extra_content') + if hasattr(extra_content, 'model_dump'): + extra_content = extra_content.model_dump() + return dict(extra_content) if isinstance(extra_content, dict) else None + async def _parse_response( self, response: Any, @@ -590,7 +602,8 @@ async def _parse_response( type="function", name=function_name, arguments=function_arguments, - index=getattr(tc, 'index', idx) + index=getattr(tc, 'index', idx), + extra_content=self._extract_tool_call_extra_content(tc), ) tool_calls.append(tool_call) @@ -754,7 +767,8 @@ def _parse_stream_chunk(self, chunk: Any) -> Optional[AssistantMessageChunk]: type="function", name=function_name, arguments=function_arguments, - index=index + index=index, + extra_content=self._extract_tool_call_extra_content(tc_delta), ) tool_calls.append(tool_call) diff --git a/openjiuwen/core/foundation/llm/schema/message.py b/openjiuwen/core/foundation/llm/schema/message.py index f737bf128..bf3c9941c 100644 --- a/openjiuwen/core/foundation/llm/schema/message.py +++ b/openjiuwen/core/foundation/llm/schema/message.py @@ -69,6 +69,7 @@ def convert_openai_tool_calls_format(cls, data: Any) -> Any: 'arguments': tc['function'].get('arguments', ''), 'index': tc.get('index'), 'response_item_id': tc.get('response_item_id'), + 'extra_content': tc.get('extra_content'), } converted_tool_calls.append(converted_tc) else: @@ -99,6 +100,8 @@ def model_dump(self, **kwargs) -> dict[str, Any]: }) if call.response_item_id is not None: tool_calls[-1]["response_item_id"] = call.response_item_id + if call.extra_content is not None: + tool_calls[-1]["extra_content"] = call.extra_content result["tool_calls"] = tool_calls if self.usage_metadata is not None: result["usage_metadata"] = self.usage_metadata.model_dump(**kwargs) diff --git a/openjiuwen/core/foundation/llm/schema/message_chunk.py b/openjiuwen/core/foundation/llm/schema/message_chunk.py index de23cb4f2..80e444b12 100644 --- a/openjiuwen/core/foundation/llm/schema/message_chunk.py +++ b/openjiuwen/core/foundation/llm/schema/message_chunk.py @@ -129,6 +129,39 @@ def merge_pydantic_models(left: Any, right: Any) -> Any: return right +def _find_tool_call_merge_index( + merged_tool_calls: list[ToolCall], + incoming: ToolCall, +) -> int | None: + """Find the accumulated tool call that owns an incoming stream fragment.""" + if incoming.id: + for position, existing in enumerate(merged_tool_calls): + if existing.id == incoming.id: + return position + + # A provider may attach the ID after the first fragment. In that case, + # only use the stream index to complete an existing call without an ID; + # a different non-empty ID represents a new call even if its index was + # reused by the provider. + if incoming.index is not None: + for position, existing in enumerate(merged_tool_calls): + if existing.index == incoming.index and not existing.id: + return position + elif merged_tool_calls and not merged_tool_calls[-1].id: + return len(merged_tool_calls) - 1 + return None + + if incoming.index is not None: + for position, existing in enumerate(merged_tool_calls): + if existing.index == incoming.index: + return position + return None + + # Some compatible providers omit both ID and index from continuation + # fragments. Preserve the existing single-call fallback in that case. + return len(merged_tool_calls) - 1 if merged_tool_calls else None + + class BaseMessageChunk(BaseMessage): model_config = ConfigDict(arbitrary_types_allowed=True, json_encoders={type(None): lambda _: None}) @@ -172,22 +205,29 @@ def __add__(self, other: Any) -> "AssistantMessageChunk": arguments=tc.arguments, index=tc.index, response_item_id=tc.response_item_id, + extra_content=tc.extra_content, )) if other.tool_calls: for incoming in other.tool_calls: - if merged_tool_calls: - last = merged_tool_calls[-1] - same_id = (last.id and incoming.id and last.id == incoming.id) or (not last.id or not incoming.id) - if (same_id and hasattr(last, 'type') and last.type == 'function' - and hasattr(incoming, 'type') and incoming.type == 'function'): - merged_tool_calls[-1] = ToolCall( - id=last.id or incoming.id, - type=last.type or incoming.type, - name=(last.name if last.name else incoming.name) or "", - arguments=(last.arguments or "") + (incoming.arguments or ""), - index=last.index, - response_item_id=last.response_item_id or incoming.response_item_id, + merge_index = _find_tool_call_merge_index(merged_tool_calls, incoming) + if merge_index is not None: + existing = merged_tool_calls[merge_index] + if existing.type == 'function' and incoming.type == 'function': + merged_tool_calls[merge_index] = ToolCall( + id=existing.id or incoming.id, + type=existing.type or incoming.type, + name=(existing.name if existing.name else incoming.name) or "", + arguments=(existing.arguments or "") + (incoming.arguments or ""), + index=existing.index if existing.index is not None else incoming.index, + response_item_id=existing.response_item_id or incoming.response_item_id, + # Provider metadata is opaque: keep the latest value + # when present instead of merging its nested fields. + extra_content=( + incoming.extra_content + if incoming.extra_content is not None + else existing.extra_content + ), ) continue # otherwise, push as a new tool_call @@ -198,6 +238,7 @@ def __add__(self, other: Any) -> "AssistantMessageChunk": arguments=incoming.arguments, index=len(merged_tool_calls), response_item_id=incoming.response_item_id, + extra_content=incoming.extra_content, )) merged_finish_reason = other.finish_reason if other.finish_reason != "null" else self.finish_reason diff --git a/openjiuwen/core/foundation/llm/schema/tool_call.py b/openjiuwen/core/foundation/llm/schema/tool_call.py index eb9a40297..2f8cd7dfe 100644 --- a/openjiuwen/core/foundation/llm/schema/tool_call.py +++ b/openjiuwen/core/foundation/llm/schema/tool_call.py @@ -1,6 +1,6 @@ # -*- coding: UTF-8 -*- # Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. -from typing import Optional +from typing import Any, Optional from pydantic import BaseModel @@ -16,6 +16,8 @@ class ToolCall(BaseModel): index: Tool call index, used to distinguish multiple tool calls response_item_id: Optional provider response item ID for protocols that distinguish response item IDs from tool call IDs. + extra_content: Optional opaque provider metadata that must be preserved + when a tool call is sent back in conversation history. """ id: Optional[str] type: str @@ -23,3 +25,4 @@ class ToolCall(BaseModel): arguments: str index: Optional[int] = None response_item_id: Optional[str] = None + extra_content: Optional[dict[str, Any]] = None diff --git a/tests/unit_tests/core/foundation/llm/test_openai_tool_call_extra_content.py b/tests/unit_tests/core/foundation/llm/test_openai_tool_call_extra_content.py new file mode 100644 index 000000000..a7c7ae5e4 --- /dev/null +++ b/tests/unit_tests/core/foundation/llm/test_openai_tool_call_extra_content.py @@ -0,0 +1,345 @@ +# coding: utf-8 +"""Regression coverage for OpenAI-compatible per-tool-call metadata.""" + +from types import SimpleNamespace +from typing import Any + +import pytest +from openai.types.chat import ChatCompletionMessageFunctionToolCall +from pydantic import BaseModel + +from openjiuwen.core.foundation.llm.model_clients.base_model_client import BaseModelClient +from openjiuwen.core.foundation.llm.model_clients.openai_model_client import OpenAIModelClient +from openjiuwen.core.foundation.llm.schema.config import ModelClientConfig, ModelRequestConfig, ProviderType +from openjiuwen.core.foundation.llm.schema.message import AssistantMessage, BaseMessage + +_ABSENT = object() + + +class _ExtraContentModel(BaseModel): + google: dict[str, str] + + +def _client() -> OpenAIModelClient: + return OpenAIModelClient( + ModelRequestConfig(model="gemini-test"), + ModelClientConfig( + client_provider=ProviderType.OpenAI, + api_key="test-key", + api_base="https://example.test/v1", + verify_ssl=False, + ), + ) + + +def _serialize(message: BaseMessage) -> dict[str, Any]: + return BaseModelClient._convert_messages_to_dict([message])[0] + + +def _raw_tool_call( + *, + call_id: str, + name: str, + arguments: str, + index: int, + extra_content: Any = _ABSENT, + via_model_extra: bool = False, +) -> SimpleNamespace: + fields: dict[str, Any] = { + "id": call_id, + "index": index, + "function": SimpleNamespace(name=name, arguments=arguments), + } + if extra_content is not _ABSENT: + if via_model_extra: + fields["model_extra"] = {"extra_content": extra_content} + else: + fields["extra_content"] = extra_content + return SimpleNamespace(**fields) + + +def _response(*tool_calls: Any) -> SimpleNamespace: + return SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace( + content=None, + tool_calls=list(tool_calls), + reasoning_content=None, + ), + logprobs=None, + token_ids=None, + ) + ], + usage=None, + prompt_token_ids=None, + ) + + +def _stream_chunk(*tool_calls: Any, finish_reason: str | None = None) -> SimpleNamespace: + return SimpleNamespace( + choices=[ + SimpleNamespace( + delta=SimpleNamespace( + content=None, + reasoning_content=None, + tool_calls=list(tool_calls), + ), + finish_reason=finish_reason, + token_ids=None, + logprobs=None, + ) + ], + usage=None, + prompt_token_ids=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("via_model_extra", [False, True]) +async def test_tool_call_extra_content_survives_history_round_trip(via_model_extra: bool) -> None: + extra_content = {"google": {"thought_signature": "signature-a"}} + raw_tool_call = _raw_tool_call( + call_id="call-1", + name="list_files", + arguments='{"path":"."}', + index=0, + extra_content=extra_content, + via_model_extra=via_model_extra, + ) + + parsed = await _client()._parse_response(_response(raw_tool_call)) + + assert parsed.tool_calls is not None + assert parsed.tool_calls[0].extra_content == extra_content + message = _serialize(parsed) + assert message["tool_calls"][0]["extra_content"] == extra_content + + +@pytest.mark.asyncio +async def test_openai_sdk_tool_call_extra_content_survives_history_round_trip() -> None: + extra_content = {"google": {"thought_signature": "signature-a"}} + raw_tool_call = ChatCompletionMessageFunctionToolCall.model_validate( + { + "id": "call-1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "extra_content": extra_content, + } + ) + + parsed = await _client()._parse_response(_response(raw_tool_call)) + + assert parsed.tool_calls is not None + assert parsed.tool_calls[0].extra_content == extra_content + message = _serialize(parsed) + assert message["tool_calls"][0]["extra_content"] == extra_content + + +@pytest.mark.asyncio +async def test_pydantic_tool_call_extra_content_is_converted_to_dict() -> None: + raw_tool_call = _raw_tool_call( + call_id="call-1", + name="list_files", + arguments="{}", + index=0, + extra_content=_ExtraContentModel( + google={"thought_signature": "signature-a"}, + ), + ) + + parsed = await _client()._parse_response(_response(raw_tool_call)) + + assert parsed.tool_calls is not None + assert parsed.tool_calls[0].extra_content == { + "google": {"thought_signature": "signature-a"} + } + + +def test_nested_openai_tool_call_preserves_extra_content_in_model_dump() -> None: + extra_content = {"google": {"thought_signature": "signature-a"}} + message = AssistantMessage.model_validate( + { + "content": "", + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "extra_content": extra_content, + } + ], + } + ) + + assert message.tool_calls is not None + assert message.tool_calls[0].extra_content == extra_content + assert message.model_dump()["tool_calls"][0]["extra_content"] == extra_content + + +def test_stream_tool_call_extra_content_survives_parse_merge_and_history() -> None: + extra_content = {"google": {"thought_signature": "signature-a"}} + first = _client()._parse_stream_chunk( + _stream_chunk( + _raw_tool_call( + call_id="call-1", + name="list_files", + arguments="{", + index=0, + extra_content=extra_content, + ) + ) + ) + second = _client()._parse_stream_chunk( + _stream_chunk( + _raw_tool_call( + call_id="call-1", + name="", + arguments='"path":"."}', + index=0, + ), + finish_reason="tool_calls", + ) + ) + + assert first is not None + assert second is not None + merged = first + second + assert merged.tool_calls is not None + assert merged.tool_calls[0].arguments == '{"path":"."}' + assert merged.tool_calls[0].extra_content == extra_content + message = _serialize(merged) + assert message["tool_calls"][0]["extra_content"] == extra_content + + +def test_stream_tool_call_explicit_empty_extra_content_replaces_previous_value() -> None: + first = _client()._parse_stream_chunk( + _stream_chunk( + _raw_tool_call( + call_id="call-1", + name="list_files", + arguments="{", + index=0, + extra_content={"google": {"thought_signature": "signature-a"}}, + ) + ) + ) + second = _client()._parse_stream_chunk( + _stream_chunk( + _raw_tool_call( + call_id="call-1", + name="", + arguments="}", + index=0, + extra_content={}, + ), + finish_reason="tool_calls", + ) + ) + + assert first is not None + assert second is not None + merged = first + second + assert merged.tool_calls is not None + assert merged.tool_calls[0].extra_content == {} + assert _serialize(merged)["tool_calls"][0]["extra_content"] == {} + + +def test_parallel_stream_tool_calls_merge_by_index_and_keep_metadata_position() -> None: + first = _client()._parse_stream_chunk( + _stream_chunk( + _raw_tool_call( + call_id="call-1", + name="weather", + arguments="{", + index=0, + ), + _raw_tool_call( + call_id="call-2", + name="weather", + arguments="{", + index=1, + ), + ) + ) + second = _client()._parse_stream_chunk( + _stream_chunk( + _raw_tool_call( + call_id="", + name="", + arguments='"city":"Paris"}', + index=0, + extra_content={"google": {"thought_signature": "signature-a"}}, + ), + _raw_tool_call( + call_id="", + name="", + arguments='"city":"London"}', + index=1, + ), + finish_reason="tool_calls", + ) + ) + + assert first is not None + assert second is not None + merged = first + second + assert merged.tool_calls is not None + assert [call.arguments for call in merged.tool_calls] == [ + '{"city":"Paris"}', + '{"city":"London"}', + ] + assert merged.tool_calls[0].extra_content == { + "google": {"thought_signature": "signature-a"} + } + assert merged.tool_calls[1].extra_content is None + + serialized_calls = _serialize(merged)["tool_calls"] + assert serialized_calls[0]["extra_content"] == { + "google": {"thought_signature": "signature-a"} + } + assert "extra_content" not in serialized_calls[1] + + +@pytest.mark.asyncio +async def test_parallel_tool_calls_keep_metadata_attached_to_original_call() -> None: + extra_content = {"google": {"thought_signature": "signature-a"}} + parsed = await _client()._parse_response( + _response( + _raw_tool_call( + call_id="call-1", + name="weather", + arguments='{"city":"Paris"}', + index=0, + extra_content=extra_content, + ), + _raw_tool_call( + call_id="call-2", + name="weather", + arguments='{"city":"London"}', + index=1, + ), + ) + ) + + message = _serialize(parsed) + serialized_calls = message["tool_calls"] + assert serialized_calls[0]["extra_content"] == extra_content + assert "extra_content" not in serialized_calls[1] + + +@pytest.mark.asyncio +async def test_standard_openai_tool_call_does_not_emit_extra_content() -> None: + parsed = await _client()._parse_response( + _response( + _raw_tool_call( + call_id="call-1", + name="list_files", + arguments="{}", + index=0, + ) + ) + ) + + message = _serialize(parsed) + assert "extra_content" not in message["tool_calls"][0]