⚡️ Speed up function get_message by 10% in PR #9765 (fix/building_splittext)

langflow/schema/schema.py

@@ -0,0 +1,119 @@
+from collections.abc import Generator
+from enum import Enum
+from typing import Literal
+
+from pandas import Series
+from pydantic import BaseModel
+from typing_extensions import TypedDict
+
+from langflow.schema.data import Data
+from langflow.schema.dataframe import DataFrame
+from langflow.schema.message import Message
+from langflow.serialization.serialization import serialize
+
+INPUT_FIELD_NAME = "input_value"
+
+InputType = Literal["chat", "text", "any"]
+OutputType = Literal["chat", "text", "any", "debug"]
+
+
+class LogType(str, Enum):
+    MESSAGE = "message"
+    DATA = "data"
+    STREAM = "stream"
+    OBJECT = "object"
+    ARRAY = "array"
+    TEXT = "text"
+    UNKNOWN = "unknown"
+
+
+class StreamURL(TypedDict):
+    location: str
+
+
+class ErrorLog(TypedDict):
+    errorMessage: str
+    stackTrace: str
+
+
+class OutputValue(BaseModel):
+    message: ErrorLog | StreamURL | dict | list | str
+    type: str
+
+
+def get_type(payload):
+    result = LogType.UNKNOWN
+    match payload:
+        case Message():
+            result = LogType.MESSAGE
+
+        case Data():
+            result = LogType.DATA
+
+        case dict():
+            result = LogType.OBJECT
+
+        case list() | DataFrame():
+            result = LogType.ARRAY
+
+        case str():
+            result = LogType.TEXT
+
+    if result == LogType.UNKNOWN and (
+        (payload and isinstance(payload, Generator))
+        or (isinstance(payload, Message) and isinstance(payload.text, Generator))
+    ):
+        result = LogType.STREAM
+
+    return result
+
+
+def get_message(payload):
+    message = getattr(payload, "data", None)
+    if message is None:
+        model_dump = getattr(payload, "model_dump", None)
+        if callable(model_dump):
+            message = model_dump()
+        elif isinstance(payload, (dict, str, Data)):
+            message = payload.data if isinstance(payload, Data) else payload
+
+    if isinstance(message, Series):
+        return message if not message.empty else payload
+
+    return message or payload
+
+
+def build_output_logs(vertex, result) -> dict:
+    outputs: dict[str, OutputValue] = {}
+    component_instance = result[0]
+    for index, output in enumerate(vertex.outputs):
+        if component_instance.status is None:
+            payload = component_instance._results
+            output_result = payload.get(output["name"])
+        else:
+            payload = component_instance._artifacts
+            output_result = payload.get(output["name"], {}).get("raw")
+        message = get_message(output_result)
+        type_ = get_type(output_result)
+
+        match type_:
+            case LogType.STREAM if "stream_url" in message:
+                message = StreamURL(location=message["stream_url"])
+
+            case LogType.STREAM:
+                message = ""
+
+            case LogType.MESSAGE if hasattr(message, "message"):
+                message = message.message
+
+            case LogType.UNKNOWN:
+                message = ""
+
+            case LogType.ARRAY:
+                if isinstance(message, DataFrame):
+                    message = message.to_dict(orient="records")
+                message = [serialize(item) for item in message]
+        name = output.get("name", f"output_{index}")
+        outputs |= {name: OutputValue(message=message, type=type_).model_dump()}
+
+    return outputs

src/backend/tests/base.py

@@ -5,11 +5,11 @@
 from uuid import uuid4
 
 import pytest
+from typing_extensions import TypedDict
+
 from langflow.custom.custom_component.component import Component
 from langflow.graph.graph.base import Graph
 from langflow.graph.vertex.base import Vertex
-from typing_extensions import TypedDict
-
 from tests.constants import SUPPORTED_VERSIONS
 from tests.integration.utils import build_component_instance_for_tests
 

src/backend/tests/conftest.py

@@ -17,6 +17,13 @@
 from dotenv import load_dotenv
 from fastapi.testclient import TestClient
 from httpx import ASGITransport, AsyncClient
+from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.orm import selectinload
+from sqlmodel import Session, SQLModel, create_engine, select
+from sqlmodel.ext.asyncio.session import AsyncSession
+from sqlmodel.pool import StaticPool
+from typer.testing import CliRunner
+
 from langflow.components.input_output import ChatInput
 from langflow.graph import Graph
 from langflow.initial_setup.constants import STARTER_FOLDER_NAME
@@ -31,13 +38,6 @@
 from langflow.services.database.models.vertex_builds.crud import delete_vertex_builds_by_flow_id
 from langflow.services.database.utils import session_getter
 from langflow.services.deps import get_db_service, session_scope
