Fix: _handle_failure method getting called 2 times

This commit is contained in:
Sameer Kankute 2026-01-19 10:19:31 +05:30
parent 14679cff3f
commit fbf2d83375
2 changed files with 49 additions and 2 deletions

View file

@ -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):

View file

@ -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()