diff --git a/py/src/braintrust/prompt_cache/test_disk_cache.py b/py/src/braintrust/prompt_cache/test_disk_cache.py index 7f49986e..64af065a 100644 --- a/py/src/braintrust/prompt_cache/test_disk_cache.py +++ b/py/src/braintrust/prompt_cache/test_disk_cache.py @@ -202,6 +202,36 @@ def test_serializer_handles_complex_objects(self): self.assertEqual(result.as_dict(), test_prompt.as_dict()) + def test_prompt_with_string_function_call_round_trips(self): + cache = disk_cache.DiskCache[prompt.PromptSchema]( + cache_dir=self.cache_dir, + serializer=lambda value: value.as_dict(), + deserializer=prompt.PromptSchema.from_dict_deep, + ) + test_prompt = prompt.PromptSchema( + id="prompt-id", + project_id="project-id", + _xact_id="version", + name="templated function call", + slug="templated-call", + description=None, + tags=None, + prompt_data=prompt.PromptData( + prompt=prompt.PromptChatBlock( + messages=[prompt.PromptMessage(role="assistant", content="", function_call="{{function_call}}")] + ) + ), + ) + cache.set("test-key", test_prompt) + + result = cache.get("test-key") + + self.assertEqual(result.as_dict(), test_prompt.as_dict()) + block = result.prompt_data.prompt + assert isinstance(block, prompt.PromptChatBlock) + self.assertIsInstance(block.messages[0], prompt.PromptMessage) + self.assertEqual(block.messages[0].function_call, "{{function_call}}") + def test_throw_on_deserializer_error(self): """Test that deserializer errors are propagated.""" diff --git a/py/src/braintrust/serializable_data_class.py b/py/src/braintrust/serializable_data_class.py index ba32ecca..84af4839 100644 --- a/py/src/braintrust/serializable_data_class.py +++ b/py/src/braintrust/serializable_data_class.py @@ -45,7 +45,7 @@ def from_dict_deep(cls, d: dict): if t == type(None) and v is None: filtered[k] = None break - if isinstance(t, type) and issubclass(t, SerializableDataClass) and v is not None: + if isinstance(t, type) and issubclass(t, SerializableDataClass) and isinstance(v, dict): try: filtered[k] = t.from_dict_deep(v) break diff --git a/py/src/braintrust/test_serializable_data_class.py b/py/src/braintrust/test_serializable_data_class.py index e31e2078..fc97cc1f 100644 --- a/py/src/braintrust/test_serializable_data_class.py +++ b/py/src/braintrust/test_serializable_data_class.py @@ -1,6 +1,7 @@ import unittest from dataclasses import dataclass +from .prompt import FunctionCall, PromptMessage from .serializable_data_class import SerializableDataClass @@ -23,6 +24,24 @@ class PromptSchema(SerializableDataClass): class TestSerializableDataClass(unittest.TestCase): + def test_union_preserves_string_and_none_values(self): + for function_call in ("{{function_call}}", None): + with self.subTest(function_call=function_call): + message = PromptMessage.from_dict_deep( + {"role": "assistant", "content": "", "function_call": function_call} + ) + self.assertEqual(message.function_call, function_call) + + def test_union_deserializes_nested_dataclass(self): + message = PromptMessage.from_dict_deep( + { + "role": "assistant", + "content": "", + "function_call": {"name": "lookup", "arguments": "{}"}, + } + ) + self.assertEqual(message.function_call, FunctionCall(name="lookup", arguments="{}")) + def test_from_dict_deep_with_none_values(self): """Test that from_dict_deep correctly handles None values in nested objects.""" test_dict = {