From f1f1f75782bb33aaf4034f54f0b7d2e198edf82b Mon Sep 17 00:00:00 2001 From: Jacob Lee Date: Thu, 11 Jul 2024 12:45:30 -0700 Subject: [PATCH] community[patch]: Make AzureML endpoint return AI messages for type assistant (#24085) --- .../langchain_community/chat_models/azureml_endpoint.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/libs/community/langchain_community/chat_models/azureml_endpoint.py b/libs/community/langchain_community/chat_models/azureml_endpoint.py index 2629c5a0ba2..c42d8a610a9 100644 --- a/libs/community/langchain_community/chat_models/azureml_endpoint.py +++ b/libs/community/langchain_community/chat_models/azureml_endpoint.py @@ -141,9 +141,8 @@ class CustomOpenAIChatContentFormatter(ContentFormatterBase): except (KeyError, IndexError, TypeError) as e: raise ValueError(self.format_error_msg.format(api_type=api_type)) from e return ChatGeneration( - message=BaseMessage( + message=AIMessage( content=choice.strip(), - type="assistant", ), generation_info=None, ) @@ -158,7 +157,9 @@ class CustomOpenAIChatContentFormatter(ContentFormatterBase): except (KeyError, IndexError, TypeError) as e: raise ValueError(self.format_error_msg.format(api_type=api_type)) from e return ChatGeneration( - message=BaseMessage( + message=AIMessage(content=choice["message"]["content"].strip()) + if choice["message"]["role"] == "assistant" + else BaseMessage( content=choice["message"]["content"].strip(), type=choice["message"]["role"], ),