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.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-23 21:51:38 -03:00
commit 2b0765d41b
6 changed files with 127 additions and 3 deletions

View file

@ -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 ###

View file

@ -119,7 +119,11 @@ class LCModelComponent(Component):
return status_message return status_message
def get_chat_result( 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]] = [] messages: list[Union[BaseMessage]] = []
if not input_value and not system_message: if not input_value and not system_message:

View file

@ -1,7 +1,8 @@
from .api_key import ApiKey from .api_key import ApiKey
from .flow import Flow from .flow import Flow
from .folder import Folder from .folder import Folder
from .message import MessageTable
from .user import User from .user import User
from .variable import Variable from .variable import Variable
__all__ = ["Flow", "User", "ApiKey", "Variable", "Folder"] __all__ = ["Flow", "User", "ApiKey", "Variable", "Folder", "MessageTable"]

View file

@ -3,7 +3,7 @@
import re import re
import warnings import warnings
from datetime import datetime, timezone 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 from uuid import UUID, uuid4
import emoji import emoji
@ -17,6 +17,7 @@ from langflow.schema import Data
if TYPE_CHECKING: if TYPE_CHECKING:
from langflow.services.database.models.folder import Folder from langflow.services.database.models.folder import Folder
from langflow.services.database.models.message import MessageTable
from langflow.services.database.models.user import User from langflow.services.database.models.user import User
@ -141,6 +142,7 @@ class Flow(FlowBase, table=True):
user: "User" = Relationship(back_populates="flows") user: "User" = Relationship(back_populates="flows")
folder_id: Optional[UUID] = Field(default=None, foreign_key="folder.id", nullable=True, index=True) folder_id: Optional[UUID] = Field(default=None, foreign_key="folder.id", nullable=True, index=True)
folder: Optional["Folder"] = Relationship(back_populates="flows") folder: Optional["Folder"] = Relationship(back_populates="flows")
messages: List["MessageTable"] = Relationship(back_populates="flow")
def to_data(self): def to_data(self):
serialized = self.model_dump() serialized = self.model_dump()

View file

@ -0,0 +1,3 @@
from .model import MessageTable, MessageCreate, MessageRead, MessageUpdate
__all__ = ["MessageTable", "MessageCreate", "MessageRead", "MessageUpdate"]

View file

@ -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