scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(langfuse): Handle null usage values to prevent validation errors (#16396)

Browse source 
* langfuse null validation fix

* formatting
This commit is contained in:
Alan Ponnachan 2025-11-11 09:04:56 +05:30 • committed by GitHub
parent 7b292ccdf5
commit 5f12e4be1e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 137 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
litellm/integrations/langfuse/langfuse.py
View file

@ -683,23 +683,33 @@ class LangFuseLogger:
_usage_obj = getattr(response_obj, "usage", None)
if _usage_obj:
# Safely get usage values, defaulting None to 0 for Langfuse compatibility.
# Some providers may return null for token counts.
prompt_tokens = getattr(_usage_obj, "prompt_tokens", None) or 0
completion_tokens = (
getattr(_usage_obj, "completion_tokens", None) or 0
)
total_tokens = getattr(_usage_obj, "total_tokens", None) or 0
cache_creation_input_tokens = (
_usage_obj.get("cache_creation_input_tokens") or 0
)
cache_read_input_tokens = (
_usage_obj.get("cache_read_input_tokens") or 0
)
usage = {
"prompt_tokens": _usage_obj.prompt_tokens,
"completion_tokens": _usage_obj.completion_tokens,
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_cost": cost if self._supports_costs() else None,
}
cache_read_input_tokens = _usage_obj.get(
"cache_read_input_tokens", 0
)
# According to langfuse documentation: "the input value must be reduced by the number of cache_read_input_tokens"
input_tokens = _usage_obj.prompt_tokens - cache_read_input_tokens
input_tokens = prompt_tokens - cache_read_input_tokens
usage_details = LangfuseUsageDetails(
input=input_tokens,
output=_usage_obj.completion_tokens,
total=_usage_obj.total_tokens,
cache_creation_input_tokens=_usage_obj.get(
"cache_creation_input_tokens", 0
),
output=completion_tokens,
total=total_tokens,
cache_creation_input_tokens=cache_creation_input_tokens,
cache_read_input_tokens=cache_read_input_tokens,
)

138
tests/test_litellm/integrations/test_langfuse.py
View file

@ -13,18 +13,22 @@ from litellm.integrations.langfuse.langfuse import LangFuseLogger
sys.path.insert(0, os.path.abspath("../.."))
from litellm.integrations.langfuse.langfuse import LangFuseLogger
# Import LangfuseUsageDetails directly from the module where it's defined
from litellm.types.integrations.langfuse import *
class TestLangfuseUsageDetails(unittest.TestCase):
class TestLangfuseUsageDetails(unittest.TestCase):
def setUp(self):
# Set up environment variables for testing
self.env_patcher = patch.dict('os.environ', {
'LANGFUSE_SECRET_KEY': 'test-secret-key',
'LANGFUSE_PUBLIC_KEY': 'test-public-key',
'LANGFUSE_HOST': 'https://test.langfuse.com'
})
self.env_patcher = patch.dict(
"os.environ",
{
"LANGFUSE_SECRET_KEY": "test-secret-key",
"LANGFUSE_PUBLIC_KEY": "test-public-key",
"LANGFUSE_HOST": "https://test.langfuse.com",
},
)
self.env_patcher.start()
# Create mock objects
@ -37,21 +41,25 @@ class TestLangfuseUsageDetails(unittest.TestCase):
self.mock_langfuse_client.trace.return_value = self.mock_langfuse_trace
# Mock the langfuse module that's imported locally in methods
self.langfuse_module_patcher = patch.dict('sys.modules', {'langfuse': MagicMock()})
self.langfuse_module_patcher = patch.dict(
"sys.modules", {"langfuse": MagicMock()}
)
self.mock_langfuse_module = self.langfuse_module_patcher.start()
# Create a mock for the langfuse module with version
self.mock_langfuse = MagicMock()
self.mock_langfuse.version = MagicMock()
self.mock_langfuse.version.__version__ = "3.0.0" # Set a version that supports all features
self.mock_langfuse.version.__version__ = (
"3.0.0" # Set a version that supports all features
)
# Mock the Langfuse class
self.mock_langfuse_class = MagicMock()
self.mock_langfuse_class.return_value = self.mock_langfuse_client
# Set up the sys.modules['langfuse'] mock
sys.modules['langfuse'] = self.mock_langfuse
sys.modules['langfuse'].Langfuse = self.mock_langfuse_class
sys.modules["langfuse"] = self.mock_langfuse
sys.modules["langfuse"].Langfuse = self.mock_langfuse_class
# Mock the Langfuse client
self.mock_langfuse_client = MagicMock()
@ -71,7 +79,16 @@ class TestLangfuseUsageDetails(unittest.TestCase):
self.logger = LangFuseLogger()
# Add the log_event_on_langfuse method to the instance
def log_event_on_langfuse(self, kwargs, response_obj, start_time=None, end_time=None, user_id=None, level="DEFAULT", status_message=None):
def log_event_on_langfuse(
self,
kwargs,
response_obj,
start_time=None,
end_time=None,
user_id=None,
level="DEFAULT",
status_message=None,
):
# This implementation calls _log_langfuse_v2 directly
return self._log_langfuse_v2(
user_id=user_id,
@ -86,12 +103,15 @@ class TestLangfuseUsageDetails(unittest.TestCase):
response_obj=response_obj,
level=level,
litellm_call_id=kwargs.get("litellm_call_id", None),
print_verbose=True # Add the missing parameter
print_verbose=True, # Add the missing parameter
)
# Bind the method to the instance
import types
self.logger.log_event_on_langfuse = types.MethodType(log_event_on_langfuse, self.logger)
self.logger.log_event_on_langfuse = types.MethodType(
log_event_on_langfuse, self.logger
)
# Make sure _is_langfuse_v2 returns True
def mock_is_langfuse_v2(self):
@ -111,7 +131,7 @@ class TestLangfuseUsageDetails(unittest.TestCase):
"output": 20,
"total": 30,
"cache_creation_input_tokens": 5,
"cache_read_input_tokens": 3
"cache_read_input_tokens": 3,
}
# Verify all fields are present
@ -127,7 +147,7 @@ class TestLangfuseUsageDetails(unittest.TestCase):
"output": 20,
"total": 30,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0
"cache_read_input_tokens": 0,
}
self.assertEqual(minimal_usage_details["input"], 10)
@ -144,9 +164,9 @@ class TestLangfuseUsageDetails(unittest.TestCase):
# Add the cache token attributes using get method
def mock_get(key, default=None):
if key == 'cache_creation_input_tokens':
if key == "cache_creation_input_tokens":
return 7
elif key == 'cache_read_input_tokens':
elif key == "cache_read_input_tokens":
return 4
return default
@ -156,7 +176,7 @@ class TestLangfuseUsageDetails(unittest.TestCase):
kwargs = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hello"}],
"litellm_params": {"metadata": {}}
"litellm_params": {"metadata": {}},
}
# Create start and end times
@ -164,12 +184,12 @@ class TestLangfuseUsageDetails(unittest.TestCase):
end_time = start_time + datetime.timedelta(seconds=1)
# Call the log_event method
with patch.object(self.logger, '_log_langfuse_v2') as mock_log_langfuse_v2:
with patch.object(self.logger, "_log_langfuse_v2") as mock_log_langfuse_v2:
self.logger.log_event_on_langfuse(
kwargs=kwargs,
response_obj=response_obj,
start_time=start_time,
end_time=end_time
end_time=end_time,
)
# Check if _log_langfuse_v2 was called
@ -189,7 +209,7 @@ class TestLangfuseUsageDetails(unittest.TestCase):
"output": 20,
"total": 30,
"cache_creation_input_tokens": None,
"cache_read_input_tokens": None
"cache_read_input_tokens": None,
}
# Verify fields can be None
@ -210,7 +230,7 @@ class TestLangfuseUsageDetails(unittest.TestCase):
"output": 25,
"total": 40,
"cache_creation_input_tokens": 7,
"cache_read_input_tokens": 4
"cache_read_input_tokens": 4,
}
# Verify the structure matches what we expect
@ -227,6 +247,80 @@ class TestLangfuseUsageDetails(unittest.TestCase):
self.assertEqual(usage_details["cache_creation_input_tokens"], 7)
self.assertEqual(usage_details["cache_read_input_tokens"], 4)
def test_log_langfuse_v2_handles_null_usage_values(self):
"""
Test that _log_langfuse_v2 correctly handles None values in the usage object
by converting them to 0, preventing validation errors.
"""
with patch(
"litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
side_effect=lambda generation_params, **kwargs: generation_params,
) as mock_add_prompt_params:
# Create a mock response object with usage information containing None values
response_obj = MagicMock()
response_obj.usage = MagicMock()
response_obj.usage.prompt_tokens = None
response_obj.usage.completion_tokens = None
response_obj.usage.total_tokens = None
# Mock the .get() method to return None for cache-related fields
def mock_get(key, default=None):
if key in ["cache_creation_input_tokens", "cache_read_input_tokens"]:
return None
return default
response_obj.usage.get = mock_get
# Prepare standard kwargs for the call
kwargs = {
"model": "gpt-4-null-usage",
"messages": [{"role": "user", "content": "Test"}],
"litellm_params": {"metadata": {}},
"optional_params": {},
"litellm_call_id": "test-call-id-null-usage",
"standard_logging_object": None,
"response_cost": 0.0,
}
# Call the method under test
self.logger._log_langfuse_v2(
user_id="test-user",
metadata={},
litellm_params=kwargs["litellm_params"],
output={"role": "assistant", "content": "Response"},
start_time=datetime.datetime.now(),
end_time=datetime.datetime.now(),
kwargs=kwargs,
optional_params=kwargs["optional_params"],
input={"messages": kwargs["messages"]},
response_obj=response_obj,
level="DEFAULT",
litellm_call_id=kwargs["litellm_call_id"],
)
# Check the arguments passed to the mocked langfuse generation call
self.mock_langfuse_trace.generation.assert_called_once()
call_args, call_kwargs = self.mock_langfuse_trace.generation.call_args
# Inspect the usage and usage_details dictionaries
usage_arg = call_kwargs.get("usage")
usage_details_arg = call_kwargs.get("usage_details")
self.assertIsNotNone(usage_arg)
self.assertIsNotNone(usage_details_arg)
# Verify that None values were converted to 0
self.assertEqual(usage_arg["prompt_tokens"], 0)
self.assertEqual(usage_arg["completion_tokens"], 0)
self.assertEqual(usage_details_arg["input"], 0)
self.assertEqual(usage_details_arg["output"], 0)
self.assertEqual(usage_details_arg["total"], 0)
self.assertEqual(usage_details_arg["cache_creation_input_tokens"], 0)
self.assertEqual(usage_details_arg["cache_read_input_tokens"], 0)
mock_add_prompt_params.assert_called_once()
def test_max_langfuse_clients_limit():
"""
Test that the max langfuse clients limit is respected when initializing multiple clients