From 7eac47f4c87e95780766dfea7a74b3ab84664e01 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 26 Jan 2024 17:08:10 -0300 Subject: [PATCH] Add support for message conversion in ChatOutputResponse class --- src/backend/langflow/utils/schemas.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/utils/schemas.py b/src/backend/langflow/utils/schemas.py index 8f0372d8f..2a411981c 100644 --- a/src/backend/langflow/utils/schemas.py +++ b/src/backend/langflow/utils/schemas.py @@ -1,10 +1,18 @@ +from typing import Dict, List, Optional, Union + +from langchain_core.messages import BaseMessage from pydantic import BaseModel -from typing import Optional class ChatOutputResponse(BaseModel): """Chat output response schema.""" - message: str + message: Union[str, List[Union[str, Dict]]] sender: Optional[str] = "Machine" sender_name: Optional[str] = "AI" + + @classmethod + def from_message(cls, message: BaseMessage, sender: Optional[str] = "Machine", sender_name: Optional[str] = "AI"): + """Build chat output response from message.""" + content = message.content + return cls(message=content, sender=sender, sender_name=sender_name)