@@ -253,18 +253,6 @@ def test_main_function_serialization_format(self, sample_messages):
253253 assert "role" in msg
254254 assert "content" in msg
255255
256- def test_main_function_respects_custom_limit (self , large_messages ):
257- """Test that the main function respects custom byte limits"""
258- custom_limit = MAX_GEN_AI_MESSAGE_BYTES // 250 # 2KB limit
259- result = truncate_and_serialize_messages (large_messages , max_bytes = custom_limit )
260-
261- # Should return AnnotatedValue due to truncation
262- assert isinstance (result , AnnotatedValue )
263-
264- # Should respect the custom limit
265- result_size = len (result .value .encode ("utf-8" ))
266- assert result_size <= custom_limit
267-
268256 def test_main_function_default_limit (self , sample_messages ):
269257 """Test that the main function uses the default limit correctly"""
270258 result = truncate_and_serialize_messages (sample_messages )
@@ -297,10 +285,10 @@ def test_messages_with_special_characters(self):
297285 ]
298286
299287 result = truncate_and_serialize_messages (messages )
300- assert result [ "serialized_data" ] is not None
288+ assert result is not None
301289
302290 # Should be valid JSON
303- parsed = json .loads (result [ "serialized_data" ] )
291+ parsed = json .loads (result )
304292 assert len (parsed ) == 2
305293 assert "🌍" in parsed [0 ]["content" ]
306294
@@ -320,26 +308,17 @@ def test_messages_with_nested_structures(self):
320308 ]
321309
322310 result = truncate_and_serialize_messages (messages )
323- assert result [ "serialized_data" ] is not None
311+ assert result is not None
324312
325313 # Should preserve the structure
326- parsed = json .loads (result ["serialized_data" ])
314+ # Handle both string and AnnotatedValue return types
315+ if isinstance (result , AnnotatedValue ):
316+ parsed = json .loads (result .value )
317+ else :
318+ parsed = json .loads (result )
327319 assert "metadata" in parsed [0 ]
328320 assert "tool_calls" in parsed [1 ]
329321
330- def test_very_small_limit (self , sample_messages ):
331- """Test behavior with extremely small size limit"""
332- tiny_limit = 10 # 10 bytes - extremely small limit
333- result = truncate_and_serialize_messages (sample_messages , max_bytes = tiny_limit )
334-
335- # With such a small limit, likely all messages will be removed
336- if result ["serialized_data" ] is None :
337- assert result ["metadata" ]["truncated_count" ] == 0
338- else :
339- # If any data remains, it should be under the limit
340- size = len (result ["serialized_data" ].encode ("utf-8" ))
341- assert size <= tiny_limit
342-
343322 def test_messages_with_none_values (self ):
344323 """Test messages containing None values"""
345324 messages = [
@@ -348,10 +327,14 @@ def test_messages_with_none_values(self):
348327 ]
349328
350329 result = truncate_and_serialize_messages (messages )
351- assert result [ "serialized_data" ] is not None
330+ assert result is not None
352331
353332 # Should handle None values gracefully
354- parsed = json .loads (result ["serialized_data" ])
333+ # Handle both string and AnnotatedValue return types
334+ if isinstance (result , AnnotatedValue ):
335+ parsed = json .loads (result .value )
336+ else :
337+ parsed = json .loads (result )
355338 assert len (parsed ) == 2
356339
357340 def test_truncation_keeps_most_recent (self ):
@@ -369,8 +352,9 @@ def test_truncation_keeps_most_recent(self):
369352 small_limit = MAX_GEN_AI_MESSAGE_BYTES // 500 # 1KB limit to force truncation
370353 result = truncate_and_serialize_messages (messages , max_bytes = small_limit )
371354
372- if result ["serialized_data" ]:
373- parsed = json .loads (result ["serialized_data" ])
355+ if result :
356+ assert isinstance (result , AnnotatedValue )
357+ parsed = json .loads (result .value )
374358 if parsed :
375359 # The last remaining message should be from the end of the original list
376360 last_kept_content = parsed [- 1 ]["content" ]
0 commit comments