From 427f2ee1faaa6eed4ee51a57fe68be6beeb6547c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 11 Jun 2024 16:27:07 -0700 Subject: [PATCH] feat - otel log /team exceptions --- litellm/integrations/opentelemetry.py | 43 +++++++++++++++++++ .../http_parsing_utils.py | 0 .../management_endpoint_utils.py | 31 ++++++++++++- litellm/proxy/proxy_server.py | 6 ++- 4 files changed, 77 insertions(+), 3 deletions(-) rename litellm/proxy/{utils => common_utils}/http_parsing_utils.py (100%) rename litellm/proxy/{utils => common_utils}/management_endpoint_utils.py (60%) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 222639758..e41272c24 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -613,3 +613,46 @@ class OpenTelemetry(CustomLogger): management_endpoint_span.set_attribute(f"response.{key}", value) management_endpoint_span.set_status(Status(StatusCode.OK)) management_endpoint_span.end(end_time=_end_time_ns) + + async def async_management_endpoint_failure_hook( + self, + logging_payload: ManagementEndpointLoggingPayload, + parent_otel_span: Optional[Span] = None, + ): + from opentelemetry import trace + from datetime import datetime + from opentelemetry.trace import Status, StatusCode + + _start_time_ns = logging_payload.start_time + _end_time_ns = logging_payload.end_time + + start_time = logging_payload.start_time + end_time = logging_payload.end_time + + if isinstance(start_time, float): + _start_time_ns = int(int(start_time) * 1e9) + else: + _start_time_ns = self._to_ns(start_time) + + if isinstance(end_time, float): + _end_time_ns = int(int(end_time) * 1e9) + else: + _end_time_ns = self._to_ns(end_time) + + if parent_otel_span is not None: + _span_name = logging_payload.route + management_endpoint_span = self.tracer.start_span( + name=_span_name, + context=trace.set_span_in_context(parent_otel_span), + start_time=_start_time_ns, + ) + + _request_data = logging_payload.request_data + if _request_data is not None: + for key, value in _request_data.items(): + management_endpoint_span.set_attribute(f"request.{key}", value) + + _exception = logging_payload.exception + management_endpoint_span.set_attribute(f"exception", str(_exception)) + management_endpoint_span.set_status(Status(StatusCode.ERROR)) + management_endpoint_span.end(end_time=_end_time_ns) diff --git a/litellm/proxy/utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py similarity index 100% rename from litellm/proxy/utils/http_parsing_utils.py rename to litellm/proxy/common_utils/http_parsing_utils.py diff --git a/litellm/proxy/utils/management_endpoint_utils.py b/litellm/proxy/common_utils/management_endpoint_utils.py similarity index 60% rename from litellm/proxy/utils/management_endpoint_utils.py rename to litellm/proxy/common_utils/management_endpoint_utils.py index 33d996992..e8b00e6a8 100644 --- a/litellm/proxy/utils/management_endpoint_utils.py +++ b/litellm/proxy/common_utils/management_endpoint_utils.py @@ -1,7 +1,7 @@ from datetime import datetime from functools import wraps from litellm.proxy._types import UserAPIKeyAuth, ManagementEndpointLoggingPayload -from http_parsing_utils import _read_request_body +from common_utils.http_parsing_utils import _read_request_body from fastapi import Request @@ -20,6 +20,7 @@ def management_endpoint_wrapper(func): try: result = await func(*args, **kwargs) end_time = datetime.now() + if kwargs is None: kwargs = {} user_api_key_dict: UserAPIKeyAuth = kwargs.get("user_api_key_dict") @@ -52,6 +53,34 @@ def management_endpoint_wrapper(func): return result except Exception as e: end_time = datetime.now() + + if kwargs is None: + kwargs = {} + user_api_key_dict: UserAPIKeyAuth = kwargs.get("user_api_key_dict") + parent_otel_span = user_api_key_dict.parent_otel_span + if parent_otel_span is not None: + from litellm.proxy.proxy_server import open_telemetry_logger + + if open_telemetry_logger is not None: + _http_request: Request = kwargs.get("http_request") + _route = _http_request.url.path + _request_body: dict = await _read_request_body( + request=_http_request + ) + logging_payload = ManagementEndpointLoggingPayload( + route=_route, + request_data=_request_body, + response=None, + start_time=start_time, + end_time=end_time, + exception=e, + ) + + await open_telemetry_logger.async_management_endpoint_failure_hook( + logging_payload=logging_payload, + parent_otel_span=parent_otel_span, + ) + raise e return wrapper diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 640f09c73..ab7b61ca6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -114,7 +114,7 @@ from litellm.proxy.utils import ( _to_ns, get_error_message_str, ) -from litellm.proxy.utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm import ( CreateBatchRequest, @@ -164,7 +164,9 @@ from litellm.proxy.auth.auth_checks import ( get_actual_routes, log_to_opentelemetry, ) -from litellm.proxy.utils.management_endpoint_utils import management_endpoint_wrapper +from litellm.proxy.common_utils.management_endpoint_utils import ( + management_endpoint_wrapper, +) from litellm.llms.custom_httpx.httpx_handler import HTTPHandler from litellm.exceptions import RejectedRequestError from litellm.integrations.slack_alerting import SlackAlertingArgs, SlackAlerting