-from sqlalchemy.ext.asyncio import create_async_engine
-from sqlalchemy.orm import selectinload
-from sqlmodel import Session, SQLModel, create_engine, select
-from sqlmodel.ext.asyncio.session import AsyncSession
-from sqlmodel.pool import StaticPool
-from typer.testing import CliRunner
-
 from tests.api_keys import get_openai_api_key
 
 load_dotenv()

src/backend/tests/integration/backward_compatibility/test_starter_projects.py

@@ -1,6 +1,6 @@
 import pytest
-from langflow.schema.message import Message
 
+from langflow.schema.message import Message
 from tests.api_keys import get_openai_api_key
 from tests.integration.utils import download_flow_from_github, run_json_flow
 

src/backend/tests/integration/components/assistants/test_assistants_components.py

@@ -1,12 +1,12 @@
 import pytest
+
 from langflow.components.datastax import (
     AssistantsCreateAssistant,
     AssistantsCreateThread,
     AssistantsGetAssistantName,
     AssistantsListAssistants,
     AssistantsRun,
 )
-
 from tests.integration.utils import run_single_component
 
 

src/backend/tests/integration/components/astra/test_astra_component.py

@@ -4,10 +4,10 @@
 from astrapy import DataAPIClient
 from langchain_astradb import AstraDBVectorStore, VectorServiceOptions
 from langchain_core.documents import Document
+
 from langflow.components.datastax import AstraDBVectorStoreComponent
 from langflow.components.openai.openai import OpenAIEmbeddingsComponent
 from langflow.schema.data import Data
-
 from tests.api_keys import get_astradb_api_endpoint, get_astradb_application_token, get_openai_api_key
 from tests.integration.components.mock_components import TextToData
 from tests.integration.utils import ComponentInputHandle, run_single_component

src/backend/tests/integration/components/helpers/test_parse_json_data.py

@@ -1,7 +1,6 @@
 from langflow.components.input_output import ChatInput
 from langflow.components.processing.parse_json_data import ParseJSONDataComponent
 from langflow.schema import Data
-
 from tests.integration.components.mock_components import TextToData
 from tests.integration.utils import ComponentInputHandle, pyleak_marker, run_single_component
 

src/backend/tests/integration/components/inputs/test_chat_input.py

@@ -1,7 +1,6 @@
 from langflow.components.input_output import ChatInput
 from langflow.memory import aget_messages
 from langflow.schema.message import Message
-
 from tests.integration.utils import pyleak_marker, run_single_component
 
 pytestmark = pyleak_marker()

src/backend/tests/integration/components/inputs/test_text_input.py

@@ -1,6 +1,5 @@
 from langflow.components.input_output import TextInputComponent
 from langflow.schema.message import Message
-
 from tests.integration.utils import pyleak_marker, run_single_component
 
 pytestmark = pyleak_marker()

src/backend/tests/integration/components/mcp/test_mcp_memory_leak.py

@@ -13,10 +13,11 @@
 
 import psutil
 import pytest
-from langflow.base.mcp.util import MCPSessionManager
 from loguru import logger
 from mcp import StdioServerParameters
 
