From fe1304afc4da990e29ddc5490a76983f1f8c1623 Mon Sep 17 00:00:00 2001 From: ccurme Date: Fri, 26 Apr 2024 11:02:19 -0400 Subject: [PATCH] openai: add unit test (#20931) Test a helper function that was added earlier. --- .../tests/unit_tests/chat_models/test_base.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/libs/partners/openai/tests/unit_tests/chat_models/test_base.py b/libs/partners/openai/tests/unit_tests/chat_models/test_base.py index e4cc2cfddf..1bbaac4d6e 100644 --- a/libs/partners/openai/tests/unit_tests/chat_models/test_base.py +++ b/libs/partners/openai/tests/unit_tests/chat_models/test_base.py @@ -19,6 +19,7 @@ from langchain_openai import ChatOpenAI from langchain_openai.chat_models.base import ( _convert_dict_to_message, _convert_message_to_dict, + _format_message_content, ) @@ -287,3 +288,36 @@ def test_custom_token_counting() -> None: llm = ChatOpenAI(custom_get_token_ids=token_encoder) assert llm.get_token_ids("foo") == [1, 2, 3] + + +def test_format_message_content() -> None: + content: Any = "hello" + assert content == _format_message_content(content) + + content = None + assert content == _format_message_content(content) + + content = [] + assert content == _format_message_content(content) + + content = [ + {"type": "text", "text": "What is in this image?"}, + { + "type": "image_url", + "image_url": { + "url": "url.com", + }, + }, + ] + assert content == _format_message_content(content) + + content = [ + {"type": "text", "text": "hello"}, + { + "type": "tool_use", + "id": "toolu_01A09q90qw90lq917835lq9", + "name": "get_weather", + "input": {"location": "San Francisco, CA", "unit": "celsius"}, + }, + ] + assert [{"type": "text", "text": "hello"}] == _format_message_content(content)