diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 9576f4742c6..6495a18c358 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1,4 +1,4 @@ -import asyncio, base64, importlib, uuid +import asyncio, base64, hashlib, importlib, uuid import json import logging import os @@ -83,7 +83,12 @@ def test_function_call_prompt_lets_the_model_answer_after_a_function_result() -> @pytest.mark.parametrize( ("thought_signature", "expected"), [ - ("encoded-signature", "call_123__thought__encoded-signature"), + ( + "encoded-signature", + "call_123__checksum__" + + hashlib.sha256(b"encoded-signature").hexdigest() + + "__thought__encoded-signature", + ), (None, "call_123"), ("", "call_123"), ], @@ -101,8 +106,12 @@ def test_encode_tool_call_id_with_signature(thought_signature, expected): "function-signature", ), ( - {"id": encode_tool_call_id_with_signature("call_123", "embedded-signature")}, - "embedded-signature", + { + "id": encode_tool_call_id_with_signature( + "call_123", "ZW1iZWRkZWQtc2lnbmF0dXJl" + ) + }, + "ZW1iZWRkZWQtc2lnbmF0dXJl", ), ({}, None), ], diff --git a/tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py b/tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py index 30e934d9a13..9a2660ccfdf 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py +++ b/tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py @@ -322,13 +322,13 @@ def test_get_thought_signature_rejects_decodable_tampering(): ) assert _is_valid_thought_signature(decodable_but_wrong) is True - assert _get_thought_signature_from_tool({"id": tampered_id, "type": "function"}) is None + assert get_thought_signature_from_tool({"id": tampered_id, "type": "function"}) is None def test_get_thought_signature_rejects_unsigned_embedded_signature(): unsigned_id = f"call_abc123{THOUGHT_SIGNATURE_SEPARATOR}{REAL_SIGNATURE}" - assert _get_thought_signature_from_tool({"id": unsigned_id, "type": "function"}) is None + assert get_thought_signature_from_tool({"id": unsigned_id, "type": "function"}) is None def test_get_thought_signature_drops_client_mangled_id_suffix():