From bc53a9dec1a1f6943d8dc6edb4c8df11aa1001c7 Mon Sep 17 00:00:00 2001 From: Sylvain DEPARTE Date: Fri, 20 Dec 2024 11:39:59 +0100 Subject: [PATCH] feat: allow to set Prefix in AIMessage (for MistralAI) --- libs/partners/mistralai/langchain_mistralai/chat_models.py | 2 ++ libs/partners/mistralai/tests/unit_tests/test_chat_models.py | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/libs/partners/mistralai/langchain_mistralai/chat_models.py b/libs/partners/mistralai/langchain_mistralai/chat_models.py index be973f3b9ec78..c9b6368bf3636 100644 --- a/libs/partners/mistralai/langchain_mistralai/chat_models.py +++ b/libs/partners/mistralai/langchain_mistralai/chat_models.py @@ -345,6 +345,8 @@ def _convert_message_to_mistral_chat_message( message_dict["content"] = "" else: message_dict["content"] = message.content + if "prefix" in message.additional_kwargs and message.additional_kwargs["prefix"] == True: + message_dict["prefix"] = True return message_dict elif isinstance(message, SystemMessage): return dict(role="system", content=message.content) diff --git a/libs/partners/mistralai/tests/unit_tests/test_chat_models.py b/libs/partners/mistralai/tests/unit_tests/test_chat_models.py index f90507c687723..27bd73f5debda 100644 --- a/libs/partners/mistralai/tests/unit_tests/test_chat_models.py +++ b/libs/partners/mistralai/tests/unit_tests/test_chat_models.py @@ -92,6 +92,10 @@ def test_mistralai_initialization_baseurl_env(env_var_name: str) -> None: AIMessage(content="Hello"), dict(role="assistant", content="Hello"), ), + ( + AIMessage(content="{", additional_kwargs={"prefix": True}), + dict(role="assistant", content="{", prefix=True), + ), ( ChatMessage(role="assistant", content="Hello"), dict(role="assistant", content="Hello"),