From 541b2c0a8513d867b6f4a593bb3bc2189ed6e878 Mon Sep 17 00:00:00 2001
From: Hassieb Pakzad <68423100+hassiebp@users.noreply.github.com>
Date: Wed, 21 May 2025 17:51:12 +0200
Subject: [PATCH] fix(environments): correctly propagate client attribute

---
 langfuse/_client/client.py | 16 ++++++++++++++--
 langfuse/_client/span.py   | 22 ++++++++++++++++++++--
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/langfuse/_client/client.py b/langfuse/_client/client.py
index 17745b816..b2c82a30a 100644
--- a/langfuse/_client/client.py
+++ b/langfuse/_client/client.py
@@ -297,6 +297,7 @@ def start_span(
                         input=input,
                         output=output,
                         metadata=metadata,
+                        environment=self._environment,
                     )
 
         otel_span = self._otel_tracer.start_span(name=name, attributes=attributes)
@@ -307,6 +308,7 @@ def start_span(
             input=input,
             output=output,
             metadata=metadata,
+            environment=self._environment,
         )
 
     def start_as_current_span(
@@ -697,6 +699,7 @@ def _start_as_current_otel_span_with_processed_media(
                     input=input,
                     output=output,
                     metadata=metadata,
+                    environment=self._environment,
                 )
                 if as_type == "span"
                 else LangfuseGeneration(
@@ -705,6 +708,7 @@ def _start_as_current_otel_span_with_processed_media(
                     input=input,
                     output=output,
                     metadata=metadata,
+                    environment=self._environment,
                 )
             )
 
@@ -849,7 +853,11 @@ def update_current_span(
         current_otel_span = self._get_current_otel_span()
 
         if current_otel_span is not None:
-            span = LangfuseSpan(otel_span=current_otel_span, langfuse_client=self)
+            span = LangfuseSpan(
+                otel_span=current_otel_span,
+                langfuse_client=self,
+                environment=self._environment,
+            )
 
             span.update(
                 input=input,
@@ -919,7 +927,11 @@ def update_current_trace(
         current_otel_span = self._get_current_otel_span()
 
         if current_otel_span is not None:
-            span = LangfuseSpan(otel_span=current_otel_span, langfuse_client=self)
+            span = LangfuseSpan(
+                otel_span=current_otel_span,
+                langfuse_client=self,
+                environment=self._environment,
+            )
 
             span.update_trace(
                 name=name,
diff --git a/langfuse/_client/span.py b/langfuse/_client/span.py
index a83791086..99327d763 100644
--- a/langfuse/_client/span.py
+++ b/langfuse/_client/span.py
@@ -68,6 +68,7 @@ def __init__(
         input: Optional[Any] = None,
         output: Optional[Any] = None,
         metadata: Optional[Any] = None,
+        environment: Optional[str] = None,
     ):
         """Initialize a new Langfuse span wrapper.
 
@@ -78,6 +79,7 @@ def __init__(
             input: Input data for the span (any JSON-serializable object)
             output: Output data from the span (any JSON-serializable object)
             metadata: Additional metadata to associate with the span
+            environment: The tracing environment
         """
         self._otel_span = otel_span
         self._otel_span.set_attribute(
@@ -88,6 +90,12 @@ def __init__(
         self.trace_id = self._langfuse_client._get_otel_trace_id(otel_span)
         self.id = self._langfuse_client._get_otel_span_id(otel_span)
 
+        self._environment = environment
+        if self._environment is not None:
+            self._otel_span.set_attribute(
+                LangfuseOtelSpanAttributes.ENVIRONMENT, self._environment
+            )
+
         # Handle media only if span is sampled
         if self._otel_span.is_recording:
             media_processed_input = self._process_media_and_apply_mask(
@@ -490,6 +498,7 @@ def __init__(
         input: Optional[Any] = None,
         output: Optional[Any] = None,
         metadata: Optional[Any] = None,
+        environment: Optional[str] = None,
     ):
         """Initialize a new LangfuseSpan.
 
@@ -499,6 +508,7 @@ def __init__(
             input: Input data for the span (any JSON-serializable object)
             output: Output data from the span (any JSON-serializable object)
             metadata: Additional metadata to associate with the span
+            environment: The tracing environment
         """
         super().__init__(
             otel_span=otel_span,
@@ -507,6 +517,7 @@ def __init__(
             input=input,
             output=output,
             metadata=metadata,
+            environment=environment,
         )
 
     def update(
@@ -643,7 +654,9 @@ def start_span(
                 )
 
         return LangfuseSpan(
-            otel_span=new_otel_span, langfuse_client=self._langfuse_client
+            otel_span=new_otel_span,
+            langfuse_client=self._langfuse_client,
+            environment=self._environment,
         )
 
     def start_as_current_span(
@@ -818,7 +831,9 @@ def start_generation(
                 )
 
         return LangfuseGeneration(
-            otel_span=new_otel_span, langfuse_client=self._langfuse_client
+            otel_span=new_otel_span,
+            langfuse_client=self._langfuse_client,
+            environment=self._environment,
         )
 
     def start_as_current_generation(
@@ -936,6 +951,7 @@ def __init__(
         input: Optional[Any] = None,
         output: Optional[Any] = None,
         metadata: Optional[Any] = None,
+        environment: Optional[str] = None,
     ):
         """Initialize a new LangfuseGeneration span.
 
@@ -945,6 +961,7 @@ def __init__(
             input: Input data for the generation (e.g., prompts)
             output: Output from the generation (e.g., completions)
             metadata: Additional metadata to associate with the generation
+            environment: The tracing environment
         """
         super().__init__(
             otel_span=otel_span,
@@ -953,6 +970,7 @@ def __init__(
             input=input,
             output=output,
             metadata=metadata,
+            environment=environment,
         )
 
     def update(