diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index f514f233b0a..ea19befccc6 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -9,6 +9,7 @@ from typing import Any, Final, ParamSpec, Protocol, TypeVar from fastapi import HTTPException, Request from pydantic import BaseModel, TypeAdapter +from starlette.exceptions import HTTPException as StarletteHTTPException import litellm from litellm._internal_context import with_service_target @@ -791,6 +792,10 @@ def _caller_auth(endpoint_kwargs: Mapping[str, object]) -> UserAPIKeyAuth: return caller if isinstance(caller, UserAPIKeyAuth) else UserAPIKeyAuth() +def _proxy_exception_handler_ends_server_span(exception: Exception) -> bool: + return not isinstance(exception, StarletteHTTPException) + + def _wrap_management_endpoint( endpoint: Callable[_P, Awaitable[_R]], *, log_payload: bool ) -> Callable[_P, Awaitable[_R]]: @@ -835,7 +840,7 @@ def _wrap_management_endpoint( user_api_key_dict: UserAPIKeyAuth = _caller_auth(kwargs) parent_otel_span = getattr(user_api_key_dict, "parent_otel_span", None) - if parent_otel_span is not None: + if parent_otel_span is not None and (log_payload or not _proxy_exception_handler_ends_server_span(e)): try: await _emit_management_endpoint_otel_span( func=endpoint, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index eb490a1cf09..f5d62c1405e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12629,6 +12629,7 @@ async def model_list( dependencies=[Depends(user_api_key_auth)], tags=["model management"], ) +@management_endpoint_span_wrapper async def model_info( model_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), diff --git a/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py b/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py index 778845f6b62..1965fba94c6 100644 --- a/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py +++ b/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py @@ -6,9 +6,11 @@ from datetime import datetime from typing import Final from unittest.mock import MagicMock +import httpx import pytest from fastapi import Request +from litellm import Router from litellm.proxy._types import ( ManagementEndpointLoggingPayload, UserAPIKeyAuth, @@ -345,16 +347,19 @@ def test_management_wrapper_swallows_post_success_errors(server_span_factory, ot TEAM_NEW_PATH: Final = "/team/new" +MODEL_LIST_PATH: Final = "/v1/models" +MODEL_RETRIEVE_PATH: Final = "/v1/models/gpt-4o" +INBOUND_TRACEPARENT: Final = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" -def _json_post(path: str, body: bytes) -> Request: +def _request(method: str, path: str, body: bytes = b"") -> Request: async def receive() -> dict[str, object]: return {"type": "http.request", "body": body, "more_body": False} return Request( { "type": "http", - "method": "POST", + "method": method, "path": path, "root_path": "", "query_string": b"", @@ -364,6 +369,15 @@ def _json_post(path: str, body: bytes) -> Request: ) +async def _answer_like_the_proxy(request: Request, exception: Exception) -> None: + from litellm.proxy.proxy_server import app + + handler: Final = next( + app.exception_handlers[cls] for cls in type(exception).__mro__ if cls in app.exception_handlers + ) + await handler(request, exception) + + @pytest.mark.parametrize( ("wrapper_name", "exports_payload"), [("management_endpoint_wrapper", True), ("management_endpoint_span_wrapper", False)], @@ -384,7 +398,7 @@ def test_span_wrapper_omits_payload(server_span_factory, otel_with_exporter, mon asyncio.run( fake_fn( - http_request=_json_post(TEAM_NEW_PATH, b'{"team_alias": "t"}'), + http_request=_request("POST", TEAM_NEW_PATH, b'{"team_alias": "t"}'), user_api_key_dict=_real_user_api_key_dict(server_span), ) ) @@ -396,8 +410,107 @@ def test_span_wrapper_omits_payload(server_span_factory, otel_with_exporter, mon assert any(key.startswith("response.") for key in attribute_keys) is exports_payload -MODEL_LIST_PATH: Final = "/v1/models" -INBOUND_TRACEPARENT: Final = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" +@pytest.mark.parametrize( + ("wrapper_name", "exports_request_body"), + [("management_endpoint_wrapper", True), ("management_endpoint_span_wrapper", False)], +) +def test_span_wrapper_omits_request_body_on_failure( + server_span_factory, otel_with_exporter, monkeypatch, wrapper_name, exports_request_body +): + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_helpers import utils as mgmt_utils + + otel, exporter = otel_with_exporter + monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False) + + server_span: Final = server_span_factory(TEAM_NEW_PATH) + wrap: Final = getattr(mgmt_utils, wrapper_name) + + @wrap + async def fake_fn(http_request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]: + raise make_fastapi_http_exception(400, "team_alias taken") + + with pytest.raises(Exception, match="team_alias taken"): + asyncio.run( + fake_fn( + http_request=_request("POST", TEAM_NEW_PATH, b'{"team_alias": "t"}'), + user_api_key_dict=_real_user_api_key_dict(server_span), + ) + ) + + spans: Final = [s for s in exporter.get_finished_spans() if s.name == TEAM_NEW_PATH] + assert len(spans) == 1 + assert any(key.startswith("request.") for key in spans[0].attributes) is exports_request_body + + +@pytest.mark.parametrize( + ("raised", "status_code"), + [ + (RuntimeError("router settings lookup failed"), 500), + (httpx.ConnectError("All connection attempts failed"), 503), + (make_fastapi_http_exception(400, "Invalid scope parameter"), 400), + ], +) +def test_span_wrapper_failure_keeps_the_answered_status(otel_with_exporter, monkeypatch, raised, status_code): + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_helpers import utils as mgmt_utils + + otel, exporter = otel_with_exporter + monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False) + + server_span: Final = _server_span(otel, MODEL_LIST_PATH, headers={}) + request: Final = _request("GET", MODEL_LIST_PATH) + request.state.parent_otel_span = server_span + + @mgmt_utils.management_endpoint_span_wrapper + async def failing_model_list(request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]: + raise raised + + async def call_and_answer() -> None: + try: + await failing_model_list(request=request, user_api_key_dict=_real_user_api_key_dict(server_span)) + except Exception as exc: + await _answer_like_the_proxy(request, exc) + + asyncio.run(call_and_answer()) + + assert_server_span_attrs( + exporter, + expected_status=status_code, + expected_url_path=MODEL_LIST_PATH, + where=f"model list {type(raised).__name__}", + ) + exported: Final = get_server_span(exporter) + assert exported is not None + assert exported.attributes.get("error.code") == str(status_code) + + +def _server_span(otel, path: str, headers: dict[str, str]): + server_span: Final = otel.create_litellm_proxy_request_started_span(start_time=datetime.now(), headers=headers) + otel.set_proxy_request_route_attributes(server_span, url_path=path, http_route=path) + return server_span + + +def _assert_ended_in_inbound_trace(exporter, path: str, handler_name: str) -> None: + assert_server_span_attrs(exporter, expected_status=200, expected_url_path=path, where=f"{path} success") + exported: Final = get_server_span(exporter) + assert exported is not None + assert exported.context.trace_id == int(INBOUND_TRACEPARENT.split("-")[1], 16) + + endpoint_spans: Final = [s for s in exporter.get_finished_spans() if s.name == handler_name] + assert len(endpoint_spans) == 1 + assert endpoint_spans[0].parent is not None + assert endpoint_spans[0].parent.span_id == exported.context.span_id + assert not any(key.startswith("response.") for key in endpoint_spans[0].attributes) + + +def _install_proxy_state(monkeypatch, router: Router | None) -> None: + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) def test_model_list_ends_server_span_in_inbound_trace(otel_with_exporter, monkeypatch): @@ -405,31 +518,29 @@ def test_model_list_ends_server_span_in_inbound_trace(otel_with_exporter, monkey otel, exporter = otel_with_exporter monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False) + _install_proxy_state(monkeypatch, None) + server_span: Final = _server_span(otel, MODEL_LIST_PATH, headers={"traceparent": INBOUND_TRACEPARENT}) - server_span: Final = otel.create_litellm_proxy_request_started_span( - start_time=datetime.now(), headers={"traceparent": INBOUND_TRACEPARENT} - ) - otel.set_proxy_request_route_attributes(server_span, url_path=MODEL_LIST_PATH, http_route=MODEL_LIST_PATH) - - result: Final = asyncio.run( - proxy_server.model_list( - user_api_key_dict=_real_user_api_key_dict(server_span), - ) - ) + result: Final = asyncio.run(proxy_server.model_list(user_api_key_dict=_real_user_api_key_dict(server_span))) assert result == {"data": [], "object": "list"} - assert_server_span_attrs( - exporter, - expected_status=200, - expected_url_path=MODEL_LIST_PATH, - where="model list success", - ) - exported: Final = get_server_span(exporter) - assert exported is not None - assert exported.context.trace_id == int(INBOUND_TRACEPARENT.split("-")[1], 16) + _assert_ended_in_inbound_trace(exporter, MODEL_LIST_PATH, "model_list") - endpoint_spans: Final = [s for s in exporter.get_finished_spans() if s.name == "model_list"] - assert len(endpoint_spans) == 1 - assert endpoint_spans[0].parent is not None - assert endpoint_spans[0].parent.span_id == exported.context.span_id - assert not any(key.startswith("response.") for key in endpoint_spans[0].attributes) + +def test_model_retrieve_ends_server_span_in_inbound_trace(otel_with_exporter, monkeypatch): + import litellm.proxy.proxy_server as proxy_server + + otel, exporter = otel_with_exporter + monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False) + _install_proxy_state( + monkeypatch, + Router(model_list=[{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-x"}}]), + ) + server_span: Final = _server_span(otel, MODEL_RETRIEVE_PATH, headers={"traceparent": INBOUND_TRACEPARENT}) + + result: Final = asyncio.run( + proxy_server.model_info(model_id="gpt-4o", user_api_key_dict=_real_user_api_key_dict(server_span)) + ) + + assert result["id"] == "gpt-4o" + _assert_ended_in_inbound_trace(exporter, MODEL_RETRIEVE_PATH, "model_info")