community[patch]: Make AzureML endpoint return AI messages for type a…

…ssistant (langchain-ai#24085)
tibor-reiss · Jul 11, 2024 · f1f1f75 · f1f1f75
1 parent 4ba14ad
commit f1f1f75
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/libs/community/langchain_community/chat_models/azureml_endpoint.py b/libs/community/langchain_community/chat_models/azureml_endpoint.py
@@ -141,9 +141,8 @@ def format_response_payload(
             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 @@ def format_response_payload(
             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"],
                 ),