diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 7bf97665fd2..e9df54b7651 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/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, ) diff --git a/tests/test_litellm/test_langfuse_environment.py b/tests/test_litellm/test_langfuse_environment.py new file mode 100644 index 00000000000..65e9921419c --- /dev/null +++ b/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}"