Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/rerank_endpoints/endpoints.py: 66%

1#### Rerank Endpoints ##### 

2 

3import asyncio 

4from typing import Final 

5 

6import orjson 

7from fastapi import APIRouter, Depends, HTTPException, Request, Response, status 

8from fastapi.responses import ORJSONResponse 

9 

10from litellm.proxy._types import * 

11from litellm.proxy.auth.user_api_key_auth import user_api_key_auth 

12from litellm.proxy.common_request_processing import ( 

13 ProxyBaseLLMRequestProcessing, 

14 log_llm_api_exception, 

15 resolve_litellm_call_id, 

16) 

17from litellm.proxy.common_utils.openai_error_payload import ( 

18 error_status_code, 

19 litellm_call_id_headers, 

20 openai_error_param, 

21 openai_error_type, 

22) 

23 

24router: Final = APIRouter() 

25 

26 

27@router.post( 

28 "/v2/rerank", 

29 dependencies=[Depends(user_api_key_auth)], 

30 response_class=ORJSONResponse, 

31 tags=["rerank"], 

32) 

33@router.post( 

34 "/v1/rerank", 

35 dependencies=[Depends(user_api_key_auth)], 

36 response_class=ORJSONResponse, 

37 tags=["rerank"], 

38) 

39@router.post( 

40 "/rerank", 

41 dependencies=[Depends(user_api_key_auth)], 

42 response_class=ORJSONResponse, 

43 tags=["rerank"], 

44) 

45async def rerank( 

46 request: Request, 

47 fastapi_response: Response, 

48 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 

49): 

50 from litellm.proxy.proxy_server import ( 

51 add_litellm_data_to_request, 

52 general_settings, 

53 llm_router, 

54 proxy_config, 

55 proxy_logging_obj, 

56 route_request, 

57 user_model, 

58 version, 

59 ) 

60 

61 litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) 

62 data = {"litellm_call_id": litellm_call_id} 

63 try: 

64 body: Final = await request.body() 

65 data = orjson.loads(body) | data 

66 

67 # Include original request and headers in the data 

68 data = await add_litellm_data_to_request( 

69 data=data, 

70 request=request, 

71 general_settings=general_settings, 

72 user_api_key_dict=user_api_key_dict, 

73 version=version, 

74 proxy_config=proxy_config, 

75 ) 

76 

77 ### CALL HOOKS ### - modify incoming data / reject request before calling the model 

78 data = await proxy_logging_obj.pre_call_hook(user_api_key_dict=user_api_key_dict, data=data, call_type="rerank") 

79 

80 ## ROUTE TO CORRECT ENDPOINT ## 

81 llm_call: Final = await route_request( 

82 data=data, 

83 route_type="arerank", 

84 llm_router=llm_router, 

85 user_model=user_model, 

86 ) 

87 response: Final = await llm_call 

88 

89 ### ALERTING ### 

90 asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) 

91 

92 ### RESPONSE HEADERS ### 

93 hidden_params: Final = getattr(response, "_hidden_params", {}) or {} 

94 model_id: Final = hidden_params.get("model_id", None) or "" 

95 cache_key: Final = hidden_params.get("cache_key", None) or "" 

96 api_base: Final = hidden_params.get("api_base", None) or "" 

97 additional_headers: Final = hidden_params.get("additional_headers", None) or {} 

98 fastapi_response.headers.update( 

99 ProxyBaseLLMRequestProcessing.get_custom_headers( 

100 user_api_key_dict=user_api_key_dict, 

101 call_id=hidden_params.get("litellm_call_id", None) or litellm_call_id, 

102 model_id=model_id, 

103 cache_key=cache_key, 

104 api_base=api_base, 

105 version=version, 

106 response_cost=hidden_params.get("response_cost", None), 

107 model_region=getattr(user_api_key_dict, "allowed_model_region", ""), 

108 request_data=data, 

109 hidden_params=hidden_params, 

110 **additional_headers, 

111 ) 

112 ) 

113 

114 return response 

115 except Exception as e: 

116 await proxy_logging_obj.post_call_failure_hook( 

117 user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data 

118 ) 

119 log_llm_api_exception(e, litellm_call_id) 

120 if isinstance(e, HTTPException): 120 ↛ 121line 120 didn't jump to line 121 because the condition on line 120 was never true

121 raise ProxyException( 

122 message=getattr(e, "message", str(e)), 

123 type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), 

124 param=openai_error_param(e), 

125 headers=litellm_call_id_headers(litellm_call_id), 

126 code=error_status_code(e, status.HTTP_400_BAD_REQUEST), 

127 ) 

128 else: 

129 error_msg: Final = f"{e}" 

130 raise ProxyException( 

131 message=getattr(e, "message", error_msg), 

132 type=openai_error_type(e, error_status_code(e, 500)), 

133 param=openai_error_param(e), 

134 headers=litellm_call_id_headers(litellm_call_id), 

135 code=error_status_code(e, 500), 

136 )