feat: add EventManager to centralize callbacks

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

@@ -31,6 +31,7 @@
     VertexBuildResponse,
     VerticesOrderResponse,
 )
+from langflow.events.event_manager import EventManager, create_default_event_manager
 from langflow.exceptions.component import ComponentBuildException
 from langflow.graph.graph.base import Graph
 from langflow.graph.utils import log_vertex_build
@@ -204,7 +205,7 @@ async def build_graph_and_get_order() -> tuple[list[str], list[str], "Graph"]:
             logger.exception(exc)
             raise HTTPException(status_code=500, detail=str(exc)) from exc
 
-    async def _build_vertex(vertex_id: str, graph: "Graph") -> VertexBuildResponse:
+    async def _build_vertex(vertex_id: str, graph: "Graph", event_manager: "EventManager") -> VertexBuildResponse:
         flow_id_str = str(flow_id)
 
         next_runnable_vertices = []
@@ -222,6 +223,7 @@ async def _build_vertex(vertex_id: str, graph: "Graph") -> VertexBuildResponse:
                     files=files,
                     get_cache=chat_service.get_cache,
                     set_cache=chat_service.set_cache,
+                    event_manager=event_manager,
                 )
                 result_dict = vertex_build_result.result_dict
                 params = vertex_build_result.params
@@ -316,17 +318,13 @@ async def _build_vertex(vertex_id: str, graph: "Graph") -> VertexBuildResponse:
             message = parse_exception(exc)
             raise HTTPException(status_code=500, detail=message) from exc
 
-    def send_event(event_type: str, value: dict, queue: asyncio.Queue) -> None:
-        json_data = {"event": event_type, "data": value}
-        event_id = uuid.uuid4()
-        logger.debug(f"sending event {event_id}: {event_type}")
-        str_data = json.dumps(json_data) + "\n\n"
-        queue.put_nowait((event_id, str_data.encode("utf-8"), time.time()))
-
     async def build_vertices(
-        vertex_id: str, graph: "Graph", queue: asyncio.Queue, client_consumed_queue: asyncio.Queue
+        vertex_id: str,
+        graph: "Graph",
+        client_consumed_queue: asyncio.Queue,
+        event_manager: "EventManager",
     ) -> None:
-        build_task = asyncio.create_task(await asyncio.to_thread(_build_vertex, vertex_id, graph))
+        build_task = asyncio.create_task(await asyncio.to_thread(_build_vertex, vertex_id, graph, event_manager))
         try:
             await build_task
         except asyncio.CancelledError:
