partners: allow to set Prefix in AIMessage (for MistralAI) (#28846)

**Description:**

Added ability to set `prefix` attribute to prevent error : 
```
httpx.HTTPStatusError: Error response 400 while fetching https://api.mistral.ai/v1/chat/completions: {"object":"error","message":"Expected last role User or Tool (or Assistant with prefix True) for serving but got assistant","type":"invalid_request_error","param":null,"code":null}
```

Co-authored-by: Sylvain DEPARTE <sylvain.departe@wizbii.com>
This commit is contained in:
Sylvain DEPARTE 2024-12-20 17:09:45 +01:00 committed by GitHub
parent 6d81137325
commit fcba567a77
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 6 additions and 0 deletions

View File

@ -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:
message_dict["prefix"] = message.additional_kwargs["prefix"]
return message_dict
elif isinstance(message, SystemMessage):
return dict(role="system", content=message.content)

View File

@ -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"),