|
21 | 21 | convert_to_openai_messages,
|
22 | 22 | count_tokens_approximately,
|
23 | 23 | filter_messages,
|
| 24 | + get_buffer_string, |
24 | 25 | merge_message_runs,
|
25 | 26 | trim_messages,
|
26 | 27 | )
|
@@ -1395,3 +1396,64 @@ def test_count_tokens_approximately_mixed_content_types() -> None:
|
1395 | 1396 |
|
1396 | 1397 | # Ensure that count is consistent if we do one message at a time
|
1397 | 1398 | assert sum(count_tokens_approximately([m]) for m in messages) == token_count
|
| 1399 | + |
| 1400 | + |
| 1401 | +def test_get_buffer_string_with_structured_content() -> None: |
| 1402 | + """Test get_buffer_string with structured content in messages.""" |
| 1403 | + messages = [ |
| 1404 | + HumanMessage(content=[{"type": "text", "text": "Hello, world!"}]), |
| 1405 | + AIMessage(content=[{"type": "text", "text": "Hi there!"}]), |
| 1406 | + SystemMessage(content=[{"type": "text", "text": "System message"}]), |
| 1407 | + ] |
| 1408 | + expected = "Human: Hello, world!\nAI: Hi there!\nSystem: System message" |
| 1409 | + actual = get_buffer_string(messages) |
| 1410 | + assert actual == expected |
| 1411 | + |
| 1412 | + |
| 1413 | +def test_get_buffer_string_with_mixed_content() -> None: |
| 1414 | + """Test get_buffer_string with mixed content types in messages.""" |
| 1415 | + messages = [ |
| 1416 | + HumanMessage(content="Simple text"), |
| 1417 | + AIMessage(content=[{"type": "text", "text": "Structured text"}]), |
| 1418 | + SystemMessage(content=[{"type": "text", "text": "Another structured text"}]), |
| 1419 | + ] |
| 1420 | + expected = ( |
| 1421 | + "Human: Simple text\nAI: Structured text\nSystem: Another structured text" |
| 1422 | + ) |
| 1423 | + actual = get_buffer_string(messages) |
| 1424 | + assert actual == expected |
| 1425 | + |
| 1426 | + |
| 1427 | +def test_get_buffer_string_with_function_call() -> None: |
| 1428 | + """Test get_buffer_string with function call in additional_kwargs.""" |
| 1429 | + messages = [ |
| 1430 | + HumanMessage(content="Hello"), |
| 1431 | + AIMessage( |
| 1432 | + content="Hi", |
| 1433 | + additional_kwargs={ |
| 1434 | + "function_call": { |
| 1435 | + "name": "test_function", |
| 1436 | + "arguments": '{"arg": "value"}', |
| 1437 | + } |
| 1438 | + }, |
| 1439 | + ), |
| 1440 | + ] |
| 1441 | + # TODO: consider changing this |
| 1442 | + expected = ( |
| 1443 | + "Human: Hello\n" |
| 1444 | + "AI: Hi{'name': 'test_function', 'arguments': '{\"arg\": \"value\"}'}" |
| 1445 | + ) |
| 1446 | + actual = get_buffer_string(messages) |
| 1447 | + assert actual == expected |
| 1448 | + |
| 1449 | + |
| 1450 | +def test_get_buffer_string_with_empty_content() -> None: |
| 1451 | + """Test get_buffer_string with empty content in messages.""" |
| 1452 | + messages = [ |
| 1453 | + HumanMessage(content=[]), |
| 1454 | + AIMessage(content=""), |
| 1455 | + SystemMessage(content=[]), |
| 1456 | + ] |
| 1457 | + expected = "Human: \nAI: \nSystem: " |
| 1458 | + actual = get_buffer_string(messages) |
| 1459 | + assert actual == expected |
0 commit comments