@@ -340,13 +338,15 @@ async def build_vertices(
             build_data = json.loads(vertex_build_response_json)
         except Exception as exc:
             raise ValueError(f"Error serializing vertex build response: {exc}") from exc
-        send_event("end_vertex", {"build_data": build_data}, queue)
+        event_manager.on_end_vertex(data={"build_data": build_data})
         await client_consumed_queue.get()
         if vertex_build_response.valid:
             if vertex_build_response.next_vertices_ids:
                 tasks = []
                 for next_vertex_id in vertex_build_response.next_vertices_ids:
-                    task = asyncio.create_task(build_vertices(next_vertex_id, graph, queue, client_consumed_queue))
+                    task = asyncio.create_task(
+                        build_vertices(next_vertex_id, graph, client_consumed_queue, event_manager)
+                    )
                     tasks.append(task)
                 try:
                     await asyncio.gather(*tasks)
@@ -355,7 +355,7 @@ async def build_vertices(
                         task.cancel()
                     return
 
-    async def event_generator(queue: asyncio.Queue, client_consumed_queue: asyncio.Queue) -> None:
+    async def event_generator(event_manager: EventManager, client_consumed_queue: asyncio.Queue) -> None:
         if not data:
             # using another thread since the DB query is I/O bound
             vertices_task = asyncio.create_task(await asyncio.to_thread(build_graph_and_get_order))
@@ -366,9 +366,9 @@ async def event_generator(queue: asyncio.Queue, client_consumed_queue: asyncio.Q
                 return
             except Exception as e:
                 if isinstance(e, HTTPException):
-                    send_event("error", {"error": str(e.detail), "statusCode": e.status_code}, queue)
+                    event_manager.on_error(data={"error": str(e.detail), "statusCode": e.status_code})
                     raise e
-                send_event("error", {"error": str(e)}, queue)
+                event_manager.on_error(data={"error": str(e)})
                 raise e
 
             ids, vertices_to_run, graph = vertices_task.result()
@@ -377,16 +377,16 @@ async def event_generator(queue: asyncio.Queue, client_consumed_queue: asyncio.Q
                 ids, vertices_to_run, graph = await build_graph_and_get_order()
             except Exception as e:
                 if isinstance(e, HTTPException):
-                    send_event("error", {"error": str(e.detail), "statusCode": e.status_code}, queue)
+                    event_manager.on_error(data={"error": str(e.detail), "statusCode": e.status_code})
                     raise e
-                send_event("error", {"error": str(e)}, queue)
+                event_manager.on_error(data={"error": str(e)})
                 raise e
-        send_event("vertices_sorted", {"ids": ids, "to_run": vertices_to_run}, queue)
+        event_manager.on_vertices_sorted(data={"ids": ids, "to_run": vertices_to_run})
         await client_consumed_queue.get()
 
         tasks = []
         for vertex_id in ids:
-            task = asyncio.create_task(build_vertices(vertex_id, graph, queue, client_consumed_queue))
+            task = asyncio.create_task(build_vertices(vertex_id, graph, client_consumed_queue, event_manager))
             tasks.append(task)
         try:
             await asyncio.gather(*tasks)
@@ -395,8 +395,8 @@ async def event_generator(queue: asyncio.Queue, client_consumed_queue: asyncio.Q
             for task in tasks:
                 task.cancel()
             return
-        send_event("end", {}, queue)
-        await queue.put((None, None, time.time))
+        event_manager.on_end(data={})
+        await event_manager.queue.put((None, None, time.time))
 
     async def consume_and_yield(queue: asyncio.Queue, client_consumed_queue: asyncio.Queue) -> typing.AsyncGenerator:
         while True:
@@ -413,7 +413,8 @@ async def consume_and_yield(queue: asyncio.Queue, client_consumed_queue: asyncio
 
     asyncio_queue: asyncio.Queue = asyncio.Queue()
     asyncio_queue_client_consumed: asyncio.Queue = asyncio.Queue()
-    main_task = asyncio.create_task(event_generator(asyncio_queue, asyncio_queue_client_consumed))
+    event_manager = create_default_event_manager(queue=asyncio_queue)
+    main_task = asyncio.create_task(event_generator(event_manager, asyncio_queue_client_consumed))
 
     def on_disconnect():
         logger.debug("Client disconnected, closing tasks")
@@ -639,6 +640,7 @@ async def build_vertex_stream(
         flow_id_str = str(flow_id)
 
         async def stream_vertex():
+            graph = None
             try:
                 cache = await chat_service.get_cache(flow_id_str)
                 if not cache:
@@ -692,7 +694,8 @@ async def stream_vertex():
                 yield str(StreamData(event="error", data={"error": exc_message}))
             finally:
                 logger.debug("Closing stream")
-                await chat_service.set_cache(flow_id_str, graph)
+                if graph:
+                    await chat_service.set_cache(flow_id_str, graph)
                 yield str(StreamData(event="close", data={"message": "Stream closed"}))
 
         return StreamingResponse(stream_vertex(), media_type="text/event-stream")

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

@@ -1,67 +1,70 @@
-from typing import Optional, Union
+from typing import AsyncIterator, Iterator, Optional, Union
 
-from langflow.base.data.utils import IMG_FILE_TYPES, TEXT_FILE_TYPES
 from langflow.custom import Component
 from langflow.memory import store_message
 from langflow.schema import Data
 from langflow.schema.message import Message
-from langflow.utils.constants import MESSAGE_SENDER_USER, MESSAGE_SENDER_AI
+from langflow.services.database.models.message.crud import update_message
+from langflow.utils.async_helpers import run_until_complete
 
 
 class ChatComponent(Component):
     display_name = "Chat Component"
     description = "Use as base for chat components."
 
-    def build_config(self):
-        return {
-            "input_value": {
-                "input_types": ["Text"],
-                "display_name": "Text",
-                "multiline": True,
-            },
-            "sender": {
-                "options": [MESSAGE_SENDER_AI, MESSAGE_SENDER_USER],
-                "display_name": "Sender Type",
-                "advanced": True,
-            },
-            "sender_name": {"display_name": "Sender Name", "advanced": True},
-            "session_id": {
-                "display_name": "Session ID",
-                "info": "If provided, the message will be stored in the memory.",
-                "advanced": True,
-            },
-            "return_message": {
-                "display_name": "Return Message",
-                "info": "Return the message as a Message containing the sender, sender_name, and session_id.",
-                "advanced": True,
-            },
-            "data_template": {
-                "display_name": "Data Template",
-                "multiline": True,
-                "info": "In case of Message being a Data, this template will be used to convert it to text.",
-                "advanced": True,
-            },
-            "files": {
-                "field_type": "file",
-                "display_name": "Files",
-                "file_types": TEXT_FILE_TYPES + IMG_FILE_TYPES,
-                "info": "Files to be sent with the message.",
-                "advanced": True,
-            },
-        }
-
     # Keep this method for backward compatibility
     def store_message(
         self,
         message: Message,
-    ) -> list[Message]:
+    ) -> Message:
         messages = store_message(
             message,
             flow_id=self.graph.flow_id,
         )
+        if len(messages) > 1:
+            raise ValueError("Only one message can be stored at a time.")
+        stored_message = messages[0]
+        if hasattr(self, "_event_manager") and self._event_manager and stored_message.id:
+            if not isinstance(message.text, str):
+                complete_message = self._stream_message(message, stored_message.id)
+                message_table = update_message(message_id=stored_message.id, message=dict(text=complete_message))
+                stored_message = Message(**message_table.model_dump())
+                self.vertex._added_message = stored_message
+        self.status = stored_message
+        return stored_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)
+        return complete_message
+
+    async def _handle_async_iterator(self, iterator: AsyncIterator, message: Message, message_id: str) -> str:
+        complete_message = ""
+        async for chunk in iterator:
+            complete_message = self._process_chunk(chunk.content, complete_message, message, message_id)
+        return complete_message
+
+    def _stream_message(self, message: Message, message_id: str) -> str:
+        iterator = message.text
+        if not isinstance(iterator, (AsyncIterator, Iterator)):
+            raise ValueError("The message must be an iterator or an async iterator.")
+
+        if isinstance(iterator, AsyncIterator):
+            return run_until_complete(self._handle_async_iterator(iterator, message, message_id))
+
+        complete_message = ""
+        for chunk in iterator:
+            complete_message = self._process_chunk(chunk.content, complete_message, message, message_id)
 
-        self.status = messages
-        return messages
+        return complete_message
 
     def build_with_data(
         self,

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

@@ -3,7 +3,7 @@
 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_NAME_AI, MESSAGE_SENDER_USER, MESSAGE_SENDER_AI
+from langflow.utils.constants import MESSAGE_SENDER_AI, MESSAGE_SENDER_NAME_AI, MESSAGE_SENDER_USER
 
 
 class ChatOutput(ChatComponent):

src/backend/base/langflow/custom/custom_component/component.py

@@ -1,17 +1,19 @@
 import inspect
+from collections.abc import Callable
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, ClassVar, get_type_hints
-from collections.abc import Callable
 from uuid import UUID
 
 import nanoid  # type: ignore
 import yaml
 from pydantic import BaseModel
 
+from langflow.events.event_manager import EventManager
 from langflow.graph.state.model import create_state_model
 from langflow.helpers.custom import format_type
 from langflow.schema.artifact import get_artifact_type, post_process_raw
 from langflow.schema.data import Data
+from langflow.schema.log import LoggableType
 from langflow.schema.message import Message
 from langflow.services.tracing.schema import Log
 from langflow.template.field.base import UNDEFINED, Input, Output
@@ -35,6 +37,7 @@ class Component(CustomComponent):
     outputs: list[Output] = []
     code_class_base_inheritance: ClassVar[str] = "Component"
     _output_logs: dict[str, Log] = {}
+    _current_output: str = ""
 
     def __init__(self, **kwargs):
         # if key starts with _ it is a config
@@ -56,6 +59,8 @@ def __init__(self, **kwargs):
         self._parameters = inputs or {}
         self._edges: list[EdgeData] = []
         self._components: list[Component] = []
+        self._current_output = ""
+        self._event_manager: EventManager | None = None
         self._state_model = None
         self.set_attributes(self._parameters)
         self._output_logs = {}
@@ -77,6 +82,9 @@ def __init__(self, **kwargs):
         self._set_output_types()
         self.set_class_code()
 
+    def set_event_manager(self, event_manager: EventManager | None = None):
+        self._event_manager = event_manager
+
     def _reset_all_output_values(self):
         for output in self.outputs:
             setattr(output, "value", UNDEFINED)
@@ -597,7 +605,9 @@ async def _build_with_tracing(self):
     async def _build_without_tracing(self):
         return await self._build_results()
 
-    async def build_results(self):
+    async def build_results(
+        self,
+    ):
         if self._tracing_service:
             return await self._build_with_tracing()
         return await self._build_without_tracing()
@@ -616,6 +626,7 @@ async def _build_results(self):
                 ):
                     if output.method is None:
                         raise ValueError(f"Output {output.name} does not have a method defined.")
+                    self._current_output = output.name
                     method: Callable = getattr(self, output.method)
                     if output.cache and output.value != UNDEFINED:
                         _results[output.name] = output.value
@@ -634,6 +645,7 @@ async def _build_results(self):
                             result.set_flow_id(self._vertex.graph.flow_id)
                         _results[output.name] = result
                         output.value = result
+
                     custom_repr = self.custom_repr()
                     if custom_repr is None and isinstance(result, (dict, Data, str)):
                         custom_repr = result
@@ -661,6 +673,7 @@ async def _build_results(self):
                     _artifacts[output.name] = artifact
                     self._output_logs[output.name] = self._logs
                     self._logs = []
+                    self._current_output = ""
         self._artifacts = _artifacts
         self._results = _results
         if self._tracing_service:
@@ -716,6 +729,25 @@ def to_tool(self):
         return ComponentTool(component=self)
 
     def get_project_name(self):
-        if hasattr(self, "_tracing_service"):
+        if hasattr(self, "_tracing_service") and self._tracing_service:
             return self._tracing_service.project_name
         return "Langflow"
+
+    def log(self, message: LoggableType | list[LoggableType], name: str | None = None):
+        """
+        Logs a message.
+
+        Args:
+            message (LoggableType | list[LoggableType]): The message to log.
+        """
+        if name is None:
+            name = f"Log {len(self._logs) + 1}"
+        log = Log(message=message, type=get_artifact_type(message), name=name)
+        self._logs.append(log)
+        if self._tracing_service and self._vertex:
+            self._tracing_service.add_log(trace_name=self.trace_name, log=log)
+        if self._event_manager is not None and self._current_output:
+            data = log.model_dump()
+            data["output"] = self._current_output
+            data["component_id"] = self._id
+            self._event_manager.on_log(data=data)