Python: fix(python/google): preserve thought_signature in Gemini function call parts (#13609)

### Motivation and Context Fixes #13480 When using Gemini models (2.5 Pro / 3 Pro) with thinking enabled, the API returns a `thought_signature` on function call `Part` objects. Per [Google's Thought Signatures documentation](https://ai.google.dev/gemini-api/docs/thought-signatures), this signature **must** be included in subsequent requests. Without it, follow-up turns with tool results fail with: ``` HTTP 400: Function call is missing a thought_signature in functionCall parts ``` The .NET connector was already fixed in #13418. This PR applies the equivalent fix to the Python connector. ### Description **Response parsing** (capture `thought_signature`): - `google_ai_chat_completion.py`: When creating `FunctionCallContent` from response parts, capture `part.thought_signature` into `FunctionCallContent.metadata` (both non-streaming and streaming paths) - `vertex_ai_chat_completion.py`: Same capture via `part_dict.get("thought_signature")` for Vertex AI **Request building** (restore `thought_signature`): - `google_ai/services/utils.py`: In `format_assistant_message()`, when a `FunctionCallContent` has `thought_signature` in its metadata, construct the `Part` directly with the `thought_signature` parameter instead of using `Part.from_function_call()` - `vertex_ai/services/utils.py`: In `format_assistant_message()`, include `thought_signature` in the `Part.from_dict()` dictionary when present **Backward compatible**: When `thought_signature` is not present (thinking disabled or older models), behavior is identical to before. ### Test Coverage Added 4 new tests: - `test_format_assistant_message_with_thought_signature` (Google AI) — verifies round-trip preservation - `test_format_assistant_message_without_thought_signature` (Google AI) — verifies backward compatibility - `test_format_assistant_message_with_thought_signature` (Vertex AI) — verifies round-trip preservation - `test_format_assistant_message_without_thought_signature` (Vertex AI) — verifies backward compatibility All 110 existing Google connector tests continue to pass. ### Contribution Checklist - [x] The code builds clean without any errors or warnings - [x] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) - [x] All unit tests pass, and I have added new tests - [x] I didn't break anyone 😄 --------- Co-authored-by: giulio-leone <[email protected]> Co-authored-by: Copilot <[email protected]>
2026-10-12 02:19:06 +00:00 · 2026-03-13 13:06:41 +00:00
parent b3cf64adf8
commit cc83889787
11 changed files with 480 additions and 16 deletions
@@ -303,6 +303,10 @@ class GoogleAIChatCompletion(GoogleAIBase, ChatCompletionClientBase):
                if part.text:
                    items.append(TextContent(text=part.text, inner_content=response, metadata=response_metadata))
                elif part.function_call:
+                    fc_metadata: dict[str, Any] = {}
+                    thought_sig = getattr(part, "thought_signature", None)
+                    if thought_sig:
+                        fc_metadata["thought_signature"] = thought_sig
                    items.append(
                        FunctionCallContent(
                            id=f"{part.function_call.name}_{idx!s}",
@@ -310,6 +314,7 @@ class GoogleAIChatCompletion(GoogleAIBase, ChatCompletionClientBase):
                                part.function_call.name  # type: ignore[arg-type]
                            ),
                            arguments={k: v for k, v in part.function_call.args.items()},  # type: ignore
+                            metadata=fc_metadata if fc_metadata else None,
                        )
                    )

@@ -360,6 +365,10 @@ class GoogleAIChatCompletion(GoogleAIBase, ChatCompletionClientBase):
                        )
                    )
                elif part.function_call:
+                    fc_metadata: dict[str, Any] = {}
+                    thought_sig = getattr(part, "thought_signature", None)
+                    if thought_sig:
+                        fc_metadata["thought_signature"] = thought_sig
                    items.append(
                        FunctionCallContent(
                            id=f"{part.function_call.name}_{idx!s}",
@@ -367,6 +376,7 @@ class GoogleAIChatCompletion(GoogleAIBase, ChatCompletionClientBase):
                                part.function_call.name  # type: ignore[arg-type]
                            ),
                            arguments={k: v for k, v in part.function_call.args.items()},  # type: ignore
+                            metadata=fc_metadata if fc_metadata else None,
                        )
                    )

