From 2f2bc4008fdfb8f589da8354346981354eef790b Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sun, 23 Jun 2024 21:51:38 -0300 Subject: [PATCH] feat: Add message table to the database This commit adds a new table called "message" to the database. The table includes columns for timestamp, sender, sender_name, session_id, text, id, flow_id, and files. The "message" table is created using Alembic migration. This addition allows for storing and retrieving messages in the application. --- .../d066bfd22890_add_message_table.py | 52 ++++++++++++++++ .../base/langflow/base/models/model.py | 6 +- .../services/database/models/__init__.py | 3 +- .../services/database/models/flow/model.py | 4 +- .../database/models/message/__init__.py | 3 + .../services/database/models/message/model.py | 62 +++++++++++++++++++ 6 files changed, 127 insertions(+), 3 deletions(-) create mode 100644 src/backend/base/langflow/alembic/versions/d066bfd22890_add_message_table.py create mode 100644 src/backend/base/langflow/services/database/models/message/__init__.py create mode 100644 src/backend/base/langflow/services/database/models/message/model.py diff --git a/src/backend/base/langflow/alembic/versions/d066bfd22890_add_message_table.py b/src/backend/base/langflow/alembic/versions/d066bfd22890_add_message_table.py new file mode 100644 index 000000000..dd63398b7 --- /dev/null +++ b/src/backend/base/langflow/alembic/versions/d066bfd22890_add_message_table.py @@ -0,0 +1,52 @@ +"""Add message table + +Revision ID: 325180f0c4e1 +Revises: 631faacf5da2 +Create Date: 2024-06-23 21:29:28.220100 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +import sqlmodel +from alembic import op + +from langflow.utils import migration + +# revision identifiers, used by Alembic. +revision: str = "325180f0c4e1" +down_revision: Union[str, None] = "631faacf5da2" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + conn = op.get_bind() + # ### commands auto generated by Alembic - please adjust! ### + if not migration.table_exists("message", conn): + op.create_table( + "message", + sa.Column("timestamp", sa.DateTime(), nullable=False), + sa.Column("sender", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("sender_name", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("session_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("text", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column("flow_id", sqlmodel.sql.sqltypes.GUID(), nullable=True), + sa.Column("files", sa.JSON(), nullable=True), + sa.ForeignKeyConstraint( + ["flow_id"], + ["flow.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + conn = op.get_bind() + # ### commands auto generated by Alembic - please adjust! ### + if migration.table_exists("message", conn): + op.drop_table("message") + # ### end Alembic commands ### diff --git a/src/backend/base/langflow/base/models/model.py b/src/backend/base/langflow/base/models/model.py index 81f124fd7..bea9c88d8 100644 --- a/src/backend/base/langflow/base/models/model.py +++ b/src/backend/base/langflow/base/models/model.py @@ -119,7 +119,11 @@ class LCModelComponent(Component): return status_message def get_chat_result( - self, runnable: LanguageModel, stream: bool, input_value: str | Message, system_message: Optional[str] = None + self, + runnable: LanguageModel, + stream: bool, + input_value: str | Message, + system_message: Optional[str] = None, ): messages: list[Union[BaseMessage]] = [] if not input_value and not system_message: diff --git a/src/backend/base/langflow/services/database/models/__init__.py b/src/backend/base/langflow/services/database/models/__init__.py index ce12a6fce..6e1f09fe3 100644 --- a/src/backend/base/langflow/services/database/models/__init__.py +++ b/src/backend/base/langflow/services/database/models/__init__.py @@ -1,7 +1,8 @@ from .api_key import ApiKey from .flow import Flow from .folder import Folder +from .message import MessageTable from .user import User from .variable import Variable -__all__ = ["Flow", "User", "ApiKey", "Variable", "Folder"] +__all__ = ["Flow", "User", "ApiKey", "Variable", "Folder", "MessageTable"] diff --git a/src/backend/base/langflow/services/database/models/flow/model.py b/src/backend/base/langflow/services/database/models/flow/model.py index 624ea0543..6d3e4aea8 100644 --- a/src/backend/base/langflow/services/database/models/flow/model.py +++ b/src/backend/base/langflow/services/database/models/flow/model.py @@ -3,7 +3,7 @@ import re import warnings from datetime import datetime, timezone -from typing import TYPE_CHECKING, Dict, Optional +from typing import TYPE_CHECKING, Dict, List, Optional from uuid import UUID, uuid4 import emoji @@ -17,6 +17,7 @@ from langflow.schema import Data if TYPE_CHECKING: from langflow.services.database.models.folder import Folder + from langflow.services.database.models.message import MessageTable from langflow.services.database.models.user import User @@ -141,6 +142,7 @@ class Flow(FlowBase, table=True): user: "User" = Relationship(back_populates="flows") folder_id: Optional[UUID] = Field(default=None, foreign_key="folder.id", nullable=True, index=True) folder: Optional["Folder"] = Relationship(back_populates="flows") + messages: List["MessageTable"] = Relationship(back_populates="flow") def to_data(self): serialized = self.model_dump() diff --git a/src/backend/base/langflow/services/database/models/message/__init__.py b/src/backend/base/langflow/services/database/models/message/__init__.py new file mode 100644 index 000000000..8cfb2ff4f --- /dev/null +++ b/src/backend/base/langflow/services/database/models/message/__init__.py @@ -0,0 +1,3 @@ +from .model import MessageTable, MessageCreate, MessageRead, MessageUpdate + +__all__ = ["MessageTable", "MessageCreate", "MessageRead", "MessageUpdate"] diff --git a/src/backend/base/langflow/services/database/models/message/model.py b/src/backend/base/langflow/services/database/models/message/model.py new file mode 100644 index 000000000..87da2ef7d --- /dev/null +++ b/src/backend/base/langflow/services/database/models/message/model.py @@ -0,0 +1,62 @@ +from datetime import datetime, timezone +from typing import TYPE_CHECKING, List, Optional +from uuid import UUID, uuid4 + +from sqlmodel import JSON, Column, Field, Relationship, SQLModel + +if TYPE_CHECKING: + from langflow.schema.message import Message + from langflow.services.database.models.flow.model import Flow + + +class MessageBase(SQLModel): + timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + sender: str + sender_name: str + session_id: str + text: str + files: list[str] = Field(default_factory=list) + + @classmethod + def from_message(cls, message: "Message", flow_id: str | None = None): + # first check if the record has all the required fields + if message.text is None or not message.sender or not message.sender_name: + raise ValueError("The message does not have the required fields (text, sender, sender_name).") + return cls( + sender=message.sender, + sender_name=message.sender_name, + text=message.text, + session_id=message.session_id, + files=message.files or [], + timestamp=message.timestamp, + flow_id=flow_id, + ) + + +class MessageTable(MessageBase, table=True): + __tablename__ = "message" + id: UUID = Field(default_factory=uuid4, primary_key=True) + flow_id: Optional[UUID] = Field(default=None, foreign_key="flow.id") + flow: "Flow" = Relationship(back_populates="messages") + files: List[str] = Field(sa_column=Column(JSON)) + + # Needed for Column(JSON) + class Config: + arbitrary_types_allowed = True + + +class MessageRead(MessageBase): + id: UUID + flow_id: Optional[UUID] = Field() + + +class MessageCreate(MessageBase): + pass + + +class MessageUpdate(SQLModel): + text: Optional[str] = None + sender: Optional[str] = None + sender_name: Optional[str] = None + session_id: Optional[str] = None + files: Optional[list[str]] = None