From 3f3ac0400f4506849bd805c4547ade03fabad892 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Mon, 4 Mar 2024 15:19:17 -0300
Subject: [PATCH] Update schemas and vertex classes

---
 src/backend/langflow/api/v1/schemas.py     |  2 +-
 src/backend/langflow/graph/schema.py       |  3 ++-
 src/backend/langflow/graph/vertex/base.py  | 29 ++++++++++++++++++++++
 src/backend/langflow/graph/vertex/types.py | 15 +++++------
 4 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/src/backend/langflow/api/v1/schemas.py b/src/backend/langflow/api/v1/schemas.py
index 94f60d822..7a91473e6 100644
--- a/src/backend/langflow/api/v1/schemas.py
+++ b/src/backend/langflow/api/v1/schemas.py
@@ -232,7 +232,7 @@ class VertexBuildResponse(BaseModel):
     next_vertices_ids: Optional[List[str]] = None
     inactive_vertices: Optional[List[str]] = None
     valid: bool
-    params: Optional[str]
+    params: Optional[Any] = Field(default_factory=dict)
     """JSON string of the params."""
     data: ResultDataResponse
     """Mapping of vertex ids to result dict containing the param name and result value."""
diff --git a/src/backend/langflow/graph/schema.py b/src/backend/langflow/graph/schema.py
index f53a0833f..a9f06ac1e 100644
--- a/src/backend/langflow/graph/schema.py
+++ b/src/backend/langflow/graph/schema.py
@@ -4,12 +4,13 @@ from typing import Any, Optional
 from pydantic import BaseModel, Field, field_serializer
 
 from langflow.graph.utils import serialize_field
-from langflow.utils.schemas import ContainsEnumMeta
+from langflow.utils.schemas import ChatOutputResponse, ContainsEnumMeta
 
 
 class ResultData(BaseModel):
     results: Optional[Any] = Field(default_factory=dict)
     artifacts: Optional[Any] = Field(default_factory=dict)
+    messages: Optional[list[ChatOutputResponse]] = Field(default_factory=list)
     timedelta: Optional[float] = None
     duration: Optional[str] = None
 
diff --git a/src/backend/langflow/graph/vertex/base.py b/src/backend/langflow/graph/vertex/base.py
index 4a00f74ea..739748f7c 100644
--- a/src/backend/langflow/graph/vertex/base.py
+++ b/src/backend/langflow/graph/vertex/base.py
@@ -28,6 +28,7 @@ from langflow.interface.initialize import loading
 from langflow.interface.listing import lazy_load_dict
 from langflow.services.deps import get_storage_service
 from langflow.utils.constants import DIRECT_TYPES
+from langflow.utils.schemas import ChatOutputResponse
 from langflow.utils.util import sync_to_async
 
 if TYPE_CHECKING:
@@ -413,15 +414,43 @@ class Vertex:
 
         self._built = True
 
+    def extract_messages_from_artifacts(self, artifacts: Dict[str, Any]) -> List[str]:
+        """
+        Extracts messages from the artifacts.
+
+        Args:
+            artifacts (Dict[str, Any]): The artifacts to extract messages from.
+
+        Returns:
+            List[str]: The extracted messages.
+        """
+        messages = []
+        for key, artifact in artifacts.items():
+            if not isinstance(artifact, dict):
+                continue
+            if "message" in artifact:
+                chat_output_response = ChatOutputResponse(
+                    message=artifact["message"],
+                    sender=artifact.get("sender"),
+                    sender_name=artifact.get("sender_name"),
+                    session_id=artifact.get("session_id"),
+                    component_id=self.id,
+                )
+                messages.append(chat_output_response.model_dump(exclude_none=True))
+
+        return messages
+
     def _finalize_build(self):
         result_dict = self.get_built_result()
         # We need to set the artifacts to pass information
         # to the frontend
         self.set_artifacts()
         artifacts = self.artifacts
+        messages = self.extract_messages_from_artifacts(artifacts)
         result_dict = ResultData(
             results=result_dict,
             artifacts=artifacts,
+            messages=messages,
         )
         self.set_result(result_dict)
 
diff --git a/src/backend/langflow/graph/vertex/types.py b/src/backend/langflow/graph/vertex/types.py
index b01443fac..1c93729f1 100644
--- a/src/backend/langflow/graph/vertex/types.py
+++ b/src/backend/langflow/graph/vertex/types.py
@@ -1,6 +1,7 @@
 import ast
 import json
-from typing import AsyncIterator, Callable, Dict, Iterator, List, Optional, Union
+from typing import (AsyncIterator, Callable, Dict, Iterator, List, Optional,
+                    Union)
 
 import yaml
 from langchain_core.messages import AIMessage
@@ -124,13 +125,13 @@ class DocumentLoaderVertex(StatefulVertex):
 
         if not isinstance(self._built_object, UnbuiltObject):
             avg_length = sum(
-                len(doc.page_content)
-                for doc in self._built_object
-                if hasattr(doc, "page_content")
+                len(record.text)
+                for record in self._built_object
+                if hasattr(record, "text")
             ) / len(self._built_object)
-            return f"""{self.display_name}({len(self._built_object)} documents)
-            \nAvg. Document Length (characters): {int(avg_length)}
-            Documents: {self._built_object[:3]}..."""
+            return f"""{self.display_name}({len(self._built_object)} records)
+            \nAvg. Record Length (characters): {int(avg_length)}
+            Records: {self._built_object[:3]}..."""
         return f"{self.vertex_type}()"