
    hl=                     L   d dl Z d dlZd dlmZ d dlmZmZmZmZm	Z	m
Z
mZ d dlZd dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d d	lmZ d d
lmZ d dlmZmZ d dlm Z m!Z! d dl" d dl#m$Z$m%Z%m&Z&  e       Z' e       Z( e       Z)e%	 	 	 	 	 dde*de*deee*ee*ef   f      de
e	d      de
e+   de
ee*      de
e,   de
e+   dee!eeee!f   f   fd       Z-e%	 	 	 	 	 	 dde*de*deee*ee*ef   f      de
e	d      de
e+   de
ee*      de
e,   de
e+   de
e+   dee!eeee!f   f   fd       Z.y)    N)partial)Any	CoroutineDictListLiteralOptionalUnion)verbose_logger)Logging)BaseRerankConfig)BedrockRerankHandler)BaseLLMHTTPHandler)TogetherAIRerank)get_optional_rerank_params)
get_secretget_secret_str)OptionalRerankParamsRerankResponse)*)ProviderConfigManagerclientexception_typemodelquery	documentscustom_llm_provider)coheretogether_aitop_nrank_fieldsreturn_documentsmax_chunks_per_docreturnc                 n  K   	 t        j                         }	d|d<   t        t        | |||||||f	i |}
t	        j
                         }t        |j                  |
      }|	j                  d|       d{   }t        j                  |      r| d{   }|S |}|S 7 )7 # t        $ r}|d}~ww xY ww)zR
    Async: Reranks a list of documents based on their relevance to the query
    TarerankN)
asyncioget_event_loopr   rerankcontextvarscopy_contextrunrun_in_executoriscoroutine	Exception)r   r   r   r   r    r!   r"   r#   kwargsloopfuncctxfunc_with_contextinit_responseresponsees                   S/var/www/Befach/backend/env/lib/python3.12/site-packages/litellm/rerank_api/main.pyr&   r&      s     %%' y
 
 &&(#CGGT2"2249JKK}-**H  %H L +  sY   B5A0B" 4B5B" B B" B5B" B5B"  B" "	B2+B--B22B5)r   r   azure_aiinfinitylitellm_proxyhosted_vllmmax_tokens_per_docc	                    |	j                  d      }
|	j                  d      }|	j                  dd      }|	j                  dd      }|	j                  dd      }|	j                  di       }|	j                  dd      }|	j                  d	d      }	 |	j                  d
d      du }t        d2i |	}||d}|j                         D cg c]
  \  }}|	| }}}t	        j
                  | ||j                  |j                        \  } }}}t        j                  | t	        j                  |      |j                  |      }t        || |	j                  d      xs t        j                  xs d|||||||||	      }t        |j                  t              rt!        |j                        |_        t#               }|j%                  | |t'        |      ||||di d|j)                  d      |       |dk(  s|dk(  r|xs |j                  xs t        j                  }|xs/ |j                  xs! t        j                  xs t+        d      xs d}|t-        d      t.        j1                  | |||||j                  ||||
xs t        j2                  xs i ||      } | S |dk(  r|xs+ |j                  xs t        j                  xs t+        d      }t.        j1                  | |||||j                  |xs |j                  |||
xs t        j2                  xs i ||      } | S |dk(  r|xs |j                  xs t        j                  }|xs+ |j                  xs t        j                  xs t5        d      }|t-        d       t.        j1                  | |||||j                  |xs |j                  |||
xs t        j2                  xs i ||      } | S |d!k(  rn|xs= |j                  xs/ t        j6                  xs t+        d"      xs t        j                  }|t9        d#      t:        j1                  | ||||||||$	      } | S |d%k(  r|t9        d&      |xs+ |j                  xs t        j                  xs t+        d'      }t.        j1                  | |||||j                  |xs |j                  |||
xs t        j2                  xs i ||(      } | S |d)k(  rb|xs+ |j                  xs t        j                  xs t+        d'      }t<        j1                  | ||||||||j)                  d      |||*      } | S |d+k(  r|xs |j                  xs t5        d,      }|xs |j                  xs t5        d-      }|t9        d.      t.        j1                  | |||||j                  ||||
xs t        j2                  xs i ||      } | S t        |t        j>                        st        |t        j@                        r|dk7  r|dk7  rt9        d/|       t.        j1                  | |||||j                  |xs |j                  |xs |j                  ||
xs t        j2                  xs i ||      } | S c c}}w # t,        $ r4}!tC        jD                  d0t        |!              tG        | ||!1      d}!~!ww xY w)3zK
    Reranks a list of documents based on their relevance to the query
    headerslitellm_logging_objlitellm_call_idNproxy_server_request
