phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix - include litellm_parent_otel_span

Browse source
This commit is contained in:
Ishaan Jaff 2024-06-07 09:25:07 -07:00
parent 1f5c3b5448
commit c7dda09773
1 changed files with 24 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

29
litellm/proxy/proxy_server.py
View file

@ -549,7 +549,10 @@ async def user_api_key_auth(
litellm_proxy_roles=jwt_handler.litellm_jwtauth, litellm_proxy_roles=jwt_handler.litellm_jwtauth,
) )
if is_allowed: if is_allowed:
return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) return UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
parent_otel_span=parent_otel_span,
)
else: else:
allowed_routes = ( allowed_routes = (
jwt_handler.litellm_jwtauth.admin_allowed_routes jwt_handler.litellm_jwtauth.admin_allowed_routes
@ -666,7 +669,6 @@ async def user_api_key_auth(
user_info=user_info, user_info=user_info,
) )
) )
# get the request body # get the request body
request_data = await _read_request_body(request=request) request_data = await _read_request_body(request=request)
@ -695,15 +697,21 @@ async def user_api_key_auth(
user_role=LitellmUserRoles.INTERNAL_USER, user_role=LitellmUserRoles.INTERNAL_USER,
user_id=user_id, user_id=user_id,
org_id=org_id, org_id=org_id,
parent_otel_span=parent_otel_span,
) )
#### ELSE #### #### ELSE ####
if master_key is None: if master_key is None:
if isinstance(api_key, str): if isinstance(api_key, str):
return UserAPIKeyAuth( return UserAPIKeyAuth(
api_key=api_key, user_role=LitellmUserRoles.PROXY_ADMIN api_key=api_key,
user_role=LitellmUserRoles.PROXY_ADMIN,
parent_otel_span=parent_otel_span,
) )
else: else:
return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) return UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
parent_otel_span=parent_otel_span,
)
elif api_key is None: # only require api key if master key is set elif api_key is None: # only require api key if master key is set
raise Exception("No api key passed in.") raise Exception("No api key passed in.")
elif api_key == "": elif api_key == "":
@ -781,6 +789,7 @@ async def user_api_key_auth(
valid_token.allowed_model_region = end_user_params.get( valid_token.allowed_model_region = end_user_params.get(
"allowed_model_region" "allowed_model_region"
) )
valid_token.parent_otel_span = parent_otel_span
return valid_token return valid_token
@ -807,6 +816,7 @@ async def user_api_key_auth(
api_key=master_key, api_key=master_key,
user_role=LitellmUserRoles.PROXY_ADMIN, user_role=LitellmUserRoles.PROXY_ADMIN,
user_id=litellm_proxy_admin_name, user_id=litellm_proxy_admin_name,
parent_otel_span=parent_otel_span,
**end_user_params, **end_user_params,
) )
await user_api_key_cache.async_set_cache( await user_api_key_cache.async_set_cache(
@ -1452,6 +1462,7 @@ async def user_api_key_auth(
return UserAPIKeyAuth( return UserAPIKeyAuth(
api_key=api_key, api_key=api_key,
user_role=LitellmUserRoles.PROXY_ADMIN, user_role=LitellmUserRoles.PROXY_ADMIN,
parent_otel_span=parent_otel_span,
**valid_token_dict, **valid_token_dict,
) )
elif ( elif (
@ -1459,7 +1470,10 @@ async def user_api_key_auth(
and route in LiteLLMRoutes.sso_only_routes.value and route in LiteLLMRoutes.sso_only_routes.value
): ):
return UserAPIKeyAuth( return UserAPIKeyAuth(
api_key=api_key, user_role="app_owner", **valid_token_dict api_key=api_key,
user_role="app_owner",
parent_otel_span=parent_otel_span,
**valid_token_dict,
) )
else: else:
raise Exception( raise Exception(
@ -1475,18 +1489,21 @@ async def user_api_key_auth(
return UserAPIKeyAuth( return UserAPIKeyAuth(
api_key=api_key, api_key=api_key,
user_role=LitellmUserRoles.PROXY_ADMIN, user_role=LitellmUserRoles.PROXY_ADMIN,
parent_otel_span=parent_otel_span,
**valid_token_dict, **valid_token_dict,
) )
elif _has_user_setup_sso() and route in LiteLLMRoutes.sso_only_routes.value: elif _has_user_setup_sso() and route in LiteLLMRoutes.sso_only_routes.value:
return UserAPIKeyAuth( return UserAPIKeyAuth(
api_key=api_key, api_key=api_key,
user_role=LitellmUserRoles.INTERNAL_USER, user_role=LitellmUserRoles.INTERNAL_USER,
parent_otel_span=parent_otel_span,
**valid_token_dict, **valid_token_dict,
) )
else: else:
return UserAPIKeyAuth( return UserAPIKeyAuth(
api_key=api_key, api_key=api_key,
user_role=LitellmUserRoles.INTERNAL_USER, user_role=LitellmUserRoles.INTERNAL_USER,
parent_otel_span=parent_otel_span,
**valid_token_dict, **valid_token_dict,
) )
else: else:
@ -4193,6 +4210,8 @@ async def chat_completion(
) # do not store the original `sk-..` api key in the db ) # do not store the original `sk-..` api key in the db
data["metadata"]["headers"] = _headers data["metadata"]["headers"] = _headers
data["metadata"]["endpoint"] = str(request.url) data["metadata"]["endpoint"] = str(request.url)
# Add the OTEL Parent Trace before sending it LiteLLM
data["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span
### TEAM-SPECIFIC PARAMS ### ### TEAM-SPECIFIC PARAMS ###
if user_api_key_dict.team_id is not None: if user_api_key_dict.team_id is not None: