forked from BerriAI/litellm
-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request BerriAI#4138 from BerriAI/litellm_log_management_e…
…ndpoint_logs_otel [FEAT] log management endpoint logs to otel
- Loading branch information
Showing
7 changed files
with
265 additions
and
35 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,31 @@ | ||
| from typing import Optional | ||
| from fastapi import Request | ||
| import ast, json | ||
|
|
||
|
|
||
| async def _read_request_body(request: Optional[Request]) -> dict: | ||
| """ | ||
| Asynchronous function to read the request body and parse it as JSON or literal data. | ||
| Parameters: | ||
| - request: The request object to read the body from | ||
| Returns: | ||
| - dict: Parsed request data as a dictionary | ||
| """ | ||
| try: | ||
| request_data: dict = {} | ||
| if request is None: | ||
| return request_data | ||
| body = await request.body() | ||
|
|
||
| if body == b"" or body is None: | ||
| return request_data | ||
| body_str = body.decode() | ||
| try: | ||
| request_data = ast.literal_eval(body_str) | ||
| except: | ||
| request_data = json.loads(body_str) | ||
| return request_data | ||
| except: | ||
| return {} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,90 @@ | ||
| from datetime import datetime | ||
| from functools import wraps | ||
| from litellm.proxy._types import UserAPIKeyAuth, ManagementEndpointLoggingPayload | ||
| from litellm.proxy.common_utils.http_parsing_utils import _read_request_body | ||
| from fastapi import Request | ||
|
|
||
|
|
||
| def management_endpoint_wrapper(func): | ||
| """ | ||
| This wrapper does the following: | ||
| 1. Log I/O, Exceptions to OTEL | ||
| 2. Create an Audit log for success calls | ||
| """ | ||
|
|
||
| @wraps(func) | ||
| async def wrapper(*args, **kwargs): | ||
| start_time = datetime.now() | ||
|
|
||
| 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") or UserAPIKeyAuth() | ||
| ) | ||
| 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 | ||
| ) | ||
| _response = dict(result) if result is not None else None | ||
|
|
||
| logging_payload = ManagementEndpointLoggingPayload( | ||
| route=_route, | ||
| request_data=_request_body, | ||
| response=_response, | ||
| start_time=start_time, | ||
| end_time=end_time, | ||
| ) | ||
|
|
||
| await open_telemetry_logger.async_management_endpoint_success_hook( | ||
| logging_payload=logging_payload, | ||
| parent_otel_span=parent_otel_span, | ||
| ) | ||
|
|
||
| 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") or UserAPIKeyAuth() | ||
| ) | ||
| 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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.