From b551462518d30adf26474f86a2add434fb067aa2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 6 Jun 2024 19:49:18 -0700 Subject: [PATCH] feat -v0 parent_otel_span in basic db reads --- litellm/proxy/auth/auth_checks.py | 5 +++++ litellm/proxy/proxy_server.py | 20 +++++++++++++++++--- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e4b8e6c8a8f..f184589f4f6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -21,6 +21,7 @@ from typing import Optional, Literal, Union from litellm.proxy.utils import PrismaClient from litellm.caching import DualCache import litellm +from opentelemetry.trace import Span all_routes = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value @@ -186,6 +187,7 @@ async def get_end_user_object( end_user_id: Optional[str], prisma_client: Optional[PrismaClient], user_api_key_cache: DualCache, + parent_otel_span: Optional[Span] = None, ) -> Optional[LiteLLM_EndUserTable]: """ Returns end user object, if in db. @@ -250,6 +252,7 @@ async def get_user_object( prisma_client: Optional[PrismaClient], user_api_key_cache: DualCache, user_id_upsert: bool, + parent_otel_span: Optional[Span] = None, ) -> Optional[LiteLLM_UserTable]: """ - Check if user id in proxy User Table @@ -300,6 +303,7 @@ async def get_team_object( team_id: str, prisma_client: Optional[PrismaClient], user_api_key_cache: DualCache, + parent_otel_span: Optional[Span] = None, ) -> LiteLLM_TeamTable: """ - Check if team id in proxy Team Table @@ -342,6 +346,7 @@ async def get_org_object( org_id: str, prisma_client: Optional[PrismaClient], user_api_key_cache: DualCache, + parent_otel_span: Optional[Span] = None, ): """ - Check if org id in proxy Org Table diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8cf2fa118d2..05eb5157391 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7,6 +7,7 @@ import secrets, subprocess import hashlib, uuid import warnings import importlib +from opentelemetry.trace import Span def showwarning(message, category, filename, lineno, file=None, line=None): @@ -398,6 +399,7 @@ disable_spend_logs = False jwt_handler = JWTHandler() prompt_injection_detection_obj: Optional[_OPTIONAL_PromptInjectionDetection] = None store_model_in_db: bool = False +open_telemetry_logger = None ### INITIALIZE GLOBAL LOGGING OBJECT ### proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) ### REDIS QUEUE ### @@ -498,7 +500,12 @@ async def user_api_key_auth( if isinstance(api_key, str): passed_in_key = api_key api_key = _get_bearer_token(api_key=api_key) - + parent_otel_span: Optional[Span] = None + if open_telemetry_logger is not None: + parent_otel_span = open_telemetry_logger.tracer.start_span( + name="Received Proxy Server Request", + start_time=time.time(), + ) ### USER-DEFINED AUTH FUNCTION ### if user_custom_auth is not None: response = await user_custom_auth(request=request, api_key=api_key) @@ -580,6 +587,7 @@ async def user_api_key_auth( team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, ) # [OPTIONAL] track spend for an org id - `LiteLLM_OrganizationTable` @@ -591,6 +599,7 @@ async def user_api_key_auth( org_id=org_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, ) # [OPTIONAL] track spend against an internal employee - `LiteLLM_UserTable` user_object = None @@ -604,6 +613,7 @@ async def user_api_key_auth( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=jwt_handler.is_upsert_user_id(), + parent_otel_span=parent_otel_span, ) # [OPTIONAL] track spend against an external user - `LiteLLM_EndUserTable` @@ -617,6 +627,7 @@ async def user_api_key_auth( end_user_id=end_user_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, ) global_proxy_spend = None @@ -715,6 +726,7 @@ async def user_api_key_auth( end_user_id=request_data["user"], prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, ) if _end_user_object is not None: end_user_params["allowed_model_region"] = ( @@ -2306,7 +2318,7 @@ class ProxyConfig: """ Load config values into proxy global state """ - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, use_queue, custom_db_client, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, use_queue, custom_db_client, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger # Load existing config config = await self.get_config(config_file_path=config_file_path) @@ -2430,7 +2442,9 @@ class ProxyConfig: OpenTelemetry, ) - imported_list.append(OpenTelemetry()) + open_telemetry_logger = OpenTelemetry() + + imported_list.append(open_telemetry_logger) elif isinstance(callback, str) and callback == "presidio": from litellm.proxy.hooks.presidio_pii_masking import ( _OPTIONAL_PresidioPIIMasking,