Add support for returning messages as a Record in ChatInput and ChatOutput

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-18 17:28:26 -03:00
commit cb95de9e1f
3 changed files with 220 additions and 209 deletions

View file

@ -1,7 +1,8 @@
from typing import Optional
from typing import Optional, Union
from langflow import CustomComponent
from langflow.field_typing import Text
from langflow.schema import Record
class ChatInput(CustomComponent):
@ -11,17 +12,29 @@ class ChatInput(CustomComponent):
def build_config(self):
return {
"message": {"input_types": ["Text"], "display_name": "Message"},
"sender": {"options": ["Machine", "User"], "display_name": "Sender Type"},
"sender_type": {
"options": ["Machine", "User"],
"display_name": "Sender Type",
},
"sender_name": {"display_name": "Sender Name"},
"as_record": {
"display_name": "As Record",
"info": "If true, the message will be returned as a Record.",
},
}
def build(
self,
sender: Optional[str] = "User",
sender_type: Optional[str] = "User",
sender_name: Optional[str] = "You",
message: Optional[str] = None,
) -> Text:
self.repr_value = message
as_record: Optional[bool] = False,
) -> Union[Text, Record]:
self.status = message
if as_record:
return Record(
text=message, data={"sender": sender_type, "sender_name": sender_name}
)
if not message:
message = ""
return message

View file

@ -2,6 +2,7 @@ from typing import Optional
from langflow import CustomComponent
from langflow.field_typing import Text
from langflow.schema import Record
class ChatOutput(CustomComponent):
@ -17,7 +18,10 @@ class ChatOutput(CustomComponent):
def build_config(self):
return {
"message": {"input_types": ["Text"], "display_name": "Message"},
"sender": {"options": ["Machine", "User"], "display_name": "Sender Type"},
"sender_type": {
"options": ["Machine", "User"],
"display_name": "Sender Type",
},
"sender_name": {"display_name": "Sender Name"},
"session_id": {
"display_name": "Session ID",
@ -28,12 +32,22 @@ class ChatOutput(CustomComponent):
def build(
self,
sender: Optional[str] = "Machine",
sender_type: Optional[str] = "Machine",
sender_name: Optional[str] = "AI",
session_id: Optional[str] = None,
message: Optional[str] = None,
as_record: Optional[bool] = False,
) -> Text:
self.repr_value = message
self.status = message
if as_record:
return Record(
text=message,
data={
"sender": sender_type,
"sender_name": sender_name,
"session_id": session_id,
},
)
if not message:
message = ""
return message