@@ -91,12 +91,24 @@ def format_assistant_message(message: ChatMessageContent) -> list[Part]:
            if item.text:
                parts.append(Part.from_text(text=item.text))
        elif isinstance(item, FunctionCallContent):
-            parts.append(
-                Part.from_function_call(
-                    name=item.name,  # type: ignore[arg-type]
-                    args=json.loads(item.arguments) if isinstance(item.arguments, str) else item.arguments,  # type: ignore[arg-type]
+            thought_signature = item.metadata.get("thought_signature") if item.metadata else None
+            if thought_signature:
+                parts.append(
+                    Part(
+                        function_call={
+                            "name": item.name,  # type: ignore[arg-type]
+                            "args": json.loads(item.arguments) if isinstance(item.arguments, str) else item.arguments,
+                        },
+                        thought_signature=thought_signature,
+                    )
+                )
+            else:
+                parts.append(
+                    Part.from_function_call(
+                        name=item.name,  # type: ignore[arg-type]
+                        args=json.loads(item.arguments) if isinstance(item.arguments, str) else item.arguments,  # type: ignore[arg-type]
+                    )
                )
-            )
        elif isinstance(item, ImageContent):
            parts.append(_create_image_part(item))
        else:
@@ -2,7 +2,7 @@

 import json
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any

 from google.cloud.aiplatform_v1beta1.types.content import Candidate
 from vertexai.generative_models import FunctionDeclaration, Part, Tool, ToolConfig
@@ -89,14 +89,16 @@ def format_assistant_message(message: ChatMessageContent) -> list[Part]:
            if item.text:
                parts.append(Part.from_text(item.text))
        elif isinstance(item, FunctionCallContent):
-            parts.append(
-                Part.from_dict({
-                    "function_call": {
-                        "name": item.name,
-                        "args": json.loads(item.arguments) if isinstance(item.arguments, str) else item.arguments,
-                    }
-                })
-            )
+            part_dict: dict[str, Any] = {
+                "function_call": {
+                    "name": item.name,  # type: ignore[arg-type]
+                    "args": json.loads(item.arguments) if isinstance(item.arguments, str) else item.arguments,
+                }
+            }
+            thought_signature = item.metadata.get("thought_signature") if item.metadata else None
+            if thought_signature:
+                part_dict["thought_signature"] = thought_signature
+            parts.append(Part.from_dict(part_dict))
        elif isinstance(item, ImageContent):
            parts.append(_create_image_part(item))
        else:
@@ -252,6 +252,10 @@ class VertexAIChatCompletion(VertexAIBase, ChatCompletionClientBase):
            if "text" in part_dict:
                items.append(TextContent(text=part.text, inner_content=response, metadata=response_metadata))
            elif "function_call" in part_dict:
+                fc_metadata: dict[str, Any] = {}
+                thought_sig = part_dict.get("thought_signature")
+                if thought_sig:
+                    fc_metadata["thought_signature"] = thought_sig
                items.append(
                    FunctionCallContent(
                        id=f"{part.function_call.name}_{idx!s}",
@@ -259,6 +263,7 @@ class VertexAIChatCompletion(VertexAIBase, ChatCompletionClientBase):
                            part.function_call.name
                        ),
                        arguments={k: v for k, v in part.function_call.args.items()},
+                        metadata=fc_metadata if fc_metadata else None,
                    )
                )

@@ -309,6 +314,10 @@ class VertexAIChatCompletion(VertexAIBase, ChatCompletionClientBase):
                    )
                )
            elif "function_call" in part_dict:
