mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
Fix: _handle_failure method getting called 2 times
This commit is contained in:
parent
14679cff3f
commit
fbf2d83375
2 changed files with 49 additions and 2 deletions
|
|
@ -170,7 +170,6 @@ class BaseResponsesAPIStreamingIterator:
|
|||
return None
|
||||
except Exception as e:
|
||||
# Ensure failures trigger failure hooks
|
||||
self._handle_failure(e)
|
||||
raise
|
||||
|
||||
def _handle_logging_completed_response(self):
|
||||
|
|
|
|||
|
|
@ -246,4 +246,52 @@ class TestBaseResponsesAPIStreamingIterator:
|
|||
|
||||
# Test with None chunk
|
||||
result = iterator._process_chunk(None)
|
||||
assert result is None
|
||||
assert result is None
|
||||
|
||||
def test_process_chunk_exception_does_not_call_handle_failure(self):
|
||||
"""
|
||||
Test that _process_chunk raises exceptions without calling _handle_failure.
|
||||
|
||||
This ensures _handle_failure is only called once in the outer exception handler
|
||||
(in __next__ or __anext__), preventing duplicate failure logging.
|
||||
|
||||
Previously, _handle_failure was called both in _process_chunk and in the outer
|
||||
exception handler, causing duplicate logs. This test verifies the fix.
|
||||
"""
|
||||
# Mock dependencies
|
||||
mock_response = Mock()
|
||||
mock_response.headers = {}
|
||||
mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
|
||||
mock_logging_obj.model_call_details = {"litellm_params": {}}
|
||||
mock_config = Mock(spec=BaseResponsesAPIConfig)
|
||||
|
||||
# Set up the mock transform method to raise an exception
|
||||
test_exception = ValueError("Test exception in transform")
|
||||
mock_config.transform_streaming_response.side_effect = test_exception
|
||||
|
||||
# Create the iterator instance
|
||||
iterator = BaseResponsesAPIStreamingIterator(
|
||||
response=mock_response,
|
||||
model="gpt-4",
|
||||
responses_api_provider_config=mock_config,
|
||||
logging_obj=mock_logging_obj
|
||||
)
|
||||
|
||||
# Mock _handle_failure to track if it's called
|
||||
with patch.object(iterator, '_handle_failure') as mock_handle_failure:
|
||||
# Prepare valid JSON chunk that will trigger transform_streaming_response
|
||||
test_chunk_data = {
|
||||
"type": "response.output_text.delta",
|
||||
"delta": "Hello"
|
||||
}
|
||||
|
||||
# _process_chunk should raise the exception without calling _handle_failure
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
iterator._process_chunk(json.dumps(test_chunk_data))
|
||||
|
||||
# Verify the exception was raised
|
||||
assert str(exc_info.value) == "Test exception in transform"
|
||||
|
||||
# Verify _handle_failure was NOT called in _process_chunk
|
||||
# It should only be called by the outer exception handler in __next__/__anext__
|
||||
mock_handle_failure.assert_not_called()
|
||||
Loading…
Add table
Reference in a new issue