chore: feedback

Author: Steffen911

langfuse/_client/client.py

@@ -23,12 +23,15 @@
     Union,
     cast,
     overload,
+    Generator,
 )
 
 import backoff
 import httpx
 from opentelemetry import (
+    baggage as otel_baggage_api,
     trace as otel_trace_api,
+    context as otel_context_api,
 )
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
@@ -111,10 +114,15 @@
     TextPromptClient,
 )
 from langfuse.types import MaskFunction, ScoreDataType, SpanLevel, TraceContext
-from langfuse._client.context_propagation import LangfuseContextPropagationMixin
 
 
-class Langfuse(LangfuseContextPropagationMixin):
+# 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.
 
     This class provides an interface for creating and managing traces, spans,
@@ -355,6 +363,108 @@ def start_span(
             status_message=status_message,
         )
 
+    @_agnosticcontextmanager
+    def with_attributes(
+        self,
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        metadata: Optional[dict[str, str]] = None,
+        as_baggage: bool = False,
+    ) -> Generator[None, None, None]:
+        """Creates a context manager that propagates the given attributes to all spans created within the context.
+
+        Args:
+            session_id (str): Session identifier.
+            user_id (str): User identifier.
+            metadata (dict): Additional metadata to associate with all spans in the context. Values must be strings and are truncated to 200 characters.
+            as_baggage (bool, optional): If True, stores the values in OpenTelemetry baggage
+                for cross-service propagation. If False, stores only in local context for
+                current-service propagation. Defaults to False.
+
+        Returns:
+            Context manager that sets values on all spans created within its scope.
+
+        Warning:
+            When as_baggage=True, the values will be included in HTTP headers of any
+            outbound requests made within this context. Only use this for non-sensitive
+            identifiers that are safe to transmit across service boundaries.
+
+        Example:
+           ```python
+            # Local context only (default)
+            with langfuse.with_attributes(session_id="session_123"):
+                with langfuse.start_as_current_span(name="process-request") as span:
+                    # This span and all its children will have session_id="session_123"
+                    child_span = langfuse.start_span(name="child-operation")
+
+            # Cross-service propagation (use with caution)
+            with langfuse.with_attributes(session_id="session_123", as_baggage=True):
+                # session_id will be propagated to external service calls
+                response = requests.get("https://api.example.com/data")
+            ```
+        """
+        current_context = otel_context_api.get_current()
+        current_span = otel_trace_api.get_current_span()
+
+        # Process session_id
+        if session_id is not None:
+            current_context = otel_context_api.set_value(
+                LANGFUSE_CTX_SESSION_ID, session_id, current_context
+            )
+            if current_span is not None and current_span.is_recording():
+                current_span.set_attribute("session.id", session_id)
+            if as_baggage:
+                current_context = otel_baggage_api.set_baggage(
+                    "session.id", session_id, current_context
+                )
+
+        # Process user_id
+        if user_id is not None:
+            current_context = otel_context_api.set_value(
+                LANGFUSE_CTX_USER_ID, user_id, current_context
+            )
+            if current_span is not None and current_span.is_recording():
+                current_span.set_attribute("user.id", user_id)
+            if as_baggage:
+                current_context = otel_baggage_api.set_baggage(
+                    "user.id", user_id, current_context
+                )
+
+        # Process metadata
+        if metadata is not None:
+            # Truncate values with size > 200 to 200 characters and emit warning including the ky
+            for k, v in metadata.items():
+                if not isinstance(v, str):
+                    # Ignore unreachable mypy warning as this runtime guard should make sense either way
+                    warnings.warn(  # type: ignore[unreachable]
+                        f"Metadata values must be strings, got {type(v)} for key '{k}'"
+                    )
+                    del metadata[k]
+                if len(v) > 200:
+                    warnings.warn(
+                        f"Metadata value for key '{k}' exceeds 200 characters and will be truncated."
+                    )
+                    metadata[k] = v[:200]
+
+            current_context = otel_context_api.set_value(
+                LANGFUSE_CTX_METADATA, metadata, current_context
+            )
+            if current_span is not None and current_span.is_recording():
+                for k, v in metadata.items():
+                    current_span.set_attribute(f"langfuse.metadata.{k}", v)
+            if as_baggage:
+                for k, v in metadata.items():
+                    current_context = otel_baggage_api.set_baggage(
+                        f"langfuse.metadata.{k}", str(v), current_context
+                    )
+
+        # Activate context, execute, and detach context
+        token = otel_context_api.attach(current_context)
+        try:
+            yield
+        finally:
+            otel_context_api.detach(token)
+
     def start_as_current_span(
         self,
         *,
@@ -1673,7 +1783,7 @@ def update_current_trace(
             ```
         """
         warnings.warn(
-            "update_current_trace is deprecated and will be removed in a future version. ",
+            "update_current_trace is deprecated and will be removed in a future version.  Use `with langfuse.with_attributes(...)` instead. ",
             DeprecationWarning,
             stacklevel=2,
         )