You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
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:
committed by
Copybara-Service
parent
a61c7e3880
commit
f9c104faf7
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user