fix: Preserve thought_signature in FunctionCall conversions between GenAI and A2A

Close: https://github.com/google/adk-python/issues/4311

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 877465519
This commit is contained in:
Xuan Yang
2026-03-02 10:28:46 -08:00
committed by Copybara-Service
parent a61c7e3880
commit f9c104faf7
2 changed files with 229 additions and 8 deletions
@@ -104,10 +104,25 @@ def convert_a2a_part_to_genai_part(
part.metadata[_get_adk_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)]
== A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
):
# Restore thought_signature if present
thought_signature = None
thought_sig_key = _get_adk_metadata_key('thought_signature')
if thought_sig_key in part.metadata:
sig_value = part.metadata[thought_sig_key]
if isinstance(sig_value, bytes):
thought_signature = sig_value
elif isinstance(sig_value, str):
try:
thought_signature = base64.b64decode(sig_value)
except Exception:
logger.warning(
'Failed to decode thought_signature: %s', sig_value
)
return genai_types.Part(
function_call=genai_types.FunctionCall.model_validate(
part.data, by_alias=True
)
),
thought_signature=thought_signature,
)
if (
part.metadata[_get_adk_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)]
@@ -214,16 +229,22 @@ def convert_genai_part_to_a2a_part(
# TODO once A2A defined how to service such information, migrate below
# logic accordingly
if part.function_call:
fc_metadata = {
_get_adk_metadata_key(
A2A_DATA_PART_METADATA_TYPE_KEY
): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
}
# Preserve thought_signature if present
if part.thought_signature is not None:
fc_metadata[_get_adk_metadata_key('thought_signature')] = (
base64.b64encode(part.thought_signature).decode('utf-8')
)
return a2a_types.Part(
root=a2a_types.DataPart(
data=part.function_call.model_dump(
by_alias=True, exclude_none=True
),
metadata={
_get_adk_metadata_key(
A2A_DATA_PART_METADATA_TYPE_KEY
): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
},
metadata=fc_metadata,
)
)
@@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import base64
import json
from unittest.mock import Mock
from unittest.mock import patch
@@ -74,7 +75,6 @@ class TestConvertA2aPartToGenaiPart:
# Arrange
test_bytes = b"test file content"
# A2A FileWithBytes expects base64-encoded string
import base64
base64_encoded = base64.b64encode(test_bytes).decode("utf-8")
a2a_part = a2a_types.Part(
@@ -328,7 +328,6 @@ class TestConvertGenaiPartToA2aPart:
assert isinstance(result.root, a2a_types.FilePart)
assert isinstance(result.root.file, a2a_types.FileWithBytes)
# A2A FileWithBytes now stores base64-encoded bytes to ensure round-trip compatibility
import base64
expected_base64 = base64.b64encode(test_bytes).decode("utf-8")
assert result.root.file.bytes == expected_base64
@@ -841,3 +840,204 @@ class TestNewConstants:
assert result.executable_code is not None
assert result.executable_code.language == genai_types.Language.PYTHON
assert result.executable_code.code == "print('Hello, World!')"
class TestThoughtSignaturePreservation:
"""Tests for thought_signature preservation in function call conversions."""
def test_genai_function_call_with_thought_signature_to_a2a(self):
"""Test that thought_signature is preserved when converting GenAI to A2A."""
# Arrange
function_call = genai_types.FunctionCall(
id="fc_gemini3",
name="my_tool",
args={"document": "test content"},
)
genai_part = genai_types.Part(
function_call=function_call,
thought_signature=b"gemini3_signature_bytes",
)
# Act
result = convert_genai_part_to_a2a_part(genai_part)
# Assert
assert result is not None
assert isinstance(result.root, a2a_types.DataPart)
assert (
result.root.metadata[
_get_adk_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)
]
== A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
)
# thought_signature should be base64 encoded in metadata
thought_sig_key = _get_adk_metadata_key("thought_signature")
assert thought_sig_key in result.root.metadata
assert (
base64.b64decode(result.root.metadata[thought_sig_key])
== b"gemini3_signature_bytes"
)
def test_genai_function_call_without_thought_signature_to_a2a(self):
"""Test function call without thought_signature doesn't add metadata key."""
# Arrange
function_call = genai_types.FunctionCall(
id="fc_regular",
name="regular_tool",
args={},
)
genai_part = genai_types.Part(function_call=function_call)
# Act
result = convert_genai_part_to_a2a_part(genai_part)
# Assert
assert result is not None
assert isinstance(result.root, a2a_types.DataPart)
# thought_signature key should not be present
thought_sig_key = _get_adk_metadata_key("thought_signature")
assert thought_sig_key not in result.root.metadata
def test_a2a_function_call_with_thought_signature_to_genai(self):
"""Test that thought_signature is restored when converting A2A to GenAI."""
# Arrange
a2a_part = a2a_types.Part(
root=a2a_types.DataPart(
data={
"id": "fc_gemini3",
"name": "my_tool",
"args": {"document": "test content"},
},
metadata={
_get_adk_metadata_key(
A2A_DATA_PART_METADATA_TYPE_KEY
): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL,
_get_adk_metadata_key("thought_signature"): (
base64.b64encode(b"restored_signature").decode("utf-8")
),
},
)
)
# Act
result = convert_a2a_part_to_genai_part(a2a_part)
# Assert
assert result is not None
assert result.function_call is not None
assert result.function_call.name == "my_tool"
# thought_signature should be decoded back to bytes
assert result.thought_signature == b"restored_signature"
def test_a2a_function_call_without_thought_signature_to_genai(self):
"""Test function call without thought_signature returns None for it."""
# Arrange
a2a_part = a2a_types.Part(
root=a2a_types.DataPart(
data={
"id": "fc_regular",
"name": "regular_tool",
"args": {},
},
metadata={
_get_adk_metadata_key(
A2A_DATA_PART_METADATA_TYPE_KEY
): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL,
},
)
)
# Act
result = convert_a2a_part_to_genai_part(a2a_part)
# Assert
assert result is not None
assert result.function_call is not None
assert result.function_call.name == "regular_tool"
# thought_signature should be None
assert result.thought_signature is None
def test_function_call_with_thought_signature_round_trip(self):
"""Test thought_signature is preserved in GenAI -> A2A -> GenAI round trip."""
# Arrange
original_signature = b"round_trip_signature_test"
function_call = genai_types.FunctionCall(
id="fc_round_trip",
name="round_trip_tool",
args={"key": "value"},
)
original_part = genai_types.Part(
function_call=function_call,
thought_signature=original_signature,
)
# Act - Convert GenAI -> A2A -> GenAI
a2a_part = convert_genai_part_to_a2a_part(original_part)
restored_part = convert_a2a_part_to_genai_part(a2a_part)
# Assert
assert restored_part is not None
assert restored_part.function_call is not None
assert restored_part.function_call.name == "round_trip_tool"
assert restored_part.thought_signature == original_signature
def test_a2a_function_call_with_bytes_thought_signature_to_genai(self):
"""Test that bytes thought_signature is used directly without decoding."""
# Arrange - metadata contains raw bytes (not base64 encoded)
a2a_part = a2a_types.Part(
root=a2a_types.DataPart(
data={
"id": "fc_bytes",
"name": "bytes_tool",
"args": {},
},
metadata={
_get_adk_metadata_key(
A2A_DATA_PART_METADATA_TYPE_KEY
): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL,
_get_adk_metadata_key(
"thought_signature"
): b"raw_bytes_signature",
},
)
)
# Act
result = convert_a2a_part_to_genai_part(a2a_part)
# Assert
assert result is not None
assert result.function_call is not None
# bytes should be used directly
assert result.thought_signature == b"raw_bytes_signature"
def test_a2a_function_call_with_invalid_base64_thought_signature(self):
"""Test that invalid base64 thought_signature logs warning and returns None."""
# Arrange - metadata contains invalid base64 string
a2a_part = a2a_types.Part(
root=a2a_types.DataPart(
data={
"id": "fc_invalid",
"name": "invalid_sig_tool",
"args": {},
},
metadata={
_get_adk_metadata_key(
A2A_DATA_PART_METADATA_TYPE_KEY
): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL,
_get_adk_metadata_key(
"thought_signature"
): "not_valid_base64!!!",
},
)
)
# Act
result = convert_a2a_part_to_genai_part(a2a_part)
# Assert
assert result is not None
assert result.function_call is not None
assert result.function_call.name == "invalid_sig_tool"
# thought_signature should be None due to decode failure
assert result.thought_signature is None