diff --git a/src/backend/base/langflow/graph/utils.py b/src/backend/base/langflow/graph/utils.py index 06b7ca90a..83e2177b1 100644 --- a/src/backend/base/langflow/graph/utils.py +++ b/src/backend/base/langflow/graph/utils.py @@ -1,12 +1,9 @@ -from enum import Enum -from typing import Any, Generator, Union +from typing import Any, Union from langchain_core.documents import Document from pydantic import BaseModel from langflow.interface.utils import extract_input_variables_from_prompt -from langflow.schema import Record -from langflow.schema.message import Message class UnbuiltObject: @@ -17,16 +14,6 @@ class UnbuiltResult: pass -class ArtifactType(str, Enum): - TEXT = "text" - RECORD = "record" - OBJECT = "object" - ARRAY = "array" - STREAM = "stream" - UNKNOWN = "unknown" - MESSAGE = "message" - - def validate_prompt(prompt: str): """Validate prompt.""" if extract_input_variables_from_prompt(prompt): @@ -63,38 +50,3 @@ def serialize_field(value): elif isinstance(value, str): return {"result": value} return value - - -def get_artifact_type(custom_component, build_result) -> str: - result = ArtifactType.UNKNOWN - value = custom_component.repr_value - match value: - case Record(): - result = ArtifactType.RECORD - - case str(): - result = ArtifactType.TEXT - - case dict(): - result = ArtifactType.OBJECT - - case list(): - result = ArtifactType.ARRAY - - case Message(): - result = ArtifactType.MESSAGE - - if result == ArtifactType.UNKNOWN: - if isinstance(build_result, Generator): - result = ArtifactType.STREAM - elif isinstance(value, Message) and isinstance(value.text, Generator): - result = ArtifactType.STREAM - - return result.value - - -def post_process_raw(raw, artifact_type: str): - if artifact_type == ArtifactType.STREAM.value: - raw = "" - - return raw diff --git a/src/backend/base/langflow/graph/vertex/base.py b/src/backend/base/langflow/graph/vertex/base.py index 5d92dc694..541990426 100644 --- a/src/backend/base/langflow/graph/vertex/base.py +++ b/src/backend/base/langflow/graph/vertex/base.py @@ -9,9 +9,10 @@ from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Dict, Iterator, from loguru import logger from langflow.graph.schema import INPUT_COMPONENTS, OUTPUT_COMPONENTS, InterfaceComponentTypes, ResultData -from langflow.graph.utils import ArtifactType, UnbuiltObject, UnbuiltResult +from langflow.graph.utils import UnbuiltObject, UnbuiltResult from langflow.interface.initialize import loading from langflow.interface.listing import lazy_load_dict +from langflow.schema.artifact import ArtifactType from langflow.schema.schema import INPUT_FIELD_NAME from langflow.services.deps import get_storage_service from langflow.services.monitor.utils import log_transaction diff --git a/src/backend/base/langflow/graph/vertex/types.py b/src/backend/base/langflow/graph/vertex/types.py index c2a54f418..f3bb1d58f 100644 --- a/src/backend/base/langflow/graph/vertex/types.py +++ b/src/backend/base/langflow/graph/vertex/types.py @@ -6,9 +6,10 @@ from langchain_core.messages import AIMessage, AIMessageChunk from loguru import logger from langflow.graph.schema import CHAT_COMPONENTS, RECORDS_COMPONENTS, InterfaceComponentTypes -from langflow.graph.utils import ArtifactType, UnbuiltObject, serialize_field +from langflow.graph.utils import UnbuiltObject, serialize_field from langflow.graph.vertex.base import Vertex from langflow.schema import Record +from langflow.schema.artifact import ArtifactType from langflow.schema.schema import INPUT_FIELD_NAME from langflow.services.monitor.utils import log_transaction, log_vertex_build from langflow.utils.schemas import ChatOutputResponse, RecordOutputResponse diff --git a/src/backend/base/langflow/interface/initialize/loading.py b/src/backend/base/langflow/interface/initialize/loading.py index 3b9f90116..132703b88 100644 --- a/src/backend/base/langflow/interface/initialize/loading.py +++ b/src/backend/base/langflow/interface/initialize/loading.py @@ -7,8 +7,8 @@ import orjson from loguru import logger from langflow.custom.eval import eval_custom_component_code -from langflow.graph.utils import get_artifact_type, post_process_raw from langflow.schema import Record +from langflow.schema.artifact import get_artifact_type, post_process_raw if TYPE_CHECKING: from langflow.custom import Component, CustomComponent @@ -159,7 +159,7 @@ async def build_custom_component(params: dict, custom_component: "CustomComponen elif hasattr(raw, "model_dump") and raw is not None: raw = raw.model_dump() - artifact_type = get_artifact_type(custom_component, build_result) + artifact_type = get_artifact_type(custom_component.repr_value, build_result) raw = post_process_raw(raw, artifact_type) artifact = {"repr": custom_repr, "raw": raw, "type": artifact_type} return custom_component, build_result, artifact diff --git a/src/backend/base/langflow/schema/artifact.py b/src/backend/base/langflow/schema/artifact.py new file mode 100644 index 000000000..1cf0cad18 --- /dev/null +++ b/src/backend/base/langflow/schema/artifact.py @@ -0,0 +1,49 @@ +from enum import Enum +from typing import Generator + +from langflow.schema import Record +from langflow.schema.message import Message + + +class ArtifactType(str, Enum): + TEXT = "text" + RECORD = "record" + OBJECT = "object" + ARRAY = "array" + STREAM = "stream" + UNKNOWN = "unknown" + MESSAGE = "message" + + +def get_artifact_type(value, build_result=None) -> str: + result = ArtifactType.UNKNOWN + match value: + case Record(): + result = ArtifactType.RECORD + + case str(): + result = ArtifactType.TEXT + + case dict(): + result = ArtifactType.OBJECT + + case list(): + result = ArtifactType.ARRAY + + case Message(): + result = ArtifactType.MESSAGE + + if result == ArtifactType.UNKNOWN: + if build_result and isinstance(build_result, Generator): + result = ArtifactType.STREAM + elif isinstance(value, Message) and isinstance(value.text, Generator): + result = ArtifactType.STREAM + + return result.value + + +def post_process_raw(raw, artifact_type: str): + if artifact_type == ArtifactType.STREAM.value: + raw = "" + + return raw