⚡️ Speed up method NativeTracer._map_trace_type by 136% in PR #11934 (temp-branch)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #11934

If you approve this dependent PR, these changes will be merged into the original PR branch temp-branch.

This PR will be automatically closed if the original PR is merged.

---

📄 136% (1.36x) speedup for NativeTracer._map_trace_type in src/backend/base/langflow/services/tracing/native.py

⏱️ Runtime : 1.23 milliseconds → 523 microseconds (best of 295 runs)

📝 Explanation and details

Brief: The optimized version speeds up _map_trace_type and initialization by removing per-call allocations and reducing string scanning. The key wins are (1) moving the trace-type dict to a module-level constant so it is not rebuilt on every call, and (2) using str.rpartition(" - ") for flow_id extraction to avoid scanning/splitting the trace_name twice. Together these reduce CPU work, temporary allocations, and Python bytecode executed per call — producing the measured ~135% speedup.

What changed (concrete):

• Moved the mapping dict from inside _map_trace_type to a module-level _TYPE_MAP constant. The static method now does a single _TYPE_MAP.get(trace_type.lower(), SpanType.CHAIN).
• Replaced the flow_id fallback logic flow_id or (trace_name.split(" - ")[-1] if " - " in trace_name else trace_name) with flow_id or trace_name.rpartition(" - ")[-1].
• Minor reorganization of imports and annotations (no behavioral change).

Why this is faster:

• Module-level mapping: building a dict is non-trivial (creates new objects and memory each call). The original profiler shows significant time spent on those dict literal lines every call. By creating the dict once at import time, each call to _map_trace_type only does a lower() and a dict lookup (both cheap), eliminating repeated allocations and GC churn.
• rpartition vs "in"+"split": the original code did an "in" test and a split (or used split[-1]) which can scan the string twice and allocate a list. rpartition scans once and returns the parts without creating a list of arbitrary length; that reduces CPU and allocations when building flow_id.
• Fewer Python-level operations: fewer bytecode instructions and attribute lookups per call. The optimized _map_trace_type is one line, which the line profiler confirms — nearly all time becomes the .lower() + .get() cost.

How this affects workloads:

• Big benefit when _map_trace_type is called many times (hot path, loops, repeated mapping). The annotated tests that call the mapper thousands of times (large-scale deterministic tests and repeated calls) are the cases that show the largest improvements.
• If _map_trace_type is only called occasionally (e.g., once per process start), the user-exvisible effect is small. But if many tracers are created or trace types are resolved repeatedly, the improvement compounds.
• The rpartition change speeds up tracer initialization where flow_id fallback is executed; again, more beneficial if many tracer objects are constructed.

Behavioral compatibility:

• Semantics are preserved: case-insensitive matching via .lower() is unchanged, unknown values still default to SpanType.CHAIN, and rpartition returns the full trace_name when the separator is absent (same as the original split behavior).
• No change to error behavior for non-string inputs (calling .lower() on None will still raise AttributeError), so tests that assert current behavior still pass.

Tests that benefit most:

• Repeated-call and large-scale tests (1000-entry deterministic mapping, repeated calls loops) — these are the scenarios where the profiler and runtime show the greatest gains.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 2077 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import pytest  # used for our unit tests
# Import the real classes from the actual module under test.
# The provided source shows the class lives in langflow.services.tracing.native
from langflow.services.tracing.native import NativeTracer, SpanType


def test_map_known_types_lowercase():
    # Verify each canonical lowercase trace_type maps to the expected SpanType member.
    expected_map = {
        "chain": SpanType.CHAIN,
        "llm": SpanType.LLM,
        "tool": SpanType.TOOL,
        "retriever": SpanType.RETRIEVER,
        "embedding": SpanType.EMBEDDING,
        "parser": SpanType.PARSER,
        "agent": SpanType.AGENT,
    }
    # Iterate over the mapping and assert the static method returns the correct enum
    for key, expected in expected_map.items():
        codeflash_output = NativeTracer._map_trace_type(key); result = codeflash_output  # call the staticmethod directly


def test_map_known_types_mixed_case():
    # The implementation uses .lower(), so mixed-case inputs should map correctly.
    mixed_inputs = {
        "Chain": SpanType.CHAIN,
        "ChAiN": SpanType.CHAIN,
        "LLM": SpanType.LLM,
        "LlM": SpanType.LLM,
        "TOol": SpanType.TOOL,
        "Retriever": SpanType.RETRIEVER,
        "EMBEDDING": SpanType.EMBEDDING,
        "Parser": SpanType.PARSER,
        "AgEnT": SpanType.AGENT,
    }
    # Ensure case-insensitivity holds for all valid keys.
    for inp, expected in mixed_inputs.items():
        codeflash_output = NativeTracer._map_trace_type(inp)


