Refactor message.py to set flow_id for Message instances

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

@@ -1,5 +1,6 @@
 import inspect
 from typing import AsyncIterator, Awaitable, Callable, ClassVar, Generator, Iterator, List, Optional, Union
+
 from uuid import UUID
 
 import yaml
@@ -9,6 +10,7 @@ from pydantic import BaseModel
 from langflow.inputs.inputs import InputTypes
 from langflow.schema.artifact import get_artifact_type, post_process_raw
 from langflow.schema.data import Data
+from langflow.schema.message import Message
 from langflow.template.field.base import UNDEFINED, Input, Output
 
 from .custom_component import CustomComponent
@@ -96,6 +98,8 @@ class Component(CustomComponent):
                         # If the method is asynchronous, we need to await it
                         if inspect.iscoroutinefunction(method):
                             result = await result
+                        if isinstance(result, Message) and result.flow_id is None:
+                            result.set_flow_id(vertex.graph.flow_id)
                         _results[output.name] = result
                         output.value = result
                         custom_repr = self.custom_repr()

src/backend/base/langflow/schema/message.py

@@ -22,8 +22,8 @@ class Message(Data):
     # Helper class to deal with image data
     text_key: str = "text"
     text: Optional[str | AsyncIterator | Iterator] = Field(default="")
-    sender: str
+    sender: Optional[str] = None
-    sender_name: str
+    sender_name: Optional[str] = None
     files: Optional[list[str | Image]] = Field(default=[])
     session_id: Optional[str] = Field(default="")
     timestamp: Annotated[str, BeforeValidator(_timestamp_to_str)] = Field(
@@ -39,6 +39,11 @@ class Message(Data):
             else:
                 new_files.append(file)
         self.files = new_files
+        if "timestamp" not in self.data:
+            self.data["timestamp"] = self.timestamp
+
+    def set_flow_id(self, flow_id: str):
+        self.flow_id = flow_id
 
     def to_lc_message(
         self,