+from langflow.base.mcp.util import MCPSessionManager
+
 pytestmark = [
     pytest.mark.timeout(300, method="thread"),
     pytest.mark.skip(reason="Skipping all MCP memory leak integration tests for now."),

src/backend/tests/integration/components/mcp/test_mcp_superuser_flow.py

@@ -1,4 +1,5 @@
 import pytest
+
 from langflow.services.auth.utils import create_user_longterm_token
 from langflow.services.deps import get_db_service, get_settings_service
 from langflow.services.utils import initialize_services

src/backend/tests/integration/components/output_parsers/test_output_parser.py

@@ -1,10 +1,10 @@
 import os
 
 import pytest
+
 from langflow.components.helpers import OutputParserComponent
 from langflow.components.openai.openai_chat_model import OpenAIModelComponent
 from langflow.components.processing import PromptComponent
-
 from tests.integration.utils import ComponentInputHandle, run_single_component
 
 

src/backend/tests/integration/components/outputs/test_chat_output.py

@@ -1,7 +1,6 @@
 from langflow.components.input_output import ChatOutput
 from langflow.memory import aget_messages
 from langflow.schema.message import Message
-
 from tests.integration.utils import run_single_component
 
 

src/backend/tests/integration/components/outputs/test_text_output.py

@@ -1,6 +1,5 @@
 from langflow.components.input_output import TextOutputComponent
 from langflow.schema.message import Message
-
 from tests.integration.utils import run_single_component
 
 

src/backend/tests/integration/components/prompts/test_prompt.py

@@ -1,6 +1,5 @@
 from langflow.components.processing import PromptComponent
 from langflow.schema.message import Message
-
 from tests.integration.utils import pyleak_marker, run_single_component
 
 pytestmark = pyleak_marker()

src/backend/tests/integration/flows/test_basic_prompting.py

@@ -2,7 +2,6 @@
 from langflow.components.processing import PromptComponent
 from langflow.graph import Graph
 from langflow.schema.message import Message
-
 from tests.integration.utils import pyleak_marker, run_flow
 
 

src/backend/tests/integration/test_dynamic_import_integration.py

@@ -8,6 +8,7 @@
 import time
 
 import pytest
+
 from langflow.components.agents import AgentComponent
 from langflow.components.data import APIRequestComponent
 from langflow.components.openai import OpenAIModelComponent

src/backend/tests/integration/test_exception_telemetry.py

@@ -4,6 +4,7 @@
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+
 from langflow.services.telemetry.service import TelemetryService
 
 

src/backend/tests/integration/test_misc.py

@@ -3,6 +3,7 @@
 import pytest
 from fastapi import status
 from httpx import AsyncClient
+
 from langflow.graph.schema import RunOutputs
 from langflow.initial_setup.setup import load_starter_projects
 from langflow.load.load import arun_flow_from_json

src/backend/tests/integration/utils.py

@@ -6,6 +6,7 @@
 import pytest
 import requests
 from astrapy.admin import parse_api_endpoint
+
 from langflow.api.v1.schemas import InputValueRequest
 from langflow.custom import Component
 from langflow.custom.eval import eval_custom_component_code

src/backend/tests/performance/test_server_init.py

@@ -1,6 +1,7 @@
 import os
 
 import pytest
+
 from langflow.services.deps import get_settings_service
 
 

src/backend/tests/unit/api/v1/test_api_schemas.py

@@ -2,11 +2,12 @@
 
 from hypothesis import HealthCheck, example, given, settings
 from hypothesis import strategies as st
+from pydantic import BaseModel
+
 from langflow.api.v1.schemas import ResultDataResponse, VertexBuildResponse
 from langflow.schema.schema import OutputValue
 from langflow.serialization import serialize
 from langflow.services.tracing.schema import Log
-from pydantic import BaseModel
 
 # Use a smaller test size for hypothesis
 TEST_TEXT_LENGTH = 50

src/backend/tests/unit/api/v1/test_endpoints.py

@@ -5,6 +5,7 @@
 from anyio import Path
 from fastapi import status
 from httpx import AsyncClient
+
 from langflow.api.v1.schemas import CustomComponentRequest, UpdateCustomComponentRequest
 from langflow.components.agents.agent import AgentComponent
 from langflow.custom.utils import build_custom_component_template

src/backend/tests/unit/api/v1/test_files.py

@@ -11,16 +11,16 @@
 import pytest
 from asgi_lifespan import LifespanManager
 from httpx import ASGITransport, AsyncClient
+from sqlalchemy.orm import selectinload
+from sqlmodel import select
+
 from langflow.main import create_app
 from langflow.services.auth.utils import get_password_hash
 from langflow.services.database.models.api_key.model import ApiKey
 from langflow.services.database.models.flow.model import Flow, FlowCreate
 from langflow.services.database.models.user.model import User, UserRead
 from langflow.services.database.utils import session_getter
 from langflow.services.deps import get_db_service
-from sqlalchemy.orm import selectinload
-from sqlmodel import select
-
 from tests.conftest import _delete_transactions_and_vertex_builds
 
 

src/backend/tests/unit/api/v1/test_flows.py

@@ -4,6 +4,7 @@
 from anyio import Path
 from fastapi import status
 from httpx import AsyncClient
+
 from langflow.services.database.models import Flow
 
 

src/backend/tests/unit/api/v1/test_mcp.py

@@ -4,6 +4,7 @@
 import pytest
 from fastapi import status
 from httpx import AsyncClient
+
 from langflow.services.auth.utils import get_password_hash
 from langflow.services.database.models.user import User
 

src/backend/tests/unit/api/v1/test_mcp_projects.py

@@ -5,6 +5,9 @@
 import pytest
 from fastapi import HTTPException, status
 from httpx import AsyncClient
+from mcp.server.sse import SseServerTransport
+from sqlmodel import select
+
 from langflow.api.v1.mcp_projects import (
     get_project_mcp_server,
     get_project_sse,
@@ -18,8 +21,6 @@
 from langflow.services.database.models.user.model import User
 from langflow.services.deps import get_db_service, get_settings_service, session_scope
 from langflow.services.utils import initialize_services
-from mcp.server.sse import SseServerTransport
-from sqlmodel import select
 
 # Mark all tests in this module as asyncio
 pytestmark = pytest.mark.asyncio