def test_empty_and_unknown_strings_default_to_chain():
    # Empty string should not match any key and therefore should default to CHAIN.
    codeflash_output = NativeTracer._map_trace_type("")
    # Completely unknown strings also default to CHAIN.
    codeflash_output = NativeTracer._map_trace_type("this-does-not-exist")
    # Numeric-like strings default to CHAIN as well.
    codeflash_output = NativeTracer._map_trace_type("12345")


def test_whitespace_and_special_characters_do_not_trim_or_match():
    # The implementation does not strip whitespace; leading/trailing whitespace will prevent a match.
    # " llm " will not equal "llm" after lower(), so should default to CHAIN.
    codeflash_output = NativeTracer._map_trace_type(" llm ")
    # Newlines or tabs included in the string should also prevent a match.
    codeflash_output = NativeTracer._map_trace_type("LLM\n")
    codeflash_output = NativeTracer._map_trace_type("\ttool")
    # Special characters appended to valid token should not match.
    codeflash_output = NativeTracer._map_trace_type("agent!")


def test_none_raises_attribute_error():
    # Passing None is a type error for the implementation because it calls .lower() on the input.
    # We assert the specific exception attribute error to document current behavior.
    with pytest.raises(AttributeError):
        NativeTracer._map_trace_type(None)  # type: ignore[arg-type]


def test_repeated_calls_are_consistent():
    # Calling the mapper repeatedly with the same input should always return the same enum value.
    for _ in range(10):
        codeflash_output = NativeTracer._map_trace_type("llm")
    # And for unknowns too.
    for _ in range(10):
        codeflash_output = NativeTracer._map_trace_type("nonsense")


