diff --git a/src/backend/base/langflow/graph/utils.py b/src/backend/base/langflow/graph/utils.py
index 332545caf..6d1dca0cb 100644
--- a/src/backend/base/langflow/graph/utils.py
+++ b/src/backend/base/langflow/graph/utils.py
@@ -5,6 +5,7 @@ from langchain_core.documents import Document
 from pydantic import BaseModel
 
 from langflow.interface.utils import extract_input_variables_from_prompt
+from langflow.schema.data import Data
 from langflow.schema.message import Message
 
 
@@ -67,7 +68,7 @@ def serialize_field(value):
 def get_artifact_type(value, build_result) -> str:
     result = ArtifactType.UNKNOWN
     match value:
-        case Record():
+        case Data():
             result = ArtifactType.RECORD
 
         case str():
diff --git a/tests/conftest.py b/tests/conftest.py
index b3f537a3c..e3948fbe8 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -23,7 +23,6 @@ from langflow.services.database.models.folder.model import Folder
 from langflow.services.database.models.user.model import User, UserCreate
 from langflow.services.database.utils import session_getter
 from langflow.services.deps import get_db_service
-from sqlmodel import Session, SQLModel, create_engine, select
 from sqlmodel.pool import StaticPool
 from typer.testing import CliRunner