diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 30aaaacad04..ea19befccc6 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -1,14 +1,15 @@ # What is this? ## Helper utils for the management endpoints (keys/users/teams) -from collections.abc import Callable, Mapping, MutableMapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence from collections.abc import Set as AbstractSet from datetime import datetime from functools import wraps from types import MappingProxyType -from typing import Any, Final, Protocol +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 @@ -687,6 +688,7 @@ async def _emit_management_endpoint_otel_span( end_time: datetime, result: Any = None, exception: Exception | None = None, + log_payload: bool = True, ) -> None: """Stamp + end the parent OTEL SERVER span for a management endpoint. @@ -716,8 +718,8 @@ async def _emit_management_endpoint_otel_span( ) route = get_request_route(http_request) - request_body: dict = await read_request_body( # rebind-ok: pre-existing rebinding on a rename-only line - request=http_request + request_body: dict = ( # rebind-ok: pre-existing rebinding on a rename-only line + await read_request_body(request=http_request) if log_payload else {} ) else: route = func.__name__ @@ -738,7 +740,7 @@ async def _emit_management_endpoint_otel_span( ) _response: dict[str, object] | None = None - if exception is None and result is not None: + if log_payload and exception is None and result is not None: try: raw: Final[Mapping[str, object]] = dict(result) _response = {k: v for k, v in raw.items() if k not in _CREDENTIAL_FIELDS} @@ -767,37 +769,60 @@ async def _emit_management_endpoint_otel_span( ) -def management_endpoint_wrapper(func): +_P: Final = ParamSpec("_P") +_R: Final = TypeVar("_R") + + +def management_endpoint_wrapper(func: Callable[_P, Awaitable[_R]]) -> Callable[_P, Awaitable[_R]]: """ This wrapper does the following: 1. Log I/O, Exceptions to OTEL 2. Create an Audit log for success calls """ + return _wrap_management_endpoint(func, log_payload=True) - @wraps(func) - async def wrapper(*args, **kwargs): + +def management_endpoint_span_wrapper(func: Callable[_P, Awaitable[_R]]) -> Callable[_P, Awaitable[_R]]: + return _wrap_management_endpoint(func, log_payload=False) + + +def _caller_auth(endpoint_kwargs: Mapping[str, object]) -> UserAPIKeyAuth: + caller: Final = endpoint_kwargs.get("user_api_key_dict") + 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]]: + @wraps(endpoint) + async def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: start_time: Final = datetime.now() try: - result: Final = await func(*args, **kwargs) + result: Final = await endpoint(*args, **kwargs) end_time = datetime.now() try: - user_api_key_dict: UserAPIKeyAuth = kwargs.get("user_api_key_dict") or UserAPIKeyAuth() + user_api_key_dict: UserAPIKeyAuth = _caller_auth(kwargs) await send_management_endpoint_alert( request_kwargs=kwargs, user_api_key_dict=user_api_key_dict, - function_name=func.__name__, + function_name=endpoint.__name__, ) parent_otel_span: Span | None = getattr(user_api_key_dict, "parent_otel_span", None) if parent_otel_span is not None: await _emit_management_endpoint_otel_span( - func=func, + func=endpoint, kwargs=kwargs, parent_otel_span=parent_otel_span, start_time=start_time, end_time=end_time, result=result, + log_payload=log_payload, ) # Delete updated/deleted info from cache @@ -813,17 +838,18 @@ def management_endpoint_wrapper(func): except Exception as e: end_time = datetime.now() - user_api_key_dict: UserAPIKeyAuth = kwargs.get("user_api_key_dict") or UserAPIKeyAuth() + 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=func, + func=endpoint, kwargs=kwargs, parent_otel_span=parent_otel_span, start_time=start_time, end_time=end_time, exception=e, + log_payload=log_payload, ) except Exception as otel_exc: # Non-Blocking Exception - never let OTEL failures swallow diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a767f5423ed..7d4a42cf025 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -181,6 +181,7 @@ from litellm.proxy.common_utils.token_count_settings import ( ) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields +from litellm.proxy.management_helpers.utils import management_endpoint_span_wrapper from litellm.repositories import database as database_runtime from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, @@ -12339,6 +12340,7 @@ async def _deployment_hidden_by_listing_callbacks(deployment: Deployment, user_a @router.get( "/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"] ) # if project requires model list +@management_endpoint_span_wrapper async def model_list( request: Request = None, # pyright: ignore[reportArgumentType] # FastAPI always injects the Request; the None default only serves direct in-process callers user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -12649,6 +12651,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 ebeb514ffff..1965fba94c6 100644 --- a/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py +++ b/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py @@ -3,10 +3,14 @@ async_management_endpoint_{success,failure}_hook integration points.""" import asyncio 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, @@ -340,3 +344,203 @@ def test_management_wrapper_swallows_post_success_errors(server_span_factory, ot result = asyncio.run(fake_fn(data={}, user_api_key_dict=_real_user_api_key_dict(server_span))) assert result == {"ok": True} + + +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 _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": method, + "path": path, + "root_path": "", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + }, + receive, + ) + + +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)], +) +def test_span_wrapper_omits_payload(server_span_factory, otel_with_exporter, monkeypatch, wrapper_name, exports_payload): + 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]: + return {"data": [{"id": "gpt-4o"}]} + + 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 + attribute_keys: Final = tuple(spans[0].attributes) + assert any(key.startswith("request.") for key in attribute_keys) is exports_payload + assert any(key.startswith("response.") for key in attribute_keys) is exports_payload + + +@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): + 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, None) + server_span: Final = _server_span(otel, MODEL_LIST_PATH, headers={"traceparent": INBOUND_TRACEPARENT}) + + 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_ended_in_inbound_trace(exporter, MODEL_LIST_PATH, "model_list") + + +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")