From 10adec5f1bc1babbd7f5cbea8290d8b1e62554ba Mon Sep 17 00:00:00 2001 From: thehunmonkgroup Date: Tue, 20 Jun 2023 11:25:55 -0400 Subject: [PATCH] add FunctionMessage support to `_convert_dict_to_message()` in OpenAI chat model (#6382) Already supported in the reverse operation in `_convert_message_to_dict()`, this just provides parity. @hwchase17 @agola11 --------- Co-authored-by: Harrison Chase --- langchain/chat_models/openai.py | 2 ++ tests/unit_tests/chat_models/test_openai.py | 25 +++++++++++++++++++++ 2 files changed, 27 insertions(+) create mode 100644 tests/unit_tests/chat_models/test_openai.py diff --git a/langchain/chat_models/openai.py b/langchain/chat_models/openai.py index c2ce0ca3..83e43ae0 100644 --- a/langchain/chat_models/openai.py +++ b/langchain/chat_models/openai.py @@ -106,6 +106,8 @@ def _convert_dict_to_message(_dict: Mapping[str, Any]) -> BaseMessage: return AIMessage(content=content, additional_kwargs=additional_kwargs) elif role == "system": return SystemMessage(content=_dict["content"]) + elif role == "function": + return FunctionMessage(content=_dict["content"], name=_dict["name"]) else: return ChatMessage(content=_dict["content"], role=role) diff --git a/tests/unit_tests/chat_models/test_openai.py b/tests/unit_tests/chat_models/test_openai.py new file mode 100644 index 00000000..9720eb98 --- /dev/null +++ b/tests/unit_tests/chat_models/test_openai.py @@ -0,0 +1,25 @@ +"""Test OpenAI Chat API wrapper.""" + +import json + +from langchain.chat_models.openai import ( + _convert_dict_to_message, +) +from langchain.schema import ( + FunctionMessage, +) + + +def test_function_message_dict_to_function_message() -> None: + content = json.dumps({"result": "Example #1"}) + name = "test_function" + result = _convert_dict_to_message( + { + "role": "function", + "name": name, + "content": content, + } + ) + assert isinstance(result, FunctionMessage) + assert result.name == name + assert result.content == content