refactor: Remove unused imports and update artifact type handling

src/backend/base/langflow/graph/utils.py

@@ -1,12 +1,9 @@
-from enum import Enum
+from typing import Any, Union
-from typing import Any, Generator, 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

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

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

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

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