feature: get messages from messages table for the playground

pyproject.toml

@@ -105,7 +105,7 @@ dependencies = [
     "jq>=1.8.0",
     "pydantic-settings==2.4.0",
     "ragstack-ai-knowledge-store>=0.2.1",
-    "duckduckgo-search>=6.3.0",
+    "duckduckgo-search>=6.3.1",
     "langchain-elasticsearch>=0.2.0",
     "opensearch-py>=2.7.1",
 ]

src/backend/base/langflow/alembic/versions/eb5e72293a8e_add_error_and_edit_flags_to_message.py

@@ -0,0 +1,49 @@
+"""Add error and edit flags to message
+
+Revision ID: eb5e72293a8e
+Revises: 5ace73a7f223
+Create Date: 2024-09-19 16:18:50.828648
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.engine.reflection import Inspector
+
+# revision identifiers, used by Alembic.
+revision: str = "eb5e72293a8e"
+down_revision: Union[str, None] = "5ace73a7f223"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    inspector = Inspector.from_engine(conn)  # type: ignore
+    table_names = inspector.get_table_names()  # noqa
+    column_names = [column["name"] for column in inspector.get_columns("message")]
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("message", schema=None) as batch_op:
+        if "error" not in column_names:
+            batch_op.add_column(sa.Column("error", sa.Boolean(), nullable=False, server_default=sa.false()))
+        if "edit" not in column_names:
+            batch_op.add_column(sa.Column("edit", sa.Boolean(), nullable=False, server_default=sa.false()))
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+    inspector = Inspector.from_engine(conn)  # type: ignore
+    table_names = inspector.get_table_names()  # noqa
+    column_names = [column["name"] for column in inspector.get_columns("message")]
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("message", schema=None) as batch_op:
+        if "edit" in column_names:
+            batch_op.drop_column("edit")
+        if "error" in column_names:
+            batch_op.drop_column("error")
+
+    # ### end Alembic commands ###

src/backend/base/langflow/api/v1/chat.py

@@ -155,6 +155,9 @@ async def build_flow(
     telemetry_service: TelemetryService = Depends(get_telemetry_service),
     session=Depends(get_session),
 ):
+    if not inputs:
+        inputs = InputValueRequest(session=str(flow_id))
+
     async def build_graph_and_get_order() -> tuple[list[str], list[str], Graph]:
         start_time = time.perf_counter()
         components_count = None

src/backend/base/langflow/api/v1/monitor.py

@@ -99,6 +99,7 @@ async def update_message(
 
     try:
         message_dict = message.model_dump(exclude_unset=True, exclude_none=True)
+        message_dict["edit"] = True
         db_message.sqlmodel_update(message_dict)
         session.add(db_message)
         session.commit()

src/backend/base/langflow/api/v1/schemas.py

@@ -297,6 +297,7 @@ class VerticesBuiltResponse(BaseModel):
 class InputValueRequest(BaseModel):
     components: list[str] | None = []
     input_value: str | None = None
+    session: str | None = None
     type: InputType | None = Field(
         "any",
         description="Defines on which components the input value should be applied. "
@@ -310,9 +311,12 @@ class InputValueRequest(BaseModel):
                 {
                     "components": ["components_id", "Component Name"],
                     "input_value": "input_value",
+                    "session": "session_id",
                 },
                 {"components": ["Component Name"], "input_value": "input_value"},
                 {"input_value": "input_value"},
+                {"components": ["Component Name"], "input_value": "input_value", "session": "session_id"},
+                {"input_value": "input_value", "session": "session_id"},
                 {"type": "chat", "input_value": "input_value"},
                 {"type": "json", "input_value": '{"key": "value"}'},
             ]

src/backend/base/langflow/base/astra_assistants/util.py

@@ -28,7 +28,7 @@ def get_patched_openai_client(shared_component_cache):
 data = json.loads(response.text)
 
 # Extract the model names into a Python list
-litellm_model_names = [model for model, _ in data.items() if model != "sample_spec"]
+litellm_model_names = [model for model in data if model != "sample_spec"]
 
 
 # To store the class names that extend ToolInterface

src/backend/base/langflow/base/io/chat.py

@@ -1,4 +1,5 @@
 from collections.abc import AsyncIterator, Iterator
+from typing import cast
 
 from langflow.custom import Component
 from langflow.memory import store_message
@@ -12,43 +13,52 @@ class ChatComponent(Component):
     display_name = "Chat Component"
     description = "Use as base for chat components."
 
-    # Keep this method for backward compatibility
-    def store_message(
-        self,
-        message: Message,
-    ) -> Message:
-        messages = store_message(
-            message,
-            flow_id=self.graph.flow_id,
-        )
-        if len(messages) > 1:
+    def store_message(self, message: Message) -> Message:
+        messages = store_message(message, flow_id=self.graph.flow_id)
+        if len(messages) != 1:
             msg = "Only one message can be stored at a time."
             raise ValueError(msg)
+
         stored_message = messages[0]
-        if (
-            hasattr(self, "_event_manager")
-            and self._event_manager
-            and stored_message.id
-            and not isinstance(message.text, str)
-        ):
+        self._send_message_event(stored_message)
+
+        if self._should_stream_message(stored_message, message):
             complete_message = self._stream_message(message, stored_message.id)
-            message_table = update_message(message_id=stored_message.id, message={"text": complete_message})
-            stored_message = Message(**message_table.model_dump())
-            self.vertex._added_message = stored_message
+            stored_message = self._update_stored_message(stored_message.id, complete_message)
+
         self.status = stored_message
         return stored_message
 
+    def _send_message_event(self, message: Message):
+        if hasattr(self, "_event_manager") and self._event_manager:
+            self._event_manager.on_message(data=message.data)
+
+    def _should_stream_message(self, stored_message: Message, original_message: Message) -> bool:
+        return bool(
+            hasattr(self, "_event_manager")
+            and self._event_manager
+            and stored_message.id
+            and not isinstance(original_message.text, str)
+        )
+
+    def _update_stored_message(self, message_id: str, complete_message: str) -> Message:
+        message_table = update_message(message_id=message_id, message={"text": complete_message})
+        updated_message = Message(**message_table.model_dump())
+        self.vertex._added_message = updated_message
+        return updated_message
+
     def _process_chunk(self, chunk: str, complete_message: str, message: Message, message_id: str) -> str:
         complete_message += chunk
-        data = {
-            "text": complete_message,
-            "chunk": chunk,
-            "sender": message.sender,
-            "sender_name": message.sender_name,
-            "id": str(message_id),
-        }
         if self._event_manager:
-            self._event_manager.on_token(data=data)
+            self._event_manager.on_token(
+                data={
+                    "text": complete_message,
+                    "chunk": chunk,
+                    "sender": message.sender,
+                    "sender_name": message.sender_name,
+                    "id": str(message_id),
+                }
+            )
         return complete_message
 
     async def _handle_async_iterator(self, iterator: AsyncIterator, message: Message, message_id: str) -> str:
@@ -69,7 +79,6 @@ def _stream_message(self, message: Message, message_id: str) -> str:
         complete_message = ""
         for chunk in iterator:
             complete_message = self._process_chunk(chunk.content, complete_message, message, message_id)
-
         return complete_message
 
     def build_with_data(
@@ -80,22 +89,25 @@ def build_with_data(
         input_value: str | Data | Message | None = None,
         files: list[str] | None = None,
         session_id: str | None = None,
-        return_message: bool | None = False,
-    ) -> Message:
-        if isinstance(input_value, Data):
-            # Update the data of the record
-            message = Message.from_data(input_value)
-        else:
-            message = Message(
-                text=input_value, sender=sender, sender_name=sender_name, files=files, session_id=session_id
-            )
+        return_message: bool = False,
+    ) -> str | Message:
+        message = self._create_message(input_value, sender, sender_name, files, session_id)
         message_text = message.text if not return_message else message
 
         self.status = message_text
         if session_id and isinstance(message, Message) and isinstance(message.text, str):
-            messages = store_message(
-                message,
-                flow_id=self.graph.flow_id,
-            )
+            messages = store_message(message, flow_id=self.graph.flow_id)
             self.status = messages
-        return message_text  # type: ignore[return-value]
+            self._send_messages_events(messages)
+
+        return cast(str | Message, message_text)
+
+    def _create_message(self, input_value, sender, sender_name, files, session_id) -> Message:
+        if isinstance(input_value, Data):
+            return Message.from_data(input_value)
+        return Message(text=input_value, sender=sender, sender_name=sender_name, files=files, session_id=session_id)
+
+    def _send_messages_events(self, messages):
+        if hasattr(self, "_event_manager") and self._event_manager:
+            for stored_message in messages:
+                self._event_manager.on_message(data=stored_message.data)

src/backend/base/langflow/components/inputs/ChatInput.py

@@ -2,7 +2,6 @@
 from langflow.base.io.chat import ChatComponent
 from langflow.inputs import BoolInput
 from langflow.io import DropdownInput, FileInput, MessageTextInput, MultilineInput, Output
-from langflow.memory import store_message
 from langflow.schema.message import Message
 from langflow.utils.constants import MESSAGE_SENDER_AI, MESSAGE_SENDER_NAME_USER, MESSAGE_SENDER_USER
 
@@ -69,18 +68,12 @@ def message_response(self) -> Message:
             session_id=self.session_id,
             files=self.files,
         )
-
-        if (
-            self.session_id
-            and isinstance(message, Message)
-            and isinstance(message.text, str)
-            and self.should_store_message
-        ):
-            store_message(
+        if self.session_id and isinstance(message, Message) and self.should_store_message:
+            stored_message = self.store_message(
                 message,
-                flow_id=self.graph.flow_id,
             )
-            self.message.value = message
+            self.message.value = stored_message
+            message = stored_message
 
         self.status = message
         return message

src/backend/base/langflow/components/outputs/ChatOutput.py

@@ -1,7 +1,6 @@
 from langflow.base.io.chat import ChatComponent
 from langflow.inputs import BoolInput
 from langflow.io import DropdownInput, MessageTextInput, Output
-from langflow.memory import store_message
 from langflow.schema.message import Message
 from langflow.utils.constants import MESSAGE_SENDER_AI, MESSAGE_SENDER_NAME_AI, MESSAGE_SENDER_USER
 
@@ -65,17 +64,12 @@ def message_response(self) -> Message:
             sender_name=self.sender_name,
             session_id=self.session_id,
         )
-        if (
-            self.session_id
-            and isinstance(message, Message)
-            and isinstance(message.text, str)
-            and self.should_store_message
-        ):
-            store_message(
+        if self.session_id and isinstance(message, Message) and self.should_store_message:
+            stored_message = self.store_message(
                 message,
-                flow_id=self.graph.flow_id,
             )
-            self.message.value = message
+            self.message.value = stored_message
+            message = stored_message
 
         self.status = message
         return message

src/backend/base/langflow/events/event_manager.py

@@ -5,6 +5,7 @@
 import uuid
 from functools import partial
 
+from fastapi.encoders import jsonable_encoder
 from typing_extensions import Protocol
 
 from langflow.schema.log import LoggableType
@@ -52,7 +53,8 @@ def register_event(self, name: str, event_type: str, callback: EventCallback | N
         self.events[name] = _callback
 
     def send_event(self, *, event_type: str, data: LoggableType):
-        json_data = {"event": event_type, "data": data}
+        jsonable_data = jsonable_encoder(data)
+        json_data = {"event": event_type, "data": jsonable_data}
         event_id = uuid.uuid4()
         str_data = json.dumps(json_data) + "\n\n"
         self.queue.put_nowait((event_id, str_data.encode("utf-8"), time.time()))

src/backend/base/langflow/graph/vertex/base.py

@@ -793,11 +793,20 @@ async def build(
                 # and we are just getting the result for the requester
                 return await self.get_requester_result(requester)
             self._reset()
-
+            # inject session_id if it is not None
+            if inputs is not None and "session" in inputs and inputs["session"] is not None and self.has_session_id:
+                session_id_value = self.get_value_from_template_dict("session_id")
+                if session_id_value == "":
+                    self.update_raw_params({"session_id": inputs["session"]}, overwrite=True)
             if self._is_chat_input() and (inputs or files):
                 chat_input = {}
-                if inputs:
-                    chat_input.update({"input_value": inputs.get(INPUT_FIELD_NAME, "")})
+                if (
+                    inputs
+                    and isinstance(inputs, dict)
+                    and "input_value" in inputs
+                    and inputs["input_value"] is not None
+                ):
+                    chat_input.update({"input_value": inputs[INPUT_FIELD_NAME]})
                 if files:
                     chat_input.update({"files": files})