def test_large_scale_deterministic_mapping_1000_entries():
    # Build a deterministic (non-random) pattern of 1000 inputs mixing valid and invalid types.
    valid_keys = ["chain", "llm", "tool", "retriever", "embedding", "parser", "agent"]
    invalid_keys = ["", "  llm  ", "UNKNOWN", "123", "agent!", "LLM\n", "some-other-type"]
    inputs = []
    # Create exactly 1000 inputs by cycling through valid and invalid lists deterministically.
    for i in range(1000):
        if i % 7 == 0:
            # pick a valid key cycling through valid_keys
            inputs.append(valid_keys[(i // 7) % len(valid_keys)])
        else:
            # pick an invalid key cycling through invalid_keys
            inputs.append(invalid_keys[i % len(invalid_keys)])

    # Now map all inputs and assert the outputs match the implementation's contract:
    # valid (case-insensitive) keys -> corresponding SpanType, everything else -> SpanType.CHAIN
    expected_map = {
        "chain": SpanType.CHAIN,
        "llm": SpanType.LLM,
        "tool": SpanType.TOOL,
        "retriever": SpanType.RETRIEVER,
        "embedding": SpanType.EMBEDDING,
        "parser": SpanType.PARSER,
        "agent": SpanType.AGENT,
    }

    # Evaluate each input deterministically and assert correctness.
    for inp in inputs:
        # If the lowercased input exactly matches a key, expect the mapped value.
        lowered = inp.lower()
        if lowered in expected_map:
            expected = expected_map[lowered]
        else:
            expected = SpanType.CHAIN
        codeflash_output = NativeTracer._map_trace_type(inp); result = codeflash_output


def test_all_known_keys_covered_and_return_enum_instances():
    # Sanity check: ensure that every expected SpanType is returned for its key and that
    # return values are actually of type SpanType (enum members).
    mapping_keys = ["chain", "llm", "tool", "retriever", "embedding", "parser", "agent"]
    for key in mapping_keys:
        codeflash_output = NativeTracer._map_trace_type(key); value = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import random  # used for deterministic large-scale generation

# imports
import pytest  # used for our unit tests
from langflow.services.tracing.native import NativeTracer


def test_basic_known_types_lowercase():
    # Basic test: verify that all known, lowercase trace_type inputs map to the expected SpanType names.
    # We compare the Enum .name to avoid importing the SpanType symbol directly.
    mapping = {
        "chain": "CHAIN",
        "llm": "LLM",
        "tool": "TOOL",
        "retriever": "RETRIEVER",
        "embedding": "EMBEDDING",
        "parser": "PARSER",
        "agent": "AGENT",
    }
    for trace_type_input, expected_name in mapping.items():
        # Call the static method directly from the real class
        codeflash_output = NativeTracer._map_trace_type(trace_type_input); result = codeflash_output


def test_case_insensitivity_and_mixed_case():
    # The mapping should be case-insensitive. Test several mixed-case variations.
    cases = [
        ("ChAiN", "CHAIN"),
        ("LlM", "LLM"),
        ("ToOl", "TOOL"),
        ("ReTrIeVeR", "RETRIEVER"),
        ("EMBedDing", "EMBEDDING"),
        ("parser", "PARSER"),
        ("AGENT", "AGENT"),
    ]
    for input_value, expected_name in cases:
        # Mixed-case input should still map correctly thanks to .lower() usage in implementation
        codeflash_output = NativeTracer._map_trace_type(input_value); result = codeflash_output


def test_unknown_and_whitespace_behavior_defaults_to_chain():
    # Unknown trace types (and types with extra whitespace) should default to CHAIN per implementation.
    # Note: the implementation does not strip whitespace, so ' chain ' is not recognized and should default.
    unknown_inputs = [
        "unknown-type",
        "",            # empty string should default
        " chain ",     # leading/trailing whitespace prevents match in current implementation
        "chains",      # similar but not exact -> default
        "123",         # numeric-looking string -> default
    ]
    for inp in unknown_inputs:
        codeflash_output = NativeTracer._map_trace_type(inp); result = codeflash_output


def test_non_string_inputs_raise_attribute_error():
    # Since the implementation calls .lower() on the input, non-string inputs should raise AttributeError.
    # This documents the current behavior (no explicit type-checking).
    non_strings = [None, 42, 3.14, object()]
    for val in non_strings:
        with pytest.raises(AttributeError):
            NativeTracer._map_trace_type(val)  # type: ignore[arg-type]


def test_idempotence_and_identity_of_enum_members():
    # Enums are singletons; repeated calls with logically equal inputs should return the same object identity.
    codeflash_output = NativeTracer._map_trace_type("chain"); first = codeflash_output
    codeflash_output = NativeTracer._map_trace_type("CHAIN"); second = codeflash_output
    codeflash_output = NativeTracer._map_trace_type("ChAiN"); third = codeflash_output


def test_large_scale_deterministic_mappings():
    # Large-scale test: generate 1000 deterministic inputs mixing known keys and unknowns,
    # and verify correct mapping for each. This checks performance at scale and deterministic behavior.
    random.seed(0)  # make generation deterministic for repeatable tests

    known_keys = ["chain", "llm", "tool", "retriever", "embedding", "parser", "agent"]
    results = []
    # Build 1000 test inputs: majority known keys, some unknowns, with varied casing
    for i in range(1000):
        if i % 10 == 0:
            # Introduce an unknown every 10th entry
            inp = f"unknown_{i}"
        else:
            # Pick a known key and randomly change case to test robustness
            base = random.choice(known_keys)
            # create a mixed-case variant deterministically
            mixed = "".join(
                ch.upper() if (j + i) % 2 == 0 else ch.lower() for j, ch in enumerate(base)
            )
            inp = mixed
        results.append(inp)

    # Now map all inputs and assert correctness
    for inp in results:
        codeflash_output = NativeTracer._map_trace_type(inp); mapped = codeflash_output
        # If the input (when lowercased) is one of the known keys exactly, expect that mapping.
        lower = inp.lower()
        if lower in known_keys:
            # Should map to the corresponding known enum name
            expected_name = lower.upper() if lower != "llm" else "LLM"  # general rule: uppercased name
        else:
            pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr11934-2026-02-27T10.49.30 and push.

[codecov]

Codecov Report

❌ Patch coverage is 33.33333% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 35.73%. Comparing base (0596d77) to head (fa03f7d).

Files with missing lines	Patch %	Lines
...c/backend/base/langflow/services/tracing/native.py	33.33%	2 Missing ⚠️

❌ Your project check has failed because the head coverage (42.40%) is below the target coverage (60.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files

@@               Coverage Diff               @@
##           temp-branch   #11937      +/-   ##
===============================================
- Coverage        35.74%   35.73%   -0.01%     
===============================================
  Files             1532     1532              
  Lines            74562    74562              
  Branches         11146    11146              
===============================================
- Hits             26651    26645       -6     
- Misses           46475    46481       +6     
  Partials          1436     1436              

Flag	Coverage Δ
backend	55.77% <33.33%> (-0.03%)	⬇️
lfx	42.40% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
...c/backend/base/langflow/services/tracing/native.py	21.33% <33.33%> (+0.66%)	⬆️

... and 7 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[codeflash-ai]

This PR has been automatically closed because the original PR #11934 by Adam-Aghili was closed.