Refactor ChatInput and ChatOutput to support returning messages as records

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-26 18:53:10 -03:00
commit 84d3bd1ba6
2 changed files with 44 additions and 32 deletions

View file

@ -1,7 +1,6 @@
from typing import Optional, Union from typing import Optional
from langflow import CustomComponent from langflow import CustomComponent
from langflow.field_typing import Text
from langflow.schema import Record from langflow.schema import Record
@ -25,6 +24,10 @@ class ChatInput(CustomComponent):
"display_name": "Session ID", "display_name": "Session ID",
"info": "Session ID of the chat history.", "info": "Session ID of the chat history.",
}, },
"return_record": {
"display_name": "Return Record",
"info": "Return the message as a record containing the sender, sender_name, and session_id.",
},
} }
def build( def build(
@ -33,21 +36,23 @@ class ChatInput(CustomComponent):
sender_name: Optional[str] = "User", sender_name: Optional[str] = "User",
message: Optional[str] = None, message: Optional[str] = None,
session_id: Optional[str] = None, session_id: Optional[str] = None,
return_record: Optional[bool] = False,
) -> Record: ) -> Record:
if isinstance(message, Record): if return_record:
# Update the data of the record if isinstance(message, Record):
message.data["sender"] = sender # Update the data of the record
message.data["sender_name"] = sender_name message.data["sender"] = sender
message.data["session_id"] = session_id message.data["sender_name"] = sender_name
else: message.data["session_id"] = session_id
message = Record( else:
text=message, message = Record(
data={ text=message,
"sender": sender, data={
"sender_name": sender_name, "sender": sender,
"session_id": session_id, "sender_name": sender_name,
}, "session_id": session_id,
) },
)
if not message: if not message:
message = "" message = ""
self.status = message self.status = message

View file

@ -1,6 +1,7 @@
from typing import Optional from typing import Optional, Union
from langflow import CustomComponent from langflow import CustomComponent
from langflow.field_typing import Text
from langflow.schema import Record from langflow.schema import Record
@ -27,6 +28,10 @@ class ChatOutput(CustomComponent):
"info": "Session ID of the chat history.", "info": "Session ID of the chat history.",
"input_types": ["Text"], "input_types": ["Text"],
}, },
"return_record": {
"display_name": "Return Record",
"info": "Return the message as a record containing the sender, sender_name, and session_id.",
},
} }
def build( def build(
@ -35,21 +40,23 @@ class ChatOutput(CustomComponent):
sender_name: Optional[str] = "AI", sender_name: Optional[str] = "AI",
session_id: Optional[str] = None, session_id: Optional[str] = None,
message: Optional[str] = None, message: Optional[str] = None,
) -> Record: return_record: Optional[bool] = False,
if isinstance(message, Record): ) -> Union[Text, Record]:
# Update the data of the record if return_record:
message.data["sender"] = sender if isinstance(message, Record):
message.data["sender_name"] = sender_name # Update the data of the record
message.data["session_id"] = session_id message.data["sender"] = sender
else: message.data["sender_name"] = sender_name
message = Record( message.data["session_id"] = session_id
text=message, else:
data={ message = Record(
"sender": sender, text=message,
"sender_name": sender_name, data={
"session_id": session_id, "sender": sender,
}, "sender_name": sender_name,
) "session_id": session_id,
},
)
if not message: if not message:
message = "" message = ""
self.status = message self.status = message