chore: update test cases and cleanup

Author: Steffen911

langfuse/_client/client.py

@@ -48,6 +48,9 @@
     ObservationTypeLiteralNoEvent,
     ObservationTypeSpanLike,
     get_observation_types_list,
+    LANGFUSE_CTX_USER_ID,
+    LANGFUSE_CTX_SESSION_ID,
+    LANGFUSE_CTX_METADATA,
 )
 from langfuse._client.datasets import DatasetClient, DatasetItemClient
 from langfuse._client.environment_variables import (
@@ -116,12 +119,6 @@
 from langfuse.types import MaskFunction, ScoreDataType, SpanLevel, TraceContext
 
 
-# Context key constants for Langfuse context propagation
-LANGFUSE_CTX_USER_ID = "langfuse.ctx.user.id"
-LANGFUSE_CTX_SESSION_ID = "langfuse.ctx.session.id"
-LANGFUSE_CTX_METADATA = "langfuse.ctx.metadata"
-
-
 class Langfuse:
     """Main client for Langfuse tracing and platform features.
 

langfuse/_client/constants.py

@@ -8,6 +8,11 @@
 
 LANGFUSE_TRACER_NAME = "langfuse-sdk"
 
+# Context key constants for Langfuse context propagation
+LANGFUSE_CTX_USER_ID = "langfuse.ctx.user.id"
+LANGFUSE_CTX_SESSION_ID = "langfuse.ctx.session.id"
+LANGFUSE_CTX_METADATA = "langfuse.ctx.metadata"
+
 
 """Note: this type is used with .__args__ / get_args in some cases and therefore must remain flat"""
 ObservationTypeGenerationLike: TypeAlias = Literal[

langfuse/_client/span_processor.py

@@ -22,7 +22,12 @@
 from opentelemetry.sdk.trace import ReadableSpan, Span
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 
-from langfuse._client.constants import LANGFUSE_TRACER_NAME
+from langfuse._client.constants import (
+    LANGFUSE_TRACER_NAME,
+    LANGFUSE_CTX_USER_ID,
+    LANGFUSE_CTX_SESSION_ID,
+    LANGFUSE_CTX_METADATA,
+)
 from langfuse._client.environment_variables import (
     LANGFUSE_FLUSH_AT,
     LANGFUSE_FLUSH_INTERVAL,
@@ -31,11 +36,6 @@
 from langfuse._client.utils import span_formatter
 from langfuse.logger import langfuse_logger
 from langfuse.version import __version__ as langfuse_version
-from langfuse._client.client import (
-    LANGFUSE_CTX_USER_ID,
-    LANGFUSE_CTX_SESSION_ID,
-    LANGFUSE_CTX_METADATA,
-)
 
 
 class LangfuseSpanProcessor(BatchSpanProcessor):
@@ -143,10 +143,11 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None
         # 1. Propagate all baggage keys as span attributes
         baggage_entries = baggage.get_all(context=current_context)
         for key, value in baggage_entries.items():
-            # Check if this baggage entry is already present as a span attribute
-            if not hasattr(span.attributes, key) or (
-                span.attributes is not None and span.attributes.get(key) != value
-            ):
+            # Only propagate user.id, session.id and langfuse.metadata.* as those are set by us on the baggage
+            if key.startswith("langfuse.metadata.") or key in [
+                "user.id",
+                "session.id",
+            ]:
                 propagated_attributes[key] = value
                 langfuse_logger.debug(
                     f"Propagated baggage key '{key}' = '{value}' to span '{span.name}'"
@@ -160,16 +161,10 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None
                 if value is not None:
                     # Convert context key to span attribute name (remove langfuse.ctx. prefix)
                     attr_key = ctx_key.replace("langfuse.ctx.", "")
-
-                    # Only propagate if not already set on span
-                    if not hasattr(span.attributes, attr_key) or (
-                        span.attributes is not None
-                        and span.attributes.get(attr_key) != value
-                    ):
-                        propagated_attributes[attr_key] = value
-                        langfuse_logger.debug(
-                            f"Propagated context key '{ctx_key}' = '{value}' to span '{span.name}'"
-                        )
+                    propagated_attributes[attr_key] = value
+                    langfuse_logger.debug(
+                        f"Propagated context key '{ctx_key}' = '{value}' to span '{span.name}'"
+                    )
             except Exception as e:
                 langfuse_logger.debug(f"Could not read context key '{ctx_key}': {e}")
 
@@ -184,24 +179,17 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None
                 for key, value in metadata_dict.items():
                     attr_key = f"langfuse.metadata.{key}"
 
-                    # Convert value to appropriate type for span attribute
-                    if isinstance(value, (str, int, float, bool)):
-                        attr_value = value
-                    else:
-                        # For complex types, convert to JSON string
-                        attr_value = json.dumps(value)
-
-                    # Only propagate if not already set or different
-                    existing_value = (
-                        span.attributes.get(attr_key)
-                        if hasattr(span, "attributes") and span.attributes is not None
-                        else None
+                    # Convert value to appropriate type for span attribute (naive or json stringify)
+                    attr_value = (
+                        value
+                        if isinstance(value, (str, int, float, bool))
+                        else json.dumps(value)
+                    )
+
+                    propagated_attributes[attr_key] = attr_value
+                    langfuse_logger.debug(
+                        f"Propagated metadata key '{key}' = '{attr_value}' to span '{span.name}'"
                     )
-                    if existing_value != attr_value:
-                        propagated_attributes[attr_key] = attr_value
-                        langfuse_logger.debug(
-                            f"Propagated metadata key '{key}' = '{attr_value}' to span '{span.name}'"
-                        )
         except Exception as e:
             langfuse_logger.debug(f"Could not read metadata from context: {e}")
 

tests/test_core_sdk.py

@@ -2071,7 +2071,7 @@ def test_context_manager_user_propagation():
     user_id = "test_user_123"
 
     with langfuse.start_as_current_span(name="parent-span") as parent_span:
-        with langfuse.user(id=user_id):
+        with langfuse.with_attributes(user_id=user_id):
             trace_id = parent_span.trace_id
 
             # Create child spans that should inherit user_id
@@ -2107,7 +2107,7 @@ def test_context_manager_session_propagation():
     session_id = "test_session_456"
 
     with langfuse.start_as_current_span(name="parent-span") as parent_span:
-        with langfuse.session(id=session_id):
+        with langfuse.with_attributes(session_id=session_id):
             trace_id = parent_span.trace_id
 
             # Create child spans that should inherit session_id
@@ -2142,8 +2142,12 @@ def test_context_manager_metadata_propagation():
     langfuse = Langfuse()
 
     with langfuse.start_as_current_span(name="parent-span") as parent_span:
-        with langfuse.metadata(
-            experiment="A/B", version="1.2.3", feature_flag="enabled"
+        with langfuse.with_attributes(
+            metadata={
+                "experiment": "A/B",
+                "version": "1.2.3",
+                "feature_flag": "enabled",
+            }
         ):
             trace_id = parent_span.trace_id
 
@@ -2185,17 +2189,18 @@ def test_context_manager_nested_contexts():
     langfuse = Langfuse()
 
     with langfuse.start_as_current_span(name="outer-span") as outer_span:
-        with langfuse.user(id="user_1"):
-            with langfuse.session(id="session_1"):
-                with langfuse.metadata(env="prod", region="us-east"):
-                    outer_trace_id = outer_span.trace_id
+        with langfuse.with_attributes(user_id="user_1", session_id="session_1"):
+            with langfuse.with_attributes(
+                metadata={"env": "prod", "region": "us-east"}
+            ):
+                outer_trace_id = outer_span.trace_id
 
-                    # Create span in outer context
-                    outer_child = langfuse.start_span(name="outer-child")
-                    outer_child.end()
+                # Create span in outer context
+                outer_child = langfuse.start_span(name="outer-child")
+                outer_child.end()
 
-                    nested_span = langfuse.start_span(name="nested-span")
-                    nested_span.end()
+                nested_span = langfuse.start_span(name="nested-span")
+                nested_span.end()
 
     langfuse.flush()
     sleep(2)
@@ -2229,8 +2234,10 @@ def test_context_manager_baggage_propagation():
 
     # Test with baggage enabled (careful with sensitive data)
     with langfuse.start_as_current_span(name="service-span") as span:
-        with langfuse.session(id="public_session_789", as_baggage=True):
-            with langfuse.metadata(as_baggage=True, service="api", version="v1.0"):
+        with langfuse.with_attributes(session_id="public_session_789", as_baggage=True):
+            with langfuse.with_attributes(
+                as_baggage=True, metadata={"service": "api", "version": "v1.0"}
+            ):
                 trace_id = span.trace_id
 
                 # Create child spans that inherit baggage context