jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

refactor: Remove unused imports and optimize get_messages function

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-25 19:50:56 -03:00
commit 0bfb702736
1 changed files with 5 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

22
src/backend/base/langflow/memory.py
View file

@ -7,7 +7,7 @@ from sqlalchemy import delete
from sqlmodel import Session, col, select from sqlmodel import Session, col, select
from langflow.schema.message import Message from langflow.schema.message import Message
from langflow.services.database.models.message.model import MessageRead, MessageTable from langflow.services.database.models.message.model import MessageTable
from langflow.services.deps import session_scope from langflow.services.deps import session_scope
@ -33,6 +33,7 @@ def get_messages(
Returns: Returns:
List[Data]: A list of Data objects representing the retrieved messages. List[Data]: A list of Data objects representing the retrieved messages.
""" """
messages_read: list[Message] = []
with session_scope() as session: with session_scope() as session:
stmt = select(MessageTable) stmt = select(MessageTable)
if sender: if sender:
@ -52,22 +53,9 @@ def get_messages(
if limit: if limit:
stmt = stmt.limit(limit) stmt = stmt.limit(limit)
messages = session.exec(stmt) messages = session.exec(stmt)
messages_read = [MessageRead.model_validate(d, from_attributes=True) for d in messages] messages_read = [Message(**d.model_dump()) for d in messages]
messages: list[Message] = [] return messages_read
for msg_read in messages_read:
msg = Message(
text=msg_read.text,
sender=msg_read.sender,
session_id=msg_read.session_id,
sender_name=msg_read.sender_name,
timestamp=msg_read.timestamp,
)
messages.append(msg)
return messages
def add_messages(messages: Message | list[Message], flow_id: Optional[str] = None): def add_messages(messages: Message | list[Message], flow_id: Optional[str] = None):
@ -102,7 +90,7 @@ def add_messagetables(messages: list[MessageTable], session: Session):
except Exception as e: except Exception as e:
logger.exception(e) logger.exception(e)
raise e raise e
return messages return [Message(**message.model_dump()) for message in messages]
def delete_messages(session_id: str): def delete_messages(session_id: str):