model_infometadatauserr   r&   FT)r#   r=   )r   r   api_baseapi_key)r   providerrF   present_version_paramsdrop_params)rerank_provider_configr   rJ   r   r   r   r    r!   r"   r#   r=   non_default_params)rA   rB   rC   rD   preset_cache_keystream_response)exclude_unset)r   rE   optional_paramslitellm_paramsr   r   r;   COHERE_API_BASEzhttps://api.cohere.comzNInvalid api base. api_base=None. Set in call or via `COHERE_API_BASE` env var.)r   r   provider_configoptional_rerank_paramslogging_objtimeoutrG   rF   	_is_asyncr?   r   model_responser9   AZURE_AI_API_BASE)r   r   rT   rS   rU   rV   rG   rF   rW   r?   r   rX   r:   INFINITY_API_BASEzPInvalid api base. api_base=None. Set in call or via `INFINITY_API_BASE` env var.r   TOGETHERAI_API_KEYzSTogetherAI API key is required, please set 'TOGETHERAI_API_KEY' in your environment)	r   r   r   r    r!   r"   r#   rG   rW   jina_aizMJina AI API key is required, please set 'JINA_AI_API_KEY' in your environmentBEDROCK_API_BASE)r   r   rT   rU   rS   rV   rG   rF   rW   r?   r   rX   bedrock)r   r   r   r    r!   r"   r#   rW   rP   rF   rU   r   r<   HOSTED_VLLM_API_KEYHOSTED_VLLM_API_BASEzbapi_base must be provided for Hosted VLLM rerank. Set in call or via HOSTED_VLLM_API_BASE env var.zUnsupported provider: zError in rerank: )r   r   original_exception )$getpopGenericLiteLLMParamsitemslitellmget_llm_providerrF   rG   r   get_provider_rerank_configLlmProvidersr   rJ   
isinstancerV   strfloatr   update_environment_variablesdict
model_dumpr   r/   base_llm_http_handlerr)   r?   r   togetherai_api_key
ValueErrortogether_rerankbedrock_rerankCohereRerankConfigCohereRerankV2Configr   errorr   )"r   r   r   r   r    r!   r"   r#   r=   r0   r?   r@   rA   rB   rC   rD   rE   r   rW   rP   unique_version_paramskvrI   _custom_llm_providerdynamic_api_keydynamic_api_baserK   rT   rX   rG   rF   r6   r7   s"                                     r8   r)   r)   H   s2   $ %jj3G-3ZZ8M-N%+ZZ0A4%HO!::&<dCL$/Jzz*b)H::fd#DZZ$'Ff
JJy%0D8	.88 #5"4!

 0557"
!Q1=A"
 "
 $$ 3$--#++	
	
  "<< --.BC(11'=	 	 8R#9

=1QW5H5HQE 4#-11%8
 o--s3&+O,C,C&DO#')88 !78#2(<($$(#% ",,4,@ !5 	9 	
"  8+/C/V  M?#:#:Mgoo 
 ! ,"++,##, /0, ,  d  -33$8 6'=/'//!#87??8b- 4 HR w "Z/  3"++3##3 12	  -33$8'= 6/'//'B?+B+B!#87??8b- 4 Hh M "Z/%S)@)@SGOOG ! 7"++7##7 ""56	  f  -33$8 6'=/'//'B?+B+B!#87??8b- 4 Hl Q "]2   #"**#--# 23# ??   i  '--#'!1#5# . 
Hr ] "Y.& c 
 ! 2"++2##2 01	  -33$8'=/ 6'//'B?+B+B!#87??8b- 4 HB g "Y.  2"++2##2 01	  &,,#'!1#5# / : : : N!/ - HV { "]2   9"**9!"78  ! :"++:!"89   x  -33$8 6'=/'//!#87??8b- 4 HT ) 173M3MN173O3OPVjnvVv  |P  Tc  |c #9:N9O!PQQ,33$8 6'=/'//'B?+B+B)E_-E-E#87??8b- 4 H  s"
t  
0Q9:-@UV
 	

sb   7Z 
ZZGZ "BZ 'B3Z A2Z BZ  A&Z BZ B*Z Z 	[/Z??[)NNNNN)NNNTNN)/r'   r*   	functoolsr   typingr   r   r   r   r   r	   r
   rg   litellm._loggingr   *litellm.litellm_core_utils.litellm_loggingr   LiteLLMLoggingObj+litellm.llms.base_llm.rerank.transformationr   #litellm.llms.bedrock.rerank.handlerr   *litellm.llms.custom_httpx.llm_http_handlerr   'litellm.llms.together_ai.rerank.handlerr   litellm.rerank_api.rerank_utilsr   litellm.secret_managers.mainr   r   litellm.types.rerankr   r   litellm.types.routerlitellm.utilsr   r   r   rt   ru   rq   rl   intboolr&   r)   rb       r8   <module>r      s      G G G  + S H D I D F C E " G G #$%'*,  
 GK'+'+(,))) E#tCH~-./) "'*A"BC	)
 C=) $s)$) tn) !) >9S#~%=>>?) )X  	'+'+(,(,


 E#tCH~-./
 "_`	
 C=
 $s)$
 tn
 !
 !
 >9S#~%=>>?
 
r   