Merge 7fad1f5bbe into a392a39627

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

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),

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")