feat: allow to set Prefix in AIMessage (for MistralAI)

langchain-ai · Dec 20, 2024 · e42a764 · e42a764
1 parent 5a351a1
commit e42a764
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git 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"]:
+            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
@@ -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"),