diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index c7b59509eb0..6dc738afe83 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -1,4 +1,6 @@ +import asyncio import json +import traceback from collections.abc import Callable from typing import Final @@ -298,7 +300,13 @@ class OpenAITextCompletion(BaseLLM): else: openai_client = client - raw_response: Final = await openai_client.completions.with_raw_response.create(**data) + try: + raw_response: Final = await openai_client.completions.with_raw_response.create(**data) + except Exception as e: + asyncio.create_task( + logging_obj.dispatch_failure_handlers(e, traceback.format_exc(), prefer_async_handlers=True) + ) + raise response: Final = raw_response.parse() streamwrapper: Final = CustomStreamWrapper( completion_stream=response, diff --git a/tests/unit/llms/openai/completion/test_completion_handler.py b/tests/unit/llms/openai/completion/test_completion_handler.py index 329956605ab..4af8c3ad3c4 100644 --- a/tests/unit/llms/openai/completion/test_completion_handler.py +++ b/tests/unit/llms/openai/completion/test_completion_handler.py @@ -6,6 +6,9 @@ Regression tests for https://github.com/BerriAI/litellm/issues/27410 """ +import asyncio + +import httpx import pytest import respx from httpx import Response @@ -13,6 +16,7 @@ from httpx import Response import litellm from litellm import atext_completion, text_completion +from litellm.integrations.custom_logger import CustomLogger @pytest.fixture(autouse=True) @@ -88,3 +92,42 @@ async def test_acompletion_forwards_client_headers_to_provider( request_headers = mock_completions_endpoint.calls.last.request.headers assert request_headers["x-mycorp-llmcall-id"] == "abc-123" + +class _FailureRecorder(CustomLogger): + def __init__(self): + self.payloads = [] + self.logged = asyncio.Event() + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + self.payloads.append(kwargs["standard_logging_object"]) + self.logged.set() + + +@pytest.mark.parametrize( + ("provider_response", "expected_error"), + [ + pytest.param(httpx.ConnectError("connection refused"), litellm.APIConnectionError, id="connection-refused"), + pytest.param(Response(500, json={"error": {"message": "boom"}}), litellm.InternalServerError, id="5xx-before-first-byte"), + ], +) +@respx.mock +async def test_astream_failing_before_first_byte_logs_one_failure(provider_response, expected_error, monkeypatch): + respx.post("https://api.openai.com/v1/completions").mock(side_effect=provider_response) + recorder = _FailureRecorder() + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + + response = await atext_completion( + model="gpt-3.5-turbo-instruct", prompt="hello", max_tokens=5, stream=True, max_retries=0 + ) + with pytest.raises(expected_error): + async for _ in response: + pass + await asyncio.wait_for(recorder.logged.wait(), timeout=5) + + assert len(recorder.payloads) == 1 + payload = recorder.payloads[0] + assert payload["status"] == "failure" + assert payload["custom_llm_provider"] == "text-completion-openai" + assert payload["model"] == "gpt-3.5-turbo-instruct"