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