From ea5139dd7aae2e4dbdb00707aa82fcf7134fde55 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Tue, 25 Jun 2024 12:13:56 -0300 Subject: [PATCH] Refactor add_messages function to separate message addition and commit logic --- src/backend/base/langflow/memory.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/src/backend/base/langflow/memory.py b/src/backend/base/langflow/memory.py index c824310d9..2835645ba 100644 --- a/src/backend/base/langflow/memory.py +++ b/src/backend/base/langflow/memory.py @@ -2,6 +2,7 @@ import warnings from typing import List, Optional from loguru import logger +from sqlmodel import Session from langflow.schema.message import Message from langflow.services.database.models.message.model import MessageTable @@ -75,22 +76,25 @@ def add_messages(messages: Message | list[Message], flow_id: Optional[str] = Non for msg in messages: messages_models.append(MessageTable.from_message(msg, flow_id=flow_id)) with session_scope() as session: - for message_model in messages_models: - try: - session.add(message_model) - session.commit() - session.refresh(message_model) - except Exception as e: - logger.error(f"Error adding message to monitor service: {e}") - logger.exception(e) - raise e - + messages_models = add_messagetables(messages_models, session) return messages_models except Exception as e: logger.exception(e) raise e +def add_messagetables(messages: list[MessageTable], session: Session): + for message in messages: + try: + session.add(message) + session.commit() + session.refresh(message) + except Exception as e: + logger.exception(e) + raise e + return messages + + def delete_messages(session_id: str): """ Delete messages from the monitor service based on the provided session ID.