From 13cb49688a4469c53e175c1b86ae44ec506a1458 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 11 Jun 2024 16:09:11 -0700 Subject: [PATCH 01/10] feat - log management endpoint logs to otel --- litellm/integrations/opentelemetry.py | 51 ++++++++++++++++++++ litellm/proxy/_types.py | 9 ++++ litellm/proxy/management_endpoint_utils.py | 56 ++++++++++++++++++++++ litellm/proxy/proxy_server.py | 20 +++++++- litellm/proxy/utils.py | 5 +- 5 files changed, 138 insertions(+), 3 deletions(-) create mode 100644 litellm/proxy/management_endpoint_utils.py diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index e18cadce23c..22263975800 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -6,17 +6,23 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm._logging import verbose_logger from litellm.types.services import ServiceLoggerPayload +from functools import wraps from typing import Union, Optional, TYPE_CHECKING, Any if TYPE_CHECKING: from opentelemetry.trace import Span as _Span from litellm.proxy.proxy_server import UserAPIKeyAuth as _UserAPIKeyAuth + from litellm.proxy._types import ( + ManagementEndpointLoggingPayload as _ManagementEndpointLoggingPayload, + ) Span = _Span UserAPIKeyAuth = _UserAPIKeyAuth + ManagementEndpointLoggingPayload = _ManagementEndpointLoggingPayload else: Span = Any UserAPIKeyAuth = Any + ManagementEndpointLoggingPayload = Any LITELLM_TRACER_NAME = os.getenv("OTEL_TRACER_NAME", "litellm") @@ -562,3 +568,48 @@ class OpenTelemetry(CustomLogger): self.OTEL_EXPORTER, ) return BatchSpanProcessor(ConsoleSpanExporter()) + + async def async_management_endpoint_success_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) + + _response = logging_payload.response + if _response is not None: + for key, value in _response.items(): + 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) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0d29f794f69..72065e5b970 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1561,3 +1561,12 @@ class SpanAttributes(str, enum.Enum): LLM_OPENAI_API_BASE = "gen_ai.openai.api_base" LLM_OPENAI_API_VERSION = "gen_ai.openai.api_version" LLM_OPENAI_API_TYPE = "gen_ai.openai.api_type" + + +class ManagementEndpointLoggingPayload(LiteLLMBase): + route: str + request_data: dict + response: Optional[dict] = None + exception: Optional[Any] = None + start_time: Optional[datetime] = None + end_time: Optional[datetime] = None diff --git a/litellm/proxy/management_endpoint_utils.py b/litellm/proxy/management_endpoint_utils.py new file mode 100644 index 00000000000..aabf6e5d3bb --- /dev/null +++ b/litellm/proxy/management_endpoint_utils.py @@ -0,0 +1,56 @@ +from datetime import datetime +from functools import wraps +from litellm.proxy._types import UserAPIKeyAuth, ManagementEndpointLoggingPayload +from litellm.proxy.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() + user_api_key_dict: UserAPIKeyAuth = kwargs["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 + ) + _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() + raise e + + return wrapper diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f76bf225793..247751f0f5d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -163,6 +163,7 @@ from litellm.proxy.auth.auth_checks import ( get_actual_routes, log_to_opentelemetry, ) +from litellm.proxy.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 @@ -9912,8 +9913,10 @@ async def delete_end_user( dependencies=[Depends(user_api_key_auth)], response_model=LiteLLM_TeamTable, ) +@management_endpoint_wrapper async def new_team( data: NewTeamRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header( None, @@ -10147,6 +10150,7 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): @router.post( "/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) +@management_endpoint_wrapper async def update_team( data: UpdateTeamRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -10252,8 +10256,10 @@ async def update_team( tags=["team management"], dependencies=[Depends(user_api_key_auth)], ) +@management_endpoint_wrapper async def team_member_add( data: TeamMemberAddRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10345,8 +10351,10 @@ async def team_member_add( tags=["team management"], dependencies=[Depends(user_api_key_auth)], ) +@management_endpoint_wrapper async def team_member_delete( data: TeamMemberDeleteRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10450,8 +10458,10 @@ async def team_member_delete( @router.post( "/team/delete", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) +@management_endpoint_wrapper async def delete_team( data: DeleteTeamRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header( None, @@ -10535,10 +10545,12 @@ async def delete_team( @router.get( "/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) +@management_endpoint_wrapper async def team_info( + http_request: Request, team_id: str = fastapi.Query( default=None, description="Team ID in the request parameters" - ) + ), ): """ get info on team + related keys @@ -10622,8 +10634,10 @@ async def team_info( @router.post( "/team/block", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) +@management_endpoint_wrapper async def block_team( data: BlockTeamRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10644,8 +10658,10 @@ async def block_team( @router.post( "/team/unblock", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) +@management_endpoint_wrapper async def unblock_team( data: BlockTeamRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10666,7 +10682,9 @@ async def unblock_team( @router.get( "/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) +@management_endpoint_wrapper async def list_team( + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6647881e292..3530b93347b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -48,6 +48,7 @@ from datetime import datetime, timedelta from litellm.integrations.slack_alerting import SlackAlerting from typing_extensions import overload from functools import wraps +from fastapi import Request if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -2595,7 +2596,7 @@ async def update_spend( raise e -async def _read_request_body(request): +async def _read_request_body(request: Optional[Request]) -> dict: """ Asynchronous function to read the request body and parse it as JSON or literal data. @@ -2608,7 +2609,7 @@ async def _read_request_body(request): import ast, json try: - request_data = {} + request_data: dict = {} if request is None: return request_data body = await request.body() From f2abac748e72b9cb5400272bdd2bca74f035377d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 11 Jun 2024 16:16:10 -0700 Subject: [PATCH 02/10] fix refactor management endpoint utils --- litellm/proxy/proxy_server.py | 5 +-- litellm/proxy/utils.py | 30 ------------------ litellm/proxy/utils/http_parsing_utils.py | 31 +++++++++++++++++++ .../{ => utils}/management_endpoint_utils.py | 7 +++-- 4 files changed, 38 insertions(+), 35 deletions(-) create mode 100644 litellm/proxy/utils/http_parsing_utils.py rename litellm/proxy/{ => utils}/management_endpoint_utils.py (90%) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 247751f0f5d..640f09c7313 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -103,7 +103,6 @@ from litellm.proxy.utils import ( hash_token, html_form, missing_keys_html_form, - _read_request_body, _is_valid_team_configs, _is_user_proxy_admin, _get_user_role, @@ -115,6 +114,8 @@ from litellm.proxy.utils import ( _to_ns, get_error_message_str, ) +from litellm.proxy.utils.http_parsing_utils import _read_request_body + from litellm import ( CreateBatchRequest, RetrieveBatchRequest, @@ -163,7 +164,7 @@ from litellm.proxy.auth.auth_checks import ( get_actual_routes, log_to_opentelemetry, ) -from litellm.proxy.management_endpoint_utils import management_endpoint_wrapper +from litellm.proxy.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 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3530b93347b..93790e76413 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2596,36 +2596,6 @@ async def update_spend( raise e -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 - """ - import ast, json - - 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 {} - - def _is_projected_spend_over_limit( current_spend: float, soft_budget_limit: Optional[float] ): diff --git a/litellm/proxy/utils/http_parsing_utils.py b/litellm/proxy/utils/http_parsing_utils.py new file mode 100644 index 00000000000..cc24713919c --- /dev/null +++ b/litellm/proxy/utils/http_parsing_utils.py @@ -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 {} diff --git a/litellm/proxy/management_endpoint_utils.py b/litellm/proxy/utils/management_endpoint_utils.py similarity index 90% rename from litellm/proxy/management_endpoint_utils.py rename to litellm/proxy/utils/management_endpoint_utils.py index aabf6e5d3bb..33d9969923a 100644 --- a/litellm/proxy/management_endpoint_utils.py +++ b/litellm/proxy/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 litellm.proxy.utils import _read_request_body +from http_parsing_utils import _read_request_body from fastapi import Request @@ -20,8 +20,9 @@ def management_endpoint_wrapper(func): try: result = await func(*args, **kwargs) end_time = datetime.now() - user_api_key_dict: UserAPIKeyAuth = kwargs["user_api_key_dict"] - + 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 From 427f2ee1faaa6eed4ee51a57fe68be6beeb6547c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 11 Jun 2024 16:27:07 -0700 Subject: [PATCH 03/10] 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 22263975800..e41272c2437 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 33d9969923a..e8b00e6a815 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 640f09c7313..ab7b61ca6ed 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 From 30f74aa4712b879602793c01cac25a7bc443c148 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 11 Jun 2024 17:17:57 -0700 Subject: [PATCH 04/10] fix import error --- litellm/proxy/common_utils/management_endpoint_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/common_utils/management_endpoint_utils.py b/litellm/proxy/common_utils/management_endpoint_utils.py index e8b00e6a815..3b5340df340 100644 --- a/litellm/proxy/common_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 common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from fastapi import Request From 479a39cebd9bd9ae5acf2a44edde9dcc5b227eae Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 11 Jun 2024 17:48:11 -0700 Subject: [PATCH 05/10] fix test new team --- litellm/proxy/proxy_server.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ab7b61ca6ed..0c444086468 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9919,7 +9919,7 @@ async def delete_end_user( @management_endpoint_wrapper async def new_team( data: NewTeamRequest, - http_request: Request, + http_request: Optional[Request] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header( None, @@ -10262,7 +10262,7 @@ async def update_team( @management_endpoint_wrapper async def team_member_add( data: TeamMemberAddRequest, - http_request: Request, + http_request: Optional[Request] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10357,7 +10357,7 @@ async def team_member_add( @management_endpoint_wrapper async def team_member_delete( data: TeamMemberDeleteRequest, - http_request: Request, + http_request: Optional[Request] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10464,7 +10464,7 @@ async def team_member_delete( @management_endpoint_wrapper async def delete_team( data: DeleteTeamRequest, - http_request: Request, + http_request: Optional[Request] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header( None, @@ -10550,7 +10550,7 @@ async def delete_team( ) @management_endpoint_wrapper async def team_info( - http_request: Request, + http_request: Optional[Request] = None, team_id: str = fastapi.Query( default=None, description="Team ID in the request parameters" ), @@ -10640,7 +10640,7 @@ async def team_info( @management_endpoint_wrapper async def block_team( data: BlockTeamRequest, - http_request: Request, + http_request: Optional[Request] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10664,7 +10664,7 @@ async def block_team( @management_endpoint_wrapper async def unblock_team( data: BlockTeamRequest, - http_request: Request, + http_request: Optional[Request] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10687,7 +10687,7 @@ async def unblock_team( ) @management_endpoint_wrapper async def list_team( - http_request: Request, + http_request: Optional[Request] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ From 51b78422df818518bbc1102ea44779ac26ec9fe7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 11 Jun 2024 17:57:50 -0700 Subject: [PATCH 06/10] fix http_request --- litellm/proxy/proxy_server.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0c444086468..ab7b61ca6ed 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9919,7 +9919,7 @@ async def delete_end_user( @management_endpoint_wrapper async def new_team( data: NewTeamRequest, - http_request: Optional[Request] = None, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header( None, @@ -10262,7 +10262,7 @@ async def update_team( @management_endpoint_wrapper async def team_member_add( data: TeamMemberAddRequest, - http_request: Optional[Request] = None, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10357,7 +10357,7 @@ async def team_member_add( @management_endpoint_wrapper async def team_member_delete( data: TeamMemberDeleteRequest, - http_request: Optional[Request] = None, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10464,7 +10464,7 @@ async def team_member_delete( @management_endpoint_wrapper async def delete_team( data: DeleteTeamRequest, - http_request: Optional[Request] = None, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header( None, @@ -10550,7 +10550,7 @@ async def delete_team( ) @management_endpoint_wrapper async def team_info( - http_request: Optional[Request] = None, + http_request: Request, team_id: str = fastapi.Query( default=None, description="Team ID in the request parameters" ), @@ -10640,7 +10640,7 @@ async def team_info( @management_endpoint_wrapper async def block_team( data: BlockTeamRequest, - http_request: Optional[Request] = None, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10664,7 +10664,7 @@ async def block_team( @management_endpoint_wrapper async def unblock_team( data: BlockTeamRequest, - http_request: Optional[Request] = None, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10687,7 +10687,7 @@ async def unblock_team( ) @management_endpoint_wrapper async def list_team( - http_request: Optional[Request] = None, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ From 7ad981b85c7631536e493cf7f3887efc969980bb Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 11 Jun 2024 18:05:10 -0700 Subject: [PATCH 07/10] fix - log mgmt data to OTEL --- litellm/proxy/proxy_server.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ab7b61ca6ed..32f0d3ae114 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9919,7 +9919,6 @@ async def delete_end_user( @management_endpoint_wrapper async def new_team( data: NewTeamRequest, - http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header( None, @@ -10262,7 +10261,6 @@ async def update_team( @management_endpoint_wrapper async def team_member_add( data: TeamMemberAddRequest, - http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10357,7 +10355,6 @@ async def team_member_add( @management_endpoint_wrapper async def team_member_delete( data: TeamMemberDeleteRequest, - http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10464,7 +10461,6 @@ async def team_member_delete( @management_endpoint_wrapper async def delete_team( data: DeleteTeamRequest, - http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header( None, @@ -10550,7 +10546,6 @@ async def delete_team( ) @management_endpoint_wrapper async def team_info( - http_request: Request, team_id: str = fastapi.Query( default=None, description="Team ID in the request parameters" ), @@ -10640,7 +10635,6 @@ async def team_info( @management_endpoint_wrapper async def block_team( data: BlockTeamRequest, - http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10664,7 +10658,6 @@ async def block_team( @management_endpoint_wrapper async def unblock_team( data: BlockTeamRequest, - http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10687,7 +10680,6 @@ async def unblock_team( ) @management_endpoint_wrapper async def list_team( - http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ From fba7d2e6233088bfd39129f4d1c76a6787e3779b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 11 Jun 2024 18:10:20 -0700 Subject: [PATCH 08/10] fix team endpoints --- litellm/proxy/proxy_server.py | 13 ++++++++++++- litellm/tests/test_key_generate_prisma.py | 11 ++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 32f0d3ae114..cd552385164 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8883,7 +8883,10 @@ async def new_user(data: NewUserRequest): role="user", user_email=data_json.get("user_email", None), ), - ) + ), + http_request=Request( + scope={"type": "http"}, + ), ) if data.send_invite_email is True: @@ -9919,6 +9922,7 @@ async def delete_end_user( @management_endpoint_wrapper async def new_team( data: NewTeamRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header( None, @@ -10261,6 +10265,7 @@ async def update_team( @management_endpoint_wrapper async def team_member_add( data: TeamMemberAddRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10355,6 +10360,7 @@ async def team_member_add( @management_endpoint_wrapper async def team_member_delete( data: TeamMemberDeleteRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10461,6 +10467,7 @@ async def team_member_delete( @management_endpoint_wrapper async def delete_team( data: DeleteTeamRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header( None, @@ -10546,6 +10553,7 @@ async def delete_team( ) @management_endpoint_wrapper async def team_info( + http_request: Request, team_id: str = fastapi.Query( default=None, description="Team ID in the request parameters" ), @@ -10635,6 +10643,7 @@ async def team_info( @management_endpoint_wrapper async def block_team( data: BlockTeamRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10658,6 +10667,7 @@ async def block_team( @management_endpoint_wrapper async def unblock_team( data: BlockTeamRequest, + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -10680,6 +10690,7 @@ async def unblock_team( ) @management_endpoint_wrapper async def list_team( + http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 277d9718f1a..aaef9d6ca86 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -137,6 +137,7 @@ async def test_new_user_response(prisma_client): NewTeamRequest( team_id=_team_id, ), + http_request=Request(scope={"type": "http"}), user_api_key_dict=UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", @@ -368,6 +369,7 @@ async def test_call_with_valid_model_using_all_models(prisma_client): new_team_response = await new_team( data=team_request, user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + http_request=Request(scope={"type": "http"}), ) print("new_team_response", new_team_response) created_team_id = new_team_response["team_id"] @@ -1086,6 +1088,7 @@ def test_generate_and_update_key(prisma_client): api_key="sk-1234", user_id="1234", ), + http_request=Request(scope={"type": "http"}), ) _team_2 = "ishaan-special-team_{}".format(uuid.uuid4()) @@ -1098,6 +1101,7 @@ def test_generate_and_update_key(prisma_client): api_key="sk-1234", user_id="1234", ), + http_request=Request(scope={"type": "http"}), ) request = NewUserRequest( @@ -2050,6 +2054,7 @@ async def test_master_key_hashing(prisma_client): api_key="sk-1234", user_id="1234", ), + http_request=Request(scope={"type": "http"}), ) _response = await new_user( @@ -2183,6 +2188,7 @@ async def test_create_update_team(prisma_client): tpm_limit=20, rpm_limit=20, ), + http_request=Request(scope={"type": "http"}), user_api_key_dict=UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", @@ -2232,7 +2238,10 @@ async def test_create_update_team(prisma_client): ) # now hit team_info - response = await team_info(team_id=_team_id) + response = await team_info( + team_id=_team_id, + http_request=Request(scope={"type": "http"}), + ) print("RESPONSE from team_info", response) From b10ac0d1f3a80ef41e5c096711b22bf1e60c7ac8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 11 Jun 2024 18:14:06 -0700 Subject: [PATCH 09/10] fix linting error --- litellm/proxy/proxy_server.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cd552385164..382c58388a8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8190,7 +8190,9 @@ async def _get_spend_report_for_time_range( return response, spend_per_tag except Exception as e: - verbose_proxy_logger.error("Exception in _get_daily_spend_reports", e) # noqa + verbose_proxy_logger.error( + "Exception in _get_daily_spend_reports {}".format(str(e)) + ) # noqa @router.post( From 995b1e73157aec39c12fde0ec84ed0997ae63f51 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 11 Jun 2024 18:33:15 -0700 Subject: [PATCH 10/10] fix when user_api_key_dict is None --- litellm/proxy/common_utils/management_endpoint_utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_utils/management_endpoint_utils.py b/litellm/proxy/common_utils/management_endpoint_utils.py index 3b5340df340..006f6aaae39 100644 --- a/litellm/proxy/common_utils/management_endpoint_utils.py +++ b/litellm/proxy/common_utils/management_endpoint_utils.py @@ -23,7 +23,9 @@ def management_endpoint_wrapper(func): if kwargs is None: kwargs = {} - user_api_key_dict: UserAPIKeyAuth = kwargs.get("user_api_key_dict") + 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 @@ -56,7 +58,9 @@ def management_endpoint_wrapper(func): if kwargs is None: kwargs = {} - user_api_key_dict: UserAPIKeyAuth = kwargs.get("user_api_key_dict") + 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