Update schemas and vertex classes

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-04 15:19:17 -03:00
commit 3f3ac0400f
4 changed files with 40 additions and 9 deletions

View file

@ -232,7 +232,7 @@ class VertexBuildResponse(BaseModel):
next_vertices_ids: Optional[List[str]] = None next_vertices_ids: Optional[List[str]] = None
inactive_vertices: Optional[List[str]] = None inactive_vertices: Optional[List[str]] = None
valid: bool valid: bool
params: Optional[str] params: Optional[Any] = Field(default_factory=dict)
"""JSON string of the params.""" """JSON string of the params."""
data: ResultDataResponse data: ResultDataResponse
"""Mapping of vertex ids to result dict containing the param name and result value.""" """Mapping of vertex ids to result dict containing the param name and result value."""

View file

@ -4,12 +4,13 @@ from typing import Any, Optional
from pydantic import BaseModel, Field, field_serializer from pydantic import BaseModel, Field, field_serializer
from langflow.graph.utils import serialize_field from langflow.graph.utils import serialize_field
from langflow.utils.schemas import ContainsEnumMeta from langflow.utils.schemas import ChatOutputResponse, ContainsEnumMeta
class ResultData(BaseModel): class ResultData(BaseModel):
results: Optional[Any] = Field(default_factory=dict) results: Optional[Any] = Field(default_factory=dict)
artifacts: 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 timedelta: Optional[float] = None
duration: Optional[str] = None duration: Optional[str] = None

View file

@ -28,6 +28,7 @@ from langflow.interface.initialize import loading
from langflow.interface.listing import lazy_load_dict from langflow.interface.listing import lazy_load_dict
from langflow.services.deps import get_storage_service from langflow.services.deps import get_storage_service
from langflow.utils.constants import DIRECT_TYPES from langflow.utils.constants import DIRECT_TYPES
from langflow.utils.schemas import ChatOutputResponse
from langflow.utils.util import sync_to_async from langflow.utils.util import sync_to_async
if TYPE_CHECKING: if TYPE_CHECKING:
@ -413,15 +414,43 @@ class Vertex:
self._built = True 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): def _finalize_build(self):
result_dict = self.get_built_result() result_dict = self.get_built_result()
# We need to set the artifacts to pass information # We need to set the artifacts to pass information
# to the frontend # to the frontend
self.set_artifacts() self.set_artifacts()
artifacts = self.artifacts artifacts = self.artifacts
messages = self.extract_messages_from_artifacts(artifacts)
result_dict = ResultData( result_dict = ResultData(
results=result_dict, results=result_dict,
artifacts=artifacts, artifacts=artifacts,
messages=messages,
) )
self.set_result(result_dict) self.set_result(result_dict)

View file

@ -1,6 +1,7 @@
import ast import ast
import json import json
from typing import AsyncIterator, Callable, Dict, Iterator, List, Optional, Union from typing import (AsyncIterator, Callable, Dict, Iterator, List, Optional,
Union)
import yaml import yaml
from langchain_core.messages import AIMessage from langchain_core.messages import AIMessage
@ -124,13 +125,13 @@ class DocumentLoaderVertex(StatefulVertex):
if not isinstance(self._built_object, UnbuiltObject): if not isinstance(self._built_object, UnbuiltObject):
avg_length = sum( avg_length = sum(
len(doc.page_content) len(record.text)
for doc in self._built_object for record in self._built_object
if hasattr(doc, "page_content") if hasattr(record, "text")
) / len(self._built_object) ) / len(self._built_object)
return f"""{self.display_name}({len(self._built_object)} documents) return f"""{self.display_name}({len(self._built_object)} records)
\nAvg. Document Length (characters): {int(avg_length)} \nAvg. Record Length (characters): {int(avg_length)}
Documents: {self._built_object[:3]}...""" Records: {self._built_object[:3]}..."""
return f"{self.vertex_type}()" return f"{self.vertex_type}()"