fix(langfuse.py): populate environment tag on provider/guardrail spans (Fixes #19926)

litellm/integrations/langfuse/langfuse.py

@@ -1044,6 +1044,10 @@ class LangFuseLogger:
                 "Not logging guardrail information as span because standard_logging_object is None"
             )
             return
+            
+        tags = []
+        if os.getenv("LANGFUSE_TRACING_ENVIRONMENT") is not None:
+            tags.append(f"environment:{os.getenv('LANGFUSE_TRACING_ENVIRONMENT')}")
 
         guardrail_information = standard_logging_object.get(
             "guardrail_information", None
@@ -1082,6 +1086,7 @@ class LangFuseLogger:
                 },
                 start_time=guardrail_entry.get("start_time", None),  # type: ignore
                 end_time=guardrail_entry.get("end_time", None),  # type: ignore
+                tags=tags,
             )
 
             verbose_logger.debug(f"Logged guardrail information as span: {span}")
@@ -1196,6 +1201,10 @@ def log_provider_specific_information_as_span(
     if _hidden_params is None:
         return
 
+    tags = []
+    if os.getenv("LANGFUSE_TRACING_ENVIRONMENT") is not None:
+        tags.append(f"environment:{os.getenv('LANGFUSE_TRACING_ENVIRONMENT')}")
+
     vertex_ai_grounding_metadata = _hidden_params.get(
         "vertex_ai_grounding_metadata", None
     )
@@ -1208,16 +1217,19 @@ def log_provider_specific_information_as_span(
                         trace.span(
                             name=key,
                             input=value,
+                            tags=tags,
                         )
                 else:
                     trace.span(
                         name="vertex_ai_grounding_metadata",
                         input=elem,
+                        tags=tags,
                     )
         else:
             trace.span(
                 name="vertex_ai_grounding_metadata",
                 input=vertex_ai_grounding_metadata,
+                tags=tags,
             )
 
 

tests/test_litellm/test_langfuse_environment.py

@@ -0,0 +1,85 @@
+
+import pytest
+from unittest.mock import MagicMock, patch
+import os
+from litellm.integrations.langfuse.langfuse import (
+    log_provider_specific_information_as_span,
+    LangFuseLogger
+)
+
+class TestLangfuseEnvironment:
+    
+    @patch.dict(os.environ, {"LANGFUSE_TRACING_ENVIRONMENT": "production"}, clear=True)
+    def test_log_provider_specific_information_as_span_includes_environment(self):
+        """
+        Verify that log_provider_specific_information_as_span includes the 
+        LANGFUSE_TRACING_ENVIRONMENT in the span tags.
+        """
+        trace = MagicMock()
+        clean_metadata = {
+            "hidden_params": {
+                "vertex_ai_grounding_metadata": {"foo": "bar"}
+            }
+        }
+        
+        # Call the function
+        log_provider_specific_information_as_span(trace, clean_metadata)
+        
+        # Check if trace.span was called
+        assert trace.span.called
+        
+        # Verify call arguments
+        # We expect tags=["environment:production"] to be passed
+        # Currently, this test should FAIL because tags are not passed.
+        
+        _, kwargs = trace.span.call_args
+        tags = kwargs.get("tags", [])
+        
+        # This assertions mirrors what we WANT to see
+        # If the bug exists, tags will be empty or None, or missing this specific tag.
+        assert "environment:production" in tags, f"Expected 'environment:production' in tags, got {tags}"
+
+    @patch.dict(os.environ, {"LANGFUSE_TRACING_ENVIRONMENT": "staging"}, clear=True)
+    def test_log_guardrail_information_as_span_includes_environment(self):
+        """
+        Verify that _log_guardrail_information_as_span includes the
+        LANGFUSE_TRACING_ENVIRONMENT in the span tags.
+        """
+        trace = MagicMock()
+        
+        # Mock sys.modules to simulate langfuse installation
+        langfuse_mock = MagicMock()
+        langfuse_mock.version.__version__ = "2.0.0"
+        
+        with patch.dict("sys.modules", {
+            "langfuse": langfuse_mock, 
+            "langfuse.Langfuse": MagicMock(),
+            "langfuse.version": langfuse_mock.version
+        }):
+            # Setup mock logger wrapper to access the method
+            # We need to bypass __init__ logic that requires valid keys/host if we don't supply them
+            # or just supply dummy values.
+            logger = LangFuseLogger(langfuse_public_key="pk", langfuse_secret="sk")
+            logger._is_langfuse_v2 = MagicMock(return_value=True)
+            
+            standard_logging_object = {
+                "guardrail_information": [
+                    {
+                        "guardrail_name": "test_guard",
+                        "guardrail_request": "input",
+                        "guardrail_response": "output"
+                    }
+                ]
+            }
+            
+            # Call the private method
+            logger._log_guardrail_information_as_span(trace, standard_logging_object)
+            
+            # Check if trace.span was called
+            assert trace.span.called
+            
+            # Verify call arguments
+            _, kwargs = trace.span.call_args
+            tags = kwargs.get("tags", [])
+            
+            assert "environment:staging" in tags, f"Expected 'environment:staging' in tags, got {tags}"