+                fc_metadata_s: dict[str, Any] = {}
+                thought_sig_s = part_dict.get("thought_signature")
+                if thought_sig_s:
+                    fc_metadata_s["thought_signature"] = thought_sig_s
                items.append(
                    FunctionCallContent(
                        id=f"{part.function_call.name}_{idx!s}",
@@ -316,6 +325,7 @@ class VertexAIChatCompletion(VertexAIBase, ChatCompletionClientBase):
                            part.function_call.name
                        ),
                        arguments={k: v for k, v in part.function_call.args.items()},
+                        metadata=fc_metadata_s if fc_metadata_s else None,
                    )
                )

@@ -440,3 +440,197 @@ def test_google_ai_chat_completion_parse_chat_history_correctly(google_ai_unit_t
    assert parsed_chat_history[0].parts[0].text == "test_user_message"
    assert parsed_chat_history[1].role == "model"
    assert parsed_chat_history[1].parts[0].text == "test_assistant_message"
+
+
+# region deserialization (Part → FunctionCallContent round-trip)
+
+
+def test_create_chat_message_content_with_thought_signature(google_ai_unit_test_env) -> None:
+    """Test that thought_signature from a Part is deserialized into FunctionCallContent.metadata."""
+    from google.genai.types import (
+        Candidate,
+        Content,
+        GenerateContentResponse,
+        GenerateContentResponseUsageMetadata,
+        Part,
+    )
+    from google.genai.types import (
+        FinishReason as GFinishReason,
+    )
+
+    from semantic_kernel.contents.function_call_content import FunctionCallContent
+
+    thought_sig_value = b"test-thought-sig-bytes"
+    part = Part.from_function_call(name="test_function", args={"key": "value"})
+    part.thought_signature = thought_sig_value
+
+    candidate = Candidate()
+    candidate.index = 0
+    candidate.content = Content(role="user", parts=[part])
+    candidate.finish_reason = GFinishReason.STOP
+
+    response = GenerateContentResponse()
+    response.candidates = [candidate]
+    response.usage_metadata = GenerateContentResponseUsageMetadata(
+        prompt_token_count=0, cached_content_token_count=0, candidates_token_count=0, total_token_count=0
+    )
+
+    completion = GoogleAIChatCompletion()
+    result = completion._create_chat_message_content(response, candidate)
+
+    fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
+    assert len(fc_items) == 1
+    assert fc_items[0].metadata is not None
+    assert fc_items[0].metadata["thought_signature"] == thought_sig_value
+
+
+def test_create_chat_message_content_without_thought_signature(google_ai_unit_test_env) -> None:
+    """Test that FunctionCallContent works when Part has no thought_signature."""
+    from google.genai.types import (
+        Candidate,
+        Content,
+        GenerateContentResponse,
+        GenerateContentResponseUsageMetadata,
+        Part,
+    )
+    from google.genai.types import (
+        FinishReason as GFinishReason,
+    )
+
+    from semantic_kernel.contents.function_call_content import FunctionCallContent
+
+    part = Part.from_function_call(name="test_function", args={"key": "value"})
+
+    candidate = Candidate()
+    candidate.index = 0
+    candidate.content = Content(role="user", parts=[part])
+    candidate.finish_reason = GFinishReason.STOP
+
+    response = GenerateContentResponse()
+    response.candidates = [candidate]
+    response.usage_metadata = GenerateContentResponseUsageMetadata(
+        prompt_token_count=0, cached_content_token_count=0, candidates_token_count=0, total_token_count=0
+    )
+
+    completion = GoogleAIChatCompletion()
+    result = completion._create_chat_message_content(response, candidate)
+
+    fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
+    assert len(fc_items) == 1
+    assert "thought_signature" not in fc_items[0].metadata
+
+
+def test_create_streaming_chat_message_content_with_thought_signature(google_ai_unit_test_env) -> None:
+    """Test that thought_signature from a Part is deserialized in streaming path."""
+    from google.genai.types import (
+        Candidate,
+        Content,
+        GenerateContentResponse,
+        GenerateContentResponseUsageMetadata,
+        Part,
+    )
+    from google.genai.types import (
+        FinishReason as GFinishReason,
+    )
+
+    from semantic_kernel.contents.function_call_content import FunctionCallContent
+
+    thought_sig_value = b"streaming-thought-sig"
+    part = Part.from_function_call(name="stream_func", args={"a": "b"})
+    part.thought_signature = thought_sig_value
+
+    candidate = Candidate()
+    candidate.index = 0
+    candidate.content = Content(role="user", parts=[part])
+    candidate.finish_reason = GFinishReason.STOP
+
+    chunk = GenerateContentResponse()
+    chunk.candidates = [candidate]
+    chunk.usage_metadata = GenerateContentResponseUsageMetadata(
+        prompt_token_count=0, cached_content_token_count=0, candidates_token_count=0, total_token_count=0
+    )
+
+    completion = GoogleAIChatCompletion()
+    result = completion._create_streaming_chat_message_content(chunk, candidate)
+
+    fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
+    assert len(fc_items) == 1
+    assert fc_items[0].metadata is not None
+    assert fc_items[0].metadata["thought_signature"] == thought_sig_value
+
+
+def test_create_streaming_chat_message_content_without_thought_signature(google_ai_unit_test_env) -> None:
+    """Test that streaming FunctionCallContent works when Part lacks thought_signature."""
+    from google.genai.types import (
+        Candidate,
+        Content,
+        GenerateContentResponse,
+        GenerateContentResponseUsageMetadata,
+        Part,
+    )
+    from google.genai.types import (
+        FinishReason as GFinishReason,
+    )
+
+    from semantic_kernel.contents.function_call_content import FunctionCallContent
+
+    part = Part.from_function_call(name="stream_func", args={"a": "b"})
+
+    candidate = Candidate()
+    candidate.index = 0
+    candidate.content = Content(role="user", parts=[part])
+    candidate.finish_reason = GFinishReason.STOP
+
+    chunk = GenerateContentResponse()
+    chunk.candidates = [candidate]
+    chunk.usage_metadata = GenerateContentResponseUsageMetadata(
+        prompt_token_count=0, cached_content_token_count=0, candidates_token_count=0, total_token_count=0
+    )
+
+    completion = GoogleAIChatCompletion()
+    result = completion._create_streaming_chat_message_content(chunk, candidate)
+
+    fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
+    assert len(fc_items) == 1
+    assert "thought_signature" not in fc_items[0].metadata
+
+
+def test_create_chat_message_content_getattr_guard_on_missing_attribute(google_ai_unit_test_env) -> None:
+    """Test that getattr guard handles SDK versions where thought_signature doesn't exist on Part."""
+    from unittest.mock import MagicMock
+
+    from google.genai.types import (
+        GenerateContentResponse,
+        GenerateContentResponseUsageMetadata,
+    )
+
+    from semantic_kernel.contents.function_call_content import FunctionCallContent
+
+    # Create a mock Part that lacks 'thought_signature' attribute entirely
+    mock_part = MagicMock()
+    mock_part.text = None
+    mock_part.function_call.name = "test_func"
+    mock_part.function_call.args = {"x": "y"}
+    del mock_part.thought_signature  # simulate older SDK without the field
+
+    # Use a fully-mocked candidate to avoid Content pydantic validation
+    mock_candidate = MagicMock()
+    mock_candidate.index = 0
+    mock_candidate.content.parts = [mock_part]
+    mock_candidate.finish_reason = 1  # STOP
+
+    response = GenerateContentResponse()
+    response.candidates = [mock_candidate]
+    response.usage_metadata = GenerateContentResponseUsageMetadata(
+        prompt_token_count=0, cached_content_token_count=0, candidates_token_count=0, total_token_count=0
+    )
+
+    completion = GoogleAIChatCompletion()
+    result = completion._create_chat_message_content(response, mock_candidate)
+
+    fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
+    assert len(fc_items) == 1
+    assert "thought_signature" not in fc_items[0].metadata
+
+
+# endregion deserialization
@@ -113,3 +113,47 @@ def test_format_assistant_message_with_unsupported_items() -> None:

    with pytest.raises(ServiceInvalidRequestError):
        format_assistant_message(assistant_message)
+
+
+def test_format_assistant_message_with_thought_signature() -> None:
+    """Test that thought_signature is preserved in function call parts."""
+    import base64
+
+    thought_sig = base64.b64encode(b"test_thought_signature_data")
+    assistant_message = ChatMessageContent(
+        role=AuthorRole.ASSISTANT,
+        items=[
+            FunctionCallContent(
+                name="test_function",
+                arguments={"arg1": "value1"},
+                metadata={"thought_signature": thought_sig},
+            ),
+        ],
+    )
+
+    formatted = format_assistant_message(assistant_message)
+    assert len(formatted) == 1
+    assert isinstance(formatted[0], Part)
+    assert formatted[0].function_call.name == "test_function"
+    assert formatted[0].function_call.args == {"arg1": "value1"}
+    assert formatted[0].thought_signature == thought_sig
+
+
+def test_format_assistant_message_without_thought_signature() -> None:
+    """Test that function calls without thought_signature still work."""
+    assistant_message = ChatMessageContent(
+        role=AuthorRole.ASSISTANT,
+        items=[
+            FunctionCallContent(
+                name="test_function",
+                arguments={"arg1": "value1"},
+            ),
+        ],
+    )
+
+    formatted = format_assistant_message(assistant_message)
+    assert len(formatted) == 1
+    assert isinstance(formatted[0], Part)
+    assert formatted[0].function_call.name == "test_function"
+    assert formatted[0].function_call.args == {"arg1": "value1"}
+    assert not getattr(formatted[0], "thought_signature", None)
@@ -402,3 +402,144 @@ def test_vertex_ai_chat_completion_parse_chat_history_correctly(vertex_ai_unit_t
    assert parsed_chat_history[0].parts[0].text == "test_user_message"
    assert parsed_chat_history[1].role == "model"
    assert parsed_chat_history[1].parts[0].text == "test_assistant_message"
+
+
+# region thought_signature deserialization tests
+
+
+def test_create_chat_message_content_with_thought_signature(vertex_ai_unit_test_env) -> None:
+    """Test that thought_signature from a Part dict is deserialized into FunctionCallContent.metadata."""
+    from unittest.mock import MagicMock
+
+    from google.cloud.aiplatform_v1beta1.types.content import Candidate as GapicCandidate
+
+    from semantic_kernel.contents.function_call_content import FunctionCallContent
+
+    thought_sig_value = "vertex-test-thought-sig"
+
+    # Create a mock Part whose to_dict() returns thought_signature
+    mock_part = MagicMock()
+    mock_part.text = None
+    mock_part.to_dict.return_value = {
+        "function_call": {"name": "test_function", "args": {"key": "value"}},
+        "thought_signature": thought_sig_value,
+    }
+    mock_part.function_call.name = "test_function"
+    mock_part.function_call.args = {"key": "value"}
+
+    # Build a mock candidate with the mock part
+    mock_candidate = MagicMock()
+    mock_candidate.index = 0
+    mock_candidate.content.parts = [mock_part]
+    mock_candidate.finish_reason = GapicCandidate.FinishReason.STOP
+
+    # Build a mock response
+    mock_response = MagicMock()
+
+    completion = VertexAIChatCompletion()
+    result = completion._create_chat_message_content(mock_response, mock_candidate)
+
+    fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
+    assert len(fc_items) == 1
+    assert fc_items[0].metadata is not None
+    assert fc_items[0].metadata["thought_signature"] == thought_sig_value
+
+
+def test_create_chat_message_content_without_thought_signature(vertex_ai_unit_test_env) -> None:
+    """Test that FunctionCallContent works when Part dict has no thought_signature."""
+    from unittest.mock import MagicMock
+
+    from google.cloud.aiplatform_v1beta1.types.content import Candidate as GapicCandidate
+
+    from semantic_kernel.contents.function_call_content import FunctionCallContent
+
+    mock_part = MagicMock()
+    mock_part.text = None
+    mock_part.to_dict.return_value = {
+        "function_call": {"name": "test_function", "args": {"key": "value"}},
+    }
+    mock_part.function_call.name = "test_function"
+    mock_part.function_call.args = {"key": "value"}
+
+    mock_candidate = MagicMock()
+    mock_candidate.index = 0
+    mock_candidate.content.parts = [mock_part]
+    mock_candidate.finish_reason = GapicCandidate.FinishReason.STOP
+
+    mock_response = MagicMock()
+
+    completion = VertexAIChatCompletion()
+    result = completion._create_chat_message_content(mock_response, mock_candidate)
+
+    fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
+    assert len(fc_items) == 1
+    assert "thought_signature" not in fc_items[0].metadata
+
+
+def test_create_streaming_chat_message_content_with_thought_signature(vertex_ai_unit_test_env) -> None:
+    """Test that thought_signature from a Part dict is deserialized in streaming path."""
+    from unittest.mock import MagicMock
+
+    from google.cloud.aiplatform_v1beta1.types.content import Candidate as GapicCandidate
+
+    from semantic_kernel.contents.function_call_content import FunctionCallContent
+
+    thought_sig_value = "vertex-streaming-thought-sig"
+
+    mock_part = MagicMock()
+    mock_part.text = None
+    mock_part.to_dict.return_value = {
+        "function_call": {"name": "stream_func", "args": {"a": "b"}},
+        "thought_signature": thought_sig_value,
+    }
+    mock_part.function_call.name = "stream_func"
+    mock_part.function_call.args = {"a": "b"}
+
+    mock_candidate = MagicMock()
+    mock_candidate.index = 0
+    mock_candidate.content.parts = [mock_part]
+    mock_candidate.finish_reason = GapicCandidate.FinishReason.STOP
+
+    mock_chunk = MagicMock()
+
+    completion = VertexAIChatCompletion()
+    result = completion._create_streaming_chat_message_content(mock_chunk, mock_candidate, function_invoke_attempt=0)
+
+    fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
+    assert len(fc_items) == 1
+    assert fc_items[0].metadata is not None
+    assert fc_items[0].metadata["thought_signature"] == thought_sig_value
+
+
+def test_create_streaming_chat_message_content_without_thought_signature(vertex_ai_unit_test_env) -> None:
+    """Test that streaming FunctionCallContent works when Part dict lacks thought_signature."""
+    from unittest.mock import MagicMock
+
+    from google.cloud.aiplatform_v1beta1.types.content import Candidate as GapicCandidate
+
+    from semantic_kernel.contents.function_call_content import FunctionCallContent
+
+    mock_part = MagicMock()
+    mock_part.text = None
+    mock_part.to_dict.return_value = {
+        "function_call": {"name": "stream_func", "args": {"a": "b"}},
+    }
+    mock_part.function_call.name = "stream_func"
+    mock_part.function_call.args = {"a": "b"}
+
+    mock_candidate = MagicMock()
+    mock_candidate.index = 0
+    mock_candidate.content.parts = [mock_part]
+    mock_candidate.finish_reason = GapicCandidate.FinishReason.STOP
+
+    mock_chunk = MagicMock()
+
+    completion = VertexAIChatCompletion()
+    result = completion._create_streaming_chat_message_content(mock_chunk, mock_candidate, function_invoke_attempt=0)
+
+    fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
+    assert len(fc_items) == 1
+    assert "thought_signature" not in fc_items[0].metadata
+
+
+# endregion thought_signature deserialization tests
@@ -110,3 +110,50 @@ def test_format_assistant_message_with_unsupported_items() -> None:

    with pytest.raises(ServiceInvalidRequestError):
        format_assistant_message(assistant_message)
+
+
+def test_format_assistant_message_with_thought_signature() -> None:
+    """Test that thought_signature is preserved in function call parts for Vertex AI."""
+    import base64
+
+    thought_sig = base64.b64encode(b"test_thought_signature_data").decode("utf-8")
+    assistant_message = ChatMessageContent(
+        role=AuthorRole.ASSISTANT,
+        items=[
+            FunctionCallContent(
+                name="test_function",
+                arguments={"arg1": "value1"},
+                metadata={"thought_signature": thought_sig},
+            ),
+        ],
+    )
+
+    formatted = format_assistant_message(assistant_message)
+    assert len(formatted) == 1
+    assert isinstance(formatted[0], Part)
+    assert formatted[0].function_call.name == "test_function"
+    assert formatted[0].function_call.args == {"arg1": "value1"}
+    part_dict = formatted[0].to_dict()
+    assert "thought_signature" in part_dict
+    assert part_dict["thought_signature"] == thought_sig
+
+
+def test_format_assistant_message_without_thought_signature() -> None:
+    """Test that function calls without thought_signature still work for Vertex AI."""
+    assistant_message = ChatMessageContent(
+        role=AuthorRole.ASSISTANT,
+        items=[
+            FunctionCallContent(
+                name="test_function",
+                arguments={"arg1": "value1"},
+            ),
+        ],
+    )
+
+    formatted = format_assistant_message(assistant_message)
+    assert len(formatted) == 1
+    assert isinstance(formatted[0], Part)
+    assert formatted[0].function_call.name == "test_function"
+    assert formatted[0].function_call.args == {"arg1": "value1"}
+    part_dict = formatted[0].to_dict()
+    assert "thought_signature" not in part_dict
@@ -97,8 +97,8 @@ async def test_trace_streaming_chat_completion(
    chat_history,
    model_diagnostics_unit_test_env,
 ):
-    mock_span = mock_tracer.start_span.return_value
    # Setup
+    mock_span = mock_tracer.start_span.return_value
    chat_completion: ChatCompletionClientBase = MockChatCompletion(ai_model_id="ai_model_id")
    iterable = MagicMock(spec=AsyncGenerator)
    iterable.__aiter__.return_value = [mock_response]
@@ -186,8 +186,8 @@ async def test_trace_streaming_chat_completion_exception(
    chat_history,
    model_diagnostics_unit_test_env,
 ):
-    mock_span = mock_tracer.start_span.return_value
    # Setup
+    mock_span = mock_tracer.start_span.return_value
    chat_completion: ChatCompletionClientBase = MockChatCompletion(ai_model_id="ai_model_id")

    with patch.object(
@@ -80,6 +80,7 @@ async def test_trace_streaming_text_completion(
 ):
    mock_span = mock_tracer.start_span.return_value
    # Setup
+    mock_span = mock_tracer.start_span.return_value
    text_completion: TextCompletionClientBase = MockTextCompletion(ai_model_id="ai_model_id")
    iterable = MagicMock(spec=AsyncGenerator)
    iterable.__aiter__.return_value = [mock_response]
@@ -154,6 +155,7 @@ async def test_trace_streaming_text_completion_exception(
 ):
    mock_span = mock_tracer.start_span.return_value
    # Setup
+    mock_span = mock_tracer.start_span.return_value
    text_completion: TextCompletionClientBase = MockTextCompletion(ai_model_id="ai_model_id")

    with patch.object(MockTextCompletion, "_inner_get_streaming_text_contents", side_effect=ServiceResponseException()):
@@ -76,6 +76,7 @@ async def test_trace_text_completion(
 ):
    mock_span = mock_tracer.start_span.return_value
    # Setup
+    mock_span = mock_tracer.start_span.return_value
    text_completion: TextCompletionClientBase = MockTextCompletion(ai_model_id="ai_model_id")

    with patch.object(MockTextCompletion, "_inner_get_text_contents", return_value=mock_response):
@@ -147,6 +148,7 @@ async def test_trace_text_completion_exception(
 ):
    mock_span = mock_tracer.start_span.return_value
    # Setup
+    mock_span = mock_tracer.start_span.return_value
    text_completion: TextCompletionClientBase = MockTextCompletion(ai_model_id="ai_model_id")

    with patch.object(MockTextCompletion, "_inner_get_text_contents", side_effect=ServiceResponseException()):